-
Notifications
You must be signed in to change notification settings - Fork 44
/
object_track.py
53 lines (46 loc) · 4.79 KB
/
object_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
import numpy as np
import cv2
from tracker import Tracker
import time
import imageio
images = []
def createimage(w,h):
size = (w, h, 1)
img = np.ones((w,h,3),np.uint8)*255
return img
def main():
data = np.array(np.load('Detections.npy'))[0:10,0:150,0:150]
tracker = Tracker(150, 30, 5)
skip_frame_count = 0
track_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
(127, 127, 255), (255, 0, 255), (255, 127, 255),
(127, 0, 255), (127, 0, 127),(127, 10, 255), (0,255, 127)]
for i in range(data.shape[1]):
centers = data[:,i,:]
frame = createimage(512,512)
if (len(centers) > 0):
tracker.update(centers)
for j in range(len(tracker.tracks)):
if (len(tracker.tracks[j].trace) > 1):
x = int(tracker.tracks[j].trace[-1][0,0])
y = int(tracker.tracks[j].trace[-1][0,1])
tl = (x-10,y-10)
br = (x+10,y+10)
cv2.rectangle(frame,tl,br,track_colors[j],1)
cv2.putText(frame,str(tracker.tracks[j].trackId), (x-10,y-20),0, 0.5, track_colors[j],2)
for k in range(len(tracker.tracks[j].trace)):
x = int(tracker.tracks[j].trace[k][0,0])
y = int(tracker.tracks[j].trace[k][0,1])
cv2.circle(frame,(x,y), 3, track_colors[j],-1)
cv2.circle(frame,(x,y), 6, track_colors[j],-1)
cv2.circle(frame,(int(data[j,i,0]),int(data[j,i,1])), 6, (0,0,0),-1)
cv2.imshow('image',frame)
# cv2.imwrite("image"+str(i)+".jpg", frame)
# images.append(imageio.imread("image"+str(i)+".jpg"))
time.sleep(0.1)
if cv2.waitKey(1) & 0xFF == ord('q'):
cv2.destroyAllWindows()
break
# imageio.mimsave('Multi-Object-Tracking.gif', images, duration=0.08)
if __name__ == '__main__':
main()