diff --git a/hubconf.py b/hubconf.py index 85f5b77f..4ad094ef 100644 --- a/hubconf.py +++ b/hubconf.py @@ -1,4 +1,13 @@ # Optional list of dependencies required by the package dependencies = ["torch", "torchvision"] -from yolort.models import yolov5n, yolov5s, yolov5m, yolov5l, yolov5ts +from yolort.models import ( + yolov5n, + yolov5n6, + yolov5s, + yolov5s6, + yolov5m, + yolov5m6, + yolov5l, + yolov5ts, +) diff --git a/test/test_models.py b/test/test_models.py index 43d72856..49587b29 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -357,7 +357,12 @@ def test_torchscript(arch): @pytest.mark.parametrize( "arch, version, upstream_version, hash_prefix", - [("yolov5s", "r4.0", "v4.0", "9ca9a642")], + [ + ("yolov5s", "r4.0", "v4.0", "9ca9a642"), + ("yolov5n", "r6.0", "v6.0", "649e089f"), + ("yolov5s", "r6.0", "v6.0", "c3b140f3"), + ("yolov5n6", "r6.0", "v6.0", "beecbbae"), + ], ) def test_load_from_yolov5( arch: str, @@ -376,8 +381,13 @@ def test_load_from_yolov5( checkpoint_path, hash_prefix=hash_prefix, ) + score_thresh = 0.25 - model_yolov5 = YOLOv5.load_from_yolov5(checkpoint_path, version=version) + model_yolov5 = YOLOv5.load_from_yolov5( + checkpoint_path, + score_thresh=score_thresh, + version=version, + ) model_yolov5.eval() out_from_yolov5 = model_yolov5.predict(img_path) assert isinstance(out_from_yolov5[0], dict) @@ -385,7 +395,11 @@ def test_load_from_yolov5( assert isinstance(out_from_yolov5[0]["labels"], Tensor) assert isinstance(out_from_yolov5[0]["scores"], Tensor) - model = models.__dict__[arch](pretrained=True, score_thresh=0.25) + model = models.__dict__[arch]( + upstream_version=version, + pretrained=True, + score_thresh=score_thresh, + ) model.eval() out = model.predict(img_path) diff --git a/tools/yolov5_to_yolort.py b/tools/yolov5_to_yolort.py new file mode 100644 index 00000000..79f600e7 --- /dev/null +++ b/tools/yolov5_to_yolort.py @@ -0,0 +1,56 @@ +# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. +import argparse +from pathlib import Path + +from yolort.utils import convert_yolov5_to_yolort + + +def get_parser(): + parser = argparse.ArgumentParser( + "Convert checkpoints from yolov5 to yolort", add_help=True + ) + + parser.add_argument( + "--checkpoint_path", + type=str, + required=True, + help="Path of the checkpoint weights", + ) + parser.add_argument( + "--version", + type=str, + default="r6.0", + help="Upstream version released by the ultralytics/yolov5, Possible " + "values are ['r3.1', 'r4.0', 'r6.0']. Default: 'r6.0'.", + ) + # Dataset Configuration + parser.add_argument( + "--image_path", + type=str, + default="./test/assets/zidane.jpg", + help="Path of the test image", + ) + + parser.add_argument( + "--output_path", type=str, default=None, help="Path where to save" + ) + return parser + + +def cli_main(): + parser = get_parser() + args = parser.parse_args() + print(f"Command Line Args: {args}") + checkpoint_path = Path(args.checkpoint_path) + assert checkpoint_path.exists(), f"Not found checkpoint file at '{checkpoint_path}'" + + if args.output_path is None: + args.output_path = checkpoint_path.parent + output_path = Path(args.output_path) + output_path.mkdir(parents=True, exist_ok=True) + + convert_yolov5_to_yolort(checkpoint_path, output_path, version=args.version) + + +if __name__ == "__main__": + cli_main() diff --git a/yolort/models/__init__.py b/yolort/models/__init__.py index c59d32da..ba36be59 100644 --- a/yolort/models/__init__.py +++ b/yolort/models/__init__.py @@ -8,7 +8,18 @@ from .yolo import YOLO from .yolo_module import YOLOv5 -__all__ = ["YOLO", "YOLOv5", "yolov5n", "yolov5s", "yolov5m", "yolov5l", "yolov5ts"] +__all__ = [ + "YOLO", + "YOLOv5", + "yolov5n", + "yolov5n6", + "yolov5s", + "yolov5s6", + "yolov5m", + "yolov5m6", + "yolov5l", + "yolov5ts", +] def yolov5n( diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py index 6fa96a1d..222751fc 100644 --- a/yolort/models/yolo.py +++ b/yolort/models/yolo.py @@ -20,6 +20,13 @@ "yolov5_darknet_pan_s_r40", "yolov5_darknet_pan_m_r40", "yolov5_darknet_pan_l_r40", + "yolov5_darknet_pan_n_r60", + "yolov5_darknet_pan_n6_r60", + "yolov5_darknet_pan_s_r60", + "yolov5_darknet_pan_s6_r60", + "yolov5_darknet_pan_m_r60", + "yolov5_darknet_pan_m6_r60", + "yolov5_darknet_pan_l_r60", "yolov5_darknet_tan_s_r40", "build_model", ] @@ -221,21 +228,6 @@ def load_from_yolov5( return model -model_urls_root = "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0" - -model_urls = { - # Path Aggregation Network - "yolov5_darknet_pan_s_r31_coco": f"{model_urls_root}/yolov5_darknet_pan_s_r31_coco-eb728698.pt", - "yolov5_darknet_pan_m_r31_coco": f"{model_urls_root}/yolov5_darknet_pan_m_r31_coco-670dc553.pt", - "yolov5_darknet_pan_l_r31_coco": f"{model_urls_root}/yolov5_darknet_pan_l_r31_coco-4dcc8209.pt", - "yolov5_darknet_pan_s_r40_coco": f"{model_urls_root}/yolov5_darknet_pan_s_r40_coco-e3fd213d.pt", - "yolov5_darknet_pan_m_r40_coco": f"{model_urls_root}/yolov5_darknet_pan_m_r40_coco-d295cb02.pt", - "yolov5_darknet_pan_l_r40_coco": f"{model_urls_root}/yolov5_darknet_pan_l_r40_coco-4416841f.pt", - # Tranformer Attention Network - "yolov5_darknet_tan_s_r40_coco": f"{model_urls_root}/yolov5_darknet_tan_s_r40_coco-fe1069ce.pt", -} - - def build_model( backbone_name: str, depth_multiple: float, @@ -280,6 +272,34 @@ def build_model( return model +model_urls_root_r30 = ( + "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0" +) +model_urls_root_r52 = ( + "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.5.2-alpha" +) + +model_urls = { + # Path Aggregation Network 3.1 and 4.0 + "yolov5_darknet_pan_s_r31_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_s_r31_coco-eb728698.pt", + "yolov5_darknet_pan_m_r31_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_m_r31_coco-670dc553.pt", + "yolov5_darknet_pan_l_r31_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_l_r31_coco-4dcc8209.pt", + "yolov5_darknet_pan_s_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_s_r40_coco-e3fd213d.pt", + "yolov5_darknet_pan_m_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_m_r40_coco-d295cb02.pt", + "yolov5_darknet_pan_l_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_l_r40_coco-4416841f.pt", + # Path Aggregation Network 6.0 + "yolov5_darknet_pan_n_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n_r60_coco-bc15659e.pt", + "yolov5_darknet_pan_n6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n6_r60_coco-4e823e0f.pt", + "yolov5_darknet_pan_s_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s_r60_coco-9f44bf3f.pt", + "yolov5_darknet_pan_s6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s6_r60_coco-b4ff1fc2.pt", + "yolov5_darknet_pan_m_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m_r60_coco-58d32352.pt", + "yolov5_darknet_pan_m6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m6_r60_coco-cc010533.pt", + "yolov5_darknet_pan_l_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_l_r60_coco-321d8dcd.pt", + # Tranformer Attention Network + "yolov5_darknet_tan_s_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_tan_s_r40_coco-fe1069ce.pt", +} + + def yolov5_darknet_pan_s_r31( pretrained: bool = False, progress: bool = True, diff --git a/yolort/utils/__init__.py b/yolort/utils/__init__.py index 6e3c884c..ecdaf098 100644 --- a/yolort/utils/__init__.py +++ b/yolort/utils/__init__.py @@ -7,7 +7,7 @@ from .hooks import FeatureExtractor from .image_utils import cv2_imshow, get_image_from_url, read_image_to_tensor -from .update_module_state import load_from_ultralytics +from .update_module_state import convert_yolov5_to_yolort, load_from_ultralytics __all__ = [ @@ -15,6 +15,7 @@ "cv2_imshow", "get_image_from_url", "get_callable_dict", + "convert_yolov5_to_yolort", "load_from_ultralytics", "load_state_dict_from_url", "read_image_to_tensor", diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py index 9a9ea1ea..d33adc73 100644 --- a/yolort/utils/update_module_state.py +++ b/yolort/utils/update_module_state.py @@ -2,6 +2,7 @@ from functools import reduce from typing import List, Dict, Optional +import torch from torch import nn from yolort.models import yolo @@ -9,7 +10,37 @@ from .image_utils import to_numpy -def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"): +def convert_yolov5_to_yolort( + checkpoint_path: str, + output_path: str, + version: str = "r6.0", + prefix: str = "yolov5_darknet_pan", + postfix: str = "custom.pt", +): + """ + Convert model checkpoint trained with ultralytics/yolov5 to yolort. + + Args: + checkpoint_path (str): Path of the YOLOv5 checkpoint model. + output_path (str): Path of the converted yolort checkpoint model. + version (str): upstream version released by the ultralytics/yolov5, Possible + values are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0". + prefix (str): The prefix string of the saved model. Default: "yolov5_darknet_pan". + postfix (str): The postfix string of the saved model. Default: "custom.pt". + """ + model_info = load_from_ultralytics(checkpoint_path, version=version) + model_state_dict = model_info["state_dict"] + + size = model_info["size"] + use_p6 = "6" if model_info["use_p6"] else "" + output_postfix = f"{prefix}_{size}{use_p6}_{version.replace('.', '')}_{postfix}" + torch.save(model_state_dict, output_path / output_postfix) + + +def load_from_ultralytics( + checkpoint_path: str, + version: str = "r6.0", +): """ Allows the user to load model state file from the checkpoint trained from the ultralytics/yolov5. @@ -85,7 +116,7 @@ def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"): use_p6=use_p6, ) module_state_updater.updating(checkpoint_yolov5) - state_dict = module_state_updater.model.state_dict() + state_dict = module_state_updater.model.half().state_dict() size = get_yolov5_size(depth_multiple, width_multiple)