class TrackerManager(object): def __init__(self, vid_list, init_time): self.init_time = init_time video_id, camera_id, max_frames, width, height = vid_list[1:] self.max_frames = max_frames self.tracker = Tracker(init_time, video_id, max_frames, camera_id, width, height) self.postprocess_trans = get_postprocess_trans(height, width) self.prev_img = None self.n = 0 def process_output(self, dets): dets = post_process(dets, self.postprocess_trans)[0] self.tracker.step(dets) self.n += 1 def is_done(self): return self.n >= self.max_frames def finalize(self): self.tracker.finalize()
def tracker_thread_fn(q_in, init_time, path, debug=0, new_thresh=0.4, track_thresh=0.2): video_id, camera_id, max_frames, width, height = get_video_params(path) postprocess_trans = get_postprocess_trans(height, width) tracker = Tracker(init_time, video_id, max_frames, camera_id, width, height, new_thresh=new_thresh, track_thresh=track_thresh) for i in range(max_frames): dets = q_in.get() get_time = time.time() for k in dets: dets[k] = dets[k].detach().cpu().numpy() dets = post_process(dets, postprocess_trans, track_thresh=track_thresh)[0] tracker.step(dets) if debug > 0 and i % 100 == 99: frame_time = time.time() - init_time FPS = (i + 1) / frame_time print("At frame {} FPS {}".format(i + 1, FPS), file=sys.stderr) tracker.finalize()
class VideoManager(object): def __init__(self, path, vid_list, model_loading_time): init_time = time.time() - model_loading_time self.init_time = init_time vid_filename = vid_list[0] video_path = os.path.join(path, vid_filename) self.cap = cv2.VideoCapture(video_path) video_id, camera_id, max_frames, width, height = vid_list[1:] self.max_frames = max_frames self.tracker = Tracker(init_time, video_id, max_frames, camera_id, width, height) self.preprocess_function = get_img_transform(height, width, new_size=512) self.postprocess_trans = get_postprocess_trans(height, width) region_mask = get_region_mask(camera_id, height, width) self.region_mask = np.where(region_mask, 255, 0).astype(np.uint8) self.prev_img = None self.n = 0 def get_img(self): ret, frame = self.cap.read() frame = cv2.bitwise_and(frame, frame, mask=self.region_mask) img = self.preprocess_function(frame) img = torch.from_numpy(img).to(torch.device('cuda')) self.n += 1 prev_img = self.prev_img if self.prev_img is not None else img self.prev_img = img return img, prev_img def process_output(self, dets): dets = post_process(dets, self.postprocess_trans)[0] self.tracker.step(dets) def is_done(self): return self.n >= self.max_frames def finalize(self): self.tracker.finalize()
def run_single_video_serial(path, debug=0, full_precision=False): init_time = time.time() if debug >= 1: print("Starting for video: {}".format(path), file=sys.stderr) video_id, camera_id, max_frames, width, height = get_video_params(path) cap = cv2.VideoCapture(path) model = create_model() model = load_model(model, 'checkpoints/coco_tracking.pth') model.to(torch.device('cuda')) model.eval() tracker = Tracker(init_time, video_id, max_frames, camera_id, width, height, debug=debug) preprocess_function = get_img_transform(height, width, new_size=512) postprocess_trans = get_postprocess_trans(height, width) region_mask = get_region_mask(camera_id, height, width) region_mask = np.where(region_mask, 255, 0).astype(np.uint8) if debug > 2: cv2.imwrite("mask.png", region_mask) pre_img = None for i in range(max_frames): ret, frame = cap.read() if debug >= 2: cv2.imshow("Frame", frame) cv2.waitKey(1) tracker.frame = np.copy(frame) frame = cv2.bitwise_and(frame, frame, mask=region_mask) img = preprocess_function(frame) img = torch.from_numpy(img).to(torch.device('cuda')) if pre_img is None: pre_img = img with torch.no_grad(): with torch.cuda.amp.autocast(enabled=not full_precision): out = model(img, pre_img, None)[-1] out = sigmoid_output(out) dets = generic_decode(out) pre_img = img for k in dets: dets[k] = dets[k].detach().cpu().numpy() dets = post_process(dets, postprocess_trans)[0] tracker.step(dets) if debug >= 1 and i % 100 == 99: frame_time = time.time() - init_time FPS = (i + 1) / frame_time print("At frame {} FPS {}".format(i + 1, FPS), file=sys.stderr) tracker.finalize() if debug >= 1: print("Finished video: {}".format(path), file=sys.stderr)