Args: image (np.ndarray): an image of shape (H, W, C) (in BGR order). This is the format used by OpenCV. Returns: predictions (dict): the output of the model. vis_output (VisImage): the visualized image output.
(self, image, confidence_threshold)
| 35 | self.predictor = DefaultPredictor(cfg) |
| 36 | |
| 37 | def run_on_image(self, image, confidence_threshold): |
| 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 | # SparseRCNN uses RGB input as default |
| 51 | # image = image[:, :, ::-1] |
| 52 | visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode) |
| 53 | if "panoptic_seg" in predictions: |
| 54 | panoptic_seg, segments_info = predictions["panoptic_seg"] |
| 55 | vis_output = visualizer.draw_panoptic_seg_predictions( |
| 56 | panoptic_seg.to(self.cpu_device), segments_info |
| 57 | ) |
| 58 | else: |
| 59 | if "sem_seg" in predictions: |
| 60 | vis_output = visualizer.draw_sem_seg( |
| 61 | predictions["sem_seg"].argmax(dim=0).to(self.cpu_device) |
| 62 | ) |
| 63 | if "instances" in predictions: |
| 64 | instances = predictions["instances"].to(self.cpu_device) |
| 65 | instances = instances[instances.scores > confidence_threshold] |
| 66 | predictions["instances"] = instances |
| 67 | vis_output = visualizer.draw_instance_predictions(predictions=instances) |
| 68 | |
| 69 | return predictions, vis_output |
| 70 | |
| 71 | def _frame_from_video(self, video): |
| 72 | while video.isOpened(): |
no test coverage detected