ultralytics 8.1.4__tar.gz → 8.1.6__tar.gz

{ultralytics-8.1.4/ultralytics.egg-info → ultralytics-8.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics
-Version: 8.1.4
+Version: 8.1.6
 Summary: Ultralytics YOLOv8 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author: Glenn Jocher, Ayush Chaurasia, Jing Qiu
 Maintainer: Glenn Jocher, Ayush Chaurasia, Jing Qiu
@@ -55,7 +55,7 @@ Requires-Dist: mkdocs-material; extra == "dev"
 Requires-Dist: mkdocstrings[python]; extra == "dev"
 Requires-Dist: mkdocs-jupyter; extra == "dev"
 Requires-Dist: mkdocs-redirects; extra == "dev"
-Requires-Dist: mkdocs-ultralytics-plugin>=0.0.38; extra == "dev"
+Requires-Dist: mkdocs-ultralytics-plugin>=0.0.40; extra == "dev"
 Provides-Extra: export
 Requires-Dist: onnx>=1.12.0; extra == "export"
 Requires-Dist: coremltools>=7.0; platform_system != "Windows" and extra == "export"

{ultralytics-8.1.4 → ultralytics-8.1.6}/pyproject.toml

@@ -93,7 +93,7 @@ dev = [
     "mkdocstrings[python]",
     "mkdocs-jupyter", # for notebooks
     "mkdocs-redirects", # for 301 redirects
-    "mkdocs-ultralytics-plugin>=0.0.38", # for meta descriptions and images, dates and authors
+    "mkdocs-ultralytics-plugin>=0.0.40", # for meta descriptions and images, dates and authors
 ]
 export = [
     "onnx>=1.12.0", # ONNX export

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.1.4"
+__version__ = "8.1.6"
 
 from ultralytics.data.explorer.explorer import Explorer
 from ultralytics.models import RTDETR, SAM, YOLO

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/cfg/__init__.py

@@ -51,7 +51,7 @@ TASK2METRIC = {
     "segment": "metrics/mAP50-95(M)",
     "classify": "metrics/accuracy_top1",
     "pose": "metrics/mAP50-95(P)",
-    "obb": "metrics/mAP50-95(OBB)",
+    "obb": "metrics/mAP50-95(B)",
 }
 
 CLI_HELP_MSG = f"""
@@ -396,6 +396,7 @@ def handle_yolo_settings(args: List[str]) -> None:
 def handle_explorer():
     """Open the Ultralytics Explorer GUI."""
     checks.check_requirements("streamlit")
+    LOGGER.info(f"💡 Loading Explorer dashboard...")
     subprocess.run(["streamlit", "run", ROOT / "data/explorer/gui/dash.py", "--server.maxMessageSize", "2048"])
 
 

ultralytics-8.1.6/ultralytics/cfg/datasets/carparts-seg.yaml

@@ -0,0 +1,43 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Carparts-seg dataset by Ultralytics
+# Documentation: https://docs.ultralytics.com/datasets/segment/carparts-seg/
+# Example usage: yolo train data=carparts-seg.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── carparts-seg  ← downloads here (132 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: ../datasets/carparts-seg # dataset root dir
+train: train/images # train images (relative to 'path') 3516 images
+val: valid/images # val images (relative to 'path') 276 images
+test: test/images # test images (relative to 'path') 401 images
+
+# Classes
+names:
+  0: back_bumper
+  1: back_door
+  2: back_glass
+  3: back_left_door
+  4: back_left_light
+  5: back_light
+  6: back_right_door
+  7: back_right_light
+  8: front_bumper
+  9: front_door
+  10: front_glass
+  11: front_left_door
+  12: front_left_light
+  13: front_light
+  14: front_right_door
+  15: front_right_light
+  16: hood
+  17: left_mirror
+  18: object
+  19: right_mirror
+  20: tailgate
+  21: trunk
+  22: wheel
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/carparts-seg.zip

ultralytics-8.1.6/ultralytics/cfg/datasets/crack-seg.yaml

@@ -0,0 +1,21 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Crack-seg dataset by Ultralytics
+# Documentation: https://docs.ultralytics.com/datasets/segment/crack-seg/
+# Example usage: yolo train data=crack-seg.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── crack-seg  ← downloads here (91.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: ../datasets/crack-seg # dataset root dir
+train: train/images # train images (relative to 'path') 3717 images
+val: valid/images # val images (relative to 'path') 112 images
+test: test/images # test images (relative to 'path') 200 images
+
+# Classes
+names:
+  0: crack
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/crack-seg.zip

ultralytics-8.1.6/ultralytics/cfg/datasets/package-seg.yaml

@@ -0,0 +1,21 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Package-seg dataset by Ultralytics
+# Documentation: https://docs.ultralytics.com/datasets/segment/package-seg/
+# Example usage: yolo train data=package-seg.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── package-seg  ← downloads here (102 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: ../datasets/package-seg # dataset root dir
+train: images/train # train images (relative to 'path') 1920 images
+val: images/val # val images (relative to 'path') 89 images
+test: test/images # test images (relative to 'path') 188 images
+
+# Classes
+names:
+  0: package
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/package-seg.zip

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/data/build.py

@@ -150,13 +150,12 @@ def check_source(source):
     return source, webcam, screenshot, from_img, in_memory, tensor
 
 
-def load_inference_source(source=None, imgsz=640, vid_stride=1, buffer=False):
+def load_inference_source(source=None, vid_stride=1, buffer=False):
     """
     Loads an inference source for object detection and applies necessary transformations.
 
     Args:
         source (str, Path, Tensor, PIL.Image, np.ndarray): The input source for inference.
-        imgsz (int, optional): The size of the image for inference. Default is 640.
         vid_stride (int, optional): The frame interval for video sources. Default is 1.
         buffer (bool, optional): Determined whether stream frames will be buffered. Default is False.
 
@@ -172,13 +171,13 @@ def load_inference_source(source=None, imgsz=640, vid_stride=1, buffer=False):
     elif in_memory:
         dataset = source
     elif webcam:
-        dataset = LoadStreams(source, imgsz=imgsz, vid_stride=vid_stride, buffer=buffer)
+        dataset = LoadStreams(source, vid_stride=vid_stride, buffer=buffer)
     elif screenshot:
-        dataset = LoadScreenshots(source, imgsz=imgsz)
+        dataset = LoadScreenshots(source)
     elif from_img:
-        dataset = LoadPilAndNumpy(source, imgsz=imgsz)
+        dataset = LoadPilAndNumpy(source)
     else:
-        dataset = LoadImages(source, imgsz=imgsz, vid_stride=vid_stride)
+        dataset = LoadImages(source, vid_stride=vid_stride)
 
     # Attach source types to the dataset
     setattr(dataset, "source_type", source_type)

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/data/converter.py

@@ -474,3 +474,66 @@ def merge_multi_segment(segments):
                     nidx = abs(idx[1] - idx[0])
                     s.append(segments[i][nidx:])
     return s
+
+
+def yolo_bbox2segment(im_dir, save_dir=None, sam_model="sam_b.pt"):
+    """
+    Converts existing object detection dataset (bounding boxes) to segmentation dataset or oriented bounding box (OBB)
+    in YOLO format. Generates segmentation data using SAM auto-annotator as needed.
+
+    Args:
+        im_dir (str | Path): Path to image directory to convert.
+        save_dir (str | Path): Path to save the generated labels, labels will be saved
+            into `labels-segment` in the same directory level of `im_dir` if save_dir is None. Default: None.
+        sam_model (str): Segmentation model to use for intermediate segmentation data; optional.
+
+    Notes:
+        The input directory structure assumed for dataset:
+            - im_dir
+                ├─ 001.jpg
+                ├─ ..
+                ├─ NNN.jpg
+            - labels
+                ├─ 001.txt
+                ├─ ..
+                ├─ NNN.txt
+    """
+    from ultralytics.data import YOLODataset
+    from ultralytics.utils.ops import xywh2xyxy
+    from ultralytics.utils import LOGGER
+    from ultralytics import SAM
+    from tqdm import tqdm
+
+    # NOTE: add placeholder to pass class index check
+    dataset = YOLODataset(im_dir, data=dict(names=list(range(1000))))
+    if len(dataset.labels[0]["segments"]) > 0:  # if it's segment data
+        LOGGER.info("Segmentation labels detected, no need to generate new ones!")
+        return
+
+    LOGGER.info("Detection labels detected, generating segment labels by SAM model!")
+    sam_model = SAM(sam_model)
+    for l in tqdm(dataset.labels, total=len(dataset.labels), desc="Generating segment labels"):
+        h, w = l["shape"]
+        boxes = l["bboxes"]
+        if len(boxes) == 0:  # skip empty labels
+            continue
+        boxes[:, [0, 2]] *= w
+        boxes[:, [1, 3]] *= h
+        im = cv2.imread(l["im_file"])
+        sam_results = sam_model(im, bboxes=xywh2xyxy(boxes), verbose=False, save=False)
+        l["segments"] = sam_results[0].masks.xyn
+
+    save_dir = Path(save_dir) if save_dir else Path(im_dir).parent / "labels-segment"
+    save_dir.mkdir(parents=True, exist_ok=True)
+    for l in dataset.labels:
+        texts = []
+        lb_name = Path(l["im_file"]).with_suffix(".txt").name
+        txt_file = save_dir / lb_name
+        cls = l["cls"]
+        for i, s in enumerate(l["segments"]):
+            line = (int(cls[i]), *s.reshape(-1))
+            texts.append(("%g " * len(line)).rstrip() % line)
+        if texts:
+            with open(txt_file, "a") as f:
+                f.writelines(text + "\n" for text in texts)
+    LOGGER.info(f"Generated segment labels saved in {save_dir}")

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/data/explorer/gui/dash.py

@@ -9,7 +9,7 @@ from ultralytics import Explorer
 from ultralytics.utils import ROOT, SETTINGS
 from ultralytics.utils.checks import check_requirements
 
-check_requirements(("streamlit>=1.29.0", "streamlit-select>=0.2"))
+check_requirements(("streamlit>=1.29.0", "streamlit-select>=0.3"))
 
 import streamlit as st
 from streamlit_select import image_select
@@ -94,6 +94,7 @@ def find_similar_imgs(imgs):
     similar = exp.get_similar(img=imgs, limit=st.session_state.get("limit"), return_type="arrow")
     paths = similar.to_pydict()["im_file"]
     st.session_state["imgs"] = paths
+    st.session_state["res"] = similar
 
 
 def similarity_form(selected_imgs):
@@ -137,6 +138,7 @@ def run_sql_query():
         exp = st.session_state["explorer"]
         res = exp.sql_query(query, return_type="arrow")
         st.session_state["imgs"] = res.to_pydict()["im_file"]
+        st.session_state["res"] = res
 
 
 def run_ai_query():
@@ -155,6 +157,7 @@ def run_ai_query():
             st.session_state["error"] = "No results found using AI generated query. Try another query or rerun it."
             return
         st.session_state["imgs"] = res["im_file"].to_list()
+        st.session_state["res"] = res
 
 
 def reset_explorer():
@@ -195,7 +198,11 @@ def layout():
     if st.session_state.get("error"):
         st.error(st.session_state["error"])
     else:
-        imgs = st.session_state.get("imgs") or exp.table.to_lance().to_table(columns=["im_file"]).to_pydict()["im_file"]
+        if st.session_state.get("imgs"):
+            imgs = st.session_state.get("imgs")
+        else:
+            imgs = exp.table.to_lance().to_table(columns=["im_file"]).to_pydict()["im_file"]
+            st.session_state["res"] = exp.table.to_arrow()
     total_imgs, selected_imgs = len(imgs), []
     with col1:
         subcol1, subcol2, subcol3, subcol4, subcol5 = st.columns(5)
@@ -230,17 +237,30 @@ def layout():
         query_form()
         ai_query_form()
         if total_imgs:
+            labels, boxes, masks, kpts, classes = None, None, None, None, None
+            task = exp.model.task
+            if st.session_state.get("display_labels"):
+                labels = st.session_state.get("res").to_pydict()["labels"][start_idx : start_idx + num]
+                boxes = st.session_state.get("res").to_pydict()["bboxes"][start_idx : start_idx + num]
+                masks = st.session_state.get("res").to_pydict()["masks"][start_idx : start_idx + num]
+                kpts = st.session_state.get("res").to_pydict()["keypoints"][start_idx : start_idx + num]
+                classes = st.session_state.get("res").to_pydict()["cls"][start_idx : start_idx + num]
             imgs_displayed = imgs[start_idx : start_idx + num]
             selected_imgs = image_select(
                 f"Total samples: {total_imgs}",
                 images=imgs_displayed,
                 use_container_width=False,
                 # indices=[i for i in range(num)] if select_all else None,
+                labels=labels,
+                classes=classes,
+                bboxes=boxes,
+                masks=masks if task == "segment" else None,
+                kpts=kpts if task == "pose" else None,
             )
 
     with col2:
         similarity_form(selected_imgs)
-        # display_labels = st.checkbox("Labels", value=False, key="display_labels")
+        display_labels = st.checkbox("Labels", value=False, key="display_labels")
         utralytics_explorer_docs_callback()
 
 

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/data/loaders.py

@@ -38,7 +38,6 @@ class LoadStreams:
 
     Attributes:
         sources (str): The source input paths or URLs for the video streams.
-        imgsz (int): The image size for processing, defaults to 640.
         vid_stride (int): Video frame-rate stride, defaults to 1.
         buffer (bool): Whether to buffer input streams, defaults to False.
         running (bool): Flag to indicate if the streaming thread is running.
@@ -60,13 +59,12 @@ class LoadStreams:
         __len__: Return the length of the sources object.
     """
 
-    def __init__(self, sources="file.streams", imgsz=640, vid_stride=1, buffer=False):
+    def __init__(self, sources="file.streams", vid_stride=1, buffer=False):
         """Initialize instance variables and check for consistent input stream shapes."""
         torch.backends.cudnn.benchmark = True  # faster for fixed-size inference
         self.buffer = buffer  # buffer input streams
         self.running = True  # running flag for Thread
         self.mode = "stream"
-        self.imgsz = imgsz
         self.vid_stride = vid_stride  # video frame-rate stride
 
         sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
@@ -193,7 +191,6 @@ class LoadScreenshots:
 
     Attributes:
         source (str): The source input indicating which screen to capture.
-        imgsz (int): The image size for processing, defaults to 640.
         screen (int): The screen number to capture.
         left (int): The left coordinate for screen capture area.
         top (int): The top coordinate for screen capture area.
@@ -210,7 +207,7 @@ class LoadScreenshots:
         __next__: Captures the next screenshot and returns it.
     """
 
-    def __init__(self, source, imgsz=640):
+    def __init__(self, source):
         """Source = [screen_number left top width height] (pixels)."""
         check_requirements("mss")
         import mss  # noqa
@@ -223,7 +220,6 @@ class LoadScreenshots:
             left, top, width, height = (int(x) for x in params)
         elif len(params) == 5:
             self.screen, left, top, width, height = (int(x) for x in params)
-        self.imgsz = imgsz
         self.mode = "stream"
         self.frame = 0
         self.sct = mss.mss()
@@ -258,7 +254,6 @@ class LoadImages:
     various formats, including single image files, video files, and lists of image and video paths.
 
     Attributes:
-        imgsz (int): Image size, defaults to 640.
         files (list): List of image and video file paths.
         nf (int): Total number of files (images and videos).
         video_flag (list): Flags indicating whether a file is a video (True) or an image (False).
@@ -274,7 +269,7 @@ class LoadImages:
         _new_video(path): Create a new cv2.VideoCapture object for a given video path.
     """
 
-    def __init__(self, path, imgsz=640, vid_stride=1):
+    def __init__(self, path, vid_stride=1):
         """Initialize the Dataloader and raise FileNotFoundError if file not found."""
         parent = None
         if isinstance(path, str) and Path(path).suffix == ".txt":  # *.txt file with img/vid/dir on each line
@@ -298,7 +293,6 @@ class LoadImages:
         videos = [x for x in files if x.split(".")[-1].lower() in VID_FORMATS]
         ni, nv = len(images), len(videos)
 
-        self.imgsz = imgsz
         self.files = images + videos
         self.nf = ni + nv  # number of files
         self.video_flag = [False] * ni + [True] * nv
@@ -377,7 +371,6 @@ class LoadPilAndNumpy:
     Attributes:
         paths (list): List of image paths or autogenerated filenames.
         im0 (list): List of images stored as Numpy arrays.
-        imgsz (int): Image size, defaults to 640.
         mode (str): Type of data being processed, defaults to 'image'.
         bs (int): Batch size, equivalent to the length of `im0`.
         count (int): Counter for iteration, initialized at 0 during `__iter__()`.
@@ -386,13 +379,12 @@ class LoadPilAndNumpy:
         _single_check(im): Validate and format a single image to a Numpy array.
     """
 
-    def __init__(self, im0, imgsz=640):
+    def __init__(self, im0):
         """Initialize PIL and Numpy Dataloader."""
         if not isinstance(im0, list):
             im0 = [im0]
         self.paths = [getattr(im, "filename", f"image{i}.jpg") for i, im in enumerate(im0)]
         self.im0 = [self._single_check(im) for im in im0]
-        self.imgsz = imgsz
         self.mode = "image"
         # Generate fake paths
         self.bs = len(self.im0)

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/engine/model.py

@@ -259,8 +259,8 @@ class Model(nn.Module):
             x in sys.argv for x in ("predict", "track", "mode=predict", "mode=track")
         )
 
-        custom = {"conf": 0.25, "save": is_cli}  # method defaults
-        args = {**self.overrides, **custom, **kwargs, "mode": "predict"}  # highest priority args on the right
+        custom = {"conf": 0.25, "save": is_cli, "mode": "predict"}  # method defaults
+        args = {**self.overrides, **custom, **kwargs}  # highest priority args on the right
         prompts = args.pop("prompts", None)  # for SAM-type models
 
         if not self.predictor:

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/engine/predictor.py

@@ -226,7 +226,7 @@ class BasePredictor:
             else None
         )
         self.dataset = load_inference_source(
-            source=source, imgsz=self.imgsz, vid_stride=self.args.vid_stride, buffer=self.args.stream_buffer
+            source=source, vid_stride=self.args.vid_stride, buffer=self.args.stream_buffer
         )
         self.source_type = self.dataset.source_type
         if not getattr(self, "stream", True) and (

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/engine/results.py

@@ -115,7 +115,7 @@ class Results(SimpleClass):
             if v is not None:
                 return len(v)
 
-    def update(self, boxes=None, masks=None, probs=None):
+    def update(self, boxes=None, masks=None, probs=None, obb=None):
         """Update the boxes, masks, and probs attributes of the Results object."""
         if boxes is not None:
             self.boxes = Boxes(ops.clip_boxes(boxes, self.orig_shape), self.orig_shape)
@@ -123,6 +123,8 @@ class Results(SimpleClass):
             self.masks = Masks(masks, self.orig_shape)
         if probs is not None:
             self.probs = probs
+        if obb is not None:
+            self.obb = OBB(obb, self.orig_shape)
 
     def _apply(self, fn, *args, **kwargs):
         """

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/models/yolo/obb/predict.py

@@ -45,8 +45,9 @@ class OBBPredictor(DetectionPredictor):
 
         results = []
         for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
-            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape, xywh=True)
+            rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
+            rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)
             # xywh, r, conf, cls
-            obb = torch.cat([pred[:, :4], pred[:, -1:], pred[:, 4:6]], dim=-1)
+            obb = torch.cat([rboxes, pred[:, 4:6]], dim=-1)
             results.append(Results(orig_img, path=img_path, names=self.model.names, obb=obb))
         return results

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/trackers/byte_tracker.py

@@ -5,6 +5,8 @@ import numpy as np
 from .basetrack import BaseTrack, TrackState
 from .utils import matching
 from .utils.kalman_filter import KalmanFilterXYAH
+from ..utils.ops import xywh2ltwh
+from ..utils import LOGGER
 
 
 class STrack(BaseTrack):
@@ -35,18 +37,18 @@ class STrack(BaseTrack):
         activate(kalman_filter, frame_id): Activate a new tracklet.
         re_activate(new_track, frame_id, new_id): Reactivate a previously lost tracklet.
         update(new_track, frame_id): Update the state of a matched track.
-        convert_coords(tlwh): Convert bounding box to x-y-angle-height format.
+        convert_coords(tlwh): Convert bounding box to x-y-aspect-height format.
         tlwh_to_xyah(tlwh): Convert tlwh bounding box to xyah format.
-        tlbr_to_tlwh(tlbr): Convert tlbr bounding box to tlwh format.
-        tlwh_to_tlbr(tlwh): Convert tlwh bounding box to tlbr format.
     """
 
     shared_kalman = KalmanFilterXYAH()
 
-    def __init__(self, tlwh, score, cls):
+    def __init__(self, xywh, score, cls):
         """Initialize new STrack instance."""
         super().__init__()
-        self._tlwh = np.asarray(self.tlbr_to_tlwh(tlwh[:-1]), dtype=np.float32)
+        # xywh+idx or xywha+idx
+        assert len(xywh) in [5, 6], f"expected 5 or 6 values but got {len(xywh)}"
+        self._tlwh = np.asarray(xywh2ltwh(xywh[:4]), dtype=np.float32)
         self.kalman_filter = None
         self.mean, self.covariance = None, None
         self.is_activated = False
@@ -54,7 +56,8 @@ class STrack(BaseTrack):
         self.score = score
         self.tracklet_len = 0
         self.cls = cls
-        self.idx = tlwh[-1]
+        self.idx = xywh[-1]
+        self.angle = xywh[4] if len(xywh) == 6 else None
 
     def predict(self):
         """Predicts mean and covariance using Kalman filter."""
@@ -123,6 +126,7 @@ class STrack(BaseTrack):
             self.track_id = self.next_id()
         self.score = new_track.score
         self.cls = new_track.cls
+        self.angle = new_track.angle
         self.idx = new_track.idx
 
     def update(self, new_track, frame_id):
@@ -145,10 +149,11 @@ class STrack(BaseTrack):
 
         self.score = new_track.score
         self.cls = new_track.cls
+        self.angle = new_track.angle
         self.idx = new_track.idx
 
     def convert_coords(self, tlwh):
-        """Convert a bounding box's top-left-width-height format to its x-y-angle-height equivalent."""
+        """Convert a bounding box's top-left-width-height format to its x-y-aspect-height equivalent."""
         return self.tlwh_to_xyah(tlwh)
 
     @property
@@ -162,7 +167,7 @@ class STrack(BaseTrack):
         return ret
 
     @property
-    def tlbr(self):
+    def xyxy(self):
         """Convert bounding box to format (min x, min y, max x, max y), i.e., (top left, bottom right)."""
         ret = self.tlwh.copy()
         ret[2:] += ret[:2]
@@ -178,19 +183,26 @@ class STrack(BaseTrack):
         ret[2] /= ret[3]
         return ret
 
-    @staticmethod
-    def tlbr_to_tlwh(tlbr):
-        """Converts top-left bottom-right format to top-left width height format."""
-        ret = np.asarray(tlbr).copy()
-        ret[2:] -= ret[:2]
+    @property
+    def xywh(self):
+        """Get current position in bounding box format (center x, center y, width, height)."""
+        ret = np.asarray(self.tlwh).copy()
+        ret[:2] += ret[2:] / 2
         return ret
 
-    @staticmethod
-    def tlwh_to_tlbr(tlwh):
-        """Converts tlwh bounding box format to tlbr format."""
-        ret = np.asarray(tlwh).copy()
-        ret[2:] += ret[:2]
-        return ret
+    @property
+    def xywha(self):
+        """Get current position in bounding box format (center x, center y, width, height, angle)."""
+        if self.angle is None:
+            LOGGER.warning("WARNING ⚠️ `angle` attr not found, returning `xywh` instead.")
+            return self.xywh
+        return np.concatenate([self.xywh, self.angle[None]])
+
+    @property
+    def result(self):
+        """Get current tracking results."""
+        coords = self.xyxy if self.angle is None else self.xywha
+        return coords.tolist() + [self.track_id, self.score, self.cls, self.idx]
 
     def __repr__(self):
         """Return a string representation of the BYTETracker object with start and end frames and track ID."""
@@ -247,7 +259,7 @@ class BYTETracker:
         removed_stracks = []
 
         scores = results.conf
-        bboxes = results.xyxy
+        bboxes = results.xywhr if hasattr(results, "xywhr") else results.xywh
         # Add index
         bboxes = np.concatenate([bboxes, np.arange(len(bboxes)).reshape(-1, 1)], axis=-1)
         cls = results.cls
@@ -349,10 +361,8 @@ class BYTETracker:
         self.removed_stracks.extend(removed_stracks)
         if len(self.removed_stracks) > 1000:
             self.removed_stracks = self.removed_stracks[-999:]  # clip remove stracks to 1000 maximum
-        return np.asarray(
-            [x.tlbr.tolist() + [x.track_id, x.score, x.cls, x.idx] for x in self.tracked_stracks if x.is_activated],
-            dtype=np.float32,
-        )
+
+        return np.asarray([x.result for x in self.tracked_stracks if x.is_activated], dtype=np.float32)
 
     def get_kalmanfilter(self):
         """Returns a Kalman filter object for tracking bounding boxes."""

{ultralytics-8.1.4 → ultralytics-8.1.6}/ultralytics/trackers/track.py

@@ -25,8 +25,6 @@ def on_predict_start(predictor: object, persist: bool = False) -> None:
     Raises:
         AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'.
     """
-    if predictor.args.task == "obb":
-        raise NotImplementedError("ERROR ❌ OBB task does not support track mode!")
     if hasattr(predictor, "trackers") and persist:
         return
 
@@ -54,11 +52,12 @@ def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None
     bs = predictor.dataset.bs
     path, im0s = predictor.batch[:2]
 
+    is_obb = predictor.args.task == "obb"
     for i in range(bs):
         if not persist and predictor.vid_path[i] != str(predictor.save_dir / Path(path[i]).name):  # new video
             predictor.trackers[i].reset()
 
-        det = predictor.results[i].boxes.cpu().numpy()
+        det = (predictor.results[i].obb if is_obb else predictor.results[i].boxes).cpu().numpy()
         if len(det) == 0:
             continue
         tracks = predictor.trackers[i].update(det, im0s[i])
@@ -66,7 +65,10 @@ def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None
             continue
         idx = tracks[:, -1].astype(int)
         predictor.results[i] = predictor.results[i][idx]
-        predictor.results[i].update(boxes=torch.as_tensor(tracks[:, :-1]))
+
+        update_args = dict()
+        update_args["obb" if is_obb else "boxes"] = torch.as_tensor(tracks[:, :-1])
+        predictor.results[i].update(**update_args)
 
 
 def register_tracker(model: object, persist: bool) -> None: