streamtest.py

"""
Stream test:
Pull the video from the drone and display in cv2 window.
Optionally encode video and dump to file.
"""
import av
import numpy
import tellopy
import cv2


def encode(frame, ovstream, output):
    """
    convert frames to packets and write to file
    """
    try:
        pkt = ovstream.encode(frame)
    except Exception as err:
        print("encoding failed{0}".format(err))

    if pkt is not None:
        try:
            output.mux(pkt)
        except Exception:
            print('mux failed: ' + str(pkt))
    return True


def main():
    # Set up tello streaming
    drone = tellopy.Tello()
    drone.log.set_level(2)
    drone.connect()
    drone.start_video()

    # container for processing the packets into frames
    container = av.open(drone.get_video_stream())
    video_st = container.streams.video[0]

    # stream and outputfile for video
    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
    save = True
    for packet in container.demux((video_st,)):
        for frame in packet.decode():
            # convert frame to cv2 image and show
            image = cv2.cvtColor(numpy.array(
                frame.to_image()), cv2.COLOR_RGB2BGR)
            cv2.imshow('frame', image)
            key = cv2.waitKey(1) & 0xFF

            # save initial 1300 frames
            if save:
                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 > 300:
                    output.close()
                    save == False

if __name__ == '__main__':
    main()