General cleanup (#69)

Co-authored-by: ayush chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

ultralytics/yolo/v8/detect/__init__.py

@@ -1,3 +1,3 @@
-from ultralytics.yolo.v8.detect.predict import DetectionPredictor, predict
-from ultralytics.yolo.v8.detect.train import DetectionTrainer, train
-from ultralytics.yolo.v8.detect.val import DetectionValidator, val
+from .predict import DetectionPredictor, predict
+from .train import DetectionTrainer, train
+from .val import DetectionValidator, val

ultralytics/yolo/v8/detect/train.py

@@ -2,18 +2,37 @@ import hydra
 import torch
 import torch.nn as nn
 
-from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
+from ultralytics.yolo import v8
+from ultralytics.yolo.data import build_dataloader
+from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer
 from ultralytics.yolo.utils.metrics import FocalLoss, bbox_iou, smooth_BCE
 from ultralytics.yolo.utils.modeling.tasks import DetectionModel
 from ultralytics.yolo.utils.plotting import plot_images, plot_results
 from ultralytics.yolo.utils.torch_utils import de_parallel
 
-from ..segment import SegmentationTrainer
-from .val import DetectionValidator
-
 
 # BaseTrainer python usage
-class DetectionTrainer(SegmentationTrainer):
+class DetectionTrainer(BaseTrainer):
+
+    def get_dataloader(self, dataset_path, batch_size, mode="train", rank=0):
+        # TODO: manage splits differently
+        # calculate stride - check if model is initialized
+        gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
+        return build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, rank=rank, mode=mode)[0]
+
+    def preprocess_batch(self, batch):
+        batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
+        return batch
+
+    def set_model_attributes(self):
+        nl = de_parallel(self.model).model[-1].nl  # number of detection layers (to scale hyps)
+        self.args.box *= 3 / nl  # scale to layers
+        self.args.cls *= self.data["nc"] / 80 * 3 / nl  # scale to classes and layers
+        self.args.obj *= (self.args.img_size / 640) ** 2 * 3 / nl  # scale to image size and layers
+        self.model.nc = self.data["nc"]  # attach number of classes to model
+        self.model.args = self.args  # attach hyperparameters to model
+        # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
+        self.model.names = self.data["names"]
 
     def load_model(self, model_cfg=None, weights=None):
         model = DetectionModel(model_cfg or weights["model"].yaml,
@@ -27,7 +46,10 @@ class DetectionTrainer(SegmentationTrainer):
         return model
 
     def get_validator(self):
-        return DetectionValidator(self.test_loader, save_dir=self.save_dir, logger=self.console, args=self.args)
+        return v8.detect.DetectionValidator(self.test_loader,
+                                            save_dir=self.save_dir,
+                                            logger=self.console,
+                                            args=self.args)
 
     def criterion(self, preds, batch):
         head = de_parallel(self.model).model[-1]

ultralytics/yolo/v8/detect/val.py

@@ -11,7 +11,7 @@ from ultralytics.yolo.engine.validator import BaseValidator
 from ultralytics.yolo.utils import ops
 from ultralytics.yolo.utils.checks import check_file, check_requirements
 from ultralytics.yolo.utils.files import yaml_load
-from ultralytics.yolo.utils.metrics import ConfusionMatrix, Metric, ap_per_class, box_iou, fitness_detection
+from ultralytics.yolo.utils.metrics import ConfusionMatrix, DetMetrics, box_iou
 from ultralytics.yolo.utils.plotting import output_to_target, plot_images
 from ultralytics.yolo.utils.torch_utils import de_parallel
 
@@ -62,7 +62,7 @@ class DetectionValidator(BaseValidator):
         self.niou = self.iouv.numel()
         self.seen = 0
         self.confusion_matrix = ConfusionMatrix(nc=self.nc)
-        self.metrics = Metric()
+        self.metrics = DetMetrics(save_dir=self.save_dir, plot=self.args.plots, names=self.names)
         self.loss = torch.zeros(3, device=self.device)
         self.jdict = []
         self.stats = []
@@ -128,10 +128,9 @@ class DetectionValidator(BaseValidator):
     def get_stats(self):
         stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)]  # to numpy
         if len(stats) and stats[0].any():
-            results = ap_per_class(*stats, plot=self.args.plots, save_dir=self.save_dir, names=self.names)
-            self.metrics.update(results[2:])
-        self.nt_per_class = np.bincount(stats[3].astype(int), minlength=self.nc)  # number of targets per class
-        metrics = {"fitness": fitness_detection(np.array(self.metrics.mean_results()).reshape(1, -1))}
+            self.metrics.process(*stats)
+        self.nt_per_class = np.bincount(stats[-1].astype(int), minlength=self.nc)  # number of targets per class
+        metrics = {"fitness": self.metrics.fitness()}
         metrics |= zip(self.metric_keys, self.metrics.mean_results())
         return metrics
 
@@ -203,8 +202,11 @@ class DetectionValidator(BaseValidator):
     def plot_predictions(self, batch, preds, ni):
         images = batch["img"]
         paths = batch["im_file"]
-        plot_images(images, *output_to_target(preds, max_det=15), paths, self.save_dir / f'val_batch{ni}_pred.jpg',
-                    self.names)  # pred
+        plot_images(images,
+                    *output_to_target(preds, max_det=15),
+                    paths=paths,
+                    fname=self.save_dir / f'val_batch{ni}_pred.jpg',
+                    names=self.names)  # pred
 
 
 @hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name)

ultralytics/yolo/v8/segment/__init__.py

@@ -1,3 +1,3 @@
-from ultralytics.yolo.v8.segment.predict import SegmentationPredictor, predict
-from ultralytics.yolo.v8.segment.train import SegmentationTrainer, train
-from ultralytics.yolo.v8.segment.val import SegmentationValidator, val
+from .predict import SegmentationPredictor, predict
+from .train import SegmentationTrainer, train
+from .val import SegmentationValidator, val

ultralytics/yolo/v8/segment/train.py

@@ -4,27 +4,18 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from ultralytics.yolo import v8
-from ultralytics.yolo.data import build_dataloader
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer
 from ultralytics.yolo.utils.metrics import FocalLoss, bbox_iou, smooth_BCE
 from ultralytics.yolo.utils.modeling.tasks import SegmentationModel
 from ultralytics.yolo.utils.ops import crop_mask, xywh2xyxy
-from ultralytics.yolo.utils.plotting import plot_images_and_masks, plot_results_with_masks
+from ultralytics.yolo.utils.plotting import plot_images, plot_results
 from ultralytics.yolo.utils.torch_utils import de_parallel
 
+from ..detect import DetectionTrainer
+
 
 # BaseTrainer python usage
-class SegmentationTrainer(BaseTrainer):
-
-    def get_dataloader(self, dataset_path, batch_size, mode="train", rank=0):
-        # TODO: manage splits differently
-        # calculate stride - check if model is initialized
-        gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
-        return build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, rank=rank, mode=mode)[0]
-
-    def preprocess_batch(self, batch):
-        batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
-        return batch
+class SegmentationTrainer(DetectionTrainer):
 
     def load_model(self, model_cfg=None, weights=None):
         model = SegmentationModel(model_cfg or weights["model"].yaml,
@@ -37,16 +28,6 @@ class SegmentationTrainer(BaseTrainer):
             v.requires_grad = True  # train all layers
         return model
 
-    def set_model_attributes(self):
-        nl = de_parallel(self.model).model[-1].nl  # number of detection layers (to scale hyps)
-        self.args.box *= 3 / nl  # scale to layers
-        self.args.cls *= self.data["nc"] / 80 * 3 / nl  # scale to classes and layers
-        self.args.obj *= (self.args.img_size / 640) ** 2 * 3 / nl  # scale to image size and layers
-        self.model.nc = self.data["nc"]  # attach number of classes to model
-        self.model.args = self.args  # attach hyperparameters to model
-        # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
-        self.model.names = self.data["names"]
-
     def get_validator(self):
         return v8.segment.SegmentationValidator(self.test_loader,
                                                 save_dir=self.save_dir,
@@ -245,16 +226,10 @@ class SegmentationTrainer(BaseTrainer):
         bboxes = batch["bboxes"]
         paths = batch["im_file"]
         batch_idx = batch["batch_idx"]
-        plot_images_and_masks(images,
-                              batch_idx,
-                              cls,
-                              bboxes,
-                              masks,
-                              paths=paths,
-                              fname=self.save_dir / f"train_batch{ni}.jpg")
+        plot_images(images, batch_idx, cls, bboxes, masks, paths=paths, fname=self.save_dir / f"train_batch{ni}.jpg")
 
     def plot_metrics(self):
-        plot_results_with_masks(file=self.csv)  # save results.png
+        plot_results(file=self.csv, segment=True)  # save results.png
 
 
 @hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name)

ultralytics/yolo/v8/segment/val.py

@@ -7,17 +7,17 @@ import torch.nn.functional as F
 
 from ultralytics.yolo.data import build_dataloader
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
-from ultralytics.yolo.engine.validator import BaseValidator
 from ultralytics.yolo.utils import ops
 from ultralytics.yolo.utils.checks import check_file, check_requirements
 from ultralytics.yolo.utils.files import yaml_load
-from ultralytics.yolo.utils.metrics import (ConfusionMatrix, Metrics, ap_per_class_box_and_mask, box_iou,
-                                            fitness_segmentation, mask_iou)
-from ultralytics.yolo.utils.plotting import output_to_target, plot_images_and_masks
+from ultralytics.yolo.utils.metrics import ConfusionMatrix, SegmentMetrics, box_iou, mask_iou
+from ultralytics.yolo.utils.plotting import output_to_target, plot_images
 from ultralytics.yolo.utils.torch_utils import de_parallel
 
+from ..detect import DetectionValidator
 
-class SegmentationValidator(BaseValidator):
+
+class SegmentationValidator(DetectionValidator):
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, logger=None, args=None):
         super().__init__(dataloader, save_dir, pbar, logger, args)
@@ -65,7 +65,7 @@ class SegmentationValidator(BaseValidator):
         self.niou = self.iouv.numel()
         self.seen = 0
         self.confusion_matrix = ConfusionMatrix(nc=self.nc)
-        self.metrics = Metrics()
+        self.metrics = SegmentMetrics(save_dir=self.save_dir, plot=self.args.plots, names=self.names)
         self.loss = torch.zeros(4, device=self.device)
         self.jdict = []
         self.stats = []
@@ -150,16 +150,6 @@ class SegmentationValidator(BaseValidator):
             # callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
             '''
 
-    def get_stats(self):
-        stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)]  # to numpy
-        if len(stats) and stats[0].any():
-            results = ap_per_class_box_and_mask(*stats, plot=self.args.plots, save_dir=self.save_dir, names=self.names)
-            self.metrics.update(results)
-        self.nt_per_class = np.bincount(stats[4].astype(int), minlength=self.nc)  # number of targets per class
-        metrics = {"fitness": fitness_segmentation(np.array(self.metrics.mean_results()).reshape(1, -1))}
-        metrics |= zip(self.metric_keys, self.metrics.mean_results())
-        return metrics
-
     def print_results(self):
         pf = '%22s' + '%11i' * 2 + '%11.3g' * 8  # print format
         self.logger.info(pf % ("all", self.seen, self.nt_per_class.sum(), *self.metrics.mean_results()))
@@ -218,6 +208,7 @@ class SegmentationValidator(BaseValidator):
         gs = max(int(de_parallel(self.model).stride if self.model else 0), 32)
         return build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, mode="val")[0]
 
+    # TODO: probably add this to class Metrics
     @property
     def metric_keys(self):
         return [
@@ -237,23 +228,22 @@ class SegmentationValidator(BaseValidator):
         bboxes = batch["bboxes"]
         paths = batch["im_file"]
         batch_idx = batch["batch_idx"]
-        plot_images_and_masks(images,
-                              batch_idx,
-                              cls,
-                              bboxes,
-                              masks,
-                              paths=paths,
-                              fname=self.save_dir / f"val_batch{ni}_labels.jpg",
-                              names=self.names)
+        plot_images(images,
+                    batch_idx,
+                    cls,
+                    bboxes,
+                    masks,
+                    paths=paths,
+                    fname=self.save_dir / f"val_batch{ni}_labels.jpg",
+                    names=self.names)
 
     def plot_predictions(self, batch, preds, ni):
         images = batch["img"]
         paths = batch["im_file"]
         if len(self.plot_masks):
             plot_masks = torch.cat(self.plot_masks, dim=0)
-        batch_idx, cls, bboxes, conf = output_to_target(preds[0], max_det=15)
-        plot_images_and_masks(images, batch_idx, cls, bboxes, plot_masks, conf, paths,
-                              self.save_dir / f'val_batch{ni}_pred.jpg', self.names)  # pred
+        plot_images(images, *output_to_target(preds[0], max_det=15), plot_masks, paths,
+                    self.save_dir / f'val_batch{ni}_pred.jpg', self.names)  # pred
         self.plot_masks.clear()