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

ultralytics/models/yolo/model.py

@@ -12,28 +12,34 @@ class YOLO(Model):
     def task_map(self):
         """Map head to model, trainer, validator, and predictor classes."""
         return {
-            'classify': {
-                'model': ClassificationModel,
-                'trainer': yolo.classify.ClassificationTrainer,
-                'validator': yolo.classify.ClassificationValidator,
-                'predictor': yolo.classify.ClassificationPredictor, },
-            'detect': {
-                'model': DetectionModel,
-                'trainer': yolo.detect.DetectionTrainer,
-                'validator': yolo.detect.DetectionValidator,
-                'predictor': yolo.detect.DetectionPredictor, },
-            'segment': {
-                'model': SegmentationModel,
-                'trainer': yolo.segment.SegmentationTrainer,
-                'validator': yolo.segment.SegmentationValidator,
-                'predictor': yolo.segment.SegmentationPredictor, },
-            'pose': {
-                'model': PoseModel,
-                'trainer': yolo.pose.PoseTrainer,
-                'validator': yolo.pose.PoseValidator,
-                'predictor': yolo.pose.PosePredictor, },
-            'obb': {
-                'model': OBBModel,
-                'trainer': yolo.obb.OBBTrainer,
-                'validator': yolo.obb.OBBValidator,
-                'predictor': yolo.obb.OBBPredictor, }, }
+            "classify": {
+                "model": ClassificationModel,
+                "trainer": yolo.classify.ClassificationTrainer,
+                "validator": yolo.classify.ClassificationValidator,
+                "predictor": yolo.classify.ClassificationPredictor,
+            },
+            "detect": {
+                "model": DetectionModel,
+                "trainer": yolo.detect.DetectionTrainer,
+                "validator": yolo.detect.DetectionValidator,
+                "predictor": yolo.detect.DetectionPredictor,
+            },
+            "segment": {
+                "model": SegmentationModel,
+                "trainer": yolo.segment.SegmentationTrainer,
+                "validator": yolo.segment.SegmentationValidator,
+                "predictor": yolo.segment.SegmentationPredictor,
+            },
+            "pose": {
+                "model": PoseModel,
+                "trainer": yolo.pose.PoseTrainer,
+                "validator": yolo.pose.PoseValidator,
+                "predictor": yolo.pose.PosePredictor,
+            },
+            "obb": {
+                "model": OBBModel,
+                "trainer": yolo.obb.OBBTrainer,
+                "validator": yolo.obb.OBBValidator,
+                "predictor": yolo.obb.OBBPredictor,
+            },
+        }

ultralytics/models/yolo/obb/__init__.py

@@ -4,4 +4,4 @@ from .predict import OBBPredictor
 from .train import OBBTrainer
 from .val import OBBValidator
 
-__all__ = 'OBBPredictor', 'OBBTrainer', 'OBBValidator'
+__all__ = "OBBPredictor", "OBBTrainer", "OBBValidator"

ultralytics/models/yolo/obb/predict.py

@@ -23,28 +23,29 @@ class OBBPredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+        """Initializes OBBPredictor with optional model and data configuration overrides."""
         super().__init__(cfg, overrides, _callbacks)
-        self.args.task = 'obb'
+        self.args.task = "obb"
 
     def postprocess(self, preds, img, orig_imgs):
         """Post-processes predictions and returns a list of Results objects."""
-        preds = ops.non_max_suppression(preds,
-                                        self.args.conf,
-                                        self.args.iou,
-                                        agnostic=self.args.agnostic_nms,
-                                        max_det=self.args.max_det,
-                                        nc=len(self.model.names),
-                                        classes=self.args.classes,
-                                        rotated=True)
+        preds = ops.non_max_suppression(
+            preds,
+            self.args.conf,
+            self.args.iou,
+            agnostic=self.args.agnostic_nms,
+            max_det=self.args.max_det,
+            nc=len(self.model.names),
+            classes=self.args.classes,
+            rotated=True,
+        )
 
         if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
             orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
 
         results = []
-        for i, pred in enumerate(preds):
-            orig_img = orig_imgs[i]
+        for i, (pred, orig_img, img_path) in enumerate(zip(preds, orig_imgs, self.batch[0])):
             pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape, xywh=True)
-            img_path = self.batch[0][i]
             # xywh, r, conf, cls
             obb = torch.cat([pred[:, :4], pred[:, -1:], pred[:, 4:6]], dim=-1)
             results.append(Results(orig_img, path=img_path, names=self.model.names, obb=obb))

ultralytics/models/yolo/obb/train.py

@@ -25,12 +25,12 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         """Initialize a OBBTrainer object with given arguments."""
         if overrides is None:
             overrides = {}
-        overrides['task'] = 'obb'
+        overrides["task"] = "obb"
         super().__init__(cfg, overrides, _callbacks)
 
     def get_model(self, cfg=None, weights=None, verbose=True):
         """Return OBBModel initialized with specified config and weights."""
-        model = OBBModel(cfg, ch=3, nc=self.data['nc'], verbose=verbose and RANK == -1)
+        model = OBBModel(cfg, ch=3, nc=self.data["nc"], verbose=verbose and RANK == -1)
         if weights:
             model.load(weights)
 
@@ -38,5 +38,5 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
 
     def get_validator(self):
         """Return an instance of OBBValidator for validation of YOLO model."""
-        self.loss_names = 'box_loss', 'cls_loss', 'dfl_loss'
+        self.loss_names = "box_loss", "cls_loss", "dfl_loss"
         return yolo.obb.OBBValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))

ultralytics/models/yolo/obb/val.py

@@ -27,26 +27,28 @@ class OBBValidator(DetectionValidator):
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
         """Initialize OBBValidator and set task to 'obb', metrics to OBBMetrics."""
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
-        self.args.task = 'obb'
+        self.args.task = "obb"
         self.metrics = OBBMetrics(save_dir=self.save_dir, plot=True, on_plot=self.on_plot)
 
     def init_metrics(self, model):
         """Initialize evaluation metrics for YOLO."""
         super().init_metrics(model)
-        val = self.data.get(self.args.split, '')  # validation path
-        self.is_dota = isinstance(val, str) and 'DOTA' in val  # is COCO
+        val = self.data.get(self.args.split, "")  # validation path
+        self.is_dota = isinstance(val, str) and "DOTA" in val  # is COCO
 
     def postprocess(self, preds):
         """Apply Non-maximum suppression to prediction outputs."""
-        return ops.non_max_suppression(preds,
-                                       self.args.conf,
-                                       self.args.iou,
-                                       labels=self.lb,
-                                       nc=self.nc,
-                                       multi_label=True,
-                                       agnostic=self.args.single_cls,
-                                       max_det=self.args.max_det,
-                                       rotated=True)
+        return ops.non_max_suppression(
+            preds,
+            self.args.conf,
+            self.args.iou,
+            labels=self.lb,
+            nc=self.nc,
+            multi_label=True,
+            agnostic=self.args.single_cls,
+            max_det=self.args.max_det,
+            rotated=True,
+        )
 
     def _process_batch(self, detections, gt_bboxes, gt_cls):
         """
@@ -61,16 +63,17 @@ class OBBValidator(DetectionValidator):
         Returns:
             (torch.Tensor): Correct prediction matrix of shape [N, 10] for 10 IoU levels.
         """
-        iou = batch_probiou(gt_bboxes, torch.cat([detections[:, :4], detections[:, -2:-1]], dim=-1))
+        iou = batch_probiou(gt_bboxes, torch.cat([detections[:, :4], detections[:, -1:]], dim=-1))
         return self.match_predictions(detections[:, 5], gt_cls, iou)
 
     def _prepare_batch(self, si, batch):
-        idx = batch['batch_idx'] == si
-        cls = batch['cls'][idx].squeeze(-1)
-        bbox = batch['bboxes'][idx]
-        ori_shape = batch['ori_shape'][si]
-        imgsz = batch['img'].shape[2:]
-        ratio_pad = batch['ratio_pad'][si]
+        """Prepares and returns a batch for OBB validation."""
+        idx = batch["batch_idx"] == si
+        cls = batch["cls"][idx].squeeze(-1)
+        bbox = batch["bboxes"][idx]
+        ori_shape = batch["ori_shape"][si]
+        imgsz = batch["img"].shape[2:]
+        ratio_pad = batch["ratio_pad"][si]
         if len(cls):
             bbox[..., :4].mul_(torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]])  # target boxes
             ops.scale_boxes(imgsz, bbox, ori_shape, ratio_pad=ratio_pad, xywh=True)  # native-space labels
@@ -78,19 +81,23 @@ class OBBValidator(DetectionValidator):
         return prepared_batch
 
     def _prepare_pred(self, pred, pbatch):
+        """Prepares and returns a batch for OBB validation with scaled and padded bounding boxes."""
         predn = pred.clone()
-        ops.scale_boxes(pbatch['imgsz'], predn[:, :4], pbatch['ori_shape'], ratio_pad=pbatch['ratio_pad'],
-                        xywh=True)  # native-space pred
+        ops.scale_boxes(
+            pbatch["imgsz"], predn[:, :4], pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"], xywh=True
+        )  # native-space pred
         return predn
 
     def plot_predictions(self, batch, preds, ni):
         """Plots predicted bounding boxes on input images and saves the result."""
-        plot_images(batch['img'],
-                    *output_to_rotated_target(preds, max_det=self.args.max_det),
-                    paths=batch['im_file'],
-                    fname=self.save_dir / f'val_batch{ni}_pred.jpg',
-                    names=self.names,
-                    on_plot=self.on_plot)  # pred
+        plot_images(
+            batch["img"],
+            *output_to_rotated_target(preds, max_det=self.args.max_det),
+            paths=batch["im_file"],
+            fname=self.save_dir / f"val_batch{ni}_pred.jpg",
+            names=self.names,
+            on_plot=self.on_plot,
+        )  # pred
 
     def pred_to_json(self, predn, filename):
         """Serialize YOLO predictions to COCO json format."""
@@ -99,12 +106,26 @@ class OBBValidator(DetectionValidator):
         rbox = torch.cat([predn[:, :4], predn[:, -1:]], dim=-1)
         poly = ops.xywhr2xyxyxyxy(rbox).view(-1, 8)
         for i, (r, b) in enumerate(zip(rbox.tolist(), poly.tolist())):
-            self.jdict.append({
-                'image_id': image_id,
-                'category_id': self.class_map[int(predn[i, 5].item())],
-                'score': round(predn[i, 4].item(), 5),
-                'rbox': [round(x, 3) for x in r],
-                'poly': [round(x, 3) for x in b]})
+            self.jdict.append(
+                {
+                    "image_id": image_id,
+                    "category_id": self.class_map[int(predn[i, 5].item())],
+                    "score": round(predn[i, 4].item(), 5),
+                    "rbox": [round(x, 3) for x in r],
+                    "poly": [round(x, 3) for x in b],
+                }
+            )
+
+    def save_one_txt(self, predn, save_conf, shape, file):
+        """Save YOLO detections to a txt file in normalized coordinates in a specific format."""
+        gn = torch.tensor(shape)[[1, 0, 1, 0]]  # normalization gain whwh
+        for *xyxy, conf, cls, angle in predn.tolist():
+            xywha = torch.tensor([*xyxy, angle]).view(1, 5)
+            xywha[:, :4] /= gn
+            xyxyxyxy = ops.xywhr2xyxyxyxy(xywha).view(-1).tolist()  # normalized xywh
+            line = (cls, *xyxyxyxy, conf) if save_conf else (cls, *xyxyxyxy)  # label format
+            with open(file, "a") as f:
+                f.write(("%g " * len(line)).rstrip() % line + "\n")
 
     def eval_json(self, stats):
         """Evaluates YOLO output in JSON format and returns performance statistics."""
@@ -112,42 +133,43 @@ class OBBValidator(DetectionValidator):
             import json
             import re
             from collections import defaultdict
-            pred_json = self.save_dir / 'predictions.json'  # predictions
-            pred_txt = self.save_dir / 'predictions_txt'  # predictions
+
+            pred_json = self.save_dir / "predictions.json"  # predictions
+            pred_txt = self.save_dir / "predictions_txt"  # predictions
             pred_txt.mkdir(parents=True, exist_ok=True)
             data = json.load(open(pred_json))
             # Save split results
-            LOGGER.info(f'Saving predictions with DOTA format to {str(pred_txt)}...')
+            LOGGER.info(f"Saving predictions with DOTA format to {str(pred_txt)}...")
             for d in data:
-                image_id = d['image_id']
-                score = d['score']
-                classname = self.names[d['category_id']].replace(' ', '-')
+                image_id = d["image_id"]
+                score = d["score"]
+                classname = self.names[d["category_id"]].replace(" ", "-")
 
-                lines = '{} {} {} {} {} {} {} {} {} {}\n'.format(
+                lines = "{} {} {} {} {} {} {} {} {} {}\n".format(
                     image_id,
                     score,
-                    d['poly'][0],
-                    d['poly'][1],
-                    d['poly'][2],
-                    d['poly'][3],
-                    d['poly'][4],
-                    d['poly'][5],
-                    d['poly'][6],
-                    d['poly'][7],
+                    d["poly"][0],
+                    d["poly"][1],
+                    d["poly"][2],
+                    d["poly"][3],
+                    d["poly"][4],
+                    d["poly"][5],
+                    d["poly"][6],
+                    d["poly"][7],
                 )
-                with open(str(pred_txt / f'Task1_{classname}') + '.txt', 'a') as f:
+                with open(str(pred_txt / f"Task1_{classname}") + ".txt", "a") as f:
                     f.writelines(lines)
             # Save merged results, this could result slightly lower map than using official merging script,
             # because of the probiou calculation.
-            pred_merged_txt = self.save_dir / 'predictions_merged_txt'  # predictions
+            pred_merged_txt = self.save_dir / "predictions_merged_txt"  # predictions
             pred_merged_txt.mkdir(parents=True, exist_ok=True)
             merged_results = defaultdict(list)
-            LOGGER.info(f'Saving merged predictions with DOTA format to {str(pred_merged_txt)}...')
+            LOGGER.info(f"Saving merged predictions with DOTA format to {str(pred_merged_txt)}...")
             for d in data:
-                image_id = d['image_id'].split('__')[0]
-                pattern = re.compile(r'\d+___\d+')
-                x, y = (int(c) for c in re.findall(pattern, d['image_id'])[0].split('___'))
-                bbox, score, cls = d['rbox'], d['score'], d['category_id']
+                image_id = d["image_id"].split("__")[0]
+                pattern = re.compile(r"\d+___\d+")
+                x, y = (int(c) for c in re.findall(pattern, d["image_id"])[0].split("___"))
+                bbox, score, cls = d["rbox"], d["score"], d["category_id"]
                 bbox[0] += x
                 bbox[1] += y
                 bbox.extend([score, cls])
@@ -165,11 +187,11 @@ class OBBValidator(DetectionValidator):
 
                 b = ops.xywhr2xyxyxyxy(bbox[:, :5]).view(-1, 8)
                 for x in torch.cat([b, bbox[:, 5:7]], dim=-1).tolist():
-                    classname = self.names[int(x[-1])].replace(' ', '-')
+                    classname = self.names[int(x[-1])].replace(" ", "-")
                     poly = [round(i, 3) for i in x[:-2]]
                     score = round(x[-2], 3)
 
-                    lines = '{} {} {} {} {} {} {} {} {} {}\n'.format(
+                    lines = "{} {} {} {} {} {} {} {} {} {}\n".format(
                         image_id,
                         score,
                         poly[0],
@@ -181,7 +203,7 @@ class OBBValidator(DetectionValidator):
                         poly[6],
                         poly[7],
                     )
-                    with open(str(pred_merged_txt / f'Task1_{classname}') + '.txt', 'a') as f:
+                    with open(str(pred_merged_txt / f"Task1_{classname}") + ".txt", "a") as f:
                         f.writelines(lines)
 
         return stats

ultralytics/models/yolo/pose/__init__.py

@@ -4,4 +4,4 @@ from .predict import PosePredictor
 from .train import PoseTrainer
 from .val import PoseValidator
 
-__all__ = 'PoseTrainer', 'PoseValidator', 'PosePredictor'
+__all__ = "PoseTrainer", "PoseValidator", "PosePredictor"

ultralytics/models/yolo/pose/predict.py

@@ -23,20 +23,24 @@ class PosePredictor(DetectionPredictor):
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
         """Initializes PosePredictor, sets task to 'pose' and logs a warning for using 'mps' as device."""
         super().__init__(cfg, overrides, _callbacks)
-        self.args.task = 'pose'
-        if isinstance(self.args.device, str) and self.args.device.lower() == 'mps':
-            LOGGER.warning("WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. "
-                           'See https://github.com/ultralytics/ultralytics/issues/4031.')
+        self.args.task = "pose"
+        if isinstance(self.args.device, str) and self.args.device.lower() == "mps":
+            LOGGER.warning(
+                "WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. "
+                "See https://github.com/ultralytics/ultralytics/issues/4031."
+            )
 
     def postprocess(self, preds, img, orig_imgs):
         """Return detection results for a given input image or list of images."""
-        preds = ops.non_max_suppression(preds,
-                                        self.args.conf,
-                                        self.args.iou,
-                                        agnostic=self.args.agnostic_nms,
-                                        max_det=self.args.max_det,
-                                        classes=self.args.classes,
-                                        nc=len(self.model.names))
+        preds = ops.non_max_suppression(
+            preds,
+            self.args.conf,
+            self.args.iou,
+            agnostic=self.args.agnostic_nms,
+            max_det=self.args.max_det,
+            classes=self.args.classes,
+            nc=len(self.model.names),
+        )
 
         if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
             orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
@@ -49,5 +53,6 @@ class PosePredictor(DetectionPredictor):
             pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
             img_path = self.batch[0][i]
             results.append(
-                Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts))
+                Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts)
+            )
         return results

ultralytics/models/yolo/pose/train.py

@@ -26,16 +26,18 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         """Initialize a PoseTrainer object with specified configurations and overrides."""
         if overrides is None:
             overrides = {}
-        overrides['task'] = 'pose'
+        overrides["task"] = "pose"
         super().__init__(cfg, overrides, _callbacks)
 
-        if isinstance(self.args.device, str) and self.args.device.lower() == 'mps':
-            LOGGER.warning("WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. "
-                           'See https://github.com/ultralytics/ultralytics/issues/4031.')
+        if isinstance(self.args.device, str) and self.args.device.lower() == "mps":
+            LOGGER.warning(
+                "WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. "
+                "See https://github.com/ultralytics/ultralytics/issues/4031."
+            )
 
     def get_model(self, cfg=None, weights=None, verbose=True):
         """Get pose estimation model with specified configuration and weights."""
-        model = PoseModel(cfg, ch=3, nc=self.data['nc'], data_kpt_shape=self.data['kpt_shape'], verbose=verbose)
+        model = PoseModel(cfg, ch=3, nc=self.data["nc"], data_kpt_shape=self.data["kpt_shape"], verbose=verbose)
         if weights:
             model.load(weights)
 
@@ -44,32 +46,33 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
     def set_model_attributes(self):
         """Sets keypoints shape attribute of PoseModel."""
         super().set_model_attributes()
-        self.model.kpt_shape = self.data['kpt_shape']
+        self.model.kpt_shape = self.data["kpt_shape"]
 
     def get_validator(self):
         """Returns an instance of the PoseValidator class for validation."""
-        self.loss_names = 'box_loss', 'pose_loss', 'kobj_loss', 'cls_loss', 'dfl_loss'
-        return yolo.pose.PoseValidator(self.test_loader,
-                                       save_dir=self.save_dir,
-                                       args=copy(self.args),
-                                       _callbacks=self.callbacks)
+        self.loss_names = "box_loss", "pose_loss", "kobj_loss", "cls_loss", "dfl_loss"
+        return yolo.pose.PoseValidator(
+            self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
+        )
 
     def plot_training_samples(self, batch, ni):
         """Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints."""
-        images = batch['img']
-        kpts = batch['keypoints']
-        cls = batch['cls'].squeeze(-1)
-        bboxes = batch['bboxes']
-        paths = batch['im_file']
-        batch_idx = batch['batch_idx']
-        plot_images(images,
-                    batch_idx,
-                    cls,
-                    bboxes,
-                    kpts=kpts,
-                    paths=paths,
-                    fname=self.save_dir / f'train_batch{ni}.jpg',
-                    on_plot=self.on_plot)
+        images = batch["img"]
+        kpts = batch["keypoints"]
+        cls = batch["cls"].squeeze(-1)
+        bboxes = batch["bboxes"]
+        paths = batch["im_file"]
+        batch_idx = batch["batch_idx"]
+        plot_images(
+            images,
+            batch_idx,
+            cls,
+            bboxes,
+            kpts=kpts,
+            paths=paths,
+            fname=self.save_dir / f"train_batch{ni}.jpg",
+            on_plot=self.on_plot,
+        )
 
     def plot_metrics(self):
         """Plots training/val metrics."""