You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
275 lines
9.6 KiB
275 lines
9.6 KiB
import time
|
|
import datetime
|
|
import os
|
|
import numpy
|
|
import av
|
|
import cv2
|
|
from ball_tracker import ball_tracker
|
|
from myTello import myTello
|
|
from cv2 import dnn
|
|
from face_tracker import face_tracker
|
|
|
|
class TelloCV(object):
|
|
"""
|
|
TelloTracker builds keyboard controls on top of TelloPy as well
|
|
as generating images from the video stream and enabling opencv support
|
|
"""
|
|
|
|
def __init__(self,Tello:myTello):
|
|
self.prev_flight_data = None
|
|
self.record = False
|
|
self.tracking = False
|
|
self.keydown = False
|
|
self.date_fmt = '%Y-%m-%d_%H%M%S'
|
|
self.speed = 20
|
|
self.speed2 = 10
|
|
self.drone = Tello
|
|
# self.init_drone()
|
|
|
|
self.caffe_prototxt_path = "../Face_Distance/model/RFB-320.prototxt"
|
|
self.caffe_model_path = "../Face_Distance/model/RFB-320.caffemodel"
|
|
self.net = dnn.readNetFromCaffe(self.caffe_prototxt_path, self.caffe_model_path)
|
|
|
|
# container for processing the packets into frames
|
|
# self.container = av.open(self.drone.get_video_stream())
|
|
# self.vid_stream = self.drone.container.streams.video[0]
|
|
self.out_file = None
|
|
self.out_stream = None
|
|
self.out_name = None
|
|
self.start_time = time.time()
|
|
|
|
# tracking a color
|
|
green_lower = (30, 50, 50)
|
|
green_upper = (80, 255, 255)
|
|
#red_lower = (0, 50, 50)
|
|
# red_upper = (20, 255, 255)
|
|
# blue_lower = (110, 50, 50)
|
|
# upper_blue = (130, 255, 255)
|
|
self.track_cmd = ""
|
|
self.ball_tracker = ball_tracker(self.drone.height,
|
|
self.drone.width,
|
|
green_lower, green_upper)
|
|
|
|
self.face_tracker = face_tracker(self.drone.height,
|
|
self.drone.width,)
|
|
|
|
# def init_drone(self):
|
|
# """Connect, uneable streaming and subscribe to events"""
|
|
# # self.drone.log.set_level(2)
|
|
# # self.drone.connect()
|
|
# # self.drone.start_video()
|
|
# self.drone.subscribe(self.drone.EVENT_FLIGHT_DATA,
|
|
# self.flight_data_handler)
|
|
# self.drone.subscribe(self.drone.EVENT_FILE_RECEIVED,
|
|
# self.handle_flight_received)
|
|
|
|
|
|
def ball_process_frame(self, frame):
|
|
"""convert frame to cv2 image and show"""
|
|
image = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
|
# image = self.write_hud(image)
|
|
if self.record:
|
|
self.record_vid(frame)
|
|
|
|
xoff, yoff, zoff = self.ball_tracker.track(image)
|
|
image = self.ball_tracker.draw_arrows(image)
|
|
|
|
distance = 100
|
|
cmd = ""
|
|
if self.tracking:
|
|
if xoff < -distance:
|
|
cmd = "counter_clockwise"
|
|
elif xoff > distance:
|
|
cmd = "clockwise"
|
|
elif yoff < -distance:
|
|
cmd = "down"
|
|
elif yoff > distance:
|
|
cmd = "up"
|
|
elif zoff < -25:
|
|
cmd = "backward"
|
|
elif zoff > 25:
|
|
cmd = "forward"
|
|
else:
|
|
if self.track_cmd is not "":
|
|
getattr(self.drone, self.track_cmd)(0)
|
|
self.track_cmd = ""
|
|
|
|
if cmd is not self.track_cmd:
|
|
if cmd is not "":
|
|
if cmd == "forward" or cmd == "backward":
|
|
print("track command forward:", cmd)
|
|
getattr(self.drone, cmd)(self.speed2)
|
|
self.track_cmd = cmd
|
|
else:
|
|
print("track command other:", cmd)
|
|
getattr(self.drone, cmd)(self.speed)
|
|
self.track_cmd = cmd
|
|
return image
|
|
|
|
def face_process_frame(self, frame):
|
|
"""convert frame to cv2 image and show"""
|
|
image = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
|
# image = self.write_hud(image)
|
|
# cv2.imshow("test",image)
|
|
# cv2.waitKey(1)
|
|
if self.record:
|
|
self.record_vid(frame)
|
|
|
|
xoff, yoff, zoff = self.face_tracker.track(image, self.net)
|
|
# image = self.tracker_face.draw_arrows(image)
|
|
#print(zoff)
|
|
|
|
distance = 80
|
|
cmd = ""
|
|
# if count!=0:
|
|
# if self.track_cmd is not "":
|
|
# getattr(self.drone, self.track_cmd)(0)
|
|
# self.track_cmd = ""
|
|
# return image
|
|
if self.tracking:
|
|
if xoff < -distance:
|
|
cmd = "counter_clockwise"
|
|
elif xoff > distance:
|
|
cmd = "clockwise"
|
|
elif yoff < -distance:
|
|
cmd = "down"
|
|
elif yoff > distance:
|
|
cmd = "up"
|
|
elif zoff < -15:
|
|
cmd = "backward"
|
|
# self.backtimes = self.backtimes + 1
|
|
# print("backtimes={}".format(self.backtimes))
|
|
elif zoff > 15:
|
|
cmd = "forward"
|
|
# self.fortimes = self.fortimes + 1
|
|
# print("fortimes={}".format(self.fortimes))
|
|
else:
|
|
if self.track_cmd is not "":
|
|
getattr(self.drone, self.track_cmd)(0)
|
|
self.track_cmd = ""
|
|
|
|
if cmd is not self.track_cmd:
|
|
if cmd is not "":
|
|
if cmd == "forward" or cmd == "backward":
|
|
print("track command forward:", cmd)
|
|
getattr(self.drone, cmd)(self.speed2)
|
|
self.track_cmd = cmd
|
|
else:
|
|
print("track command other:", cmd)
|
|
getattr(self.drone, cmd)(self.speed)
|
|
self.track_cmd = cmd
|
|
|
|
return image
|
|
|
|
|
|
def write_hud(self, frame):
|
|
"""Draw drone info, tracking and record on frame"""
|
|
stats = self.prev_flight_data.split('|')
|
|
stats.append("Tracking:" + str(self.tracking))
|
|
if self.drone.zoom:
|
|
stats.append("VID")
|
|
else:
|
|
stats.append("PIC")
|
|
if self.record:
|
|
diff = int(time.time() - self.start_time)
|
|
mins, secs = divmod(diff, 60)
|
|
stats.append("REC {:02d}:{:02d}".format(mins, secs))
|
|
|
|
for idx, stat in enumerate(stats):
|
|
text = stat.lstrip()
|
|
cv2.putText(frame, text, (0, 30 + (idx * 30)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.0, (255, 0, 0), lineType=30)
|
|
return frame
|
|
|
|
def toggle_recording(self, speed):
|
|
"""Handle recording keypress, creates output stream and file"""
|
|
if speed == 0:
|
|
return
|
|
self.record = not self.record
|
|
|
|
if self.record:
|
|
datename = [os.getenv('HOME'), datetime.datetime.now().strftime(self.date_fmt)]
|
|
self.out_name = '{}/Pictures/tello-{}.mp4'.format(*datename)
|
|
print("Outputting video to:", self.out_name)
|
|
self.out_file = av.open(self.out_name, 'w')
|
|
self.start_time = time.time()
|
|
self.out_stream = self.out_file.add_stream(
|
|
'mpeg4', self.drone.vid_stream.rate)
|
|
self.out_stream.pix_fmt = 'yuv420p'
|
|
self.out_stream.width = self.drone.vid_stream.width
|
|
self.out_stream.height = self.drone.vid_stream.height
|
|
|
|
if not self.record:
|
|
print("Video saved to ", self.out_name)
|
|
self.out_file.close()
|
|
self.out_stream = None
|
|
|
|
def record_vid(self, frame):
|
|
"""
|
|
convert frames to packets and write to file
|
|
"""
|
|
new_frame = av.VideoFrame(
|
|
width=frame.width, height=frame.height, format=frame.format.name)
|
|
for i in range(len(frame.planes)):
|
|
new_frame.planes[i].update(frame.planes[i])
|
|
pkt = None
|
|
try:
|
|
pkt = self.out_stream.encode(new_frame)
|
|
except IOError as err:
|
|
print("encoding failed: {0}".format(err))
|
|
if pkt is not None:
|
|
try:
|
|
self.out_file.mux(pkt)
|
|
except IOError:
|
|
print('mux failed: ' + str(pkt))
|
|
|
|
def take_picture(self, speed):
|
|
"""Tell drone to take picture, image sent to file handler"""
|
|
if speed == 0:
|
|
return
|
|
self.drone.take_picture()
|
|
|
|
def palm_land(self, speed):
|
|
"""Tell drone to land"""
|
|
if speed == 0:
|
|
return
|
|
self.drone.palm_land()
|
|
|
|
def toggle_tracking(self, speed):
|
|
""" Handle tracking keypress"""
|
|
if speed == 0: # handle key up event
|
|
return
|
|
self.tracking = not self.tracking
|
|
print("tracking:", self.tracking)
|
|
return
|
|
|
|
def toggle_zoom(self, speed):
|
|
"""
|
|
In "video" mode the self.drone sends 1280x720 frames.
|
|
In "photo" mode it sends 2592x1936 (952x720) frames.
|
|
The video will always be centered in the window.
|
|
In photo mode, if we keep the window at 1280x720 that gives us ~160px on
|
|
each side for status information, which is ample.
|
|
Video mode is harder because then we need to abandon the 16:9 display size
|
|
if we want to put the HUD next to the video.
|
|
"""
|
|
if speed == 0:
|
|
return
|
|
self.drone.set_video_mode(not self.drone.zoom)
|
|
|
|
def flight_data_handler(self, event, sender, data):
|
|
"""Listener to flight data from the drone."""
|
|
text = str(data)
|
|
if self.prev_flight_data != text:
|
|
self.prev_flight_data = text
|
|
|
|
def handle_flight_received(self, event, sender, data):
|
|
"""Create a file in ~/Pictures/ to receive image from the drone"""
|
|
path = '%s/Pictures/tello-%s.jpeg' % (
|
|
os.getenv('HOME'),
|
|
datetime.datetime.now().strftime(self.date_fmt))
|
|
with open(path, 'wb') as out_file:
|
|
out_file.write(data)
|
|
print('Saved photo to %s' % path)
|