Skip to content

Commit

Permalink
examples: update Tello examples for main thread friendly macOS/Window…
Browse files Browse the repository at this point in the history 
…s, add Tello face tracker

Signed-off-by: Ron Evans <[email protected]>
  • Loading branch information
@deadprogram
deadprogram committed Jul 10, 2018
1 parent 84d4d0e commit 4409ee7
Show file tree
Hide file tree
Showing 3 changed files with 338 additions and 33 deletions.
300 changes: 300 additions & 0 deletions examples/tello_facetracker.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,300 @@
// +build example
//
// Do not build by default.

/*
You must have ffmpeg and OpenCV installed in order to run this code. It will connect to the Tello
and then open a window using OpenCV showing the streaming video.
How to run
go run examples/tello_facetracker.go ~/Downloads/res10_300x300_ssd_iter_140000.caffemodel ~/Development/opencv/samples/dnn/face_detector/deploy.prototxt
*/

package main

import (
"fmt"
"image"
"image/color"
"io"
"math"
"os"
"os/exec"
"strconv"
"time"

"gobot.io/x/gobot"
"gobot.io/x/gobot/platforms/dji/tello"
"gobot.io/x/gobot/platforms/joystick"
"gocv.io/x/gocv"
)

const (
maxJoyVal = 32768
frameX = 400
frameY = 300
frameSize = frameX * frameY * 3
)

// ffmpeg command to decode video stream from drone
var ffmpeg = exec.Command("ffmpeg", "-hwaccel", "auto", "-hwaccel_device", "opencl", "-i", "pipe:0",
"-pix_fmt", "bgr24", "-s", strconv.Itoa(frameX)+"x"+strconv.Itoa(frameY), "-f", "rawvideo", "pipe:1")
var ffmpegIn, _ = ffmpeg.StdinPipe()
var ffmpegOut, _ = ffmpeg.StdoutPipe()

// gocv
var window = gocv.NewWindow("Tello")
var net *gocv.Net
var green = color.RGBA{0, 255, 0, 0}

// tracking
var tracking = false
var detected = false
var detectSize = false
var distTolerance = 0.05 * dist(0, 0, frameX, frameY)
var refDistance float64
var left float64
var top float64
var right float64
var bottom float64

// drone
var drone = tello.NewDriver("8890")
var flightData *tello.FlightData

// joystick
var joyAdaptor = joystick.NewAdaptor()
var stick = joystick.NewDriver(joyAdaptor, "dualshock4")

func init() {
// process joystick events in main goroutine to keep SDL happy
handleJoystick()

// process drone events in separate goroutine for concurrency
go func() {
if err := ffmpeg.Start(); err != nil {
fmt.Println(err)
return
}

drone.On(tello.FlightDataEvent, func(data interface{}) {
// TODO: protect flight data from race condition
flightData = data.(*tello.FlightData)
})

drone.On(tello.ConnectedEvent, func(data interface{}) {
fmt.Println("Connected")
drone.StartVideo()
drone.SetVideoEncoderRate(tello.VideoBitRateAuto)
drone.SetExposure(0)
gobot.Every(100*time.Millisecond, func() {
drone.StartVideo()
})
})

drone.On(tello.VideoFrameEvent, func(data interface{}) {
pkt := data.([]byte)
if _, err := ffmpegIn.Write(pkt); err != nil {
fmt.Println(err)
}
})

robot := gobot.NewRobot("tello",
[]gobot.Connection{joyAdaptor},
[]gobot.Device{drone, stick},
)

robot.Start()
}()
}

func main() {
if len(os.Args) < 5 {
fmt.Println("How to run:\ngo run facetracker.go [model] [config] ([backend] [device])")
return
}

model := os.Args[1]
config := os.Args[2]
backend := gocv.NetBackendDefault
if len(os.Args) > 3 {
backend = gocv.ParseNetBackend(os.Args[3])
}

target := gocv.NetTargetCPU
if len(os.Args) > 4 {
target = gocv.ParseNetTarget(os.Args[4])
}

n := gocv.ReadNet(model, config)
if n.Empty() {
fmt.Printf("Error reading network model from : %v %v\n", model, config)
return
}
net = &n
defer net.Close()
net.SetPreferableBackend(gocv.NetBackendType(backend))
net.SetPreferableTarget(gocv.NetTargetType(target))

for {
// get next frame from stream
buf := make([]byte, frameSize)
if _, err := io.ReadFull(ffmpegOut, buf); err != nil {
fmt.Println(err)
continue
}
img, _ := gocv.NewMatFromBytes(frameY, frameX, gocv.MatTypeCV8UC3, buf)
if img.Empty() {
continue
}

trackFace(&img)

window.IMShow(img)
if window.WaitKey(10) >= 0 {
break
}
}
}

func trackFace(frame *gocv.Mat) {
W := float64(frame.Cols())
H := float64(frame.Rows())

blob := gocv.BlobFromImage(*frame, 1.0, image.Pt(300, 300), gocv.NewScalar(104, 177, 123, 0), false, false)
defer blob.Close()

net.SetInput(blob, "data")

detBlob := net.Forward("detection_out")
defer detBlob.Close()

detections := gocv.GetBlobChannel(detBlob, 0, 0)
defer detections.Close()

for r := 0; r < detections.Rows(); r++ {
confidence := detections.GetFloatAt(r, 2)
if confidence < 0.5 {
continue
}

left = float64(detections.GetFloatAt(r, 3)) * W
top = float64(detections.GetFloatAt(r, 4)) * H
right = float64(detections.GetFloatAt(r, 5)) * W
bottom = float64(detections.GetFloatAt(r, 6)) * H

left = math.Min(math.Max(0.0, left), W-1.0)
right = math.Min(math.Max(0.0, right), W-1.0)
bottom = math.Min(math.Max(0.0, bottom), H-1.0)
top = math.Min(math.Max(0.0, top), H-1.0)

detected = true
rect := image.Rect(int(left), int(top), int(right), int(bottom))
gocv.Rectangle(frame, rect, green, 3)
}

if !tracking || !detected {
return
}

if detectSize {
detectSize = false
refDistance = dist(left, top, right, bottom)
}

distance := dist(left, top, right, bottom)

// x axis
switch {
case right < W/2:
drone.CounterClockwise(50)
case left > W/2:
drone.Clockwise(50)
default:
drone.Clockwise(0)
}

// y axis
switch {
case top < H/10:
drone.Up(25)
case bottom > H-H/10:
drone.Down(25)
default:
drone.Up(0)
}

// z axis
switch {
case distance < refDistance-distTolerance:
drone.Forward(20)
case distance > refDistance+distTolerance:
drone.Backward(20)
default:
drone.Forward(0)
}
}

func dist(x1, y1, x2, y2 float64) float64 {
return math.Sqrt((x2-x1)*(x2-x1) + (y2-y1)*(y2-y1))
}

func handleJoystick() {
stick.On(joystick.CirclePress, func(data interface{}) {
drone.Forward(0)
drone.Up(0)
drone.Clockwise(0)
tracking = !tracking
if tracking {
detectSize = true
println("tracking")
} else {
detectSize = false
println("not tracking")
}
})
stick.On(joystick.SquarePress, func(data interface{}) {
fmt.Println("battery:", flightData.BatteryPercentage)
})
stick.On(joystick.TrianglePress, func(data interface{}) {
drone.TakeOff()
println("Takeoff")
})
stick.On(joystick.XPress, func(data interface{}) {
drone.Land()
println("Land")
})
stick.On(joystick.RightY, func(data interface{}) {
val := float64(data.(int16))
if val >= 0 {
drone.Backward(tello.ValidatePitch(val, maxJoyVal))
} else {
drone.Forward(tello.ValidatePitch(val, maxJoyVal))
}
})
stick.On(joystick.RightX, func(data interface{}) {
val := float64(data.(int16))
if val >= 0 {
drone.Right(tello.ValidatePitch(val, maxJoyVal))
} else {
drone.Left(tello.ValidatePitch(val, maxJoyVal))
}
})
stick.On(joystick.LeftY, func(data interface{}) {
val := float64(data.(int16))
if val >= 0 {
drone.Down(tello.ValidatePitch(val, maxJoyVal))
} else {
drone.Up(tello.ValidatePitch(val, maxJoyVal))
}
})
stick.On(joystick.LeftX, func(data interface{}) {
val := float64(data.(int16))
if val >= 0 {
drone.Clockwise(tello.ValidatePitch(val, maxJoyVal))
} else {
drone.CounterClockwise(tello.ValidatePitch(val, maxJoyVal))
}
})
}
57 changes: 31 additions & 26 deletions examples/tello_opencv.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,49 +17,35 @@ import (
"fmt"
"io"
"os/exec"
"strconv"
"time"

"gobot.io/x/gobot"
"gobot.io/x/gobot/platforms/dji/tello"
"gobot.io/x/gobot/platforms/opencv"
"gocv.io/x/gocv"
)

const (
frameSize = 960 * 720 * 3
frameX = 960
frameY = 720
frameSize = frameX * frameY * 3
)

func main() {
drone := tello.NewDriver("8890")
window := opencv.NewWindowDriver()
window := gocv.NewWindow("Tello")

ffmpeg := exec.Command("ffmpeg", "-hwaccel", "auto", "-hwaccel_device", "opencl", "-i", "pipe:0",
"-pix_fmt", "bgr24", "-s", strconv.Itoa(frameX)+"x"+strconv.Itoa(frameY), "-f", "rawvideo", "pipe:1")
ffmpegIn, _ := ffmpeg.StdinPipe()
ffmpegOut, _ := ffmpeg.StdoutPipe()

work := func() {
ffmpeg := exec.Command("ffmpeg", "-i", "pipe:0", "-pix_fmt", "bgr24", "-vcodec", "rawvideo",
"-an", "-sn", "-s", "960x720", "-f", "rawvideo", "pipe:1")
ffmpegIn, _ := ffmpeg.StdinPipe()
ffmpegOut, _ := ffmpeg.StdoutPipe()
if err := ffmpeg.Start(); err != nil {
fmt.Println(err)
return
}

go func() {
for {
buf := make([]byte, frameSize)
if _, err := io.ReadFull(ffmpegOut, buf); err != nil {
fmt.Println(err)
continue
}

img := gocv.NewMatFromBytes(720, 960, gocv.MatTypeCV8UC3, buf)
if img.Empty() {
continue
}
window.ShowImage(img)
window.WaitKey(1)
}
}()

drone.On(tello.ConnectedEvent, func(data interface{}) {
fmt.Println("Connected")
drone.StartVideo()
Expand All @@ -81,9 +67,28 @@ func main() {

robot := gobot.NewRobot("tello",
[]gobot.Connection{},
[]gobot.Device{drone, window},
[]gobot.Device{drone},
work,
)

robot.Start()
// calling Start(false) lets the Start routine return immediately without an additional blocking goroutine
robot.Start(false)

// now handle video frames from ffmpeg stream in main thread, to be macOS/Windows friendly
for {
buf := make([]byte, frameSize)
if _, err := io.ReadFull(ffmpegOut, buf); err != nil {
fmt.Println(err)
continue
}
img, _ := gocv.NewMatFromBytes(frameY, frameX, gocv.MatTypeCV8UC3, buf)
if img.Empty() {
continue
}

window.IMShow(img)
if window.WaitKey(1) >= 0 {
break
}
}
}
Loading

0 comments on commit 4409ee7

Please sign in to comment.