Connecting with TorchVision's model banks

[zhiqwang]

This release adds an example to showcase how to use the model banks provided by TorchVision to build a new YOLOv5 model, introduces Visualizer to display the predictions and fixes an incompatible bug in YOLOv5 Detect Layer.

Highlights

Build a new YOLOv5 model via TorchVision's model banks

We add an example to showcase how to construct a YOLOv5-Lite model with TorchVision's pre-trained MobileNetV3-Small FPN as the backbone.

import torch
from yolort.models.yolo_lite import yolov5_mobilenet_v3_small_fpn

model = yolov5_mobilenet_v3_small_fpn()
model = model.eval()

images = torch.rand(4, 3, 640, 640)
out = model(images)

• Use C3 following the model specification of YOLOv5 #343
• Constructs YOLOv5 model with TorchVision MobileNetV3 backbone #342

Introduce a new visualization interface

from torchvision.io import read_image
from yolort.models import yolov5n6
from yolort.utils import Visualizer

model = yolov5n6(pretrained=True, score_thresh=0.45)
model = model.eval()
img_path = "bus.jpg"
preds = model.predict(img_path)

metalabels_path = "coco.names"
image = read_image(img_path)
v = Visualizer(image, metalabels=metalabels_path)
v.draw_instance_predictions(preds[0])
v.imshow(scale=0.5)

• Add Visualizer for visualization #341

Bugfixes

• Fix YOLOv5 Detect Layer compatibility #345
• Add missing version param in export_tensorrt #335

Code Quality

• Update Visualizer in docs and PR Recommendations #346
• Bump actions/checkout from 2 to 3 #344
• [pre-commit.ci] pre-commit autoupdate #338
• Bump actions/setup-python from 2 to 3.0.0 #337

Contributors

We're grateful for our community, which helps us improving yolort by submitting issues and PRs, and providing feedback and suggestions. The following persons have contributed patches for this release:

@Luwill6, @zhiqwang

Full Changelog: v0.6.0...v0.6.1

---

This discussion was created from the release Connecting to TorchVision's model banks.