def __init__(self, gui_mode):
bg_im = None
self.gui_mode_pyqt = (args.gui_mode == 'pyqt')
self.gui_mode_opencv = (args.gui_mode == 'opencv')
if self.gui_mode_pyqt:
self.qtui = OpenTrajUI(reserve_n_agents=100)
self.agent_index = -1
class Play:
def __init__(self, gui_mode):
bg_im = None
self.gui_mode_pyqt = (args.gui_mode == 'pyqt')
self.gui_mode_opencv = (args.gui_mode == 'opencv')
if self.gui_mode_pyqt:
self.qtui = OpenTrajUI(reserve_n_agents=100)
self.agent_index = -1
def is_a_video(self, filename):
return '.mp4' in filename or '.avi' in filename
def to_image_frame(self, Hinv, loc):
"""
Given H^-1 and world coordinates, returns (u, v) in image coordinates.
"""
if loc.ndim > 1:
locHomogenous = np.hstack((loc, np.ones((loc.shape[0], 1))))
loc_tr = np.transpose(locHomogenous)
loc_tr = np.matmul(Hinv, loc_tr) # to camera frame
locXYZ = np.transpose(loc_tr /
loc_tr[2]) # to pixels (from millimeters)
return locXYZ[:, :2].astype(int)
else:
locHomogenous = np.hstack((loc, 1))
locHomogenous = np.dot(Hinv, locHomogenous) # to camera frame
locXYZ = locHomogenous / locHomogenous[
2] # to pixels (from millimeters)
return locXYZ[:2].astype(int)
def set_background_im(self, im, timestamp=-1):
self.bg_im = im.copy()
if self.gui_mode_pyqt:
self.qtui.update_im(im)
self.qtui.erase_paths()
self.qtui.erase_circles()
if timestamp >= 0:
if self.gui_mode_pyqt:
self.qtui.setTimestamp(timestamp)
elif self.gui_mode_opencv:
cv2.putText(self.bg_im, '%d' % timestamp, (30, 80),
cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
def draw_trajectory(self, id, ll, color, width):
if self.gui_mode_pyqt:
self.qtui.draw_path(ll[..., ::-1], color, [width])
elif self.gui_mode_opencv:
for tt in range(ll.shape[0] - 1):
cv2.line(self.bg_im, (ll[tt][1], ll[tt][0]),
(ll[tt + 1][1], ll[tt + 1][0]), color, width)
def draw_agent(self, id, pos, radius, color, width):
if self.gui_mode_pyqt:
self.qtui.draw_circle(pos, radius, color, width)
elif self.gui_mode_opencv:
cv2.circle(self.bg_im, (pos[1], pos[0]), radius, color, width)
def play(self, traj_dataset, Hinv, media_file):
timestamps = sorted(traj_dataset.__t_p_dict__.keys())
if os.path.exists(media_file):
if self.is_a_video(media_file):
cap = cv2.VideoCapture(media_file)
else:
ref_im = cv2.imread(media_file)
ids_t = []
# for t in range(timestamps[0], timestamps[-1]):
# for t in range(0, timestamps[-1]):
t = 0
pause = False
while True:
if self.is_a_video(media_file) and not pause:
ret, ref_im = cap.read()
ref_im_copy = np.copy(ref_im)
self.set_background_im(ref_im, t)
if t in timestamps:
xys_t = traj_dataset.__t_p_dict__[t]
ids_t = traj_dataset.__t_id_dict__[t]
for i, id in enumerate(ids_t):
xy_i = np.array(xys_t[i])
UV_i = self.to_image_frame(Hinv, xy_i)
# fetch entire trajectory
traj_i = traj_dataset.__id_p_dict__[id]
TRAJ_i = self.to_image_frame(Hinv, traj_i)
self.draw_trajectory(id, TRAJ_i, (255, 255, 0), 2)
self.draw_agent(id, (UV_i[1], UV_i[0]), 5, (0, 0, 255), 2)
# if not ids_t:
# print('No agent')
if not pause and t < timestamps[-1]:
t += 1
delay_ms = 20
if self.gui_mode_pyqt:
self.qtui.processEvents()
pause = self.qtui.pause
time.sleep(delay_ms / 1000.)
else:
# cv2.namedWindow('OpenTraj (Press ESC for exit)', cv2.WINDOW_NORMAL)
# cv2.imshow('OpenTraj (Press ESC for exit)', ref_im_copy)
key = cv2.waitKey(delay_ms * (1 - pause)) & 0xFF
if key == 27: # press ESCAPE to quit
break
elif key == 32: # press SPACE to pause/play
pause = not pause