Ball tracking via drone.

streamerwithyolo.py

@@ -0,0 +1,59 @@
+import av
+import numpy
+import tellopy
+from cv2 import cv2
+import caffe
+
+def videoFrameHandler(event, sender, data):
+    with open('tmpvid', 'wb') as w:
+        w.write(data)
+
+def encode(frame, ovstream, output):
+    try:
+        pkt = ovstream.encode(frame)
+    except Exception:
+        return False
+    if pkt is not None:
+        try:
+            output.mux(pkt)
+        except Exception:
+            print('mux failed: ' + str(pkt))
+    return True
+
+def main():
+    # drone = tellopy.Tello()
+    # drone.log.set_level(2)
+    # drone.connect()
+    # drone.start_video()
+    # drone.subscribe(drone.EVENT_VIDEO_FRAME, videoFrameHandler)
+
+    container = av.open('ball_tracking_example.mp4')
+    # container = av.open(drone.get_video_stream())
+    video_st = container.streams.video[0]
+    output = av.open('archive.mp4', 'w')
+    ovstream = output.add_stream('mpeg4', video_st.rate)
+    ovstream.pix_fmt = 'yuv420p'
+    ovstream.width = video_st.width
+    ovstream.height = video_st.height
+
+    net = caffe.Net('mobilenet-yolov3.prototxt','mobilenet_yolov3_lite_deploy.caffemodel')
+
+    counter = 0
+    for packet in container.demux((video_st,)):
+        for frame in packet.decode():
+            image = cv2.cvtColor(numpy.array(frame.to_image()), cv2.COLOR_RGB2BGR)
+            print(type(image))
+            cv2.imshow('frame', image)
+
+            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])
+            encode(new_frame, ovstream, output)
+            counter += 1
+            print("Frames encoded:", counter)
+        if counter > 500:
+            output.close()
+            break
+
+if __name__ == '__main__':
+    main()

streamtest.py

@@ -11,6 +11,7 @@ def videoFrameHandler(event, sender, data):
     with open('tmpvid', 'wb') as w:
         w.write(data)
 
+
 def main():
     videostream = io.BytesIO()
     videostream.seek(0)
@@ -21,18 +22,20 @@ def main():
     drone.start_video()
     drone.subscribe(drone.EVENT_VIDEO_FRAME, videoFrameHandler)
 
-    #container = av.open('ball_tracking_example.mp4')
-    container = av.open('tmp.mp4', mode='w')
+    #
+    in_container = av.open(drone.get_video_stream())
+    out_container = av.open('tmp.mp4', mode='w')
     # in tello.py it is set to 0x20
-    stream = container.add_stream('mpeg4', rate=32)
+    out_stream = out_container.add_stream('mpeg4', rate=32)
     while True:
-        for frame in container.decode(video=0):
-            image = cv2.cvtColor(numpy.array(frame.to_image()), cv2.COLOR_RGB2BGR)
-
-            print("image", type(image))
-            cv2.imshow('frame',image)
-            if cv2.waitKey(20) & 0xFF == ord('q'):
-                break
+        for packet in in_container.demux((in_container.streams.video[0],)):
+            for frame in packet.decode():
+                image = cv2.cvtColor(numpy.array(frame.to_image()), cv2.COLOR_RGB2BGR)
+
+                print("image", type(image))
+                cv2.imshow('frame',image)
+                if cv2.waitKey(20) & 0xFF == ord('q'):
+                    break
         print(len(videostream))
 
 if __name__ == '__main__':

tellotracker.py

@@ -23,6 +23,9 @@
 
 import time
 import sys
+
+import imutils
+import numpy
 import tellopy
 import pygame
 import pygame.display
@@ -31,11 +34,14 @@
 import pygame.font
 import os
 import datetime
+import av
+from cv2 import cv2
+
+from tracker import Tracker
 from subprocess import Popen, PIPE
 # from tellopy import logger
 
 # log = tellopy.logger.Logger('TelloUI')
-
 prev_flight_data = None
 video_player = None
 video_recorder = None
@@ -80,8 +86,52 @@ def palm_land(drone, speed):
         return
     drone.palm_land()
 
+
 def toggle_tracking(drone, speed):
     global tracking
+    container = av.open(drone.get_video_stream())
+    video_st = container.streams.video[0]
+    green_lower = (40, 100, 100)
+    green_upper = (75, 255, 255)
+
+    for packet in container.demux((video_st,)):
+        for frame in packet.decode():
+
+            image = cv2.cvtColor(numpy.array(frame.to_image()), cv2.COLOR_RGB2BGR)
+            greentracker = Tracker(image, green_lower, green_upper)
+            greentracker.track()
+            greentracker.show()
+
+            # In progress drone tracking
+            # offsets = greentracker.get_offsets()
+            # mids = greentracker.get_mids()
+            #
+            # time.sleep(0.01)
+            # for e in pygame.event.get():
+            #     if e.type == pygame.locals.KEYDOWN:
+            #         print('+' + pygame.key.name(e.key))
+            # if (mids[0]-offsets[0] < -20 and mids[1]-offsets[1] < -20):
+            #     print("moving left and going up")
+            #     # drone.clockwise(speed)
+            # elif (mids[0]-offsets[0] < -20 and mids[1]-offsets[1] > -20):
+            #     print("moving left and going down")
+            #     # drone.clockwise(speed)
+            # elif (mids[0]-offsets[0] > 20 and mids[1]-offsets[1] < -20):
+            #     print("moving right and going up")
+            #     # drone.counter_clockwise(speed)
+            # elif (mids[0] - offsets[0] > 20 and mids[1] - offsets[1] < -20):
+            #     print("moving right and going down")
+            #     # drone.counter_clockwise(speed)
+            # elif (mids[0] - offsets[0] > 20 ):
+            #     print("moving right")
+            #     # drone.counter_clockwise(speed)
+            # elif (mids[0] - offsets[0] < -20 ):
+            #     print("moving left")
+            #     # drone.counter_clockwise(speed)
+
+
+
+
     if speed == 0: # handle key up event
         return
     tracking = not(tracking)
@@ -152,7 +202,10 @@ def update(self, drone, data):
         return self._surface
 
 def flight_data_mode(drone, *args):
-    return (drone.zoom and "VID" or "PIC")
+    # error received on the below line - drone object doesn't have attribute zoom.
+    # return (drone.zoom and "VID" or "PIC")
+    return ("VID" or "PIC")
+
 
 def tracker_mode(drone, *args):
     if tracking:
@@ -204,6 +257,7 @@ def flightDataHandler(event, sender, data):
 def videoFrameHandler(event, sender, data):
     global video_player
     global video_recorder
+    global container
     # print(len(data))
     if video_player is None:
         cmd = [ 'mplayer', '-fps', '35', '-really-quiet' ]
@@ -252,9 +306,12 @@ def main():
     drone.log.set_level(2)
     drone.connect()
     drone.start_video()
+
     drone.subscribe(drone.EVENT_FLIGHT_DATA, flightDataHandler)
     drone.subscribe(drone.EVENT_VIDEO_FRAME, videoFrameHandler)
-    drone.subscribe(drone.EVENT_FILE_RECEIVED, handleFileReceived)
+    # Error received on the below line as drone object don't have EVENT_FILE_RECEIVED attribute
+    # drone.subscribe(drone.EVENT_FILE_RECEIVED, handleFileReceived)
+
     speed = 30
 
     try:

tracker.py

@@ -43,9 +43,10 @@ def main():
 
     # define the lower and upper boundaries of the "green"
     # ball in the HSV color space. NB the hue range in 
-    # opencv is 180, normally it is 360 
-    green_lower = (50, 50, 50)
-    green_upper = (70,255,255)
+    # opencv is 180, normally it is 360
+    # changed green_lower and upper values
+    green_lower = (40, 20, 25)
+    green_upper = (55,255,120)
     red_lower = (0, 50, 50) 
     red_upper = (20,255,255)
     blue_lower = (110, 50, 50)
@@ -85,21 +86,40 @@ def main():
     cv2.destroyAllWindows()
 
 class Tracker:
-    def __init__(self, vs, stream, color_lower, color_upper):
-        self.vs = vs
-        self.stream = stream
+
+    #changed the constructor for integrating with tellotracker.
+    def __init__(self, frame, color_lower, color_upper):
         self.color_lower = color_lower
         self.color_upper = color_upper
-        self.next_frame = True
-        self.frame = None
-        self.get_frame()
-        height, width, depth = self.frame.shape
-        self.width = width
-        self.height = height
-        self.midx = int(width / 2)
-        self.midy = int(height / 2)
+        self.width = frame.shape[0]
+        self.height = frame.shape[1]
+        self.midx = int(self.width / 2)
+        self.midy = int(self.height / 2)
         self.xoffset = 0
         self.yoffset = 0
+        self.frame = frame
+
+    # def __init__(self, vs, stream, color_lower, color_upper):
+    #     self.vs = vs
+    #     self.stream = stream
+    #     self.color_lower = color_lower
+    #     self.color_upper = color_upper
+    #     self.next_frame = True
+    #     self.frame = None
+    #     self.get_frame()
+    #     height, width, depth = self.frame.shape
+    #     self.width = width
+    #     self.height = height
+    #     self.midx = int(width / 2)
+    #     self.midy = int(height / 2)
+    #     self.xoffset = 0
+    #     self.yoffset = 0
+
+    def get_mids(self):
+        return (self.midx,self.midy)
+
+    def get_offsets(self):
+        return (self.xoffset, self.yoffset)
 
     def get_frame(self):
         # grab the current frame

trystream.py

@@ -0,0 +1,59 @@
+import av
+import numpy
+import tellopy
+from cv2 import cv2
+
+
+def videoFrameHandler(event, sender, data):
+    with open('tmpvid', 'wb') as w:
+        w.write(data)
+
+def encode(frame, ovstream, output):
+    try:
+        pkt = ovstream.encode(frame)
+    except Exception:
+        return False
+    if pkt is not None:
+        try:
+            output.mux(pkt)
+        except Exception:
+            print('mux failed: ' + str(pkt))
+    return True
+
+def main():
+    drone = tellopy.Tello()
+    drone.log.set_level(2)
+    drone.connect()
+    drone.start_video()
+    drone.subscribe(drone.EVENT_VIDEO_FRAME, videoFrameHandler)
+
+    # container = av.open('ball_tracking_example.mp4')
+    container = av.open(drone.get_video_stream())
+    video_st = container.streams.video[0]
+    output = av.open('archive.mp4', 'w')
+    ovstream = output.add_stream('mpeg4', video_st.rate)
+    ovstream.pix_fmt = 'yuv420p'
+    ovstream.width = video_st.width
+    ovstream.height = video_st.height
+
+    counter = 0
+    for packet in container.demux((video_st,)):
+        for frame in packet.decode():
+
+            image = cv2.cvtColor(numpy.array(frame.to_image()), cv2.COLOR_RGB2BGR)
+            print(type(image))
+            # cv2.imshow('frame', image)
+            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])
+            cv2.imshow('frame', image)
+            key = cv2.waitKey(1) & 0xFF
+            encode(new_frame, ovstream, output)
+            counter += 1
+            print("Frames encoded:", counter)
+        if counter > 1300:
+            output.close()
+            break
+
+if __name__ == '__main__':
+    main()