-
Notifications
You must be signed in to change notification settings - Fork 2
/
state_track.py
94 lines (70 loc) · 2.53 KB
/
state_track.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
84
85
86
87
88
89
90
91
92
93
94
"Tracker state as part of the PTZ tracker finite state machine"
import time
import numpy as np
import cv2
from state_id import async_id
def in_track_fn(system):
print('=== tracking')
# start tracker on identified frame
success = system.start_tracker()
# set timer
local_timer = time.time()
timeout = 0.5
# estimated moving average parameters
estimated_drone_prob = 1.0
alpha = 0.1
while success:
zoom = system.get_telemetry()[2]
move(system, zoom)
# get next frame
frame = system.get_frame()
if frame is None:
continue
# update tracker
success, system.drone_bbox = system.tracker.update(frame)
if not success:
system.camera.ptz.stop()
system.camera.ptz.zoomto(0)
break
# Draw bounding box
x, y, w, h = system.drone_bbox
cv_im = frame.copy()
p1 = (int(x), int(y))
p2 = (int(x + w), int(y + h))
cv2.rectangle(cv_im, p1, p2, (255, 0, 0), 2, 1)
system.update_gui(frame=cv_im)
# if time is left, identify object in bbox
if time.time() - local_timer > timeout:
status = async_id(system, frame)
# exponential weighted moving average
estimated_drone_prob = (
1 - alpha) * estimated_drone_prob + alpha * status
# reset timer
local_timer = time.time()
# drone is probably lost
if estimated_drone_prob < 0.5:
break
# user quit inititated
if system.gui.RESET or system.gui.ABORT:
system.gui.RESET = False
break
system.fsm.lost_track()
def move(system, zoom):
# control camera
x, y, w, h = system.drone_bbox
center = (x + w / 2, y + h / 2)
pan_error, tilt_error = system.camera.errors_pt(center,
system.camera.width,
system.camera.height)
zoom_error = system.camera.error_zoom(w, system.camera.height)
system.camera.control(pan_error=pan_error / scale_factor(zoom),
tilt_error=tilt_error / scale_factor(zoom))
system.camera.control_zoom(zoom_error)
def scale_factor(zoom_value):
factor = 0.1214 * np.exp(318.16 * 10**(-6) * zoom_value) + 1.0605
return factor * 1.3
def out_track_fn(system):
print('lost_track')
# reinitialize short-term tracker
system.tracker = cv2.TrackerCSRT_create()
system.drone_bbox = None