ultralytics 8.3.165__tar.gz → 8.3.167__tar.gz

{ultralytics-8.3.165/ultralytics.egg-info → ultralytics-8.3.167}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ultralytics
-Version: 8.3.165
+Version: 8.3.167
 Summary: Ultralytics YOLO 🚀 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.165"
+__version__ = "8.3.167"
 
 import os
 

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/Argoverse.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── Argoverse  ← downloads here (31.5 GB)
+#     └── Argoverse ← downloads here (31.5 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: Argoverse # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/DOTAv1.5.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── dota1.5  ← downloads here (2GB)
+#     └── dota1.5 ← downloads here (2GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: DOTAv1.5 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/DOTAv1.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── dota1  ← downloads here (2GB)
+#     └── dota1 ← downloads here (2GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: DOTAv1 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/GlobalWheat2020.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── GlobalWheat2020  ← downloads here (7.0 GB)
+#     └── GlobalWheat2020 ← downloads here (7.0 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: GlobalWheat2020 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/HomeObjects-3K.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── homeobjects-3K  ← downloads here (390 MB)
+#     └── homeobjects-3K ← downloads here (390 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: homeobjects-3K # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/ImageNet.yaml

@@ -7,7 +7,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── imagenet  ← downloads here (144 GB)
+#     └── imagenet ← downloads here (144 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: imagenet # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/Objects365.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── Objects365  ← downloads here (712 GB = 367G data + 345G zips)
+#     └── Objects365 ← downloads here (712 GB = 367G data + 345G zips)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: Objects365 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/SKU-110K.yaml

@@ -6,13 +6,13 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── SKU-110K  ← downloads here (13.6 GB)
+#     └── SKU-110K ← downloads here (13.6 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: SKU-110K # dataset root dir
-train: train.txt # train images (relative to 'path')  8219 images
-val: val.txt # val images (relative to 'path')  588 images
-test: test.txt # test images (optional)  2936 images
+train: train.txt # train images (relative to 'path') 8219 images
+val: val.txt # val images (relative to 'path') 588 images
+test: test.txt # test images (optional) 2936 images
 
 # Classes
 names:

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/VOC.yaml

@@ -6,16 +6,16 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── VOC  ← downloads here (2.8 GB)
+#     └── VOC ← downloads here (2.8 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: VOC
-train: # train images (relative to 'path')  16551 images
+train: # train images (relative to 'path') 16551 images
   - images/train2012
   - images/train2007
   - images/val2012
   - images/val2007
-val: # val images (relative to 'path')  4952 images
+val: # val images (relative to 'path') 4952 images
   - images/test2007
 test: # test images (optional)
   - images/test2007

ultralytics-8.3.167/ultralytics/cfg/datasets/VisDrone.yaml

@@ -0,0 +1,87 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University
+# Documentation: https://docs.ultralytics.com/datasets/detect/visdrone/
+# Example usage: yolo train data=VisDrone.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── VisDrone ← downloads here (2.3 GB)
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: VisDrone # dataset root dir
+train: images/train # train images (relative to 'path') 6471 images
+val: images/val # val images (relative to 'path') 548 images
+test: images/test # test-dev images (optional) 1610 images
+
+# Classes
+names:
+  0: pedestrian
+  1: people
+  2: bicycle
+  3: car
+  4: van
+  5: truck
+  6: tricycle
+  7: awning-tricycle
+  8: bus
+  9: motor
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import os
+  from pathlib import Path
+  import shutil
+
+  from ultralytics.utils.downloads import download
+
+
+  def visdrone2yolo(dir, split, source_name=None):
+      """Convert VisDrone annotations to YOLO format with images/{split} and labels/{split} structure."""
+      from PIL import Image
+      from tqdm import tqdm
+
+      source_dir = dir / (source_name or f"VisDrone2019-DET-{split}")
+      images_dir = dir / "images" / split
+      labels_dir = dir / "labels" / split
+      labels_dir.mkdir(parents=True, exist_ok=True)
+      
+      # Move images to new structure
+      if (source_images_dir := source_dir / "images").exists():
+          images_dir.mkdir(parents=True, exist_ok=True)
+          for img in source_images_dir.glob("*.jpg"):
+              img.rename(images_dir / img.name)
+      
+      for f in tqdm((source_dir / "annotations").glob("*.txt"), desc=f"Converting {split}"):
+          img_size = Image.open(images_dir / f.with_suffix(".jpg").name).size
+          dw, dh = 1.0 / img_size[0], 1.0 / img_size[1]
+          lines = []
+          
+          with open(f, encoding="utf-8") as file:
+              for row in [x.split(",") for x in file.read().strip().splitlines()]:
+                  if row[4] != "0":  # Skip ignored regions
+                      x, y, w, h = map(int, row[:4])
+                      cls = int(row[5]) - 1
+                      # Convert to YOLO format
+                      x_center, y_center = (x + w / 2) * dw, (y + h / 2) * dh
+                      w_norm, h_norm = w * dw, h * dh
+                      lines.append(f"{cls} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}\n")
+          
+          (labels_dir / f.name).write_text("".join(lines), encoding="utf-8")
+
+
+  # Download (ignores test-challenge split)
+  dir = Path(yaml["path"])  # dataset root dir
+  urls = [
+      "https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-train.zip",
+      "https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-val.zip",
+      "https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-test-dev.zip",
+      # "https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-test-challenge.zip",
+  ]
+  download(urls, dir=dir, curl=True, threads=4)
+
+  # Convert
+  splits = {"VisDrone2019-DET-train": "train", "VisDrone2019-DET-val": "val", "VisDrone2019-DET-test-dev": "test"}
+  for folder, split in splits.items():
+      visdrone2yolo(dir, split, folder)  # convert VisDrone annotations to YOLO labels
+      shutil.rmtree(dir / folder)  # cleanup original directory

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/african-wildlife.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── african-wildlife  ← downloads here (100 MB)
+#     └── african-wildlife ← downloads here (100 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: african-wildlife # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/brain-tumor.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── brain-tumor  ← downloads here (4.21 MB)
+#     └── brain-tumor ← downloads here (4.21 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: brain-tumor # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/carparts-seg.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── carparts-seg  ← downloads here (133 MB)
+#     └── carparts-seg ← downloads here (133 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: carparts-seg # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco-pose.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco-pose  ← downloads here (20.1 GB)
+#     └── coco-pose ← downloads here (20.1 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco-pose # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco  ← downloads here (20.1 GB)
+#     └── coco ← downloads here (20.1 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco128-seg.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco128-seg  ← downloads here (7 MB)
+#     └── coco128-seg ← downloads here (7 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco128-seg # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco128.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco128  ← downloads here (7 MB)
+#     └── coco128 ← downloads here (7 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco128 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco8-grayscale.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco8-grayscale  ← downloads here (1 MB)
+#     └── coco8-grayscale ← downloads here (1 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco8-grayscale # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco8-multispectral.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco8-multispectral  ← downloads here (20.2 MB)
+#     └── coco8-multispectral ← downloads here (20.2 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco8-multispectral # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco8-pose.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco8-pose  ← downloads here (1 MB)
+#     └── coco8-pose ← downloads here (1 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco8-pose # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco8-seg.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco8-seg  ← downloads here (1 MB)
+#     └── coco8-seg ← downloads here (1 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco8-seg # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/coco8.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── coco8  ← downloads here (1 MB)
+#     └── coco8 ← downloads here (1 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: coco8 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/crack-seg.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── crack-seg  ← downloads here (91.6 MB)
+#     └── crack-seg ← downloads here (91.6 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: crack-seg # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/dog-pose.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── dog-pose  ← downloads here (337 MB)
+#     └── dog-pose ← downloads here (337 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: dog-pose # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/dota8-multispectral.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── dota8-multispectral  ← downloads here (37.3MB)
+#     └── dota8-multispectral ← downloads here (37.3MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: dota8-multispectral # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/dota8.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── dota8  ← downloads here (1MB)
+#     └── dota8 ← downloads here (1MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: dota8 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/hand-keypoints.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── hand-keypoints  ← downloads here (369 MB)
+#     └── hand-keypoints ← downloads here (369 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: hand-keypoints # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/lvis.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── lvis  ← downloads here (20.1 GB)
+#     └── lvis ← downloads here (20.1 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: lvis # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/medical-pills.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── medical-pills  ← downloads here (8.19 MB)
+#     └── medical-pills ← downloads here (8.19 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: medical-pills # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/open-images-v7.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── open-images-v7  ← downloads here (561 GB)
+#     └── open-images-v7 ← downloads here (561 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: open-images-v7 # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/package-seg.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── package-seg  ← downloads here (103 MB)
+#     └── package-seg ← downloads here (103 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: package-seg # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/signature.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── signature  ← downloads here (11.3 MB)
+#     └── signature ← downloads here (11.3 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: signature # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/tiger-pose.yaml

@@ -6,7 +6,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── tiger-pose  ← downloads here (22.8 MB)
+#     └── tiger-pose ← downloads here (49.8 MB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: tiger-pose # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/cfg/datasets/xView.yaml

@@ -7,7 +7,7 @@
 # parent
 # ├── ultralytics
 # └── datasets
-#     └── xView  ← downloads here (20.7 GB)
+#     └── xView ← downloads here (20.7 GB)
 
 # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
 path: xView # dataset root dir

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/engine/exporter.py

@@ -142,7 +142,7 @@ def export_formats():
         ["PaddlePaddle", "paddle", "_paddle_model", True, True, ["batch"]],
         ["MNN", "mnn", ".mnn", True, True, ["batch", "half", "int8"]],
         ["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
-        ["IMX", "imx", "_imx_model", True, True, ["int8", "fraction"]],
+        ["IMX", "imx", "_imx_model", True, True, ["int8", "fraction", "nms"]],
         ["RKNN", "rknn", "_rknn_model", False, False, ["batch", "name"]],
     ]
     return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
@@ -313,8 +313,11 @@ class Exporter:
             if not self.args.int8:
                 LOGGER.warning("IMX export requires int8=True, setting int8=True.")
                 self.args.int8 = True
-            if model.task != "detect":
-                raise ValueError("IMX export only supported for detection models.")
+            if not self.args.nms:
+                LOGGER.warning("IMX export requires nms=True, setting nms=True.")
+                self.args.nms = True
+            if model.task not in {"detect", "pose"}:
+                raise ValueError("IMX export only supported for detection and pose estimation models.")
         if not hasattr(model, "names"):
             model.names = default_class_names()
         model.names = check_class_names(model.names)
@@ -428,7 +431,7 @@ class Exporter:
 
         y = None
         for _ in range(2):  # dry runs
-            y = NMSModel(model, self.args)(im) if self.args.nms and not coreml else model(im)
+            y = NMSModel(model, self.args)(im) if self.args.nms and not (coreml or imx) else model(im)
         if self.args.half and onnx and self.device.type != "cpu":
             im, model = im.half(), model.half()  # to FP16
 
@@ -1166,15 +1169,14 @@ class Exporter:
         )
         if getattr(self.model, "end2end", False):
             raise ValueError("IMX export is not supported for end2end models.")
-        check_requirements(
-            ("model-compression-toolkit>=2.3.0,<2.4.1", "sony-custom-layers>=0.3.0", "edge-mdt-tpc>=1.1.0")
-        )
+        check_requirements(("model-compression-toolkit>=2.4.1", "sony-custom-layers>=0.3.0", "edge-mdt-tpc>=1.1.0"))
         check_requirements("imx500-converter[pt]>=3.16.1")  # Separate requirements for imx500-converter
+        check_requirements("mct-quantizers>=1.6.0")  # Separate for compatibility with model-compression-toolkit
 
         import model_compression_toolkit as mct
         import onnx
         from edgemdt_tpc import get_target_platform_capabilities
-        from sony_custom_layers.pytorch import multiclass_nms
+        from sony_custom_layers.pytorch import multiclass_nms_with_indices
 
         LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")
 
@@ -1198,13 +1200,23 @@ class Exporter:
 
         bit_cfg = mct.core.BitWidthConfig()
         if "C2PSA" in self.model.__str__():  # YOLO11
-            layer_names = ["sub", "mul_2", "add_14", "cat_21"]
-            weights_memory = 2585350.2439
-            n_layers = 238  # 238 layers for fused YOLO11n
+            if self.model.task == "detect":
+                layer_names = ["sub", "mul_2", "add_14", "cat_21"]
+                weights_memory = 2585350.2439
+                n_layers = 238  # 238 layers for fused YOLO11n
+            elif self.model.task == "pose":
+                layer_names = ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"]
+                weights_memory = 2437771.67
+                n_layers = 257  # 257 layers for fused YOLO11n-pose
         else:  # YOLOv8
-            layer_names = ["sub", "mul", "add_6", "cat_17"]
-            weights_memory = 2550540.8
-            n_layers = 168  # 168 layers for fused YOLOv8n
+            if self.model.task == "detect":
+                layer_names = ["sub", "mul", "add_6", "cat_17"]
+                weights_memory = 2550540.8
+                n_layers = 168  # 168 layers for fused YOLOv8n
+            elif self.model.task == "pose":
+                layer_names = ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"]
+                weights_memory = 2482451.85
+                n_layers = 187  # 187 layers for fused YOLO11n-pose
 
         # Check if the model has the expected number of layers
         if len(list(self.model.modules())) != n_layers:
@@ -1251,6 +1263,7 @@ class Exporter:
                 score_threshold: float = 0.001,
                 iou_threshold: float = 0.7,
                 max_detections: int = 300,
+                task: str = "detect",
             ):
                 """
                 Initialize NMSWrapper with PyTorch Module and NMS parameters.
@@ -1260,34 +1273,40 @@ class Exporter:
                     score_threshold (float): Score threshold for non-maximum suppression.
                     iou_threshold (float): Intersection over union threshold for non-maximum suppression.
                     max_detections (int): The number of detections to return.
+                    task (str): Task type, either 'detect' or 'pose'.
                 """
                 super().__init__()
                 self.model = model
                 self.score_threshold = score_threshold
                 self.iou_threshold = iou_threshold
                 self.max_detections = max_detections
+                self.task = task
 
             def forward(self, images):
                 """Forward pass with model inference and NMS post-processing."""
                 # model inference
                 outputs = self.model(images)
 
-                boxes = outputs[0]
-                scores = outputs[1]
-                nms = multiclass_nms(
+                boxes, scores = outputs[0], outputs[1]
+                nms_outputs = multiclass_nms_with_indices(
                     boxes=boxes,
                     scores=scores,
                     score_threshold=self.score_threshold,
                     iou_threshold=self.iou_threshold,
                     max_detections=self.max_detections,
                 )
-                return nms
+                if self.task == "pose":
+                    kpts = outputs[2]  # (bs, max_detections, kpts 17*3)
+                    out_kpts = torch.gather(kpts, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, kpts.size(-1)))
+                    return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_kpts
+                return nms_outputs
 
         quant_model = NMSWrapper(
             model=quant_model,
             score_threshold=self.args.conf or 0.001,
             iou_threshold=self.args.iou,
             max_detections=self.args.max_det,
+            task=self.model.task,
         ).to(self.device)
 
         f = Path(str(self.file).replace(self.file.suffix, "_imx_model"))

{ultralytics-8.3.165 → ultralytics-8.3.167}/ultralytics/models/yolo/model.py

@@ -243,10 +243,6 @@ class YOLOE(Model):
         """
         super().__init__(model=model, task=task, verbose=verbose)
 
-        # Assign default COCO class names when there are no custom names
-        if not hasattr(self.model, "names"):
-            self.model.names = YAML.load(ROOT / "cfg/datasets/coco8.yaml").get("names")
-
     @property
     def task_map(self) -> Dict[str, Dict[str, Any]]:
         """Map head to model, validator, and predictor classes."""
@@ -287,7 +283,7 @@ class YOLOE(Model):
         Examples:
             >>> model = YOLOE("yoloe-11s-seg.pt")
             >>> img = torch.rand(1, 3, 640, 640)
-            >>> visual_features = model.model.backbone(img)
+            >>> visual_features = torch.rand(1, 1, 80, 80)
             >>> pe = model.get_visual_pe(img, visual_features)
         """
         assert isinstance(self.model, YOLOEModel)