| 13 | |
| 14 | |
| 15 | class VisualizationDemo(object): |
| 16 | def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False): |
| 17 | """ |
| 18 | Args: |
| 19 | cfg (CfgNode): |
| 20 | instance_mode (ColorMode): |
| 21 | parallel (bool): whether to run the model in different processes from visualization. |
| 22 | Useful since the visualization logic can be slow. |
| 23 | """ |
| 24 | self.metadata = MetadataCatalog.get( |
| 25 | cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused" |
| 26 | ) |
| 27 | self.cpu_device = torch.device("cpu") |
| 28 | self.instance_mode = instance_mode |
| 29 | |
| 30 | self.parallel = parallel |
| 31 | if parallel: |
| 32 | num_gpu = torch.cuda.device_count() |
| 33 | self.predictor = AsyncPredictor(cfg, num_gpus=num_gpu) |
| 34 | else: |
| 35 | self.predictor = DefaultPredictor(cfg) |
| 36 | |
| 37 | def run_on_image(self, image): |
| 38 | """ |
| 39 | Args: |
| 40 | image (np.ndarray): an image of shape (H, W, C) (in BGR order). |
| 41 | This is the format used by OpenCV. |
| 42 | |
| 43 | Returns: |
| 44 | predictions (dict): the output of the model. |
| 45 | vis_output (VisImage): the visualized image output. |
| 46 | """ |
| 47 | vis_output = None |
| 48 | predictions = self.predictor(image) |
| 49 | # Convert image from OpenCV BGR format to Matplotlib RGB format. |
| 50 | image = image[:, :, ::-1] |
| 51 | visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode) |
| 52 | if "panoptic_seg" in predictions: |
| 53 | panoptic_seg, segments_info = predictions["panoptic_seg"] |
| 54 | vis_output = visualizer.draw_panoptic_seg_predictions( |
| 55 | panoptic_seg.to(self.cpu_device), segments_info |
| 56 | ) |
| 57 | else: |
| 58 | if "sem_seg" in predictions: |
| 59 | vis_output = visualizer.draw_sem_seg( |
| 60 | predictions["sem_seg"].argmax(dim=0).to(self.cpu_device) |
| 61 | ) |
| 62 | if "instances" in predictions: |
| 63 | instances = predictions["instances"].to(self.cpu_device) |
| 64 | vis_output = visualizer.draw_instance_predictions(predictions=instances) |
| 65 | |
| 66 | return predictions, vis_output |
| 67 | |
| 68 | def _frame_from_video(self, video): |
| 69 | while video.isOpened(): |
| 70 | success, frame = video.read() |
| 71 | if success: |
| 72 | yield frame |