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/data/augment.py

@@ -521,23 +521,25 @@ class CopyPaste:
         instances.convert_bbox(format="xyxy")
         if self.p and len(instances.segments):
             n = len(instances)
-            h, w, _ = im.shape  # height, width, channels
+            _, w, _ = im.shape  # height, width, channels
             im_new = np.zeros(im.shape, np.uint8)
-            j = random.sample(range(n), k=round(self.p * n))
-            c, instance = cls[j], instances[j]
-            instance.fliplr(w)
-            ioa = bbox_ioa(instance.bboxes, instances.bboxes)  # intersection over area, (N, M)
-            i = (ioa < 0.30).all(1)  # (N, )
-            if i.sum():
-                cls = np.concatenate((cls, c[i]), axis=0)
-                instances = Instances.concatenate((instances, instance[i]), axis=0)
-                cv2.drawContours(im_new, instances.segments[j][i].astype(np.int32), -1, (255, 255, 255), cv2.FILLED)
 
-            result = cv2.bitwise_and(src1=im, src2=im_new)
-            result = cv2.flip(result, 1)  # augment segments (flip left-right)
-            i = result > 0  # pixels to replace
-            # i[:, :] = result.max(2).reshape(h, w, 1)  # act over ch
+            # calculate ioa first then select indexes randomly
+            ins_flip = deepcopy(instances)
+            ins_flip.fliplr(w)
+
+            ioa = bbox_ioa(ins_flip.bboxes, instances.bboxes)  # intersection over area, (N, M)
+            indexes = np.nonzero((ioa < 0.30).all(1))[0]  # (N, )
+            n = len(indexes)
+            for j in random.sample(list(indexes), k=round(self.p * n)):
+                cls = np.concatenate((cls, cls[[j]]), axis=0)
+                instances = Instances.concatenate((instances, ins_flip[[j]]), axis=0)
+                cv2.drawContours(im_new, instances.segments[[j]].astype(np.int32), -1, (1, 1, 1), cv2.FILLED)
+
+            result = cv2.flip(im, 1)  # augment segments (flip left-right)
+            i = cv2.flip(im_new, 1).astype(bool)
             im[i] = result[i]  # cv2.imwrite('debug.jpg', im)  # debug
+
         labels["img"] = im
         labels["cls"] = cls
         labels["instances"] = instances

ultralytics/yolo/engine/model.py

@@ -4,6 +4,7 @@ Top-level YOLO model interface. First principle usage example - https://github.c
 import torch
 import yaml
 
+from ultralytics import yolo
 from ultralytics.yolo.utils import LOGGER
 from ultralytics.yolo.utils.checks import check_yaml
 from ultralytics.yolo.utils.modeling import attempt_load_weights

ultralytics/yolo/engine/trainer.py

@@ -327,7 +327,7 @@ class BaseTrainer:
         metrics = self.validator(self)
         fitness = metrics.pop("fitness", -self.loss.detach().cpu().numpy())  # use loss as fitness measure if not found
         if not self.best_fitness or self.best_fitness < fitness:
-            self.best_fitness = self.fitness
+            self.best_fitness = fitness
         return metrics, fitness
 
     def log(self, text, rank=-1):

ultralytics/yolo/utils/metrics.py

@@ -263,18 +263,6 @@ class ConfusionMatrix:
             print(' '.join(map(str, self.matrix[i])))
 
 
-def fitness_detection(x):
-    # Model fitness as a weighted combination of metrics
-    w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
-    return (x[:, :4] * w).sum(1)
-
-
-def fitness_segmentation(x):
-    # Model fitness as a weighted combination of metrics
-    w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9]
-    return (x[:, :8] * w).sum(1)
-
-
 def smooth(y, f=0.05):
     # Box filter of fraction f
     nf = round(len(y) * f * 2) // 2 + 1  # number of filter elements (must be odd)
@@ -422,55 +410,6 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
     return tp, fp, p, r, f1, ap, unique_classes.astype(int)
 
 
-def ap_per_class_box_and_mask(
-        tp_m,
-        tp_b,
-        conf,
-        pred_cls,
-        target_cls,
-        plot=False,
-        save_dir=".",
-        names=(),
-):
-    """
-    Args:
-        tp_b: tp of boxes.
-        tp_m: tp of masks.
-        other arguments see `func: ap_per_class`.
-    """
-    results_boxes = ap_per_class(tp_b,
-                                 conf,
-                                 pred_cls,
-                                 target_cls,
-                                 plot=plot,
-                                 save_dir=save_dir,
-                                 names=names,
-                                 prefix="Box")[2:]
-    results_masks = ap_per_class(tp_m,
-                                 conf,
-                                 pred_cls,
-                                 target_cls,
-                                 plot=plot,
-                                 save_dir=save_dir,
-                                 names=names,
-                                 prefix="Mask")[2:]
-
-    results = {
-        "boxes": {
-            "p": results_boxes[0],
-            "r": results_boxes[1],
-            "f1": results_boxes[2],
-            "ap": results_boxes[3],
-            "ap_class": results_boxes[4]},
-        "masks": {
-            "p": results_masks[0],
-            "r": results_masks[1],
-            "f1": results_masks[2],
-            "ap": results_masks[3],
-            "ap_class": results_masks[4]}}
-    return results
-
-
 class Metric:
 
     def __init__(self) -> None:
@@ -542,6 +481,11 @@ class Metric:
             maps[c] = self.ap[i]
         return maps
 
+    def fitness(self):
+        # Model fitness as a weighted combination of metrics
+        w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
+        return (np.array(self.mean_results()) * w).sum()
+
     def update(self, results):
         """
         Args:
@@ -555,20 +499,80 @@ class Metric:
         self.ap_class_index = ap_class_index
 
 
-class Metrics:
-    """Metric for boxes and masks."""
+class DetMetrics:
 
-    def __init__(self) -> None:
+    def __init__(self, save_dir=Path("."), plot=False, names=()) -> None:
+        self.save_dir = save_dir
+        self.plot = plot
+        self.names = names
+        self.metric = Metric()
+
+    def process(self, tp, conf, pred_cls, target_cls):
+        results = ap_per_class(tp, conf, pred_cls, target_cls, plot=self.plot, save_dir=self.save_dir,
+                               names=self.names)[2:]
+        self.metric.update(results)
+
+    @property
+    def keys(self):
+        return ["metrics/precision(B)", "metrics/recall(B)", "metrics/mAP_0.5(B)", "metrics/mAP_0.5:0.95(B)"]
+
+    def mean_results(self):
+        return self.metric.mean_results()
+
+    def class_result(self, i):
+        return self.metric.class_result(i)
+
+    def get_maps(self, nc):
+        return self.metric.get_maps(nc)
+
+    def fitness(self):
+        return self.metric.fitness()
+
+    @property
+    def ap_class_index(self):
+        return self.metric.ap_class_index
+
+
+class SegmentMetrics:
+
+    def __init__(self, save_dir=Path("."), plot=False, names=()) -> None:
+        self.save_dir = save_dir
+        self.plot = plot
+        self.names = names
         self.metric_box = Metric()
         self.metric_mask = Metric()
 
-    def update(self, results):
-        """
-        Args:
-            results: Dict{'boxes': Dict{}, 'masks': Dict{}}
-        """
-        self.metric_box.update(list(results["boxes"].values()))
-        self.metric_mask.update(list(results["masks"].values()))
+    def process(self, tp_m, tp_b, conf, pred_cls, target_cls):
+        results_mask = ap_per_class(tp_m,
+                                    conf,
+                                    pred_cls,
+                                    target_cls,
+                                    plot=self.plot,
+                                    save_dir=self.save_dir,
+                                    names=self.names,
+                                    prefix="Mask")[2:]
+        self.metric_mask.update(results_mask)
+        results_box = ap_per_class(tp_b,
+                                   conf,
+                                   pred_cls,
+                                   target_cls,
+                                   plot=self.plot,
+                                   save_dir=self.save_dir,
+                                   names=self.names,
+                                   prefix="Box")[2:]
+        self.metric_box.update(results_box)
+
+    @property
+    def keys(self):
+        return [
+            "metrics/precision(B)",
+            "metrics/recall(B)",
+            "metrics/mAP_0.5(B)",
+            "metrics/mAP_0.5:0.95(B)",  # metrics
+            "metrics/precision(M)",
+            "metrics/recall(M)",
+            "metrics/mAP_0.5(M)",
+            "metrics/mAP_0.5:0.95(M)"]
 
     def mean_results(self):
         return self.metric_box.mean_results() + self.metric_mask.mean_results()
@@ -579,6 +583,9 @@ class Metrics:
     def get_maps(self, nc):
         return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc)
 
+    def fitness(self):
+        return self.metric_mask.fitness() + self.metric_box.fitness()
+
     @property
     def ap_class_index(self):
         # boxes and masks have the same ap_class_index

ultralytics/yolo/utils/plotting.py

@@ -84,7 +84,7 @@ class Annotator:
                             thickness=tf,
                             lineType=cv2.LINE_AA)
 
-    def masks(self, masks, colors, im_gpu=None, alpha=0.5):
+    def masks(self, masks, colors, im_gpu, alpha=0.5, retina_masks=False):
         """Plot masks at once.
         Args:
             masks (tensor): predicted masks on cuda, shape: [n, h, w]
@@ -95,37 +95,21 @@ class Annotator:
         if self.pil:
             # convert to numpy first
             self.im = np.asarray(self.im).copy()
-        if im_gpu is None:
-            # Add multiple masks of shape(h,w,n) with colors list([r,g,b], [r,g,b], ...)
-            if len(masks) == 0:
-                return
-            if isinstance(masks, torch.Tensor):
-                masks = torch.as_tensor(masks, dtype=torch.uint8)
-                masks = masks.permute(1, 2, 0).contiguous()
-                masks = masks.cpu().numpy()
-            # masks = np.ascontiguousarray(masks.transpose(1, 2, 0))
-            masks = scale_image(masks.shape[:2], masks, self.im.shape)
-            masks = np.asarray(masks, dtype=np.float32)
-            colors = np.asarray(colors, dtype=np.float32)  # shape(n,3)
-            s = masks.sum(2, keepdims=True).clip(0, 1)  # add all masks together
-            masks = (masks @ colors).clip(0, 255)  # (h,w,n) @ (n,3) = (h,w,3)
-            self.im[:] = masks * alpha + self.im * (1 - s * alpha)
-        else:
-            if len(masks) == 0:
-                self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
-            colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
-            colors = colors[:, None, None]  # shape(n,1,1,3)
-            masks = masks.unsqueeze(3)  # shape(n,h,w,1)
-            masks_color = masks * (colors * alpha)  # shape(n,h,w,3)
+        if len(masks) == 0:
+            self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
+        colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
+        colors = colors[:, None, None]  # shape(n,1,1,3)
+        masks = masks.unsqueeze(3)  # shape(n,h,w,1)
+        masks_color = masks * (colors * alpha)  # shape(n,h,w,3)
 
-            inv_alph_masks = (1 - masks * alpha).cumprod(0)  # shape(n,h,w,1)
-            mcs = (masks_color * inv_alph_masks).sum(0) * 2  # mask color summand shape(n,h,w,3)
+        inv_alph_masks = (1 - masks * alpha).cumprod(0)  # shape(n,h,w,1)
+        mcs = (masks_color * inv_alph_masks).sum(0) * 2  # mask color summand shape(n,h,w,3)
 
-            im_gpu = im_gpu.flip(dims=[0])  # flip channel
-            im_gpu = im_gpu.permute(1, 2, 0).contiguous()  # shape(h,w,3)
-            im_gpu = im_gpu * inv_alph_masks[-1] + mcs
-            im_mask = (im_gpu * 255).byte().cpu().numpy()
-            self.im[:] = scale_image(im_gpu.shape, im_mask, self.im.shape)
+        im_gpu = im_gpu.flip(dims=[0])  # flip channel
+        im_gpu = im_gpu.permute(1, 2, 0).contiguous()  # shape(h,w,3)
+        im_gpu = im_gpu * inv_alph_masks[-1] + mcs
+        im_mask = (im_gpu * 255).byte().cpu().numpy()
+        self.im[:] = im_mask if retina_masks else scale_image(im_gpu.shape, im_mask, self.im.shape)
         if self.pil:
             # convert im back to PIL and update draw
             self.fromarray(self.im)
@@ -186,15 +170,14 @@ def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False,
 
 
 @threaded
-def plot_images_and_masks(images,
-                          batch_idx,
-                          cls,
-                          bboxes,
-                          masks,
-                          confs=None,
-                          paths=None,
-                          fname='images.jpg',
-                          names=None):
+def plot_images(images,
+                batch_idx,
+                cls,
+                bboxes,
+                masks=np.zeros(0, dtype=np.uint8),
+                paths=None,
+                fname='images.jpg',
+                names=None):
     # Plot image grid with labels
     if isinstance(images, torch.Tensor):
         images = images.cpu().float().numpy()
@@ -242,10 +225,10 @@ def plot_images_and_masks(images,
         if len(cls) > 0:
             idx = batch_idx == i
 
-            boxes = xywh2xyxy(bboxes[idx]).T
+            boxes = xywh2xyxy(bboxes[idx, :4]).T
             classes = cls[idx].astype('int')
-            labels = confs is None  # labels if no conf column
-            conf = None if labels else confs[idx]  # check for confidence presence (label vs pred)
+            labels = bboxes.shape[1] == 4  # labels if no conf column
+            conf = None if labels else bboxes[idx, 4]  # check for confidence presence (label vs pred)
 
             if boxes.shape[1]:
                 if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
@@ -291,126 +274,15 @@ def plot_images_and_masks(images,
     annotator.im.save(fname)  # save
 
 
-def plot_results_with_masks(file="path/to/results.csv", dir="", best=True):
+def plot_results(file='path/to/results.csv', dir='', segment=False):
     # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
     save_dir = Path(file).parent if file else Path(dir)
-    fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
-    ax = ax.ravel()
-    files = list(save_dir.glob("results*.csv"))
-    assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
-    for f in files:
-        try:
-            data = pd.read_csv(f)
-            index = np.argmax(0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] +
-                              0.1 * data.values[:, 11])
-            s = [x.strip() for x in data.columns]
-            x = data.values[:, 0]
-            for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]):
-                y = data.values[:, j]
-                # y[y == 0] = np.nan  # don't show zero values
-                ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=2)
-                if best:
-                    # best
-                    ax[i].scatter(index, y[index], color="r", label=f"best:{index}", marker="*", linewidth=3)
-                    ax[i].set_title(s[j] + f"\n{round(y[index], 5)}")
-                else:
-                    # last
-                    ax[i].scatter(x[-1], y[-1], color="r", label="last", marker="*", linewidth=3)
-                    ax[i].set_title(s[j] + f"\n{round(y[-1], 5)}")
-                # if j in [8, 9, 10]:  # share train and val loss y axes
-                #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
-        except Exception as e:
-            print(f"Warning: Plotting error for {f}: {e}")
-    ax[1].legend()
-    fig.savefig(save_dir / "results.png", dpi=200)
-    plt.close()
-
-
-def output_to_target(output, max_det=300):
-    # Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting
-    targets = []
-    for i, o in enumerate(output):
-        box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
-        j = torch.full((conf.shape[0], 1), i)
-        targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
-    targets = torch.cat(targets, 0).numpy()
-    return targets[:, 0], targets[:, 1], targets[:, 2:6], targets[:, 6]
-
-
-@threaded
-def plot_images(images, batch_idx, cls, bboxes, confs=None, paths=None, fname='images.jpg', names=None):
-    # Plot image grid with labels
-    if isinstance(images, torch.Tensor):
-        images = images.cpu().float().numpy()
-    if isinstance(cls, torch.Tensor):
-        cls = cls.cpu().numpy()
-    if isinstance(bboxes, torch.Tensor):
-        bboxes = bboxes.cpu().numpy()
-    if isinstance(batch_idx, torch.Tensor):
-        batch_idx = batch_idx.cpu().numpy()
-
-    max_size = 1920  # max image size
-    max_subplots = 16  # max image subplots, i.e. 4x4
-    bs, _, h, w = images.shape  # batch size, _, height, width
-    bs = min(bs, max_subplots)  # limit plot images
-    ns = np.ceil(bs ** 0.5)  # number of subplots (square)
-    if np.max(images[0]) <= 1:
-        images *= 255  # de-normalise (optional)
-
-    # Build Image
-    mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)  # init
-    for i, im in enumerate(images):
-        if i == max_subplots:  # if last batch has fewer images than we expect
-            break
-        x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
-        im = im.transpose(1, 2, 0)
-        mosaic[y:y + h, x:x + w, :] = im
-
-    # Resize (optional)
-    scale = max_size / ns / max(h, w)
-    if scale < 1:
-        h = math.ceil(scale * h)
-        w = math.ceil(scale * w)
-        mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
-
-    # Annotate
-    fs = int((h + w) * ns * 0.01)  # font size
-    annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
-    for i in range(i + 1):
-        x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
-        annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2)  # borders
-        if paths:
-            annotator.text((x + 5, y + 5 + h), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220))  # filenames
-        if len(cls) > 0:
-            idx = batch_idx == i
-
-            boxes = xywh2xyxy(bboxes[idx]).T
-            classes = cls[idx].astype('int')
-            labels = confs is None  # labels if no conf column
-            conf = None if labels else confs[idx]  # check for confidence presence (label vs pred)
-
-            if boxes.shape[1]:
-                if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
-                    boxes[[0, 2]] *= w  # scale to pixels
-                    boxes[[1, 3]] *= h
-                elif scale < 1:  # absolute coords need scale if image scales
-                    boxes *= scale
-            boxes[[0, 2]] += x
-            boxes[[1, 3]] += y
-            for j, box in enumerate(boxes.T.tolist()):
-                c = classes[j]
-                color = colors(c)
-                c = names[c] if names else c
-                if labels or conf[j] > 0.25:  # 0.25 conf thresh
-                    label = f'{c}' if labels else f'{c} {conf[j]:.1f}'
-                    annotator.box_label(box, label, color=color)
-    annotator.im.save(fname)  # save
-
-
-def plot_results(file='path/to/results.csv', dir=''):
-    # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
-    save_dir = Path(file).parent if file else Path(dir)
-    fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
+    if segment:
+        fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
+        index = [1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]
+    else:
+        fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
+        index = [1, 2, 3, 4, 5, 8, 9, 10, 6, 7]
     ax = ax.ravel()
     files = list(save_dir.glob('results*.csv'))
     assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.'
@@ -419,7 +291,7 @@ def plot_results(file='path/to/results.csv', dir=''):
             data = pd.read_csv(f)
             s = [x.strip() for x in data.columns]
             x = data.values[:, 0]
-            for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]):
+            for i, j in enumerate(index):
                 y = data.values[:, j].astype('float')
                 # y[y == 0] = np.nan  # don't show zero values
                 ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=8)
@@ -431,3 +303,14 @@ def plot_results(file='path/to/results.csv', dir=''):
     ax[1].legend()
     fig.savefig(save_dir / 'results.png', dpi=200)
     plt.close()
+
+
+def output_to_target(output, max_det=300):
+    # Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting
+    targets = []
+    for i, o in enumerate(output):
+        box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
+        j = torch.full((conf.shape[0], 1), i)
+        targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
+    targets = torch.cat(targets, 0).numpy()
+    return targets[:, 0], targets[:, 1], targets[:, 2:]

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()