0

I am learning detectron2 for using Faster R-CNN.

I am trying to draw the bounding boxes of same class. For example, if bounding box includes "cat", it is blue while if another bounding box includes "dog", it is assigned to red. However, the below code can output only each bounding box has different colors would you please kindly advise to me?

from detectron2.utils.visualizer import ColorMode
import glob
for imageName in random.sample(glob.glob(os.path.join(test_path, '*.png')), 3):
  im = cv2.imread(imageName)
  outputs = predictor(im)
  v = Visualizer(im[:, :, ::-1],
        metadata=train_metadata,
        scale=0.8
         )
  out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
  cv2_imshow(out.get_image()[:, :, ::-1])
Seunghyeon
  • 103
  • 1
  • 2
  • 10

1 Answers1

1

I have not tried it yet but it seems a fairly simple task.

There are two functions in visualizer.py called overlay_instances and draw_box which are present in detectron2 directory.

I will write down the path below. detectron2/detectron2/utils/visualizer.py

Basically what happens is labels are been set in overlay_instances function and then calls the draw_box function to well draw boxes.}

You can play around in these two functions to set a specific color for the specific classes.

Also, FYI, the overlay_instances function is being called by draw_instance_predictions function which you are using.

out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

greybeard
  • 2,249
  • 8
  • 30
  • 66