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()