From d63ee112d49451fb55cf50e723f98479a4710c0f Mon Sep 17 00:00:00 2001 From: Laughing <61612323+Laughing-q@users.noreply.github.com> Date: Thu, 8 Dec 2022 08:28:13 -0600 Subject: [PATCH] General cleanup (#69) Co-authored-by: ayush chaurasia Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- .../tests/data/dataloader/yolodetection.py | 81 ++++--- .../tests/data/dataloader/yolosegment.py | 91 +++----- ultralytics/yolo/data/augment.py | 32 +-- ultralytics/yolo/engine/model.py | 1 + ultralytics/yolo/engine/trainer.py | 2 +- ultralytics/yolo/utils/metrics.py | 149 +++++++------ ultralytics/yolo/utils/plotting.py | 203 ++++-------------- ultralytics/yolo/v8/detect/__init__.py | 6 +- ultralytics/yolo/v8/detect/train.py | 34 ++- ultralytics/yolo/v8/detect/val.py | 18 +- ultralytics/yolo/v8/segment/__init__.py | 6 +- ultralytics/yolo/v8/segment/train.py | 37 +--- ultralytics/yolo/v8/segment/val.py | 44 ++-- 13 files changed, 268 insertions(+), 436 deletions(-) diff --git a/ultralytics/tests/data/dataloader/yolodetection.py b/ultralytics/tests/data/dataloader/yolodetection.py index 7d37a84..db6da14 100644 --- a/ultralytics/tests/data/dataloader/yolodetection.py +++ b/ultralytics/tests/data/dataloader/yolodetection.py @@ -1,8 +1,12 @@ import cv2 +import hydra import numpy as np -from omegaconf import OmegaConf from ultralytics.yolo.data import build_dataloader +from ultralytics.yolo.utils import ROOT +from ultralytics.yolo.utils.plotting import plot_images + +DEFAULT_CONFIG = ROOT / "yolo/utils/configs/default.yaml" class Colors: @@ -51,47 +55,34 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None): ) -with open("ultralytics/tests/data/dataloader/hyp_test.yaml") as f: - hyp = OmegaConf.load(f) - -dataloader, dataset = build_dataloader( - img_path="/d/dataset/COCO/coco128-seg/images", - img_size=640, - label_path=None, - cache=False, - hyp=hyp, - augment=False, - prefix="", - rect=False, - batch_size=4, - stride=32, - pad=0.5, - use_segments=True, - use_keypoints=False, -) - -for d in dataloader: - idx = 1 # show which image inside one batch - img = d["img"][idx].numpy() - img = np.ascontiguousarray(img.transpose(1, 2, 0)) - ih, iw = img.shape[:2] - # print(img.shape) - bidx = d["batch_idx"] - cls = d["cls"][bidx == idx].numpy() - bboxes = d["bboxes"][bidx == idx].numpy() - print(bboxes.shape) - bboxes[:, [0, 2]] *= iw - bboxes[:, [1, 3]] *= ih - nl = len(cls) - - for i, b in enumerate(bboxes): - x, y, w, h = b - x1 = x - w / 2 - x2 = x + w / 2 - y1 = y - h / 2 - y2 = y + h / 2 - c = int(cls[i][0]) - plot_one_box([int(x1), int(y1), int(x2), int(y2)], img, label=f"{c}", color=colors(c)) - cv2.imshow("p", img) - if cv2.waitKey(0) == ord("q"): - break +@hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name) +def test(cfg): + cfg.task = "detect" + cfg.mode = "train" + dataloader, _ = build_dataloader( + cfg=cfg, + batch_size=4, + img_path="/d/dataset/COCO/coco128-seg/images", + stride=32, + label_path=None, + mode=cfg.mode, + ) + + for d in dataloader: + images = d["img"] + cls = d["cls"].squeeze(-1) + bboxes = d["bboxes"] + paths = d["im_file"] + batch_idx = d["batch_idx"] + result = plot_images(images, batch_idx, cls, bboxes, paths=paths) + + cv2.imshow("p", result) + if cv2.waitKey(0) == ord("q"): + break + + +if __name__ == "__main__": + test() + # test(augment=True, rect=False) + # test(augment=False, rect=True) + # test(augment=False, rect=False) diff --git a/ultralytics/tests/data/dataloader/yolosegment.py b/ultralytics/tests/data/dataloader/yolosegment.py index 24c2263..cd38ba2 100644 --- a/ultralytics/tests/data/dataloader/yolosegment.py +++ b/ultralytics/tests/data/dataloader/yolosegment.py @@ -1,9 +1,11 @@ import cv2 -import numpy as np -import torch -from omegaconf import OmegaConf +import hydra from ultralytics.yolo.data import build_dataloader +from ultralytics.yolo.utils import ROOT +from ultralytics.yolo.utils.plotting import plot_images + +DEFAULT_CONFIG = ROOT / "yolo/utils/configs/default.yaml" class Colors: @@ -52,77 +54,34 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None): ) -with open("ultralytics/tests/data/dataloader/hyp_test.yaml") as f: - hyp = OmegaConf.load(f) - - -def test(augment, rect): +@hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name) +def test(cfg): + cfg.task = "segment" + cfg.mode = "train" dataloader, _ = build_dataloader( - img_path="/d/dataset/COCO/coco128-seg/images", - img_size=640, - label_path=None, - cache=False, - hyp=hyp, - augment=augment, - prefix="", - rect=rect, + cfg=cfg, batch_size=4, + img_path="/d/dataset/COCO/coco128-seg/images", stride=32, - pad=0.5, - use_segments=True, - use_keypoints=False, + label_path=None, + mode=cfg.mode, ) for d in dataloader: - # info - im_file = d["im_file"] - ori_shape = d["ori_shape"] - resize_shape = d["resized_shape"] - print(ori_shape, resize_shape) - print(im_file) - - # labels - idx = 1 # show which image inside one batch - img = d["img"][idx].numpy() - img = np.ascontiguousarray(img.transpose(1, 2, 0)) - ih, iw = img.shape[:2] - # print(img.shape) - bidx = d["batch_idx"] - cls = d["cls"][bidx == idx].numpy() - bboxes = d["bboxes"][bidx == idx].numpy() - masks = d["masks"][idx] - print(bboxes.shape) - bboxes[:, [0, 2]] *= iw - bboxes[:, [1, 3]] *= ih - nl = len(cls) - - index = torch.arange(nl).view(nl, 1, 1) + 1 - masks = masks.repeat(nl, 1, 1) - # print(masks.shape, index.shape) - masks = torch.where(masks == index, 1, 0) - masks = masks.numpy().astype(np.uint8) - print(masks.shape) - # keypoints = d["keypoints"] - - for i, b in enumerate(bboxes): - x, y, w, h = b - x1 = x - w / 2 - x2 = x + w / 2 - y1 = y - h / 2 - y2 = y + h / 2 - c = int(cls[i][0]) - # print(x1, y1, x2, y2) - plot_one_box([int(x1), int(y1), int(x2), int(y2)], img, label=f"{c}", color=colors(c)) - mask = masks[i] - mask = cv2.resize(mask, (iw, ih)) - mask = mask.astype(bool) - img[mask] = img[mask] * 0.5 + np.array(colors(c)) * 0.5 - cv2.imshow("p", img) + images = d["img"] + masks = d["masks"] + cls = d["cls"].squeeze(-1) + bboxes = d["bboxes"] + paths = d["im_file"] + batch_idx = d["batch_idx"] + result = plot_images(images, batch_idx, cls, bboxes, masks, paths=paths) + cv2.imshow("p", result) if cv2.waitKey(0) == ord("q"): break if __name__ == "__main__": - test(augment=True, rect=False) - test(augment=False, rect=True) - test(augment=False, rect=False) + test() + # test(augment=True, rect=False) + # test(augment=False, rect=True) + # test(augment=False, rect=False) diff --git a/ultralytics/yolo/data/augment.py b/ultralytics/yolo/data/augment.py index 71b9441..16640ce 100644 --- a/ultralytics/yolo/data/augment.py +++ b/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 diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 163210d..447f311 100644 --- a/ultralytics/yolo/engine/model.py +++ b/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 diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 4dcbf4c..307b5d8 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/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): diff --git a/ultralytics/yolo/utils/metrics.py b/ultralytics/yolo/utils/metrics.py index 1f80b92..f82b5c2 100644 --- a/ultralytics/yolo/utils/metrics.py +++ b/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 diff --git a/ultralytics/yolo/utils/plotting.py b/ultralytics/yolo/utils/plotting.py index cdd2fd2..89128c9 100644 --- a/ultralytics/yolo/utils/plotting.py +++ b/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) - - 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) + 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) + + 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,38 +274,34 @@ 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) + 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." + 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] + 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=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)}") + ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=8) + ax[i].set_title(s[j], fontsize=12) # 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}") + print(f'Warning: Plotting error for {f}: {e}') ax[1].legend() - fig.savefig(save_dir / "results.png", dpi=200) + fig.savefig(save_dir / 'results.png', dpi=200) plt.close() @@ -334,100 +313,4 @@ def output_to_target(output, max_det=300): 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) - 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) - 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]): - 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) - ax[i].set_title(s[j], fontsize=12) - # 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() + return targets[:, 0], targets[:, 1], targets[:, 2:] diff --git a/ultralytics/yolo/v8/detect/__init__.py b/ultralytics/yolo/v8/detect/__init__.py index e158fe8..a683006 100644 --- a/ultralytics/yolo/v8/detect/__init__.py +++ b/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 diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index 553f0d9..e967ebf 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/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] diff --git a/ultralytics/yolo/v8/detect/val.py b/ultralytics/yolo/v8/detect/val.py index f6f861f..56c98b9 100644 --- a/ultralytics/yolo/v8/detect/val.py +++ b/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) diff --git a/ultralytics/yolo/v8/segment/__init__.py b/ultralytics/yolo/v8/segment/__init__.py index a9d00cb..f299f27 100644 --- a/ultralytics/yolo/v8/segment/__init__.py +++ b/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 diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py index 0286ce1..2ec1df1 100644 --- a/ultralytics/yolo/v8/segment/train.py +++ b/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 +# BaseTrainer python usage +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) diff --git a/ultralytics/yolo/v8/segment/val.py b/ultralytics/yolo/v8/segment/val.py index 7ada26c..56bc038 100644 --- a/ultralytics/yolo/v8/segment/val.py +++ b/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()