Add Visualizer for visualization (#341)

* Add Visualizer from detectron2

* Refactor

* Move display into cv2_imshow()

* Use OpenCV to draw bounding boxes

* Cleanup

* Refactor the method in Visualizer

* Fix docstring

* Add metadata attribute in Visualizer

* Fix Visualizer._create_text_labels()

* Apply pre-commit

* Fix Visualizer.overlay_instances()

* Apply pre-commit

* Cleanup

* Add test_visualizer

* Apply pre-commit

* Add Visualizer.imshow()

test/test_utils.py

@@ -4,13 +4,15 @@
 import pytest
 import torch
 from torch import Tensor
+from torchvision.io import read_image
 from yolort import models
 from yolort.models import YOLO
 from yolort.utils import (
-    FeatureExtractor,
     get_image_from_url,
     load_from_ultralytics,
     read_image_to_tensor,
+    FeatureExtractor,
+    Visualizer,
 )
 from yolort.utils.image_utils import box_cxcywh_to_xyxy
 from yolort.v5 import (
@@ -22,6 +24,20 @@
 )
 
 
+@pytest.mark.parametrize("arch", ["yolov5n"])
+def test_visualizer(arch):
+    model = models.__dict__[arch](pretrained=True, size=(320, 320), score_thresh=0.45)
+    model = model.eval()
+    img_path = "test/assets/zidane.jpg"
+    preds = model.predict(img_path)
+
+    metalabels_path = "notebooks/assets/coco.names"
+    image = read_image(img_path)
+    v = Visualizer(image, metalabels=metalabels_path)
+    output = v.draw_instance_predictions(preds[0])
+    assert isinstance(output, np.ndarray)
+
+
 @pytest.mark.parametrize(
     "arch, version, upstream_version, hash_prefix, use_p6",
     [

yolort/data/builtin_meta.py

@@ -0,0 +1,154 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+"""
+Note:
+For your custom dataset, there is no need to hard-code metadata anywhere in the code.
+For example, for COCO-format dataset, metadata will be obtained automatically
+when calling `load_coco_json`. For other dataset, metadata may also be obtained in other ways
+during loading.
+
+However, we hard-coded metadata for a few common dataset here.
+The only goal is to allow users who don't have these dataset to use pre-trained models.
+Users don't have to download a COCO json (which contains metadata), in order to visualize a
+COCO model (with correct class names and colors).
+"""
+
+
+# All coco categories, together with their nice-looking visualization colors
+# It's from https://github.com/cocodataset/panopticapi/blob/master/panoptic_coco_categories.json
+COCO_CATEGORIES = [
+    {"id": 1, "color": [220, 20, 60], "isthing": 1, "name": "person"},
+    {"id": 2, "color": [119, 11, 32], "isthing": 1, "name": "bicycle"},
+    {"id": 3, "color": [0, 0, 142], "isthing": 1, "name": "car"},
+    {"id": 4, "color": [0, 0, 230], "isthing": 1, "name": "motorcycle"},
+    {"id": 5, "color": [106, 0, 228], "isthing": 1, "name": "airplane"},
+    {"id": 6, "color": [0, 60, 100], "isthing": 1, "name": "bus"},
+    {"id": 7, "color": [0, 80, 100], "isthing": 1, "name": "train"},
+    {"id": 8, "color": [0, 0, 70], "isthing": 1, "name": "truck"},
+    {"id": 9, "color": [0, 0, 192], "isthing": 1, "name": "boat"},
+    {"id": 10, "color": [250, 170, 30], "isthing": 1, "name": "traffic light"},
+    {"id": 11, "color": [100, 170, 30], "isthing": 1, "name": "fire hydrant"},
+    {"id": 13, "color": [220, 220, 0], "isthing": 1, "name": "stop sign"},
+    {"id": 14, "color": [175, 116, 175], "isthing": 1, "name": "parking meter"},
+    {"id": 15, "color": [250, 0, 30], "isthing": 1, "name": "bench"},
+    {"id": 16, "color": [165, 42, 42], "isthing": 1, "name": "bird"},
+    {"id": 17, "color": [255, 77, 255], "isthing": 1, "name": "cat"},
+    {"id": 18, "color": [0, 226, 252], "isthing": 1, "name": "dog"},
+    {"id": 19, "color": [182, 182, 255], "isthing": 1, "name": "horse"},
+    {"id": 20, "color": [0, 82, 0], "isthing": 1, "name": "sheep"},
+    {"id": 21, "color": [120, 166, 157], "isthing": 1, "name": "cow"},
+    {"id": 22, "color": [110, 76, 0], "isthing": 1, "name": "elephant"},
+    {"id": 23, "color": [174, 57, 255], "isthing": 1, "name": "bear"},
+    {"id": 24, "color": [199, 100, 0], "isthing": 1, "name": "zebra"},
+    {"id": 25, "color": [72, 0, 118], "isthing": 1, "name": "giraffe"},
+    {"id": 27, "color": [255, 179, 240], "isthing": 1, "name": "backpack"},
+    {"id": 28, "color": [0, 125, 92], "isthing": 1, "name": "umbrella"},
+    {"id": 31, "color": [209, 0, 151], "isthing": 1, "name": "handbag"},
+    {"id": 32, "color": [188, 208, 182], "isthing": 1, "name": "tie"},
+    {"id": 33, "color": [0, 220, 176], "isthing": 1, "name": "suitcase"},
+    {"id": 34, "color": [255, 99, 164], "isthing": 1, "name": "frisbee"},
+    {"id": 35, "color": [92, 0, 73], "isthing": 1, "name": "skis"},
+    {"id": 36, "color": [133, 129, 255], "isthing": 1, "name": "snowboard"},
+    {"id": 37, "color": [78, 180, 255], "isthing": 1, "name": "sports ball"},
+    {"id": 38, "color": [0, 228, 0], "isthing": 1, "name": "kite"},
+    {"id": 39, "color": [174, 255, 243], "isthing": 1, "name": "baseball bat"},
+    {"id": 40, "color": [45, 89, 255], "isthing": 1, "name": "baseball glove"},
+    {"id": 41, "color": [134, 134, 103], "isthing": 1, "name": "skateboard"},
+    {"id": 42, "color": [145, 148, 174], "isthing": 1, "name": "surfboard"},
+    {"id": 43, "color": [255, 208, 186], "isthing": 1, "name": "tennis racket"},
+    {"id": 44, "color": [197, 226, 255], "isthing": 1, "name": "bottle"},
+    {"id": 46, "color": [171, 134, 1], "isthing": 1, "name": "wine glass"},
+    {"id": 47, "color": [109, 63, 54], "isthing": 1, "name": "cup"},
+    {"id": 48, "color": [207, 138, 255], "isthing": 1, "name": "fork"},
+    {"id": 49, "color": [151, 0, 95], "isthing": 1, "name": "knife"},
+    {"id": 50, "color": [9, 80, 61], "isthing": 1, "name": "spoon"},
+    {"id": 51, "color": [84, 105, 51], "isthing": 1, "name": "bowl"},
+    {"id": 52, "color": [74, 65, 105], "isthing": 1, "name": "banana"},
+    {"id": 53, "color": [166, 196, 102], "isthing": 1, "name": "apple"},
+    {"id": 54, "color": [208, 195, 210], "isthing": 1, "name": "sandwich"},
+    {"id": 55, "color": [255, 109, 65], "isthing": 1, "name": "orange"},
+    {"id": 56, "color": [0, 143, 149], "isthing": 1, "name": "broccoli"},
+    {"id": 57, "color": [179, 0, 194], "isthing": 1, "name": "carrot"},
+    {"id": 58, "color": [209, 99, 106], "isthing": 1, "name": "hot dog"},
+    {"id": 59, "color": [5, 121, 0], "isthing": 1, "name": "pizza"},
+    {"id": 60, "color": [227, 255, 205], "isthing": 1, "name": "donut"},
+    {"id": 61, "color": [147, 186, 208], "isthing": 1, "name": "cake"},
+    {"id": 62, "color": [153, 69, 1], "isthing": 1, "name": "chair"},
+    {"id": 63, "color": [3, 95, 161], "isthing": 1, "name": "couch"},
+    {"id": 64, "color": [163, 255, 0], "isthing": 1, "name": "potted plant"},
+    {"id": 65, "color": [119, 0, 170], "isthing": 1, "name": "bed"},
+    {"id": 67, "color": [0, 182, 199], "isthing": 1, "name": "dining table"},
+    {"id": 70, "color": [0, 165, 120], "isthing": 1, "name": "toilet"},
+    {"id": 72, "color": [183, 130, 88], "isthing": 1, "name": "tv"},
+    {"id": 73, "color": [95, 32, 0], "isthing": 1, "name": "laptop"},
+    {"id": 74, "color": [130, 114, 135], "isthing": 1, "name": "mouse"},
+    {"id": 75, "color": [110, 129, 133], "isthing": 1, "name": "remote"},
+    {"id": 76, "color": [166, 74, 118], "isthing": 1, "name": "keyboard"},
+    {"id": 77, "color": [219, 142, 185], "isthing": 1, "name": "cell phone"},
+    {"id": 78, "color": [79, 210, 114], "isthing": 1, "name": "microwave"},
+    {"id": 79, "color": [178, 90, 62], "isthing": 1, "name": "oven"},
+    {"id": 80, "color": [65, 70, 15], "isthing": 1, "name": "toaster"},
+    {"id": 81, "color": [127, 167, 115], "isthing": 1, "name": "sink"},
+    {"id": 82, "color": [59, 105, 106], "isthing": 1, "name": "refrigerator"},
+    {"id": 84, "color": [142, 108, 45], "isthing": 1, "name": "book"},
+    {"id": 85, "color": [196, 172, 0], "isthing": 1, "name": "clock"},
+    {"id": 86, "color": [95, 54, 80], "isthing": 1, "name": "vase"},
+    {"id": 87, "color": [128, 76, 255], "isthing": 1, "name": "scissors"},
+    {"id": 88, "color": [201, 57, 1], "isthing": 1, "name": "teddy bear"},
+    {"id": 89, "color": [246, 0, 122], "isthing": 1, "name": "hair drier"},
+    {"id": 90, "color": [191, 162, 208], "isthing": 1, "name": "toothbrush"},
+    {"id": 92, "color": [255, 255, 128], "isthing": 0, "name": "banner"},
+    {"id": 93, "color": [147, 211, 203], "isthing": 0, "name": "blanket"},
+    {"id": 95, "color": [150, 100, 100], "isthing": 0, "name": "bridge"},
+    {"id": 100, "color": [168, 171, 172], "isthing": 0, "name": "cardboard"},
+    {"id": 107, "color": [146, 112, 198], "isthing": 0, "name": "counter"},
+    {"id": 109, "color": [210, 170, 100], "isthing": 0, "name": "curtain"},
+    {"id": 112, "color": [92, 136, 89], "isthing": 0, "name": "door-stuff"},
+    {"id": 118, "color": [218, 88, 184], "isthing": 0, "name": "floor-wood"},
+    {"id": 119, "color": [241, 129, 0], "isthing": 0, "name": "flower"},
+    {"id": 122, "color": [217, 17, 255], "isthing": 0, "name": "fruit"},
+    {"id": 125, "color": [124, 74, 181], "isthing": 0, "name": "gravel"},
+    {"id": 128, "color": [70, 70, 70], "isthing": 0, "name": "house"},
+    {"id": 130, "color": [255, 228, 255], "isthing": 0, "name": "light"},
+    {"id": 133, "color": [154, 208, 0], "isthing": 0, "name": "mirror-stuff"},
+    {"id": 138, "color": [193, 0, 92], "isthing": 0, "name": "net"},
+    {"id": 141, "color": [76, 91, 113], "isthing": 0, "name": "pillow"},
+    {"id": 144, "color": [255, 180, 195], "isthing": 0, "name": "platform"},
+    {"id": 145, "color": [106, 154, 176], "isthing": 0, "name": "playingfield"},
+    {"id": 147, "color": [230, 150, 140], "isthing": 0, "name": "railroad"},
+    {"id": 148, "color": [60, 143, 255], "isthing": 0, "name": "river"},
+    {"id": 149, "color": [128, 64, 128], "isthing": 0, "name": "road"},
+    {"id": 151, "color": [92, 82, 55], "isthing": 0, "name": "roof"},
+    {"id": 154, "color": [254, 212, 124], "isthing": 0, "name": "sand"},
+    {"id": 155, "color": [73, 77, 174], "isthing": 0, "name": "sea"},
+    {"id": 156, "color": [255, 160, 98], "isthing": 0, "name": "shelf"},
+    {"id": 159, "color": [255, 255, 255], "isthing": 0, "name": "snow"},
+    {"id": 161, "color": [104, 84, 109], "isthing": 0, "name": "stairs"},
+    {"id": 166, "color": [169, 164, 131], "isthing": 0, "name": "tent"},
+    {"id": 168, "color": [225, 199, 255], "isthing": 0, "name": "towel"},
+    {"id": 171, "color": [137, 54, 74], "isthing": 0, "name": "wall-brick"},
+    {"id": 175, "color": [135, 158, 223], "isthing": 0, "name": "wall-stone"},
+    {"id": 176, "color": [7, 246, 231], "isthing": 0, "name": "wall-tile"},
+    {"id": 177, "color": [107, 255, 200], "isthing": 0, "name": "wall-wood"},
+    {"id": 178, "color": [58, 41, 149], "isthing": 0, "name": "water-other"},
+    {"id": 180, "color": [183, 121, 142], "isthing": 0, "name": "window-blind"},
+    {"id": 181, "color": [255, 73, 97], "isthing": 0, "name": "window-other"},
+    {"id": 184, "color": [107, 142, 35], "isthing": 0, "name": "tree-merged"},
+    {"id": 185, "color": [190, 153, 153], "isthing": 0, "name": "fence-merged"},
+    {"id": 186, "color": [146, 139, 141], "isthing": 0, "name": "ceiling-merged"},
+    {"id": 187, "color": [70, 130, 180], "isthing": 0, "name": "sky-other-merged"},
+    {"id": 188, "color": [134, 199, 156], "isthing": 0, "name": "cabinet-merged"},
+    {"id": 189, "color": [209, 226, 140], "isthing": 0, "name": "table-merged"},
+    {"id": 190, "color": [96, 36, 108], "isthing": 0, "name": "floor-other-merged"},
+    {"id": 191, "color": [96, 96, 96], "isthing": 0, "name": "pavement-merged"},
+    {"id": 192, "color": [64, 170, 64], "isthing": 0, "name": "mountain-merged"},
+    {"id": 193, "color": [152, 251, 152], "isthing": 0, "name": "grass-merged"},
+    {"id": 194, "color": [208, 229, 228], "isthing": 0, "name": "dirt-merged"},
+    {"id": 195, "color": [206, 186, 171], "isthing": 0, "name": "paper-merged"},
+    {"id": 196, "color": [152, 161, 64], "isthing": 0, "name": "food-other-merged"},
+    {"id": 197, "color": [116, 112, 0], "isthing": 0, "name": "building-other-merged"},
+    {"id": 198, "color": [0, 114, 143], "isthing": 0, "name": "rock-merged"},
+    {"id": 199, "color": [102, 102, 156], "isthing": 0, "name": "wall-other-merged"},
+    {"id": 200, "color": [250, 141, 255], "isthing": 0, "name": "rug-merged"},
+]

yolort/utils/__init__.py

@@ -8,6 +8,7 @@
 from .hooks import FeatureExtractor
 from .image_utils import cv2_imshow, get_image_from_url, read_image_to_tensor
 from .update_module_state import convert_yolov5_to_yolort, load_from_ultralytics
+from .visualizer import Visualizer
 
 
 __all__ = [
@@ -19,6 +20,7 @@
     "load_from_ultralytics",
     "load_state_dict_from_url",
     "read_image_to_tensor",
+    "Visualizer",
 ]