ultralytics 8.0.237__py3-none-any.whl → 8.0.239__py3-none-any.whl

ultralytics/engine/results.py

@@ -98,15 +98,15 @@ class Results(SimpleClass):
         self.probs = Probs(probs) if probs is not None else None
         self.keypoints = Keypoints(keypoints, self.orig_shape) if keypoints is not None else None
         self.obb = OBB(obb, self.orig_shape) if obb is not None else None
-        self.speed = {'preprocess': None, 'inference': None, 'postprocess': None}  # milliseconds per image
+        self.speed = {"preprocess": None, "inference": None, "postprocess": None}  # milliseconds per image
         self.names = names
         self.path = path
         self.save_dir = None
-        self._keys = 'boxes', 'masks', 'probs', 'keypoints', 'obb'
+        self._keys = "boxes", "masks", "probs", "keypoints", "obb"
 
     def __getitem__(self, idx):
         """Return a Results object for the specified index."""
-        return self._apply('__getitem__', idx)
+        return self._apply("__getitem__", idx)
 
     def __len__(self):
         """Return the number of detections in the Results object."""
@@ -146,19 +146,19 @@ class Results(SimpleClass):
 
     def cpu(self):
         """Return a copy of the Results object with all tensors on CPU memory."""
-        return self._apply('cpu')
+        return self._apply("cpu")
 
     def numpy(self):
         """Return a copy of the Results object with all tensors as numpy arrays."""
-        return self._apply('numpy')
+        return self._apply("numpy")
 
     def cuda(self):
         """Return a copy of the Results object with all tensors on GPU memory."""
-        return self._apply('cuda')
+        return self._apply("cuda")
 
     def to(self, *args, **kwargs):
         """Return a copy of the Results object with tensors on the specified device and dtype."""
-        return self._apply('to', *args, **kwargs)
+        return self._apply("to", *args, **kwargs)
 
     def new(self):
         """Return a new Results object with the same image, path, and names."""
@@ -169,7 +169,7 @@ class Results(SimpleClass):
         conf=True,
         line_width=None,
         font_size=None,
-        font='Arial.ttf',
+        font="Arial.ttf",
         pil=False,
         img=None,
         im_gpu=None,
@@ -229,14 +229,20 @@ class Results(SimpleClass):
             font_size,
             font,
             pil or (pred_probs is not None and show_probs),  # Classify tasks default to pil=True
-            example=names)
+            example=names,
+        )
 
         # Plot Segment results
         if pred_masks and show_masks:
             if im_gpu is None:
                 img = LetterBox(pred_masks.shape[1:])(image=annotator.result())
-                im_gpu = torch.as_tensor(img, dtype=torch.float16, device=pred_masks.data.device).permute(
-                    2, 0, 1).flip(0).contiguous() / 255
+                im_gpu = (
+                    torch.as_tensor(img, dtype=torch.float16, device=pred_masks.data.device)
+                    .permute(2, 0, 1)
+                    .flip(0)
+                    .contiguous()
+                    / 255
+                )
             idx = pred_boxes.cls if pred_boxes else range(len(pred_masks))
             annotator.masks(pred_masks.data, colors=[colors(x, True) for x in idx], im_gpu=im_gpu)
 
@@ -244,14 +250,14 @@ class Results(SimpleClass):
         if pred_boxes is not None and show_boxes:
             for d in reversed(pred_boxes):
                 c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
-                name = ('' if id is None else f'id:{id} ') + names[c]
-                label = (f'{name} {conf:.2f}' if conf else name) if labels else None
+                name = ("" if id is None else f"id:{id} ") + names[c]
+                label = (f"{name} {conf:.2f}" if conf else name) if labels else None
                 box = d.xyxyxyxy.reshape(-1, 4, 2).squeeze() if is_obb else d.xyxy.squeeze()
                 annotator.box_label(box, label, color=colors(c, True), rotated=is_obb)
 
         # Plot Classify results
         if pred_probs is not None and show_probs:
-            text = ',\n'.join(f'{names[j] if names else j} {pred_probs.data[j]:.2f}' for j in pred_probs.top5)
+            text = ",\n".join(f"{names[j] if names else j} {pred_probs.data[j]:.2f}" for j in pred_probs.top5)
             x = round(self.orig_shape[0] * 0.03)
             annotator.text([x, x], text, txt_color=(255, 255, 255))  # TODO: allow setting colors
 
@@ -264,11 +270,11 @@ class Results(SimpleClass):
 
     def verbose(self):
         """Return log string for each task."""
-        log_string = ''
+        log_string = ""
         probs = self.probs
         boxes = self.boxes
         if len(self) == 0:
-            return log_string if probs is not None else f'{log_string}(no detections), '
+            return log_string if probs is not None else f"{log_string}(no detections), "
         if probs is not None:
             log_string += f"{', '.join(f'{self.names[j]} {probs.data[j]:.2f}' for j in probs.top5)}, "
         if boxes:
@@ -285,34 +291,35 @@ class Results(SimpleClass):
             txt_file (str): txt file path.
             save_conf (bool): save confidence score or not.
         """
-        boxes = self.boxes
+        is_obb = self.obb is not None
+        boxes = self.obb if is_obb else self.boxes
         masks = self.masks
         probs = self.probs
         kpts = self.keypoints
         texts = []
         if probs is not None:
             # Classify
-            [texts.append(f'{probs.data[j]:.2f} {self.names[j]}') for j in probs.top5]
+            [texts.append(f"{probs.data[j]:.2f} {self.names[j]}") for j in probs.top5]
         elif boxes:
             # Detect/segment/pose
             for j, d in enumerate(boxes):
                 c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
-                line = (c, *d.xywhn.view(-1))
+                line = (c, *(d.xyxyxyxyn.view(-1) if is_obb else d.xywhn.view(-1)))
                 if masks:
                     seg = masks[j].xyn[0].copy().reshape(-1)  # reversed mask.xyn, (n,2) to (n*2)
                     line = (c, *seg)
                 if kpts is not None:
                     kpt = torch.cat((kpts[j].xyn, kpts[j].conf[..., None]), 2) if kpts[j].has_visible else kpts[j].xyn
-                    line += (*kpt.reshape(-1).tolist(), )
-                line += (conf, ) * save_conf + (() if id is None else (id, ))
-                texts.append(('%g ' * len(line)).rstrip() % line)
+                    line += (*kpt.reshape(-1).tolist(),)
+                line += (conf,) * save_conf + (() if id is None else (id,))
+                texts.append(("%g " * len(line)).rstrip() % line)
 
         if texts:
             Path(txt_file).parent.mkdir(parents=True, exist_ok=True)  # make directory
-            with open(txt_file, 'a') as f:
-                f.writelines(text + '\n' for text in texts)
+            with open(txt_file, "a") as f:
+                f.writelines(text + "\n" for text in texts)
 
-    def save_crop(self, save_dir, file_name=Path('im.jpg')):
+    def save_crop(self, save_dir, file_name=Path("im.jpg")):
         """
         Save cropped predictions to `save_dir/cls/file_name.jpg`.
 
@@ -321,18 +328,23 @@ class Results(SimpleClass):
             file_name (str | pathlib.Path): File name.
         """
         if self.probs is not None:
-            LOGGER.warning('WARNING ⚠️ Classify task do not support `save_crop`.')
+            LOGGER.warning("WARNING ⚠️ Classify task do not support `save_crop`.")
+            return
+        if self.obb is not None:
+            LOGGER.warning("WARNING ⚠️ OBB task do not support `save_crop`.")
             return
         for d in self.boxes:
-            save_one_box(d.xyxy,
-                         self.orig_img.copy(),
-                         file=Path(save_dir) / self.names[int(d.cls)] / f'{Path(file_name)}.jpg',
-                         BGR=True)
+            save_one_box(
+                d.xyxy,
+                self.orig_img.copy(),
+                file=Path(save_dir) / self.names[int(d.cls)] / f"{Path(file_name)}.jpg",
+                BGR=True,
+            )
 
     def tojson(self, normalize=False):
         """Convert the object to JSON format."""
         if self.probs is not None:
-            LOGGER.warning('Warning: Classify task do not support `tojson` yet.')
+            LOGGER.warning("Warning: Classify task do not support `tojson` yet.")
             return
 
         import json
@@ -342,19 +354,19 @@ class Results(SimpleClass):
         data = self.boxes.data.cpu().tolist()
         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': row[0] / w, 'y1': row[1] / h, 'x2': row[2] / w, 'y2': row[3] / h}
+            box = {"x1": row[0] / w, "y1": row[1] / h, "x2": row[2] / w, "y2": row[3] / h}
             conf = row[-2]
             class_id = int(row[-1])
             name = self.names[class_id]
-            result = {'name': name, 'class': class_id, 'confidence': conf, 'box': box}
+            result = {"name": name, "class": class_id, "confidence": conf, "box": box}
             if self.boxes.is_track:
-                result['track_id'] = int(row[-3])  # track ID
+                result["track_id"] = int(row[-3])  # track ID
             if self.masks:
                 x, y = self.masks.xy[i][:, 0], self.masks.xy[i][:, 1]  # numpy array
-                result['segments'] = {'x': (x / w).tolist(), 'y': (y / h).tolist()}
+                result["segments"] = {"x": (x / w).tolist(), "y": (y / h).tolist()}
             if self.keypoints is not None:
                 x, y, visible = self.keypoints[i].data[0].cpu().unbind(dim=1)  # torch Tensor
-                result['keypoints'] = {'x': (x / w).tolist(), 'y': (y / h).tolist(), 'visible': visible.tolist()}
+                result["keypoints"] = {"x": (x / w).tolist(), "y": (y / h).tolist(), "visible": visible.tolist()}
             results.append(result)
 
         # Convert detections to JSON
@@ -393,7 +405,7 @@ class Boxes(BaseTensor):
         if boxes.ndim == 1:
             boxes = boxes[None, :]
         n = boxes.shape[-1]
-        assert n in (6, 7), f'expected 6 or 7 values but got {n}'  # xyxy, track_id, conf, cls
+        assert n in (6, 7), f"expected 6 or 7 values but got {n}"  # xyxy, track_id, conf, cls
         super().__init__(boxes, orig_shape)
         self.is_track = n == 7
         self.orig_shape = orig_shape
@@ -470,7 +482,8 @@ class Masks(BaseTensor):
         """Return normalized segments."""
         return [
             ops.scale_coords(self.data.shape[1:], x, self.orig_shape, normalize=True)
-            for x in ops.masks2segments(self.data)]
+            for x in ops.masks2segments(self.data)
+        ]
 
     @property
     @lru_cache(maxsize=1)
@@ -478,7 +491,8 @@ class Masks(BaseTensor):
         """Return segments in pixel coordinates."""
         return [
             ops.scale_coords(self.data.shape[1:], x, self.orig_shape, normalize=False)
-            for x in ops.masks2segments(self.data)]
+            for x in ops.masks2segments(self.data)
+        ]
 
 
 class Keypoints(BaseTensor):
@@ -606,7 +620,7 @@ class OBB(BaseTensor):
         if boxes.ndim == 1:
             boxes = boxes[None, :]
         n = boxes.shape[-1]
-        assert n in (7, 8), f'expected 7 or 8 values but got {n}'  # xywh, rotation, track_id, conf, cls
+        assert n in (7, 8), f"expected 7 or 8 values but got {n}"  # xywh, rotation, track_id, conf, cls
         super().__init__(boxes, orig_shape)
         self.is_track = n == 8
         self.orig_shape = orig_shape
@@ -637,6 +651,15 @@ class OBB(BaseTensor):
         """Return the boxes in xyxyxyxy format, (N, 4, 2)."""
         return ops.xywhr2xyxyxyxy(self.xywhr)
 
+    @property
+    @lru_cache(maxsize=2)
+    def xyxyxyxyn(self):
+        """Return the boxes in xyxyxyxy format, (N, 4, 2)."""
+        xyxyxyxyn = self.xyxyxyxy.clone() if isinstance(self.xyxyxyxy, torch.Tensor) else np.copy(self.xyxyxyxy)
+        xyxyxyxyn[..., 0] /= self.orig_shape[1]
+        xyxyxyxyn[..., 1] /= self.orig_shape[1]
+        return xyxyxyxyn
+
     @property
     @lru_cache(maxsize=2)
     def xyxy(self):