Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Visualizer for visualization #341

Merged
merged 16 commits into from
Mar 3, 2022
Merged

Add Visualizer for visualization #341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
8b40feb
Add Visualizer from detectron2
zhiqwang Oct 5, 2021
8e762cd
Refactor
zhiqwang Feb 28, 2022
dfadb53
Move display into cv2_imshow()
zhiqwang Feb 28, 2022
f35d002
Use OpenCV to draw bounding boxes
zhiqwang Mar 1, 2022
2150615
Cleanup
zhiqwang Mar 2, 2022
4806c85
Refactor the method in Visualizer
zhiqwang Mar 2, 2022
544f4de
Fix docstring
zhiqwang Mar 3, 2022
2ece383
Add metadata attribute in Visualizer
zhiqwang Mar 3, 2022
78c156b
Fix Visualizer._create_text_labels()
zhiqwang Mar 3, 2022
2cba497
Apply pre-commit
zhiqwang Mar 3, 2022
b27fcb2
Fix Visualizer.overlay_instances()
zhiqwang Mar 3, 2022
39a7469
Apply pre-commit
zhiqwang Mar 3, 2022
9284a22
Cleanup
zhiqwang Mar 3, 2022
b0e2f7c
Add test_visualizer
zhiqwang Mar 3, 2022
9386c16
Apply pre-commit
zhiqwang Mar 3, 2022
88b830a
Add Visualizer.imshow()
zhiqwang Mar 3, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion test/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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",
[
Expand Down
154 changes: 154 additions & 0 deletions yolort/data/builtin_meta.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"},
]
2 changes: 2 additions & 0 deletions yolort/utils/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__ = [
Expand All @@ -19,6 +20,7 @@
"load_from_ultralytics",
"load_state_dict_from_url",
"read_image_to_tensor",
"Visualizer",
]


Expand Down
Loading