YOLOv5 segmentation model support (#9052)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix duplicate plots.py

* Fix check_font()

* # torch.use_deterministic_algorithms(True)

* update doc detect->predict

* Resolve precommit for segment/train and segment/val

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Resolve precommit for utils/segment

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Resolve precommit min_wh

* Resolve precommit utils/segment/plots

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Resolve precommit utils/segment/general

* Align NMS-seg closer to NMS

* restore deterministic init_seeds code

* remove easydict dependency

* update

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* restore output_to_target mask

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update

* cleanup

* Remove unused ImageFont import

* Unified NMS

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* DetectMultiBackend compatibility

* segment/predict.py update

* update plot colors

* fix bbox shifted

* sort bbox by confidence

* enable overlap by default

* Merge detect/segment output_to_target() function

* Start segmentation CI

* fix plots

* Update ci-testing.yml

* fix training whitespace

* optimize process mask functions (can we merge both?)

* Update predict/detect

* Update plot_images

* Update plot_images_and_masks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix

* Add train to CI

* fix precommit

* fix precommit CI

* fix precommit pycocotools

* fix val float issues

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix masks float float issues

* suppress errors

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix no-predictions plotting bug

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add CSV Logger

* fix val len(plot_masks)

* speed up evaluation

* fix process_mask

* fix plots

* update segment/utils build_targets

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* optimize utils/segment/general crop()

* optimize utils/segment/general crop() 2

* minor updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* torch.where revert

* downsample only if different shape

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* loss cleanup

* loss cleanup 2

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* loss cleanup 3

* update project names

* Rename -seg yamls from _underscore to -dash

* prepare for yolov5n-seg.pt

* precommit space fix

* add coco128-seg.yaml

* update coco128-seg comments

* cleanup val.py

* Major val.py cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* precommit fix

* precommit fix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* optional pycocotools

* remove CI pip install pycocotools (auto-installed now)

* seg yaml fix

* optimize mask_iou() and masks_iou()

* threaded fix

* Major train.py update

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Major segments/val/process_batch() update

* yolov5/val updates from segment

* process_batch numpy/tensor fix

* opt-in to pycocotools with --save-json

* threaded pycocotools ops for 2x speed increase

* Avoid permute contiguous if possible

* Add max_det=300 argument to both val.py and segment/val.py

* fix onnx_dynamic

* speed up pycocotools ops

* faster process_mask(upsample=True) for predict

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* eliminate permutations for process_mask(upsample=True)

* eliminate permute-contiguous in crop(), use native dimension order

* cleanup comment

* Add Proto() module

* fix class count

* fix anchor order

* broadcast mask_gti in loss for speed

* Cleanup seg loss

* faster indexing

* faster indexing fix

* faster indexing fix2

* revert faster indexing

* fix validation plotting

* Loss cleanup and mxyxy simplification

* Loss cleanup and mxyxy simplification 2

* revert validation plotting

* replace missing tanh

* Eliminate last permutation

* delete unneeded .float()

* Remove MaskIOULoss and crop(if HWC)

* Final v6.3 SegmentationModel architecture updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add support for TF export

* remove debugger trace

* add call

* update

* update

* Merge master

* Merge master

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update dataloaders.py

* Restore CI

* Update dataloaders.py

* Fix TF/TFLite export for segmentation model

* Merge master

* Cleanup predict.py mask plotting

* cleanup scale_masks()

* rename scale_masks to scale_image

* cleanup/optimize plot_masks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add Annotator.masks()

* Annotator.masks() fix

* Update plots.py

* Annotator mask optimization

* Rename crop() to crop_mask()

* Do not crop in predict.py

* crop always

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Merge master

* Add vid-stride from master PR

* Update seg model outputs

* Update seg model outputs

* Add segmentation benchmarks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add segmentation benchmarks

* Add segmentation benchmarks

* Add segmentation benchmarks

* Fix DetectMultiBackend for OpenVINO

* update Annotator.masks

* fix val plot

* revert val plot

* clean up

* revert pil

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix CI error

* fix predict log

* remove upsample

* update interpolate

* fix validation plot logging

* Annotator.masks() cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove segmentation_model definition

* Restore 0.99999 decimals

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>
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: Jiacong Fang <zldrobit@126.com>

val.py

@@ -71,12 +71,12 @@ def save_one_json(predn, jdict, path, class_map):
 
 def process_batch(detections, labels, iouv):
     """
-    Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.
+    Return correct prediction matrix
     Arguments:
-        detections (Array[N, 6]), x1, y1, x2, y2, conf, class
-        labels (Array[M, 5]), class, x1, y1, x2, y2
+        detections (array[N, 6]), x1, y1, x2, y2, conf, class
+        labels (array[M, 5]), class, x1, y1, x2, y2
     Returns:
-        correct (Array[N, 10]), for 10 IoU levels
+        correct (array[N, 10]), for 10 IoU levels
     """
     correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
     iou = box_iou(labels[:, 1:], detections[:, :4])
@@ -102,6 +102,7 @@ def run(
         imgsz=640,  # inference size (pixels)
         conf_thres=0.001,  # confidence threshold
         iou_thres=0.6,  # NMS IoU threshold
+        max_det=300,  # maximum detections per image
         task='val',  # train, val, test, speed or study
         device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
         workers=8,  # max dataloader workers (per RANK in DDP mode)
@@ -187,7 +188,7 @@ def run(
     if isinstance(names, (list, tuple)):  # old format
         names = dict(enumerate(names))
     class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
-    s = ('%22s' + '%11s' * 6) % ('Class', 'Images', 'Instances', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
+    s = ('%22s' + '%11s' * 6) % ('Class', 'Images', 'Instances', 'P', 'R', 'mAP50', 'mAP50-95')
     dt, p, r, f1, mp, mr, map50, map = (Profile(), Profile(), Profile()), 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
     loss = torch.zeros(3, device=device)
     jdict, stats, ap, ap_class = [], [], [], []
@@ -205,7 +206,7 @@ def run(
 
         # Inference
         with dt[1]:
-            out, train_out = model(im) if compute_loss else (model(im, augment=augment), None)
+            preds, train_out = model(im) if compute_loss else (model(im, augment=augment), None)
 
         # Loss
         if compute_loss:
@@ -215,10 +216,16 @@ def run(
         targets[:, 2:] *= torch.tensor((width, height, width, height), device=device)  # to pixels
         lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling
         with dt[2]:
-            out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
+            preds = non_max_suppression(preds,
+                                        conf_thres,
+                                        iou_thres,
+                                        labels=lb,
+                                        multi_label=True,
+                                        agnostic=single_cls,
+                                        max_det=max_det)
 
         # Metrics
-        for si, pred in enumerate(out):
+        for si, pred in enumerate(preds):
             labels = targets[targets[:, 0] == si, 1:]
             nl, npr = labels.shape[0], pred.shape[0]  # number of labels, predictions
             path, shape = Path(paths[si]), shapes[si][0]
@@ -258,9 +265,9 @@ def run(
         # Plot images
         if plots and batch_i < 3:
             plot_images(im, targets, paths, save_dir / f'val_batch{batch_i}_labels.jpg', names)  # labels
-            plot_images(im, output_to_target(out), paths, save_dir / f'val_batch{batch_i}_pred.jpg', names)  # pred
+            plot_images(im, output_to_target(preds), paths, save_dir / f'val_batch{batch_i}_pred.jpg', names)  # pred
 
-        callbacks.run('on_val_batch_end', batch_i, im, targets, paths, shapes, out)
+        callbacks.run('on_val_batch_end', batch_i, im, targets, paths, shapes, preds)
 
     # Compute metrics
     stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy
@@ -332,11 +339,12 @@ def run(
 def parse_opt():
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
+    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
     parser.add_argument('--batch-size', type=int, default=32, help='batch size')
     parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
     parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
     parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
+    parser.add_argument('--max-det', type=int, default=300, help='maximum detections per image')
     parser.add_argument('--task', default='val', help='train, val, test, speed or study')
     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')