Exemplo n.º 1
0
 def draw_scene(self, frame, monitor=None, event_count=None):
     image_rgb = frame.image.numpy()[:, :, ::-1]
     visualizer = dt_visualizer(image_rgb, None)
     visualizer.draw_text('Frame %d' % (frame.image_id), (10, 10),
                          font_size=15, horizontal_alignment='left')
     instances = frame.instances.to('cpu')
     obj_colors = self._draw_instances(visualizer, instances, image_rgb)
     if instances.has('boxes_3d'):
         for obj_i in range(len(instances)):
             angle = instances.angles[obj_i]
             if angle < 0:
                 continue
             box_3d = instances.boxes_3d[obj_i]
             color = obj_colors[obj_i] if self.box_3d_color_sync else None
             self._draw_box_3d(visualizer, box_3d, color)
     if instances.has('boxes_3d_gt'):
         for obj_i in range(len(instances)):
             box_3d = instances.boxes_3d_gt[obj_i]
             if box_3d[0, 0, 0] < 0:
                 continue
             color = obj_colors[obj_i] if self.box_3d_color_sync else None
             self._draw_box_3d(
                 visualizer, box_3d, color, linestyle='dotted')
     if monitor is not None:
         movement_colors = self._draw_region_movements(
             visualizer, image_rgb, monitor.region, monitor.movements,
             event_count)
         self._draw_events(visualizer, instances, movement_colors)
     output = visualizer.get_output()
     visual_image = output.get_image()
     plt.close(output.fig)
     return visual_image
Exemplo n.º 2
0
 def draw_scene(self, frame, image_roi=None, mois=None):
     image_rgb = frame.image.numpy()[:, :, ::-1]
     visualizer = dt_visualizer(image_rgb, None)
     instances = frame.instances.to('cpu')
     labels, colors, masks = [], [], []
     for obj_i in range(len(instances)):
         obj_type = ObjectType(instances.pred_classes[obj_i].item())
         obj_id = obj_i
         if instances.has('track_ids'):
             obj_id = instances.track_ids[obj_i].item()
         obj_type = obj_type.name
         score = instances.scores[obj_i] * 100
         label = '%s-%s %.0f%%' % (obj_type, obj_id, score)
         labels.append(label)
         x0, y0, x1, y1 = instances.pred_boxes.tensor[obj_i].type(torch.int)
         roi = image_rgb[y0:y1, x0:x1]
         color = self.color_manager.get_color((obj_type, obj_id), roi)
         colors.append(color)
         mask = [np.array([0, 0])]
         if instances.has('pred_masks'):
             mask = instances.pred_masks[obj_i].numpy()
         if instances.has('contours'):
             contour = instances.contours[obj_i]
             if contour is not None:
                 contour = contour[contour[:, 0] >= 0]
                 mask = [contour.cpu().numpy()]
         masks.append(mask)
     visualizer.overlay_instances(masks=masks,
                                  boxes=instances.pred_boxes,
                                  labels=labels,
                                  assigned_colors=colors)
     if instances.has('boxes_3d'):
         for obj_i in range(len(instances)):
             angle = instances.angles[obj_i]
             if angle < 0:
                 continue
             box_3d = instances.boxes_3d[obj_i]
             color = colors[obj_i] if self.box_3d_color_sync else None
             self._draw_box_3d(visualizer, box_3d, color)
     if instances.has('boxes_3d_gt'):
         for obj_i in range(len(instances)):
             box_3d = instances.boxes_3d_gt[obj_i]
             if box_3d[0, 0, 0] < 0:
                 continue
             color = colors[obj_i] if self.box_3d_color_sync else None
             self._draw_box_3d(visualizer,
                               box_3d,
                               color,
                               linestyle='dotted')
     self._draw_roi_mois(visualizer, image_roi, mois)
     visual_image = visualizer.get_output().get_image()
     return visual_image