(self, video_file, camera_id)
| 441 | return results |
| 442 | |
| 443 | def predict_video(self, video_file, camera_id): |
| 444 | video_out_name = 'output.mp4' |
| 445 | if camera_id != -1: |
| 446 | capture = cv2.VideoCapture(camera_id) |
| 447 | else: |
| 448 | capture = cv2.VideoCapture(video_file) |
| 449 | video_out_name = os.path.split(video_file)[-1] |
| 450 | # Get Video info : resolution, fps, frame count |
| 451 | width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 452 | height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| 453 | fps = int(capture.get(cv2.CAP_PROP_FPS)) |
| 454 | frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 455 | print("fps: %d, frame_count: %d" % (fps, frame_count)) |
| 456 | |
| 457 | if not os.path.exists(self.output_dir): |
| 458 | os.makedirs(self.output_dir) |
| 459 | out_path = os.path.join(self.output_dir, video_out_name) |
| 460 | fourcc = cv2.VideoWriter_fourcc(*'mp4v') |
| 461 | writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) |
| 462 | index = 1 |
| 463 | while (1): |
| 464 | ret, frame = capture.read() |
| 465 | if not ret: |
| 466 | break |
| 467 | print('detect frame: %d' % (index)) |
| 468 | index += 1 |
| 469 | results = self.predict_image([frame[:, :, ::-1]], visual=False) |
| 470 | |
| 471 | im = visualize_box_mask( |
| 472 | frame, |
| 473 | results, |
| 474 | self.pred_config.labels, |
| 475 | threshold=self.threshold) |
| 476 | im = np.array(im) |
| 477 | writer.write(im) |
| 478 | if camera_id != -1: |
| 479 | cv2.imshow('Mask Detection', im) |
| 480 | if cv2.waitKey(1) & 0xFF == ord('q'): |
| 481 | break |
| 482 | writer.release() |
| 483 | |
| 484 | def save_coco_results(self, |
| 485 | image_list, |
no test coverage detected