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)