Action Predictor for action recognition.
| 18 | |
| 19 | |
| 20 | class Predictor: |
| 21 | """ |
| 22 | Action Predictor for action recognition. |
| 23 | """ |
| 24 | |
| 25 | def __init__(self, cfg, gpu_id=None): |
| 26 | """ |
| 27 | Args: |
| 28 | cfg (CfgNode): configs. Details can be found in |
| 29 | slowfast/config/defaults.py |
| 30 | gpu_id (Optional[int]): GPU id. |
| 31 | """ |
| 32 | if cfg.NUM_GPUS: |
| 33 | self.gpu_id = ( |
| 34 | torch.cuda.current_device() if gpu_id is None else gpu_id |
| 35 | ) |
| 36 | |
| 37 | # Build the video model and print model statistics. |
| 38 | self.model = build_model(cfg, gpu_id=gpu_id) |
| 39 | self.model.eval() |
| 40 | self.cfg = cfg |
| 41 | |
| 42 | if cfg.DETECTION.ENABLE: |
| 43 | # self.object_detector = Detectron2Predictor(cfg, gpu_id=self.gpu_id) |
| 44 | pass |
| 45 | |
| 46 | logger.info("Start loading model weights.") |
| 47 | cu.load_test_checkpoint(cfg, self.model) |
| 48 | logger.info("Finish loading model weights") |
| 49 | |
| 50 | def __call__(self, task): |
| 51 | """ |
| 52 | Returns the prediction results for the current task. |
| 53 | Args: |
| 54 | task (TaskInfo object): task object that contain |
| 55 | the necessary information for action prediction. (e.g. frames, boxes) |
| 56 | Returns: |
| 57 | task (TaskInfo object): the same task info object but filled with |
| 58 | prediction values (a tensor) and the corresponding boxes for |
| 59 | action detection task. |
| 60 | """ |
| 61 | if self.cfg.DETECTION.ENABLE: |
| 62 | task = self.object_detector(task) |
| 63 | |
| 64 | frames, bboxes = task.frames, task.bboxes |
| 65 | if bboxes is not None: |
| 66 | bboxes = cv2_transform.scale_boxes( |
| 67 | self.cfg.DATA.TEST_CROP_SIZE, |
| 68 | bboxes, |
| 69 | task.img_height, |
| 70 | task.img_width, |
| 71 | ) |
| 72 | if self.cfg.DEMO.INPUT_FORMAT == "BGR": |
| 73 | frames = [ |
| 74 | cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames |
| 75 | ] |
| 76 | |
| 77 | frames = [ |