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