-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
83 lines (67 loc) · 2.16 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
"""
APPLYING CV TO LIVE VIDEO STREAM
- if video error turn off public firewall
"""
import sys
import traceback
import tellopy
import av
import cv2.cv2 as cv2
import time
import torch
from Actions import *
filename = "C:/Users/McKinstryJohn/Desktop/Python/Tello/vids/output.avi"
codec = cv2.VideoWriter_fourcc("X", "V", "I", "D")
frame_rate = 30
resolution = (640, 480)
Output = cv2.VideoWriter(filename, codec, frame_rate, resolution)
trace_print = False
def main():
drone = tellopy.Tello()
start = time.time()
try:
drone.connect()
drone.wait_for_connection(60.0)
retry = 3
container = None
while container is None and 0 < retry:
retry -= 1
try:
container = av.open(drone.get_video_stream())
except av.AVError as ave:
if trace_print:
print(ave)
print('retry...')
frame_skip = 300 # skip first N frames
drone.takeoff()
while time.time() - start < 20:
for frame in container.decode(video=0):
if 0 < frame_skip:
frame_skip = frame_skip - 1
continue
start_time = time.time()
frame_ = model(frame.to_image())
# xmin ymin xmax ymax confidence class name
find(drone, frame_.pandas().xyxy[0], obj)
# press ESC to end stream
if cv2.waitKey(1) == 27:
break
if frame.time_base < 1.0 / 60:
time_base = 1.0 / 60
else:
time_base = frame.time_base
frame_skip = int((time.time() - start_time) / time_base)
except Exception as ex:
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback)
if trace_print:
print(ex)
finally:
drone.quit()
cv2.destroyAllWindows()
Output.release()
if __name__ == '__main__':
obj = 'person'
# v5x6 is their current best
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
main()