(args)
| 158 | |
| 159 | |
| 160 | def main(args): |
| 161 | cap = cv2.VideoCapture(args.video_input) |
| 162 | frame_rate = int(round(cap.get(cv2.CAP_PROP_FPS))) |
| 163 | video_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 164 | video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| 165 | |
| 166 | |
| 167 | demo_images_path = os.path.join(args.demo_output, 'demo_images') |
| 168 | if not os.path.exists(demo_images_path): |
| 169 | os.makedirs(demo_images_path) |
| 170 | |
| 171 | device = torch.device(args.device) |
| 172 | model, _, postprocessors = build_tracktest_model(args) |
| 173 | model.to(device) |
| 174 | model.eval() |
| 175 | tracker = Tracker(score_thresh=args.track_thresh) |
| 176 | |
| 177 | checkpoint = torch.load(args.resume, map_location='cpu') |
| 178 | _, _ = model.load_state_dict(checkpoint['model'], strict=False) |
| 179 | print("Model is loaded") |
| 180 | |
| 181 | |
| 182 | mean = [0.485, 0.456, 0.406] |
| 183 | std = [0.229, 0.224, 0.225] |
| 184 | color_list = colormap() |
| 185 | |
| 186 | |
| 187 | print("Starting inference") |
| 188 | count = 0 |
| 189 | tracker.reset_all() |
| 190 | pre_embed = None |
| 191 | res, img = cap.read() |
| 192 | |
| 193 | |
| 194 | while res: |
| 195 | count += 1 |
| 196 | resized_img, nh, nw = resize(img) |
| 197 | rgb_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB) |
| 198 | tensor_img = F.normalize(F.to_tensor(rgb_img), mean, std) |
| 199 | samples = nested_tensor_from_tensor_list([tensor_img]).to(device) |
| 200 | outputs, pre_embed = model(samples, pre_embed) |
| 201 | |
| 202 | orig_sizes = torch.stack([torch.as_tensor([video_height, video_width])], dim=0).to(device) |
| 203 | results = postprocessors['bbox'](outputs, orig_sizes) |
| 204 | |
| 205 | if count == 1: |
| 206 | res_track = tracker.init_track(results[0]) |
| 207 | else: |
| 208 | res_track = tracker.step(results[0]) |
| 209 | |
| 210 | for ret in res_track: |
| 211 | if ret['active'] == 0: |
| 212 | continue |
| 213 | bbox = ret['bbox'] |
| 214 | tracking_id = ret['tracking_id'] |
| 215 | |
| 216 | cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color_list[tracking_id%79].tolist(), thickness=2) |
| 217 | cv2.putText(img, "{}".format(tracking_id), (int(bbox[0]), int(bbox[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color_list[tracking_id%79].tolist(), 2) |
no test coverage detected