From 9d48190e6d624df45714b089491e304525326db1 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 6 May 2024 12:26:01 +0200 Subject: [PATCH] `ultralytics 8.2.10` add Classify and OBB Tasks to `Results.summary()` (#11653) Signed-off-by: Glenn Jocher Co-authored-by: Laughing-q <1185102784@qq.com> --- tests/test_exports.py | 1 + tests/test_python.py | 57 +++++++++++++---------------------- ultralytics/__init__.py | 2 +- ultralytics/cfg/__init__.py | 1 + ultralytics/engine/results.py | 38 +++++++++++++---------- 5 files changed, 46 insertions(+), 53 deletions(-) diff --git a/tests/test_exports.py b/tests/test_exports.py index fb829c76..a9b5d24a 100644 --- a/tests/test_exports.py +++ b/tests/test_exports.py @@ -18,6 +18,7 @@ from ultralytics.utils import ( checks, ) from ultralytics.utils.torch_utils import TORCH_1_9, TORCH_1_13 + from . import MODEL, SOURCE diff --git a/tests/test_python.py b/tests/test_python.py index 86ed2eee..3ca57d4e 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -12,7 +12,7 @@ import yaml from PIL import Image from ultralytics import RTDETR, YOLO -from ultralytics.cfg import TASK2DATA +from ultralytics.cfg import MODELS, TASK2DATA from ultralytics.data.build import load_inference_source from ultralytics.utils import ( ASSETS, @@ -76,42 +76,27 @@ def test_predict_txt(): _ = YOLO(MODEL)(source=txt_file, imgsz=32) -def test_predict_img(): +@pytest.mark.parametrize("model_name", MODELS) +def test_predict_img(model_name): """Test YOLO prediction on various types of image sources.""" - model = YOLO(MODEL) - seg_model = YOLO(WEIGHTS_DIR / "yolov8n-seg.pt") - cls_model = YOLO(WEIGHTS_DIR / "yolov8n-cls.pt") - pose_model = YOLO(WEIGHTS_DIR / "yolov8n-pose.pt") - obb_model = YOLO(WEIGHTS_DIR / "yolov8n-obb.pt") - im = cv2.imread(str(SOURCE)) + model = YOLO(WEIGHTS_DIR / model_name) + im = cv2.imread(str(SOURCE)) # uint8 numpy array assert len(model(source=Image.open(SOURCE), save=True, verbose=True, imgsz=32)) == 1 # PIL assert len(model(source=im, save=True, save_txt=True, imgsz=32)) == 1 # ndarray + assert len(model(torch.rand((2, 3, 32, 32)), imgsz=32)) == 2 # batch-size 2 Tensor, FP32 0.0-1.0 RGB order assert len(model(source=[im, im], save=True, save_txt=True, imgsz=32)) == 2 # batch assert len(list(model(source=[im, im], save=True, stream=True, imgsz=32))) == 2 # stream - assert len(model(torch.zeros(320, 640, 3).numpy(), imgsz=32)) == 1 # tensor to numpy + assert len(model(torch.zeros(320, 640, 3).numpy().astype(np.uint8), imgsz=32)) == 1 # tensor to numpy batch = [ str(SOURCE), # filename Path(SOURCE), # Path "https://ultralytics.com/images/zidane.jpg" if ONLINE else SOURCE, # URI cv2.imread(str(SOURCE)), # OpenCV Image.open(SOURCE), # PIL - np.zeros((320, 640, 3)), - ] # numpy + np.zeros((320, 640, 3), dtype=np.uint8), # numpy + ] assert len(model(batch, imgsz=32)) == len(batch) # multiple sources in a batch - # Test tensor inference - im = torch.rand((4, 3, 32, 32)) # batch-size 4, FP32 0.0-1.0 RGB order - results = model(im, imgsz=32) - assert len(results) == im.shape[0] - results = seg_model(im, imgsz=32) - assert len(results) == im.shape[0] - results = cls_model(im, imgsz=32) - assert len(results) == im.shape[0] - results = pose_model(im, imgsz=32) - assert len(results) == im.shape[0] - results = obb_model(im, imgsz=32) - assert len(results) == im.shape[0] - def test_predict_grey_and_4ch(): """Test YOLO prediction on SOURCE converted to greyscale and 4-channel images.""" @@ -236,19 +221,19 @@ def test_predict_callback_and_setup(): print(boxes) -def test_results(): +@pytest.mark.parametrize("model", MODELS) +def test_results(model): """Test various result formats for the YOLO model.""" - for m in "yolov8n-pose.pt", "yolov8n-seg.pt", "yolov8n.pt", "yolov8n-cls.pt": - results = YOLO(WEIGHTS_DIR / m)([SOURCE, SOURCE], imgsz=160) - for r in results: - r = r.cpu().numpy() - r = r.to(device="cpu", dtype=torch.float32) - r.save_txt(txt_file=TMP / "runs/tests/label.txt", save_conf=True) - r.save_crop(save_dir=TMP / "runs/tests/crops/") - r.tojson(normalize=True) - r.plot(pil=True) - r.plot(conf=True, boxes=True) - print(r, len(r), r.path) + results = YOLO(WEIGHTS_DIR / model)([SOURCE, SOURCE], imgsz=160) + for r in results: + r = r.cpu().numpy() + r = r.to(device="cpu", dtype=torch.float32) + r.save_txt(txt_file=TMP / "runs/tests/label.txt", save_conf=True) + r.save_crop(save_dir=TMP / "runs/tests/crops/") + r.tojson(normalize=True) + r.plot(pil=True) + r.plot(conf=True, boxes=True) + print(r, len(r), r.path) def test_labels_and_crops(): diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 1909d745..153bb770 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, AGPL-3.0 license -__version__ = "8.2.9" +__version__ = "8.2.10" from ultralytics.data.explorer.explorer import Explorer from ultralytics.models import RTDETR, SAM, YOLO, YOLOWorld diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index 322abeb6..e2357151 100644 --- a/ultralytics/cfg/__init__.py +++ b/ultralytics/cfg/__init__.py @@ -53,6 +53,7 @@ TASK2METRIC = { "pose": "metrics/mAP50-95(P)", "obb": "metrics/mAP50-95(B)", } +MODELS = {TASK2MODEL[task] for task in TASKS} ARGV = sys.argv or ["", ""] # sometimes sys.argv = [] CLI_HELP_MSG = f""" diff --git a/ultralytics/engine/results.py b/ultralytics/engine/results.py index ba6f2137..2e1b60f4 100644 --- a/ultralytics/engine/results.py +++ b/ultralytics/engine/results.py @@ -387,26 +387,32 @@ class Results(SimpleClass): def summary(self, normalize=False, decimals=5): """Convert the results to a summarized format.""" - if self.probs is not None: - LOGGER.warning("Warning: Classify results do not support the `summary()` method yet.") - return - # Create list of detection dictionaries results = [] - data = self.boxes.data.cpu().tolist() + if self.probs is not None: + class_id = self.probs.top1 + results.append( + { + "name": self.names[class_id], + "class": class_id, + "confidence": round(self.probs.top1conf.item(), decimals), + } + ) + return results + + data = self.boxes or self.obb + is_obb = self.obb is not None h, w = self.orig_shape if normalize else (1, 1) for i, row in enumerate(data): # xyxy, track_id if tracking, conf, class_id - box = { - "x1": round(row[0] / w, decimals), - "y1": round(row[1] / h, decimals), - "x2": round(row[2] / w, decimals), - "y2": round(row[3] / h, decimals), - } - conf = round(row[-2], decimals) - class_id = int(row[-1]) - result = {"name": self.names[class_id], "class": class_id, "confidence": conf, "box": box} - if self.boxes.is_track: - result["track_id"] = int(row[-3]) # track ID + class_id, conf = int(row.cls), round(row.conf.item(), decimals) + box = (row.xyxyxyxy if is_obb else row.xyxy).squeeze().reshape(-1, 2).tolist() + xy = {} + for i, b in enumerate(box): + xy[f"x{i + 1}"] = round(b[0] / w, decimals) + xy[f"y{i + 1}"] = round(b[1] / h, decimals) + result = {"name": self.names[class_id], "class": class_id, "confidence": conf, "box": xy} + if data.is_track: + result["track_id"] = int(row.id.item()) # track ID if self.masks: result["segments"] = { "x": (self.masks.xy[i][:, 0] / w).round(decimals).tolist(),