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