Update Actions with Lychee and GitHub Token (#12592)

* Add Ultralytics Actions

* Auto-format by Ultralytics actions

* Update format.yml

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

---------

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

.github/workflows/format.yml

@@ -0,0 +1,25 @@
+# Ultralytics 🚀 - AGPL-3.0 license
+# Ultralytics Actions https://github.com/ultralytics/actions
+# This workflow automatically formats code and documentation in PRs to official Ultralytics standards
+
+name: Ultralytics Actions
+
+on:
+  push:
+    branches: [main,master]
+  pull_request_target:
+    branches: [main,master]
+
+jobs:
+  format:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Run Ultralytics Formatting
+        uses: ultralytics/actions@main
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}  # automatically generated
+          python: true
+          docstrings: true
+          markdown: true
+          spelling: true
+          links: true

benchmarks.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Run YOLOv5 benchmarks on all supported export formats
+Run YOLOv5 benchmarks on all supported export formats.
 
 Format                      | `export.py --include`         | Model
 ---                         | ---                           | ---
@@ -50,11 +50,11 @@ from val import run as val_det
 
 
 def run(
-        weights=ROOT / 'yolov5s.pt',  # weights path
+    weights=ROOT / "yolov5s.pt",  # weights path
     imgsz=640,  # inference size (pixels)
     batch_size=1,  # batch size
-        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
-        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     half=False,  # use FP16 half-precision inference
     test=False,  # test exports only
     pt_only=False,  # test PyTorch only
@@ -65,63 +65,60 @@ def run(
     model_type = type(attempt_load(weights, fuse=False))  # DetectionModel, SegmentationModel, etc.
     for i, (name, f, suffix, cpu, gpu) in export.export_formats().iterrows():  # index, (name, file, suffix, CPU, GPU)
         try:
-            assert i not in (9, 10), 'inference not supported'  # Edge TPU and TF.js are unsupported
+            assert i not in (9, 10), "inference not supported"  # Edge TPU and TF.js are unsupported
-            assert i != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13'  # CoreML
+            assert i != 5 or platform.system() == "Darwin", "inference only supported on macOS>=10.13"  # CoreML
-            if 'cpu' in device.type:
+            if "cpu" in device.type:
-                assert cpu, 'inference not supported on CPU'
+                assert cpu, "inference not supported on CPU"
-            if 'cuda' in device.type:
+            if "cuda" in device.type:
-                assert gpu, 'inference not supported on GPU'
+                assert gpu, "inference not supported on GPU"
 
             # Export
-            if f == '-':
+            if f == "-":
                 w = weights  # PyTorch format
             else:
-                w = export.run(weights=weights,
+                w = export.run(
-                               imgsz=[imgsz],
+                    weights=weights, imgsz=[imgsz], include=[f], batch_size=batch_size, device=device, half=half
-                               include=[f],
+                )[-1]  # all others
-                               batch_size=batch_size,
+            assert suffix in str(w), "export failed"
-                               device=device,
-                               half=half)[-1]  # all others
-            assert suffix in str(w), 'export failed'
 
             # Validate
             if model_type == SegmentationModel:
-                result = val_seg(data, w, batch_size, imgsz, plots=False, device=device, task='speed', half=half)
+                result = val_seg(data, w, batch_size, imgsz, plots=False, device=device, task="speed", half=half)
                 metric = result[0][7]  # (box(p, r, map50, map), mask(p, r, map50, map), *loss(box, obj, cls))
             else:  # DetectionModel:
-                result = val_det(data, w, batch_size, imgsz, plots=False, device=device, task='speed', half=half)
+                result = val_det(data, w, batch_size, imgsz, plots=False, device=device, task="speed", half=half)
                 metric = result[0][3]  # (p, r, map50, map, *loss(box, obj, cls))
             speed = result[2][1]  # times (preprocess, inference, postprocess)
             y.append([name, round(file_size(w), 1), round(metric, 4), round(speed, 2)])  # MB, mAP, t_inference
         except Exception as e:
             if hard_fail:
-                assert type(e) is AssertionError, f'Benchmark --hard-fail for {name}: {e}'
+                assert type(e) is AssertionError, f"Benchmark --hard-fail for {name}: {e}"
-            LOGGER.warning(f'WARNING ⚠️ Benchmark failure for {name}: {e}')
+            LOGGER.warning(f"WARNING ⚠️ Benchmark failure for {name}: {e}")
             y.append([name, None, None, None])  # mAP, t_inference
         if pt_only and i == 0:
             break  # break after PyTorch
 
     # Print results
-    LOGGER.info('\n')
+    LOGGER.info("\n")
     parse_opt()
     notebook_init()  # print system info
-    c = ['Format', 'Size (MB)', 'mAP50-95', 'Inference time (ms)'] if map else ['Format', 'Export', '', '']
+    c = ["Format", "Size (MB)", "mAP50-95", "Inference time (ms)"] if map else ["Format", "Export", "", ""]
     py = pd.DataFrame(y, columns=c)
-    LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)')
+    LOGGER.info(f"\nBenchmarks complete ({time.time() - t:.2f}s)")
     LOGGER.info(str(py if map else py.iloc[:, :2]))
     if hard_fail and isinstance(hard_fail, str):
-        metrics = py['mAP50-95'].array  # values to compare to floor
+        metrics = py["mAP50-95"].array  # values to compare to floor
         floor = eval(hard_fail)  # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n
-        assert all(x > floor for x in metrics if pd.notna(x)), f'HARD FAIL: mAP50-95 < floor {floor}'
+        assert all(x > floor for x in metrics if pd.notna(x)), f"HARD FAIL: mAP50-95 < floor {floor}"
     return py
 
 
 def test(
-        weights=ROOT / 'yolov5s.pt',  # weights path
+    weights=ROOT / "yolov5s.pt",  # weights path
     imgsz=640,  # inference size (pixels)
     batch_size=1,  # batch size
-        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
-        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     half=False,  # use FP16 half-precision inference
     test=False,  # test exports only
     pt_only=False,  # test PyTorch only
@@ -131,34 +128,37 @@ def test(
     device = select_device(device)
     for i, (name, f, suffix, gpu) in export.export_formats().iterrows():  # index, (name, file, suffix, gpu-capable)
         try:
-            w = weights if f == '-' else \
+            w = (
-                export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1]  # weights
+                weights
-            assert suffix in str(w), 'export failed'
+                if f == "-"
+                else export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1]
+            )  # weights
+            assert suffix in str(w), "export failed"
             y.append([name, True])
         except Exception:
             y.append([name, False])  # mAP, t_inference
 
     # Print results
-    LOGGER.info('\n')
+    LOGGER.info("\n")
     parse_opt()
     notebook_init()  # print system info
-    py = pd.DataFrame(y, columns=['Format', 'Export'])
+    py = pd.DataFrame(y, columns=["Format", "Export"])
-    LOGGER.info(f'\nExports complete ({time.time() - t:.2f}s)')
+    LOGGER.info(f"\nExports complete ({time.time() - t:.2f}s)")
     LOGGER.info(str(py))
     return py
 
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
+    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path")
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
+    parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
-    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
+    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--test', action='store_true', help='test exports only')
+    parser.add_argument("--test", action="store_true", help="test exports only")
-    parser.add_argument('--pt-only', action='store_true', help='test PyTorch only')
+    parser.add_argument("--pt-only", action="store_true", help="test PyTorch only")
-    parser.add_argument('--hard-fail', nargs='?', const=True, default=False, help='Exception on error or < min metric')
+    parser.add_argument("--hard-fail", nargs="?", const=True, default=False, help="Exception on error or < min metric")
     opt = parser.parse_args()
     opt.data = check_yaml(opt.data)  # check YAML
     print_args(vars(opt))
@@ -169,6 +169,6 @@ def main(opt):
     test(**vars(opt)) if opt.test else run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

classify/predict.py

@@ -48,43 +48,54 @@ from ultralytics.utils.plotting import Annotator
 from models.common import DetectMultiBackend
 from utils.augmentations import classify_transforms
 from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
-from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
+from utils.general import (
-                           increment_path, print_args, strip_optimizer)
+    LOGGER,
+    Profile,
+    check_file,
+    check_img_size,
+    check_imshow,
+    check_requirements,
+    colorstr,
+    cv2,
+    increment_path,
+    print_args,
+    strip_optimizer,
+)
 from utils.torch_utils import select_device, smart_inference_mode
 
 
 @smart_inference_mode()
 def run(
-        weights=ROOT / 'yolov5s-cls.pt',  # model.pt path(s)
+    weights=ROOT / "yolov5s-cls.pt",  # model.pt path(s)
-        source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
+    source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
-        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
     imgsz=(224, 224),  # inference size (height, width)
-        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     view_img=False,  # show results
     save_txt=False,  # save results to *.txt
     nosave=False,  # do not save images/videos
     augment=False,  # augmented inference
     visualize=False,  # visualize features
     update=False,  # update all models
-        project=ROOT / 'runs/predict-cls',  # save results to project/name
+    project=ROOT / "runs/predict-cls",  # save results to project/name
-        name='exp',  # save results to project/name
+    name="exp",  # save results to project/name
     exist_ok=False,  # existing project/name ok, do not increment
     half=False,  # use FP16 half-precision inference
     dnn=False,  # use OpenCV DNN for ONNX inference
     vid_stride=1,  # video frame-rate stride
 ):
     source = str(source)
-    save_img = not nosave and not source.endswith('.txt')  # save inference images
+    save_img = not nosave and not source.endswith(".txt")  # save inference images
     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
-    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
+    is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
-    webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
+    webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
-    screenshot = source.lower().startswith('screen')
+    screenshot = source.lower().startswith("screen")
     if is_url and is_file:
         source = check_file(source)  # download
 
     # Directories
     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+    (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
     # Load model
     device = select_device(device)
@@ -127,15 +138,15 @@ def run(
             seen += 1
             if webcam:  # batch_size >= 1
                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
-                s += f'{i}: '
+                s += f"{i}: "
             else:
-                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
+                p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
 
             p = Path(p)  # to Path
             save_path = str(save_dir / p.name)  # im.jpg
-            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
+            txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
 
-            s += '%gx%g ' % im.shape[2:]  # print string
+            s += "%gx%g " % im.shape[2:]  # print string
             annotator = Annotator(im0, example=str(names), pil=True)
 
             # Print results
@@ -143,17 +154,17 @@ def run(
             s += f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top5i)}, "
 
             # Write results
-            text = '\n'.join(f'{prob[j]:.2f} {names[j]}' for j in top5i)
+            text = "\n".join(f"{prob[j]:.2f} {names[j]}" for j in top5i)
             if save_img or view_img:  # Add bbox to image
                 annotator.text([32, 32], text, txt_color=(255, 255, 255))
             if save_txt:  # Write to file
-                with open(f'{txt_path}.txt', 'a') as f:
+                with open(f"{txt_path}.txt", "a") as f:
-                    f.write(text + '\n')
+                    f.write(text + "\n")
 
             # Stream results
             im0 = annotator.result()
             if view_img:
-                if platform.system() == 'Linux' and p not in windows:
+                if platform.system() == "Linux" and p not in windows:
                     windows.append(p)
                     cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                     cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
@@ -162,7 +173,7 @@ def run(
 
             # Save results (image with detections)
             if save_img:
-                if dataset.mode == 'image':
+                if dataset.mode == "image":
                     cv2.imwrite(save_path, im0)
                 else:  # 'video' or 'stream'
                     if vid_path[i] != save_path:  # new video
@@ -175,18 +186,18 @@ def run(
                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                         else:  # stream
                             fps, w, h = 30, im0.shape[1], im0.shape[0]
-                        save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
+                        save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
-                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                     vid_writer[i].write(im0)
 
         # Print time (inference-only)
-        LOGGER.info(f'{s}{dt[1].dt * 1E3:.1f}ms')
+        LOGGER.info(f"{s}{dt[1].dt * 1E3:.1f}ms")
 
     # Print results
-    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
-    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
+    LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
     if save_txt or save_img:
-        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     if update:
         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
@@ -194,23 +205,23 @@ def run(
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)')
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-cls.pt", help="model path(s)")
-    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
+    parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[224], help='inference size h,w')
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[224], help="inference size h,w")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--view-img', action='store_true', help='show results')
+    parser.add_argument("--view-img", action="store_true", help="show results")
-    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
+    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
-    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
+    parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
-    parser.add_argument('--augment', action='store_true', help='augmented inference')
+    parser.add_argument("--augment", action="store_true", help="augmented inference")
-    parser.add_argument('--visualize', action='store_true', help='visualize features')
+    parser.add_argument("--visualize", action="store_true", help="visualize features")
-    parser.add_argument('--update', action='store_true', help='update all models')
+    parser.add_argument("--update", action="store_true", help="update all models")
-    parser.add_argument('--project', default=ROOT / 'runs/predict-cls', help='save results to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/predict-cls", help="save results to project/name")
-    parser.add_argument('--name', default='exp', help='save results to project/name')
+    parser.add_argument("--name", default="exp", help="save results to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
-    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
+    parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
     opt = parser.parse_args()
     opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
     print_args(vars(opt))
@@ -218,10 +229,10 @@ def parse_opt():
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

classify/train.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Train a YOLOv5 classifier model on a classification dataset
+Train a YOLOv5 classifier model on a classification dataset.
 
 Usage - Single-GPU training:
     $ python classify/train.py --model yolov5s-cls.pt --data imagenette160 --epochs 5 --img 224
@@ -40,33 +40,61 @@ from classify import val as validate
 from models.experimental import attempt_load
 from models.yolo import ClassificationModel, DetectionModel
 from utils.dataloaders import create_classification_dataloader
-from utils.general import (DATASETS_DIR, LOGGER, TQDM_BAR_FORMAT, WorkingDirectory, check_git_info, check_git_status,
+from utils.general import (
-                           check_requirements, colorstr, download, increment_path, init_seeds, print_args, yaml_save)
+    DATASETS_DIR,
+    LOGGER,
+    TQDM_BAR_FORMAT,
+    WorkingDirectory,
+    check_git_info,
+    check_git_status,
+    check_requirements,
+    colorstr,
+    download,
+    increment_path,
+    init_seeds,
+    print_args,
+    yaml_save,
+)
 from utils.loggers import GenericLogger
 from utils.plots import imshow_cls
-from utils.torch_utils import (ModelEMA, de_parallel, model_info, reshape_classifier_output, select_device, smart_DDP,
+from utils.torch_utils import (
-                               smart_optimizer, smartCrossEntropyLoss, torch_distributed_zero_first)
+    ModelEMA,
+    de_parallel,
+    model_info,
+    reshape_classifier_output,
+    select_device,
+    smart_DDP,
+    smart_optimizer,
+    smartCrossEntropyLoss,
+    torch_distributed_zero_first,
+)
 
-LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
+LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
-WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
+WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
 GIT_INFO = check_git_info()
 
 
 def train(opt, device):
     init_seeds(opt.seed + 1 + RANK, deterministic=True)
-    save_dir, data, bs, epochs, nw, imgsz, pretrained = \
+    save_dir, data, bs, epochs, nw, imgsz, pretrained = (
-        opt.save_dir, Path(opt.data), opt.batch_size, opt.epochs, min(os.cpu_count() - 1, opt.workers), \
+        opt.save_dir,
-        opt.imgsz, str(opt.pretrained).lower() == 'true'
+        Path(opt.data),
-    cuda = device.type != 'cpu'
+        opt.batch_size,
+        opt.epochs,
+        min(os.cpu_count() - 1, opt.workers),
+        opt.imgsz,
+        str(opt.pretrained).lower() == "true",
+    )
+    cuda = device.type != "cpu"
 
     # Directories
-    wdir = save_dir / 'weights'
+    wdir = save_dir / "weights"
     wdir.mkdir(parents=True, exist_ok=True)  # make dir
-    last, best = wdir / 'last.pt', wdir / 'best.pt'
+    last, best = wdir / "last.pt", wdir / "best.pt"
 
     # Save run settings
-    yaml_save(save_dir / 'opt.yaml', vars(opt))
+    yaml_save(save_dir / "opt.yaml", vars(opt))
 
     # Logger
     logger = GenericLogger(opt=opt, console_logger=LOGGER) if RANK in {-1, 0} else None
@@ -75,51 +103,55 @@ def train(opt, device):
     with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
         data_dir = data if data.is_dir() else (DATASETS_DIR / data)
         if not data_dir.is_dir():
-            LOGGER.info(f'\nDataset not found ⚠️, missing path {data_dir}, attempting download...')
+            LOGGER.info(f"\nDataset not found ⚠️, missing path {data_dir}, attempting download...")
             t = time.time()
-            if str(data) == 'imagenet':
+            if str(data) == "imagenet":
-                subprocess.run(['bash', str(ROOT / 'data/scripts/get_imagenet.sh')], shell=True, check=True)
+                subprocess.run(["bash", str(ROOT / "data/scripts/get_imagenet.sh")], shell=True, check=True)
             else:
-                url = f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{data}.zip'
+                url = f"https://github.com/ultralytics/yolov5/releases/download/v1.0/{data}.zip"
                 download(url, dir=data_dir.parent)
             s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n"
             LOGGER.info(s)
 
     # Dataloaders
-    nc = len([x for x in (data_dir / 'train').glob('*') if x.is_dir()])  # number of classes
+    nc = len([x for x in (data_dir / "train").glob("*") if x.is_dir()])  # number of classes
-    trainloader = create_classification_dataloader(path=data_dir / 'train',
+    trainloader = create_classification_dataloader(
+        path=data_dir / "train",
         imgsz=imgsz,
         batch_size=bs // WORLD_SIZE,
         augment=True,
         cache=opt.cache,
         rank=LOCAL_RANK,
-                                                   workers=nw)
+        workers=nw,
+    )
 
-    test_dir = data_dir / 'test' if (data_dir / 'test').exists() else data_dir / 'val'  # data/test or data/val
+    test_dir = data_dir / "test" if (data_dir / "test").exists() else data_dir / "val"  # data/test or data/val
     if RANK in {-1, 0}:
-        testloader = create_classification_dataloader(path=test_dir,
+        testloader = create_classification_dataloader(
+            path=test_dir,
             imgsz=imgsz,
             batch_size=bs // WORLD_SIZE * 2,
             augment=False,
             cache=opt.cache,
             rank=-1,
-                                                      workers=nw)
+            workers=nw,
+        )
 
     # Model
     with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
-        if Path(opt.model).is_file() or opt.model.endswith('.pt'):
+        if Path(opt.model).is_file() or opt.model.endswith(".pt"):
-            model = attempt_load(opt.model, device='cpu', fuse=False)
+            model = attempt_load(opt.model, device="cpu", fuse=False)
         elif opt.model in torchvision.models.__dict__:  # TorchVision models i.e. resnet50, efficientnet_b0
-            model = torchvision.models.__dict__[opt.model](weights='IMAGENET1K_V1' if pretrained else None)
+            model = torchvision.models.__dict__[opt.model](weights="IMAGENET1K_V1" if pretrained else None)
         else:
-            m = hub.list('ultralytics/yolov5')  # + hub.list('pytorch/vision')  # models
+            m = hub.list("ultralytics/yolov5")  # + hub.list('pytorch/vision')  # models
-            raise ModuleNotFoundError(f'--model {opt.model} not found. Available models are: \n' + '\n'.join(m))
+            raise ModuleNotFoundError(f"--model {opt.model} not found. Available models are: \n" + "\n".join(m))
         if isinstance(model, DetectionModel):
             LOGGER.warning("WARNING ⚠️ pass YOLOv5 classifier model with '-cls' suffix, i.e. '--model yolov5s-cls.pt'")
             model = ClassificationModel(model=model, nc=nc, cutoff=opt.cutoff or 10)  # convert to classification model
         reshape_classifier_output(model, nc)  # update class count
     for m in model.modules():
-        if not pretrained and hasattr(m, 'reset_parameters'):
+        if not pretrained and hasattr(m, "reset_parameters"):
             m.reset_parameters()
         if isinstance(m, torch.nn.Dropout) and opt.dropout is not None:
             m.p = opt.dropout  # set dropout
@@ -135,8 +167,8 @@ def train(opt, device):
         if opt.verbose:
             LOGGER.info(model)
         images, labels = next(iter(trainloader))
-        file = imshow_cls(images[:25], labels[:25], names=model.names, f=save_dir / 'train_images.jpg')
+        file = imshow_cls(images[:25], labels[:25], names=model.names, f=save_dir / "train_images.jpg")
-        logger.log_images(file, name='Train Examples')
+        logger.log_images(file, name="Train Examples")
         logger.log_graph(model, imgsz)  # log model
 
     # Optimizer
@@ -163,11 +195,13 @@ def train(opt, device):
     best_fitness = 0.0
     scaler = amp.GradScaler(enabled=cuda)
     val = test_dir.stem  # 'val' or 'test'
-    LOGGER.info(f'Image sizes {imgsz} train, {imgsz} test\n'
+    LOGGER.info(
+        f'Image sizes {imgsz} train, {imgsz} test\n'
         f'Using {nw * WORLD_SIZE} dataloader workers\n'
         f"Logging results to {colorstr('bold', save_dir)}\n"
         f'Starting {opt.model} training on {data} dataset with {nc} classes for {epochs} epochs...\n\n'
-                f"{'Epoch':>10}{'GPU_mem':>10}{'train_loss':>12}{f'{val}_loss':>12}{'top1_acc':>12}{'top5_acc':>12}")
+        f"{'Epoch':>10}{'GPU_mem':>10}{'train_loss':>12}{f'{val}_loss':>12}{'top1_acc':>12}{'top5_acc':>12}"
+    )
     for epoch in range(epochs):  # loop over the dataset multiple times
         tloss, vloss, fitness = 0.0, 0.0, 0.0  # train loss, val loss, fitness
         model.train()
@@ -198,15 +232,14 @@ def train(opt, device):
             if RANK in {-1, 0}:
                 # Print
                 tloss = (tloss * i + loss.item()) / (i + 1)  # update mean losses
-                mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
+                mem = "%.3gG" % (torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0)  # (GB)
-                pbar.desc = f"{f'{epoch + 1}/{epochs}':>10}{mem:>10}{tloss:>12.3g}" + ' ' * 36
+                pbar.desc = f"{f'{epoch + 1}/{epochs}':>10}{mem:>10}{tloss:>12.3g}" + " " * 36
 
                 # Test
                 if i == len(pbar) - 1:  # last batch
-                    top1, top5, vloss = validate.run(model=ema.ema,
+                    top1, top5, vloss = validate.run(
-                                                     dataloader=testloader,
+                        model=ema.ema, dataloader=testloader, criterion=criterion, pbar=pbar
-                                                     criterion=criterion,
+                    )  # test accuracy, loss
-                                                     pbar=pbar)  # test accuracy, loss
                     fitness = top1  # define fitness as top1 accuracy
 
         # Scheduler
@@ -220,26 +253,28 @@ def train(opt, device):
 
             # Log
             metrics = {
-                'train/loss': tloss,
+                "train/loss": tloss,
-                f'{val}/loss': vloss,
+                f"{val}/loss": vloss,
-                'metrics/accuracy_top1': top1,
+                "metrics/accuracy_top1": top1,
-                'metrics/accuracy_top5': top5,
+                "metrics/accuracy_top5": top5,
-                'lr/0': optimizer.param_groups[0]['lr']}  # learning rate
+                "lr/0": optimizer.param_groups[0]["lr"],
+            }  # learning rate
             logger.log_metrics(metrics, epoch)
 
             # Save model
             final_epoch = epoch + 1 == epochs
             if (not opt.nosave) or final_epoch:
                 ckpt = {
-                    'epoch': epoch,
+                    "epoch": epoch,
-                    'best_fitness': best_fitness,
+                    "best_fitness": best_fitness,
-                    'model': deepcopy(ema.ema).half(),  # deepcopy(de_parallel(model)).half(),
+                    "model": deepcopy(ema.ema).half(),  # deepcopy(de_parallel(model)).half(),
-                    'ema': None,  # deepcopy(ema.ema).half(),
+                    "ema": None,  # deepcopy(ema.ema).half(),
-                    'updates': ema.updates,
+                    "updates": ema.updates,
-                    'optimizer': None,  # optimizer.state_dict(),
+                    "optimizer": None,  # optimizer.state_dict(),
-                    'opt': vars(opt),
+                    "opt": vars(opt),
-                    'git': GIT_INFO,  # {remote, branch, commit} if a git repo
+                    "git": GIT_INFO,  # {remote, branch, commit} if a git repo
-                    'date': datetime.now().isoformat()}
+                    "date": datetime.now().isoformat(),
+                }
 
                 # Save last, best and delete
                 torch.save(ckpt, last)
@@ -249,49 +284,51 @@ def train(opt, device):
 
     # Train complete
     if RANK in {-1, 0} and final_epoch:
-        LOGGER.info(f'\nTraining complete ({(time.time() - t0) / 3600:.3f} hours)'
+        LOGGER.info(
+            f'\nTraining complete ({(time.time() - t0) / 3600:.3f} hours)'
             f"\nResults saved to {colorstr('bold', save_dir)}"
             f'\nPredict:         python classify/predict.py --weights {best} --source im.jpg'
             f'\nValidate:        python classify/val.py --weights {best} --data {data_dir}'
             f'\nExport:          python export.py --weights {best} --include onnx'
             f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{best}')"
-                    f'\nVisualize:       https://netron.app\n')
+            f'\nVisualize:       https://netron.app\n'
+        )
 
         # Plot examples
         images, labels = (x[:25] for x in next(iter(testloader)))  # first 25 images and labels
         pred = torch.max(ema.ema(images.to(device)), 1)[1]
-        file = imshow_cls(images, labels, pred, de_parallel(model).names, verbose=False, f=save_dir / 'test_images.jpg')
+        file = imshow_cls(images, labels, pred, de_parallel(model).names, verbose=False, f=save_dir / "test_images.jpg")
 
         # Log results
-        meta = {'epochs': epochs, 'top1_acc': best_fitness, 'date': datetime.now().isoformat()}
+        meta = {"epochs": epochs, "top1_acc": best_fitness, "date": datetime.now().isoformat()}
-        logger.log_images(file, name='Test Examples (true-predicted)', epoch=epoch)
+        logger.log_images(file, name="Test Examples (true-predicted)", epoch=epoch)
         logger.log_model(best, epochs, metadata=meta)
 
 
 def parse_opt(known=False):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', type=str, default='yolov5s-cls.pt', help='initial weights path')
+    parser.add_argument("--model", type=str, default="yolov5s-cls.pt", help="initial weights path")
-    parser.add_argument('--data', type=str, default='imagenette160', help='cifar10, cifar100, mnist, imagenet, ...')
+    parser.add_argument("--data", type=str, default="imagenette160", help="cifar10, cifar100, mnist, imagenet, ...")
-    parser.add_argument('--epochs', type=int, default=10, help='total training epochs')
+    parser.add_argument("--epochs", type=int, default=10, help="total training epochs")
-    parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs')
+    parser.add_argument("--batch-size", type=int, default=64, help="total batch size for all GPUs")
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='train, val image size (pixels)')
+    parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=224, help="train, val image size (pixels)")
-    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
+    parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
-    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
+    parser.add_argument("--cache", type=str, nargs="?", const="ram", help='--cache images in "ram" (default) or "disk"')
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    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)')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--project', default=ROOT / 'runs/train-cls', help='save to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/train-cls", help="save to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--pretrained', nargs='?', const=True, default=True, help='start from i.e. --pretrained False')
+    parser.add_argument("--pretrained", nargs="?", const=True, default=True, help="start from i.e. --pretrained False")
-    parser.add_argument('--optimizer', choices=['SGD', 'Adam', 'AdamW', 'RMSProp'], default='Adam', help='optimizer')
+    parser.add_argument("--optimizer", choices=["SGD", "Adam", "AdamW", "RMSProp"], default="Adam", help="optimizer")
-    parser.add_argument('--lr0', type=float, default=0.001, help='initial learning rate')
+    parser.add_argument("--lr0", type=float, default=0.001, help="initial learning rate")
-    parser.add_argument('--decay', type=float, default=5e-5, help='weight decay')
+    parser.add_argument("--decay", type=float, default=5e-5, help="weight decay")
-    parser.add_argument('--label-smoothing', type=float, default=0.1, help='Label smoothing epsilon')
+    parser.add_argument("--label-smoothing", type=float, default=0.1, help="Label smoothing epsilon")
-    parser.add_argument('--cutoff', type=int, default=None, help='Model layer cutoff index for Classify() head')
+    parser.add_argument("--cutoff", type=int, default=None, help="Model layer cutoff index for Classify() head")
-    parser.add_argument('--dropout', type=float, default=None, help='Dropout (fraction)')
+    parser.add_argument("--dropout", type=float, default=None, help="Dropout (fraction)")
-    parser.add_argument('--verbose', action='store_true', help='Verbose mode')
+    parser.add_argument("--verbose", action="store_true", help="Verbose mode")
-    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
+    parser.add_argument("--seed", type=int, default=0, help="Global training seed")
-    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
+    parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
     return parser.parse_known_args()[0] if known else parser.parse_args()
 
 
@@ -300,17 +337,17 @@ def main(opt):
     if RANK in {-1, 0}:
         print_args(vars(opt))
         check_git_status()
-        check_requirements(ROOT / 'requirements.txt')
+        check_requirements(ROOT / "requirements.txt")
 
     # DDP mode
     device = select_device(opt.device, batch_size=opt.batch_size)
     if LOCAL_RANK != -1:
-        assert opt.batch_size != -1, 'AutoBatch is coming soon for classification, please pass a valid --batch-size'
+        assert opt.batch_size != -1, "AutoBatch is coming soon for classification, please pass a valid --batch-size"
-        assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
+        assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE"
-        assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
+        assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
         torch.cuda.set_device(LOCAL_RANK)
-        device = torch.device('cuda', LOCAL_RANK)
+        device = torch.device("cuda", LOCAL_RANK)
-        dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo')
+        dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
 
     # Parameters
     opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)  # increment run
@@ -328,6 +365,6 @@ def run(**kwargs):
     return opt
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

classify/val.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Validate a trained YOLOv5 classification model on a classification dataset
+Validate a trained YOLOv5 classification model on a classification dataset.
 
 Usage:
     $ bash data/scripts/get_imagenet.sh --val  # download ImageNet val split (6.3G, 50000 images)
@@ -36,22 +36,30 @@ ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 
 from models.common import DetectMultiBackend
 from utils.dataloaders import create_classification_dataloader
-from utils.general import (LOGGER, TQDM_BAR_FORMAT, Profile, check_img_size, check_requirements, colorstr,
+from utils.general import (
-                           increment_path, print_args)
+    LOGGER,
+    TQDM_BAR_FORMAT,
+    Profile,
+    check_img_size,
+    check_requirements,
+    colorstr,
+    increment_path,
+    print_args,
+)
 from utils.torch_utils import select_device, smart_inference_mode
 
 
 @smart_inference_mode()
 def run(
-    data=ROOT / '../datasets/mnist',  # dataset dir
+    data=ROOT / "../datasets/mnist",  # dataset dir
-    weights=ROOT / 'yolov5s-cls.pt',  # model.pt path(s)
+    weights=ROOT / "yolov5s-cls.pt",  # model.pt path(s)
     batch_size=128,  # batch size
     imgsz=224,  # inference size (pixels)
-    device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     workers=8,  # max dataloader workers (per RANK in DDP mode)
     verbose=False,  # verbose output
-    project=ROOT / 'runs/val-cls',  # save to project/name
+    project=ROOT / "runs/val-cls",  # save to project/name
-    name='exp',  # save to project/name
+    name="exp",  # save to project/name
     exist_ok=False,  # existing project/name ok, do not increment
     half=False,  # use FP16 half-precision inference
     dnn=False,  # use OpenCV DNN for ONNX inference
@@ -64,7 +72,7 @@ def run(
     training = model is not None
     if training:  # called by train.py
         device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
-        half &= device.type != 'cpu'  # half precision only supported on CUDA
+        half &= device.type != "cpu"  # half precision only supported on CUDA
         model.half() if half else model.float()
     else:  # called directly
         device = select_device(device, batch_size=batch_size)
@@ -84,25 +92,22 @@ def run(
             device = model.device
             if not (pt or jit):
                 batch_size = 1  # export.py models default to batch-size 1
-                LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
+                LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
 
         # Dataloader
         data = Path(data)
-        test_dir = data / 'test' if (data / 'test').exists() else data / 'val'  # data/test or data/val
+        test_dir = data / "test" if (data / "test").exists() else data / "val"  # data/test or data/val
-        dataloader = create_classification_dataloader(path=test_dir,
+        dataloader = create_classification_dataloader(
-                                                      imgsz=imgsz,
+            path=test_dir, imgsz=imgsz, batch_size=batch_size, augment=False, rank=-1, workers=workers
-                                                      batch_size=batch_size,
+        )
-                                                      augment=False,
-                                                      rank=-1,
-                                                      workers=workers)
 
     model.eval()
     pred, targets, loss, dt = [], [], 0, (Profile(device=device), Profile(device=device), Profile(device=device))
     n = len(dataloader)  # number of batches
-    action = 'validating' if dataloader.dataset.root.stem == 'val' else 'testing'
+    action = "validating" if dataloader.dataset.root.stem == "val" else "testing"
-    desc = f'{pbar.desc[:-36]}{action:>36}' if pbar else f'{action}'
+    desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}"
     bar = tqdm(dataloader, desc, n, not training, bar_format=TQDM_BAR_FORMAT, position=0)
-    with torch.cuda.amp.autocast(enabled=device.type != 'cpu'):
+    with torch.cuda.amp.autocast(enabled=device.type != "cpu"):
         for images, labels in bar:
             with dt[0]:
                 images, labels = images.to(device, non_blocking=True), labels.to(device)
@@ -123,19 +128,19 @@ def run(
     top1, top5 = acc.mean(0).tolist()
 
     if pbar:
-        pbar.desc = f'{pbar.desc[:-36]}{loss:>12.3g}{top1:>12.3g}{top5:>12.3g}'
+        pbar.desc = f"{pbar.desc[:-36]}{loss:>12.3g}{top1:>12.3g}{top5:>12.3g}"
     if verbose:  # all classes
         LOGGER.info(f"{'Class':>24}{'Images':>12}{'top1_acc':>12}{'top5_acc':>12}")
         LOGGER.info(f"{'all':>24}{targets.shape[0]:>12}{top1:>12.3g}{top5:>12.3g}")
         for i, c in model.names.items():
             acc_i = acc[targets == i]
             top1i, top5i = acc_i.mean(0).tolist()
-            LOGGER.info(f'{c:>24}{acc_i.shape[0]:>12}{top1i:>12.3g}{top5i:>12.3g}')
+            LOGGER.info(f"{c:>24}{acc_i.shape[0]:>12}{top1i:>12.3g}{top5i:>12.3g}")
 
         # Print results
-        t = tuple(x.t / len(dataloader.dataset.samples) * 1E3 for x in dt)  # speeds per image
+        t = tuple(x.t / len(dataloader.dataset.samples) * 1e3 for x in dt)  # speeds per image
         shape = (1, 3, imgsz, imgsz)
-        LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}' % t)
+        LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}" % t)
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
 
     return top1, top5, loss
@@ -143,28 +148,28 @@ def run(
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--data', type=str, default=ROOT / '../datasets/mnist', help='dataset path')
+    parser.add_argument("--data", type=str, default=ROOT / "../datasets/mnist", help="dataset path")
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model.pt path(s)')
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-cls.pt", help="model.pt path(s)")
-    parser.add_argument('--batch-size', type=int, default=128, help='batch size')
+    parser.add_argument("--batch-size", type=int, default=128, help="batch size")
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='inference size (pixels)')
+    parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=224, help="inference size (pixels)")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    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)')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--verbose', nargs='?', const=True, default=True, help='verbose output')
+    parser.add_argument("--verbose", nargs="?", const=True, default=True, help="verbose output")
-    parser.add_argument('--project', default=ROOT / 'runs/val-cls', help='save to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/val-cls", help="save to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
     opt = parser.parse_args()
     print_args(vars(opt))
     return opt
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

detect.py

@@ -47,21 +47,35 @@ from ultralytics.utils.plotting import Annotator, colors, save_one_box
 
 from models.common import DetectMultiBackend
 from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
-from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
+from utils.general import (
-                           increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
+    LOGGER,
+    Profile,
+    check_file,
+    check_img_size,
+    check_imshow,
+    check_requirements,
+    colorstr,
+    cv2,
+    increment_path,
+    non_max_suppression,
+    print_args,
+    scale_boxes,
+    strip_optimizer,
+    xyxy2xywh,
+)
 from utils.torch_utils import select_device, smart_inference_mode
 
 
 @smart_inference_mode()
 def run(
-        weights=ROOT / 'yolov5s.pt',  # model path or triton URL
+    weights=ROOT / "yolov5s.pt",  # model path or triton URL
-        source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
+    source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
-        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
     imgsz=(640, 640),  # inference size (height, width)
     conf_thres=0.25,  # confidence threshold
     iou_thres=0.45,  # NMS IOU threshold
     max_det=1000,  # maximum detections per image
-        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     view_img=False,  # show results
     save_txt=False,  # save results to *.txt
     save_csv=False,  # save results in CSV format
@@ -73,8 +87,8 @@ def run(
     augment=False,  # augmented inference
     visualize=False,  # visualize features
     update=False,  # update all models
-        project=ROOT / 'runs/detect',  # save results to project/name
+    project=ROOT / "runs/detect",  # save results to project/name
-        name='exp',  # save results to project/name
+    name="exp",  # save results to project/name
     exist_ok=False,  # existing project/name ok, do not increment
     line_thickness=3,  # bounding box thickness (pixels)
     hide_labels=False,  # hide labels
@@ -84,17 +98,17 @@ def run(
     vid_stride=1,  # video frame-rate stride
 ):
     source = str(source)
-    save_img = not nosave and not source.endswith('.txt')  # save inference images
+    save_img = not nosave and not source.endswith(".txt")  # save inference images
     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
-    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
+    is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
-    webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
+    webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
-    screenshot = source.lower().startswith('screen')
+    screenshot = source.lower().startswith("screen")
     if is_url and is_file:
         source = check_file(source)  # download
 
     # Directories
     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+    (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
     # Load model
     device = select_device(device)
@@ -148,12 +162,12 @@ def run(
         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
 
         # Define the path for the CSV file
-        csv_path = save_dir / 'predictions.csv'
+        csv_path = save_dir / "predictions.csv"
 
         # Create or append to the CSV file
         def write_to_csv(image_name, prediction, confidence):
-            data = {'Image Name': image_name, 'Prediction': prediction, 'Confidence': confidence}
+            data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
-            with open(csv_path, mode='a', newline='') as f:
+            with open(csv_path, mode="a", newline="") as f:
                 writer = csv.DictWriter(f, fieldnames=data.keys())
                 if not csv_path.is_file():
                     writer.writeheader()
@@ -164,14 +178,14 @@ def run(
             seen += 1
             if webcam:  # batch_size >= 1
                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
-                s += f'{i}: '
+                s += f"{i}: "
             else:
-                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
+                p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
 
             p = Path(p)  # to Path
             save_path = str(save_dir / p.name)  # im.jpg
-            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
+            txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
-            s += '%gx%g ' % im.shape[2:]  # print string
+            s += "%gx%g " % im.shape[2:]  # print string
             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
             imc = im0.copy() if save_crop else im0  # for save_crop
             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
@@ -187,9 +201,9 @@ def run(
                 # Write results
                 for *xyxy, conf, cls in reversed(det):
                     c = int(cls)  # integer class
-                    label = names[c] if hide_conf else f'{names[c]}'
+                    label = names[c] if hide_conf else f"{names[c]}"
                     confidence = float(conf)
-                    confidence_str = f'{confidence:.2f}'
+                    confidence_str = f"{confidence:.2f}"
 
                     if save_csv:
                         write_to_csv(p.name, label, confidence_str)
@@ -197,20 +211,20 @@ def run(
                     if save_txt:  # Write to file
                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                         line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
-                        with open(f'{txt_path}.txt', 'a') as f:
+                        with open(f"{txt_path}.txt", "a") as f:
-                            f.write(('%g ' * len(line)).rstrip() % line + '\n')
+                            f.write(("%g " * len(line)).rstrip() % line + "\n")
 
                     if save_img or save_crop or view_img:  # Add bbox to image
                         c = int(cls)  # integer class
-                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
+                        label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}")
                         annotator.box_label(xyxy, label, color=colors(c, True))
                     if save_crop:
-                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
+                        save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True)
 
             # Stream results
             im0 = annotator.result()
             if view_img:
-                if platform.system() == 'Linux' and p not in windows:
+                if platform.system() == "Linux" and p not in windows:
                     windows.append(p)
                     cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                     cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
@@ -219,7 +233,7 @@ def run(
 
             # Save results (image with detections)
             if save_img:
-                if dataset.mode == 'image':
+                if dataset.mode == "image":
                     cv2.imwrite(save_path, im0)
                 else:  # 'video' or 'stream'
                     if vid_path[i] != save_path:  # new video
@@ -232,18 +246,18 @@ def run(
                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                         else:  # stream
                             fps, w, h = 30, im0.shape[1], im0.shape[0]
-                        save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
+                        save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
-                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                     vid_writer[i].write(im0)
 
         # Print time (inference-only)
         LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
 
     # Print results
-    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
-    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
+    LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
     if save_txt or save_img:
-        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     if update:
         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
@@ -251,34 +265,34 @@ def run(
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL')
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
-    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
+    parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
-    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
+    parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
-    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
+    parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
-    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
+    parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--view-img', action='store_true', help='show results')
+    parser.add_argument("--view-img", action="store_true", help="show results")
-    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
+    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
-    parser.add_argument('--save-csv', action='store_true', help='save results in CSV format')
+    parser.add_argument("--save-csv", action="store_true", help="save results in CSV format")
-    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
+    parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
-    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
+    parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
-    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
+    parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
-    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
+    parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
-    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
+    parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
-    parser.add_argument('--augment', action='store_true', help='augmented inference')
+    parser.add_argument("--augment", action="store_true", help="augmented inference")
-    parser.add_argument('--visualize', action='store_true', help='visualize features')
+    parser.add_argument("--visualize", action="store_true", help="visualize features")
-    parser.add_argument('--update', action='store_true', help='update all models')
+    parser.add_argument("--update", action="store_true", help="update all models")
-    parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/detect", help="save results to project/name")
-    parser.add_argument('--name', default='exp', help='save results to project/name')
+    parser.add_argument("--name", default="exp", help="save results to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
+    parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
-    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
+    parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
-    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
+    parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
-    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
+    parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
     opt = parser.parse_args()
     opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
     print_args(vars(opt))
@@ -286,10 +300,10 @@ def parse_opt():
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

export.py

@@ -64,30 +64,42 @@ FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLOv5 root directory
 if str(ROOT) not in sys.path:
     sys.path.append(str(ROOT))  # add ROOT to PATH
-if platform.system() != 'Windows':
+if platform.system() != "Windows":
     ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 
 from models.experimental import attempt_load
 from models.yolo import ClassificationModel, Detect, DetectionModel, SegmentationModel
 from utils.dataloaders import LoadImages
-from utils.general import (LOGGER, Profile, check_dataset, check_img_size, check_requirements, check_version,
+from utils.general import (
-                           check_yaml, colorstr, file_size, get_default_args, print_args, url2file, yaml_save)
+    LOGGER,
+    Profile,
+    check_dataset,
+    check_img_size,
+    check_requirements,
+    check_version,
+    check_yaml,
+    colorstr,
+    file_size,
+    get_default_args,
+    print_args,
+    url2file,
+    yaml_save,
+)
 from utils.torch_utils import select_device, smart_inference_mode
 
-MACOS = platform.system() == 'Darwin'  # macOS environment
+MACOS = platform.system() == "Darwin"  # macOS environment
 
 
 class iOSModel(torch.nn.Module):
-
     def __init__(self, model, im):
         super().__init__()
         b, c, h, w = im.shape  # batch, channel, height, width
         self.model = model
         self.nc = model.nc  # number of classes
         if w == h:
-            self.normalize = 1. / w
+            self.normalize = 1.0 / w
         else:
-            self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h])  # broadcast (slower, smaller)
+            self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h])  # broadcast (slower, smaller)
             # np = model(im)[0].shape[1]  # number of points
             # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4)  # explicit (faster, larger)
 
@@ -99,19 +111,20 @@ class iOSModel(torch.nn.Module):
 def export_formats():
     # YOLOv5 export formats
     x = [
-        ['PyTorch', '-', '.pt', True, True],
+        ["PyTorch", "-", ".pt", True, True],
-        ['TorchScript', 'torchscript', '.torchscript', True, True],
+        ["TorchScript", "torchscript", ".torchscript", True, True],
-        ['ONNX', 'onnx', '.onnx', True, True],
+        ["ONNX", "onnx", ".onnx", True, True],
-        ['OpenVINO', 'openvino', '_openvino_model', True, False],
+        ["OpenVINO", "openvino", "_openvino_model", True, False],
-        ['TensorRT', 'engine', '.engine', False, True],
+        ["TensorRT", "engine", ".engine", False, True],
-        ['CoreML', 'coreml', '.mlmodel', True, False],
+        ["CoreML", "coreml", ".mlmodel", True, False],
-        ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],
+        ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True],
-        ['TensorFlow GraphDef', 'pb', '.pb', True, True],
+        ["TensorFlow GraphDef", "pb", ".pb", True, True],
-        ['TensorFlow Lite', 'tflite', '.tflite', True, False],
+        ["TensorFlow Lite", "tflite", ".tflite", True, False],
-        ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
+        ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False],
-        ['TensorFlow.js', 'tfjs', '_web_model', False, False],
+        ["TensorFlow.js", "tfjs", "_web_model", False, False],
-        ['PaddlePaddle', 'paddle', '_paddle_model', True, True], ]
+        ["PaddlePaddle", "paddle", "_paddle_model", True, True],
-    return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])
+    ]
+    return pd.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"])
 
 
 def try_export(inner_func):
@@ -119,28 +132,28 @@ def try_export(inner_func):
     inner_args = get_default_args(inner_func)
 
     def outer_func(*args, **kwargs):
-        prefix = inner_args['prefix']
+        prefix = inner_args["prefix"]
         try:
             with Profile() as dt:
                 f, model = inner_func(*args, **kwargs)
-            LOGGER.info(f'{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)')
+            LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)")
             return f, model
         except Exception as e:
-            LOGGER.info(f'{prefix} export failure ❌ {dt.t:.1f}s: {e}')
+            LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
             return None, None
 
     return outer_func
 
 
 @try_export
-def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
+def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
     # YOLOv5 TorchScript model export
-    LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
-    f = file.with_suffix('.torchscript')
+    f = file.with_suffix(".torchscript")
 
     ts = torch.jit.trace(model, im, strict=False)
-    d = {'shape': im.shape, 'stride': int(max(model.stride)), 'names': model.names}
+    d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
-    extra_files = {'config.txt': json.dumps(d)}  # torch._C.ExtraFilesMap()
+    extra_files = {"config.txt": json.dumps(d)}  # torch._C.ExtraFilesMap()
     if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
         optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
     else:
@@ -149,22 +162,22 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
 
 
 @try_export
-def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')):
+def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
     # YOLOv5 ONNX export
-    check_requirements('onnx>=1.12.0')
+    check_requirements("onnx>=1.12.0")
     import onnx
 
-    LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...")
-    f = str(file.with_suffix('.onnx'))
+    f = str(file.with_suffix(".onnx"))
 
-    output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
+    output_names = ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"]
     if dynamic:
-        dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
+        dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
         if isinstance(model, SegmentationModel):
-            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
+            dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
-            dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
+            dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
         elif isinstance(model, DetectionModel):
-            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
+            dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
 
     torch.onnx.export(
         model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
@@ -173,16 +186,17 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX
         verbose=False,
         opset_version=opset,
         do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
-        input_names=['images'],
+        input_names=["images"],
         output_names=output_names,
-        dynamic_axes=dynamic or None)
+        dynamic_axes=dynamic or None,
+    )
 
     # Checks
     model_onnx = onnx.load(f)  # load onnx model
     onnx.checker.check_model(model_onnx)  # check onnx model
 
     # Metadata
-    d = {'stride': int(max(model.stride)), 'names': model.names}
+    d = {"stride": int(max(model.stride)), "names": model.names}
     for k, v in d.items():
         meta = model_onnx.metadata_props.add()
         meta.key, meta.value = k, str(v)
@@ -192,36 +206,37 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX
     if simplify:
         try:
             cuda = torch.cuda.is_available()
-            check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
+            check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1"))
             import onnxsim
 
-            LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
+            LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...")
             model_onnx, check = onnxsim.simplify(model_onnx)
-            assert check, 'assert check failed'
+            assert check, "assert check failed"
             onnx.save(model_onnx, f)
         except Exception as e:
-            LOGGER.info(f'{prefix} simplifier failure: {e}')
+            LOGGER.info(f"{prefix} simplifier failure: {e}")
     return f, model_onnx
 
 
 @try_export
-def export_openvino(file, metadata, half, int8, data, prefix=colorstr('OpenVINO:')):
+def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
     # YOLOv5 OpenVINO export
-    check_requirements('openvino-dev>=2023.0')  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+    check_requirements("openvino-dev>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
     import openvino.runtime as ov  # noqa
     from openvino.tools import mo  # noqa
 
-    LOGGER.info(f'\n{prefix} starting export with openvino {ov.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
-    f = str(file).replace(file.suffix, f'_openvino_model{os.sep}')
+    f = str(file).replace(file.suffix, f"_openvino_model{os.sep}")
-    f_onnx = file.with_suffix('.onnx')
+    f_onnx = file.with_suffix(".onnx")
-    f_ov = str(Path(f) / file.with_suffix('.xml').name)
+    f_ov = str(Path(f) / file.with_suffix(".xml").name)
     if int8:
-        check_requirements('nncf>=2.4.0')  # requires at least version 2.4.0 to use the post-training quantization
+        check_requirements("nncf>=2.4.0")  # requires at least version 2.4.0 to use the post-training quantization
         import nncf
         import numpy as np
         from openvino.runtime import Core
 
         from utils.dataloaders import create_dataloader
+
         core = Core()
         onnx_model = core.read_model(f_onnx)  # export
 
@@ -233,24 +248,21 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr('OpenVINO:
                 input_tensor = np.expand_dims(input_tensor, 0)
             return input_tensor
 
-        def gen_dataloader(yaml_path, task='train', imgsz=640, workers=4):
+        def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4):
             data_yaml = check_yaml(yaml_path)
             data = check_dataset(data_yaml)
-            dataloader = create_dataloader(data[task],
+            dataloader = create_dataloader(
-                                           imgsz=imgsz,
+                data[task], imgsz=imgsz, batch_size=1, stride=32, pad=0.5, single_cls=False, rect=False, workers=workers
-                                           batch_size=1,
+            )[0]
-                                           stride=32,
-                                           pad=0.5,
-                                           single_cls=False,
-                                           rect=False,
-                                           workers=workers)[0]
             return dataloader
 
         # noqa: F811
 
         def transform_fn(data_item):
             """
-            Quantization transform function. Extracts and preprocess input data from dataloader item for quantization.
+            Quantization transform function.
+
+            Extracts and preprocess input data from dataloader item for quantization.
             Parameters:
                data_item: Tuple with data item produced by DataLoader during iteration
             Returns:
@@ -264,77 +276,77 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr('OpenVINO:
         quantization_dataset = nncf.Dataset(ds, transform_fn)
         ov_model = nncf.quantize(onnx_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED)
     else:
-        ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework='onnx', compress_to_fp16=half)  # export
+        ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half)  # export
 
     ov.serialize(ov_model, f_ov)  # save
-    yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata)  # add metadata.yaml
+    yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
     return f, None
 
 
 @try_export
-def export_paddle(model, im, file, metadata, prefix=colorstr('PaddlePaddle:')):
+def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
     # YOLOv5 Paddle export
-    check_requirements(('paddlepaddle', 'x2paddle'))
+    check_requirements(("paddlepaddle", "x2paddle"))
     import x2paddle
     from x2paddle.convert import pytorch2paddle
 
-    LOGGER.info(f'\n{prefix} starting export with X2Paddle {x2paddle.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
-    f = str(file).replace('.pt', f'_paddle_model{os.sep}')
+    f = str(file).replace(".pt", f"_paddle_model{os.sep}")
 
-    pytorch2paddle(module=model, save_dir=f, jit_type='trace', input_examples=[im])  # export
+    pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im])  # export
-    yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata)  # add metadata.yaml
+    yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
     return f, None
 
 
 @try_export
-def export_coreml(model, im, file, int8, half, nms, prefix=colorstr('CoreML:')):
+def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
     # YOLOv5 CoreML export
-    check_requirements('coremltools')
+    check_requirements("coremltools")
     import coremltools as ct
 
-    LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
-    f = file.with_suffix('.mlmodel')
+    f = file.with_suffix(".mlmodel")
 
     if nms:
         model = iOSModel(model, im)
     ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
-    ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
+    ct_model = ct.convert(ts, inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
-    bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None)
+    bits, mode = (8, "kmeans_lut") if int8 else (16, "linear") if half else (32, None)
     if bits < 32:
         if MACOS:  # quantization only supported on macOS
             with warnings.catch_warnings():
-                warnings.filterwarnings('ignore', category=DeprecationWarning)  # suppress numpy==1.20 float warning
+                warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress numpy==1.20 float warning
                 ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
         else:
-            print(f'{prefix} quantization only supported on macOS, skipping...')
+            print(f"{prefix} quantization only supported on macOS, skipping...")
     ct_model.save(f)
     return f, ct_model
 
 
 @try_export
-def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
+def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")):
     # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
-    assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
+    assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
     try:
         import tensorrt as trt
     except Exception:
-        if platform.system() == 'Linux':
+        if platform.system() == "Linux":
-            check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com')
+            check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com")
         import tensorrt as trt
 
-    if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
+    if trt.__version__[0] == "7":  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
         grid = model.model[-1].anchor_grid
         model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
         export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
         model.model[-1].anchor_grid = grid
     else:  # TensorRT >= 8
-        check_version(trt.__version__, '8.0.0', hard=True)  # require tensorrt>=8.0.0
+        check_version(trt.__version__, "8.0.0", hard=True)  # require tensorrt>=8.0.0
         export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
-    onnx = file.with_suffix('.onnx')
+    onnx = file.with_suffix(".onnx")
 
-    LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
-    assert onnx.exists(), f'failed to export ONNX file: {onnx}'
+    assert onnx.exists(), f"failed to export ONNX file: {onnx}"
-    f = file.with_suffix('.engine')  # TensorRT engine file
+    f = file.with_suffix(".engine")  # TensorRT engine file
     logger = trt.Logger(trt.Logger.INFO)
     if verbose:
         logger.min_severity = trt.Logger.Severity.VERBOSE
@@ -344,11 +356,11 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
     config.max_workspace_size = workspace * 1 << 30
     # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)  # fix TRT 8.4 deprecation notice
 
-    flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
     network = builder.create_network(flag)
     parser = trt.OnnxParser(network, logger)
     if not parser.parse_from_file(str(onnx)):
-        raise RuntimeError(f'failed to load ONNX file: {onnx}')
+        raise RuntimeError(f"failed to load ONNX file: {onnx}")
 
     inputs = [network.get_input(i) for i in range(network.num_inputs)]
     outputs = [network.get_output(i) for i in range(network.num_outputs)]
@@ -359,22 +371,23 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
 
     if dynamic:
         if im.shape[0] <= 1:
-            LOGGER.warning(f'{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument')
+            LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
         profile = builder.create_optimization_profile()
         for inp in inputs:
             profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
         config.add_optimization_profile(profile)
 
-    LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}')
+    LOGGER.info(f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}")
     if builder.platform_has_fast_fp16 and half:
         config.set_flag(trt.BuilderFlag.FP16)
-    with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
+    with builder.build_engine(network, config) as engine, open(f, "wb") as t:
         t.write(engine.serialize())
     return f, None
 
 
 @try_export
-def export_saved_model(model,
+def export_saved_model(
+    model,
     im,
     file,
     dynamic,
@@ -385,7 +398,8 @@ def export_saved_model(model,
     iou_thres=0.45,
     conf_thres=0.25,
     keras=False,
-                       prefix=colorstr('TensorFlow SavedModel:')):
+    prefix=colorstr("TensorFlow SavedModel:"),
+):
     # YOLOv5 TensorFlow SavedModel export
     try:
         import tensorflow as tf
@@ -396,13 +410,13 @@ def export_saved_model(model,
 
     from models.tf import TFModel
 
-    LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
-    if tf.__version__ > '2.13.1':
+    if tf.__version__ > "2.13.1":
-        helper_url = 'https://github.com/ultralytics/yolov5/issues/12489'
+        helper_url = "https://github.com/ultralytics/yolov5/issues/12489"
         LOGGER.info(
-            f'WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}'
+            f"WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}"
         )  # handling issue https://github.com/ultralytics/yolov5/issues/12489
-    f = str(file).replace('.pt', '_saved_model')
+    f = str(file).replace(".pt", "_saved_model")
     batch_size, ch, *imgsz = list(im.shape)  # BCHW
 
     tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
@@ -414,7 +428,7 @@ def export_saved_model(model,
     keras_model.trainable = False
     keras_model.summary()
     if keras:
-        keras_model.save(f, save_format='tf')
+        keras_model.save(f, save_format="tf")
     else:
         spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
         m = tf.function(lambda x: keras_model(x))  # full model
@@ -423,21 +437,24 @@ def export_saved_model(model,
         tfm = tf.Module()
         tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec])
         tfm.__call__(im)
-        tf.saved_model.save(tfm,
+        tf.saved_model.save(
+            tfm,
             f,
-                            options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version(
+            options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
-                                tf.__version__, '2.6') else tf.saved_model.SaveOptions())
+            if check_version(tf.__version__, "2.6")
+            else tf.saved_model.SaveOptions(),
+        )
     return f, keras_model
 
 
 @try_export
-def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):
+def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
     # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
     import tensorflow as tf
     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
-    LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
-    f = file.with_suffix('.pb')
+    f = file.with_suffix(".pb")
 
     m = tf.function(lambda x: keras_model(x))  # full model
     m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
@@ -448,14 +465,15 @@ def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):
 
 
 @try_export
-def export_tflite(keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms,
+def export_tflite(
-                  prefix=colorstr('TensorFlow Lite:')):
+    keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")
+):
     # YOLOv5 TensorFlow Lite export
     import tensorflow as tf
 
-    LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
     batch_size, ch, *imgsz = list(im.shape)  # BCHW
-    f = str(file).replace('.pt', '-fp16.tflite')
+    f = str(file).replace(".pt", "-fp16.tflite")
 
     converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
     converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
@@ -463,7 +481,8 @@ def export_tflite(keras_model, im, file, int8, per_tensor, data, nms, agnostic_n
     converter.optimizations = [tf.lite.Optimize.DEFAULT]
     if int8:
         from models.tf import representative_dataset_gen
-        dataset = LoadImages(check_dataset(check_yaml(data))['train'], img_size=imgsz, auto=False)
+
+        dataset = LoadImages(check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False)
         converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
         converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
         converter.target_spec.supported_types = []
@@ -472,77 +491,87 @@ def export_tflite(keras_model, im, file, int8, per_tensor, data, nms, agnostic_n
         converter.experimental_new_quantizer = True
         if per_tensor:
             converter._experimental_disable_per_channel = True
-        f = str(file).replace('.pt', '-int8.tflite')
+        f = str(file).replace(".pt", "-int8.tflite")
     if nms or agnostic_nms:
         converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
 
     tflite_model = converter.convert()
-    open(f, 'wb').write(tflite_model)
+    open(f, "wb").write(tflite_model)
     return f, None
 
 
 @try_export
-def export_edgetpu(file, prefix=colorstr('Edge TPU:')):
+def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
     # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
-    cmd = 'edgetpu_compiler --version'
+    cmd = "edgetpu_compiler --version"
-    help_url = 'https://coral.ai/docs/edgetpu/compiler/'
+    help_url = "https://coral.ai/docs/edgetpu/compiler/"
-    assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}'
+    assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}"
-    if subprocess.run(f'{cmd} > /dev/null 2>&1', shell=True).returncode != 0:
+    if subprocess.run(f"{cmd} > /dev/null 2>&1", shell=True).returncode != 0:
-        LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
+        LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
-        sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0  # sudo installed on system
+        sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0  # sudo installed on system
         for c in (
-                'curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
+            "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -",
             'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
-                'sudo apt-get update', 'sudo apt-get install edgetpu-compiler'):
+            "sudo apt-get update",
-            subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True)
+            "sudo apt-get install edgetpu-compiler",
+        ):
+            subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True)
     ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
 
-    LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
+    LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
-    f = str(file).replace('.pt', '-int8_edgetpu.tflite')  # Edge TPU model
+    f = str(file).replace(".pt", "-int8_edgetpu.tflite")  # Edge TPU model
-    f_tfl = str(file).replace('.pt', '-int8.tflite')  # TFLite model
+    f_tfl = str(file).replace(".pt", "-int8.tflite")  # TFLite model
 
-    subprocess.run([
+    subprocess.run(
-        'edgetpu_compiler',
+        [
-        '-s',
+            "edgetpu_compiler",
-        '-d',
+            "-s",
-        '-k',
+            "-d",
-        '10',
+            "-k",
-        '--out_dir',
+            "10",
+            "--out_dir",
             str(file.parent),
-        f_tfl, ], check=True)
+            f_tfl,
+        ],
+        check=True,
+    )
     return f, None
 
 
 @try_export
-def export_tfjs(file, int8, prefix=colorstr('TensorFlow.js:')):
+def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
     # YOLOv5 TensorFlow.js export
-    check_requirements('tensorflowjs')
+    check_requirements("tensorflowjs")
     import tensorflowjs as tfjs
 
-    LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
+    LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
-    f = str(file).replace('.pt', '_web_model')  # js dir
+    f = str(file).replace(".pt", "_web_model")  # js dir
-    f_pb = file.with_suffix('.pb')  # *.pb path
+    f_pb = file.with_suffix(".pb")  # *.pb path
-    f_json = f'{f}/model.json'  # *.json path
+    f_json = f"{f}/model.json"  # *.json path
 
     args = [
-        'tensorflowjs_converter',
+        "tensorflowjs_converter",
-        '--input_format=tf_frozen_model',
+        "--input_format=tf_frozen_model",
-        '--quantize_uint8' if int8 else '',
+        "--quantize_uint8" if int8 else "",
-        '--output_node_names=Identity,Identity_1,Identity_2,Identity_3',
+        "--output_node_names=Identity,Identity_1,Identity_2,Identity_3",
         str(f_pb),
-        str(f), ]
+        str(f),
+    ]
     subprocess.run([arg for arg in args if arg], check=True)
 
     json = Path(f_json).read_text()
-    with open(f_json, 'w') as j:  # sort JSON Identity_* in ascending order
+    with open(f_json, "w") as j:  # sort JSON Identity_* in ascending order
         subst = re.sub(
             r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
             r'"Identity.?.?": {"name": "Identity.?.?"}, '
             r'"Identity.?.?": {"name": "Identity.?.?"}, '
-            r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '
+            r'"Identity.?.?": {"name": "Identity.?.?"}}}',
+            r'{"outputs": {"Identity": {"name": "Identity"}, '
             r'"Identity_1": {"name": "Identity_1"}, '
             r'"Identity_2": {"name": "Identity_2"}, '
-            r'"Identity_3": {"name": "Identity_3"}}}', json)
+            r'"Identity_3": {"name": "Identity_3"}}}',
+            json,
+        )
         j.write(subst)
     return f, None
 
@@ -555,8 +584,8 @@ def add_tflite_metadata(file, metadata, num_outputs):
         from tflite_support import metadata as _metadata
         from tflite_support import metadata_schema_py_generated as _metadata_fb
 
-        tmp_file = Path('/tmp/meta.txt')
+        tmp_file = Path("/tmp/meta.txt")
-        with open(tmp_file, 'w') as meta_f:
+        with open(tmp_file, "w") as meta_f:
             meta_f.write(str(metadata))
 
         model_meta = _metadata_fb.ModelMetadataT()
@@ -580,22 +609,22 @@ def add_tflite_metadata(file, metadata, num_outputs):
         tmp_file.unlink()
 
 
-def pipeline_coreml(model, im, file, names, y, prefix=colorstr('CoreML Pipeline:')):
+def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
     # YOLOv5 CoreML pipeline
     import coremltools as ct
     from PIL import Image
 
-    print(f'{prefix} starting pipeline with coremltools {ct.__version__}...')
+    print(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
     batch_size, ch, h, w = list(im.shape)  # BCHW
     t = time.time()
 
     # YOLOv5 Output shapes
     spec = model.get_spec()
     out0, out1 = iter(spec.description.output)
-    if platform.system() == 'Darwin':
+    if platform.system() == "Darwin":
-        img = Image.new('RGB', (w, h))  # img(192 width, 320 height)
+        img = Image.new("RGB", (w, h))  # img(192 width, 320 height)
         # img = torch.zeros((*opt.img_size, 3)).numpy()  # img size(320,192,3) iDetection
-        out = model.predict({'image': img})
+        out = model.predict({"image": img})
         out0_shape, out1_shape = out[out0.name].shape, out[out1.name].shape
     else:  # linux and windows can not run model.predict(), get sizes from pytorch output y
         s = tuple(y[0].shape)
@@ -605,7 +634,7 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr('CoreML Pipeline:
     nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height
     na, nc = out0_shape
     # na, nc = out0.type.multiArrayType.shape  # number anchors, classes
-    assert len(names) == nc, f'{len(names)} names found for nc={nc}'  # check
+    assert len(names) == nc, f"{len(names)} names found for nc={nc}"  # check
 
     # Define output shapes (missing)
     out0.type.multiArrayType.shape[:] = out0_shape  # (3780, 80)
@@ -639,8 +668,8 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr('CoreML Pipeline:
         nms_spec.description.output.add()
         nms_spec.description.output[i].ParseFromString(decoder_output)
 
-    nms_spec.description.output[0].name = 'confidence'
+    nms_spec.description.output[0].name = "confidence"
-    nms_spec.description.output[1].name = 'coordinates'
+    nms_spec.description.output[1].name = "coordinates"
 
     output_sizes = [nc, 4]
     for i in range(2):
@@ -656,10 +685,10 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr('CoreML Pipeline:
     nms = nms_spec.nonMaximumSuppression
     nms.confidenceInputFeatureName = out0.name  # 1x507x80
     nms.coordinatesInputFeatureName = out1.name  # 1x507x4
-    nms.confidenceOutputFeatureName = 'confidence'
+    nms.confidenceOutputFeatureName = "confidence"
-    nms.coordinatesOutputFeatureName = 'coordinates'
+    nms.coordinatesOutputFeatureName = "coordinates"
-    nms.iouThresholdInputFeatureName = 'iouThreshold'
+    nms.iouThresholdInputFeatureName = "iouThreshold"
-    nms.confidenceThresholdInputFeatureName = 'confidenceThreshold'
+    nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
     nms.iouThreshold = 0.45
     nms.confidenceThreshold = 0.25
     nms.pickTop.perClass = True
@@ -667,10 +696,14 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr('CoreML Pipeline:
     nms_model = ct.models.MLModel(nms_spec)
 
     # 4. Pipeline models together
-    pipeline = ct.models.pipeline.Pipeline(input_features=[('image', ct.models.datatypes.Array(3, ny, nx)),
+    pipeline = ct.models.pipeline.Pipeline(
-                                                           ('iouThreshold', ct.models.datatypes.Double()),
+        input_features=[
-                                                           ('confidenceThreshold', ct.models.datatypes.Double())],
+            ("image", ct.models.datatypes.Array(3, ny, nx)),
-                                           output_features=['confidence', 'coordinates'])
+            ("iouThreshold", ct.models.datatypes.Double()),
+            ("confidenceThreshold", ct.models.datatypes.Double()),
+        ],
+        output_features=["confidence", "coordinates"],
+    )
     pipeline.add_model(model)
     pipeline.add_model(nms_model)
 
@@ -681,36 +714,40 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr('CoreML Pipeline:
 
     # Update metadata
     pipeline.spec.specificationVersion = 5
-    pipeline.spec.description.metadata.versionString = 'https://github.com/ultralytics/yolov5'
+    pipeline.spec.description.metadata.versionString = "https://github.com/ultralytics/yolov5"
-    pipeline.spec.description.metadata.shortDescription = 'https://github.com/ultralytics/yolov5'
+    pipeline.spec.description.metadata.shortDescription = "https://github.com/ultralytics/yolov5"
-    pipeline.spec.description.metadata.author = 'glenn.jocher@ultralytics.com'
+    pipeline.spec.description.metadata.author = "glenn.jocher@ultralytics.com"
-    pipeline.spec.description.metadata.license = 'https://github.com/ultralytics/yolov5/blob/master/LICENSE'
+    pipeline.spec.description.metadata.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE"
-    pipeline.spec.description.metadata.userDefined.update({
+    pipeline.spec.description.metadata.userDefined.update(
-        'classes': ','.join(names.values()),
+        {
-        'iou_threshold': str(nms.iouThreshold),
+            "classes": ",".join(names.values()),
-        'confidence_threshold': str(nms.confidenceThreshold)})
+            "iou_threshold": str(nms.iouThreshold),
+            "confidence_threshold": str(nms.confidenceThreshold),
+        }
+    )
 
     # Save the model
-    f = file.with_suffix('.mlmodel')  # filename
+    f = file.with_suffix(".mlmodel")  # filename
     model = ct.models.MLModel(pipeline.spec)
-    model.input_description['image'] = 'Input image'
+    model.input_description["image"] = "Input image"
-    model.input_description['iouThreshold'] = f'(optional) IOU Threshold override (default: {nms.iouThreshold})'
+    model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})"
-    model.input_description['confidenceThreshold'] = \
+    model.input_description[
-        f'(optional) Confidence Threshold override (default: {nms.confidenceThreshold})'
+        "confidenceThreshold"
-    model.output_description['confidence'] = 'Boxes × Class confidence (see user-defined metadata "classes")'
+    ] = f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})"
-    model.output_description['coordinates'] = 'Boxes × [x, y, width, height] (relative to image size)'
+    model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")'
+    model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)"
     model.save(f)  # pipelined
-    print(f'{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)')
+    print(f"{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)")
 
 
 @smart_inference_mode()
 def run(
-        data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
+    data=ROOT / "data/coco128.yaml",  # 'dataset.yaml path'
-        weights=ROOT / 'yolov5s.pt',  # weights path
+    weights=ROOT / "yolov5s.pt",  # weights path
     imgsz=(640, 640),  # image (height, width)
     batch_size=1,  # batch size
-        device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="cpu",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
-        include=('torchscript', 'onnx'),  # include formats
+    include=("torchscript", "onnx"),  # include formats
     half=False,  # FP16 half-precision export
     inplace=False,  # set YOLOv5 Detect() inplace=True
     keras=False,  # use Keras
@@ -731,23 +768,23 @@ def run(
 ):
     t = time.time()
     include = [x.lower() for x in include]  # to lowercase
-    fmts = tuple(export_formats()['Argument'][1:])  # --include arguments
+    fmts = tuple(export_formats()["Argument"][1:])  # --include arguments
     flags = [x in include for x in fmts]
-    assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}'
+    assert sum(flags) == len(include), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}"
     jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
-    file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)  # PyTorch weights
+    file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights)  # PyTorch weights
 
     # Load PyTorch model
     device = select_device(device)
     if half:
-        assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0'
+        assert device.type != "cpu" or coreml, "--half only compatible with GPU export, i.e. use --device 0"
-        assert not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both'
+        assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both"
     model = attempt_load(weights, device=device, inplace=True, fuse=True)  # load FP32 model
 
     # Checks
     imgsz *= 2 if len(imgsz) == 1 else 1  # expand
     if optimize:
-        assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
+        assert device.type == "cpu", "--optimize not compatible with cuda devices, i.e. use --device cpu"
 
     # Input
     gs = int(max(model.stride))  # grid size (max stride)
@@ -767,12 +804,12 @@ def run(
     if half and not coreml:
         im, model = im.half(), model.half()  # to FP16
     shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
-    metadata = {'stride': int(max(model.stride)), 'names': model.names}  # model metadata
+    metadata = {"stride": int(max(model.stride)), "names": model.names}  # model metadata
     LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
 
     # Exports
-    f = [''] * len(fmts)  # exported filenames
+    f = [""] * len(fmts)  # exported filenames
-    warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning)  # suppress TracerWarning
+    warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning)  # suppress TracerWarning
     if jit:  # TorchScript
         f[0], _ = export_torchscript(model, im, file, optimize)
     if engine:  # TensorRT required before ONNX
@@ -786,9 +823,10 @@ def run(
         if nms:
             pipeline_coreml(ct_model, im, file, model.names, y)
     if any((saved_model, pb, tflite, edgetpu, tfjs)):  # TensorFlow formats
-        assert not tflite or not tfjs, 'TFLite and TF.js models must be exported separately, please pass only one type.'
+        assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type."
-        assert not isinstance(model, ClassificationModel), 'ClassificationModel export to TF formats not yet supported.'
+        assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported."
-        f[5], s_model = export_saved_model(model.cpu(),
+        f[5], s_model = export_saved_model(
+            model.cpu(),
             im,
             file,
             dynamic,
@@ -798,18 +836,14 @@ def run(
             topk_all=topk_all,
             iou_thres=iou_thres,
             conf_thres=conf_thres,
-                                           keras=keras)
+            keras=keras,
+        )
         if pb or tfjs:  # pb prerequisite to tfjs
             f[6], _ = export_pb(s_model, file)
         if tflite or edgetpu:
-            f[7], _ = export_tflite(s_model,
+            f[7], _ = export_tflite(
-                                    im,
+                s_model, im, file, int8 or edgetpu, per_tensor, data=data, nms=nms, agnostic_nms=agnostic_nms
-                                    file,
+            )
-                                    int8 or edgetpu,
-                                    per_tensor,
-                                    data=data,
-                                    nms=nms,
-                                    agnostic_nms=agnostic_nms)
             if edgetpu:
                 f[8], _ = export_edgetpu(file)
             add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs))
@@ -823,58 +857,66 @@ def run(
     if any(f):
         cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel))  # type
         det &= not seg  # segmentation models inherit from SegmentationModel(DetectionModel)
-        dir = Path('segment' if seg else 'classify' if cls else '')
+        dir = Path("segment" if seg else "classify" if cls else "")
-        h = '--half' if half else ''  # --half FP16 inference arg
+        h = "--half" if half else ""  # --half FP16 inference arg
-        s = '# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference' if cls else \
+        s = (
-            '# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference' if seg else ''
+            "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference"
-        LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)'
+            if cls
+            else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference"
+            if seg
+            else ""
+        )
+        LOGGER.info(
+            f'\nExport complete ({time.time() - t:.1f}s)'
             f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
             f"\nDetect:          python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}"
             f"\nValidate:        python {dir / 'val.py'} --weights {f[-1]} {h}"
             f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')  {s}"
-                    f'\nVisualize:       https://netron.app')
+            f'\nVisualize:       https://netron.app'
+        )
     return f  # return list of exported files/dirs
 
 
 def parse_opt(known=False):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
+    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.pt path(s)")
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)")
-    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
+    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
-    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
+    parser.add_argument("--half", action="store_true", help="FP16 half-precision export")
-    parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
+    parser.add_argument("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True")
-    parser.add_argument('--keras', action='store_true', help='TF: use Keras')
+    parser.add_argument("--keras", action="store_true", help="TF: use Keras")
-    parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
+    parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile")
-    parser.add_argument('--int8', action='store_true', help='CoreML/TF/OpenVINO INT8 quantization')
+    parser.add_argument("--int8", action="store_true", help="CoreML/TF/OpenVINO INT8 quantization")
-    parser.add_argument('--per-tensor', action='store_true', help='TF per-tensor quantization')
+    parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization")
-    parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes')
+    parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes")
-    parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
+    parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model")
-    parser.add_argument('--opset', type=int, default=17, help='ONNX: opset version')
+    parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version")
-    parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
+    parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log")
-    parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
+    parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)")
-    parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
+    parser.add_argument("--nms", action="store_true", help="TF: add NMS to model")
-    parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
+    parser.add_argument("--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model")
-    parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
+    parser.add_argument("--topk-per-class", type=int, default=100, help="TF.js NMS: topk per class to keep")
-    parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
+    parser.add_argument("--topk-all", type=int, default=100, help="TF.js NMS: topk for all classes to keep")
-    parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
+    parser.add_argument("--iou-thres", type=float, default=0.45, help="TF.js NMS: IoU threshold")
-    parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
+    parser.add_argument("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold")
     parser.add_argument(
-        '--include',
+        "--include",
-        nargs='+',
+        nargs="+",
-        default=['torchscript'],
+        default=["torchscript"],
-        help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle')
+        help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle",
+    )
     opt = parser.parse_known_args()[0] if known else parser.parse_args()
     print_args(vars(opt))
     return opt
 
 
 def main(opt):
-    for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
+    for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
         run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

hubconf.py

@@ -14,7 +14,8 @@ import torch
 
 
 def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
-    """Creates or loads a YOLOv5 model
+    """
+    Creates or loads a YOLOv5 model.
 
     Arguments:
         name (str): model name 'yolov5s' or path 'path/to/best.pt'
@@ -39,9 +40,9 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
 
     if not verbose:
         LOGGER.setLevel(logging.WARNING)
-    check_requirements(ROOT / 'requirements.txt', exclude=('opencv-python', 'tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("opencv-python", "tensorboard", "thop"))
     name = Path(name)
-    path = name.with_suffix('.pt') if name.suffix == '' and not name.is_dir() else name  # checkpoint path
+    path = name.with_suffix(".pt") if name.suffix == "" and not name.is_dir() else name  # checkpoint path
     try:
         device = select_device(device)
         if pretrained and channels == 3 and classes == 80:
@@ -49,91 +50,95 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
                 model = DetectMultiBackend(path, device=device, fuse=autoshape)  # detection model
                 if autoshape:
                     if model.pt and isinstance(model.model, ClassificationModel):
-                        LOGGER.warning('WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. '
+                        LOGGER.warning(
-                                       'You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224).')
+                            "WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. "
+                            "You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224)."
+                        )
                     elif model.pt and isinstance(model.model, SegmentationModel):
-                        LOGGER.warning('WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. '
+                        LOGGER.warning(
-                                       'You will not be able to run inference with this model.')
+                            "WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. "
+                            "You will not be able to run inference with this model."
+                        )
                     else:
                         model = AutoShape(model)  # for file/URI/PIL/cv2/np inputs and NMS
             except Exception:
                 model = attempt_load(path, device=device, fuse=False)  # arbitrary model
         else:
-            cfg = list((Path(__file__).parent / 'models').rglob(f'{path.stem}.yaml'))[0]  # model.yaml path
+            cfg = list((Path(__file__).parent / "models").rglob(f"{path.stem}.yaml"))[0]  # model.yaml path
             model = DetectionModel(cfg, channels, classes)  # create model
             if pretrained:
                 ckpt = torch.load(attempt_download(path), map_location=device)  # load
-                csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
+                csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
-                csd = intersect_dicts(csd, model.state_dict(), exclude=['anchors'])  # intersect
+                csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"])  # intersect
                 model.load_state_dict(csd, strict=False)  # load
-                if len(ckpt['model'].names) == classes:
+                if len(ckpt["model"].names) == classes:
-                    model.names = ckpt['model'].names  # set class names attribute
+                    model.names = ckpt["model"].names  # set class names attribute
         if not verbose:
             LOGGER.setLevel(logging.INFO)  # reset to default
         return model.to(device)
 
     except Exception as e:
-        help_url = 'https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading'
+        help_url = "https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading"
-        s = f'{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help.'
+        s = f"{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help."
         raise Exception(s) from e
 
 
-def custom(path='path/to/model.pt', autoshape=True, _verbose=True, device=None):
+def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
     # YOLOv5 custom or local model
     return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
 
 
 def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-nano model https://github.com/ultralytics/yolov5
-    return _create('yolov5n', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-small model https://github.com/ultralytics/yolov5
-    return _create('yolov5s', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-medium model https://github.com/ultralytics/yolov5
-    return _create('yolov5m', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-large model https://github.com/ultralytics/yolov5
-    return _create('yolov5l', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-xlarge model https://github.com/ultralytics/yolov5
-    return _create('yolov5x', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-nano-P6 model https://github.com/ultralytics/yolov5
-    return _create('yolov5n6', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
-    return _create('yolov5s6', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
-    return _create('yolov5m6', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
-    return _create('yolov5l6', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
     # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
-    return _create('yolov5x6', pretrained, channels, classes, autoshape, _verbose, device)
+    return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     import argparse
     from pathlib import Path
 
@@ -144,7 +149,7 @@ if __name__ == '__main__':
 
     # Argparser
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', type=str, default='yolov5s', help='model name')
+    parser.add_argument("--model", type=str, default="yolov5s", help="model name")
     opt = parser.parse_args()
     print_args(vars(opt))
 
@@ -154,12 +159,13 @@ if __name__ == '__main__':
 
     # Images
     imgs = [
-        'data/images/zidane.jpg',  # filename
+        "data/images/zidane.jpg",  # filename
-        Path('data/images/zidane.jpg'),  # Path
+        Path("data/images/zidane.jpg"),  # Path
-        'https://ultralytics.com/images/zidane.jpg',  # URI
+        "https://ultralytics.com/images/zidane.jpg",  # URI
-        cv2.imread('data/images/bus.jpg')[:, :, ::-1],  # OpenCV
+        cv2.imread("data/images/bus.jpg")[:, :, ::-1],  # OpenCV
-        Image.open('data/images/bus.jpg'),  # PIL
+        Image.open("data/images/bus.jpg"),  # PIL
-        np.zeros((320, 640, 3))]  # numpy
+        np.zeros((320, 640, 3)),
+    ]  # numpy
 
     # Inference
     results = model(imgs, size=320)  # batched inference

models/common.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Common modules."""
-Common modules
-"""
 
 import ast
 import contextlib
@@ -28,20 +26,34 @@ from torch.cuda import amp
 try:
     import ultralytics
 
-    assert hasattr(ultralytics, '__version__')  # verify package is not directory
+    assert hasattr(ultralytics, "__version__")  # verify package is not directory
 except (ImportError, AssertionError):
     import os
 
-    os.system('pip install -U ultralytics')
+    os.system("pip install -U ultralytics")
     import ultralytics
 
 from ultralytics.utils.plotting import Annotator, colors, save_one_box
 
 from utils import TryExcept
 from utils.dataloaders import exif_transpose, letterbox
-from utils.general import (LOGGER, ROOT, Profile, check_requirements, check_suffix, check_version, colorstr,
+from utils.general import (
-                           increment_path, is_jupyter, make_divisible, non_max_suppression, scale_boxes, xywh2xyxy,
+    LOGGER,
-                           xyxy2xywh, yaml_load)
+    ROOT,
+    Profile,
+    check_requirements,
+    check_suffix,
+    check_version,
+    colorstr,
+    increment_path,
+    is_jupyter,
+    make_divisible,
+    non_max_suppression,
+    scale_boxes,
+    xywh2xyxy,
+    xyxy2xywh,
+    yaml_load,
+)
 from utils.torch_utils import copy_attr, smart_inference_mode
 
 
@@ -223,7 +235,7 @@ class SPP(nn.Module):
     def forward(self, x):
         x = self.cv1(x)
         with warnings.catch_warnings():
-            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
+            warnings.simplefilter("ignore")  # suppress torch 1.9.0 max_pool2d() warning
             return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
 
 
@@ -239,7 +251,7 @@ class SPPF(nn.Module):
     def forward(self, x):
         x = self.cv1(x)
         with warnings.catch_warnings():
-            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
+            warnings.simplefilter("ignore")  # suppress torch 1.9.0 max_pool2d() warning
             y1 = self.m(x)
             y2 = self.m(y1)
             return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
@@ -278,9 +290,11 @@ class GhostBottleneck(nn.Module):
         self.conv = nn.Sequential(
             GhostConv(c1, c_, 1, 1),  # pw
             DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),  # dw
-            GhostConv(c_, c2, 1, 1, act=False))  # pw-linear
+            GhostConv(c_, c2, 1, 1, act=False),
-        self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
+        )  # pw-linear
-                                                                            act=False)) if s == 2 else nn.Identity()
+        self.shortcut = (
+            nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
+        )
 
     def forward(self, x):
         return self.conv(x) + self.shortcut(x)
@@ -326,7 +340,7 @@ class Concat(nn.Module):
 
 class DetectMultiBackend(nn.Module):
     # YOLOv5 MultiBackend class for python inference on various backends
-    def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
+    def __init__(self, weights="yolov5s.pt", device=torch.device("cpu"), dnn=False, data=None, fp16=False, fuse=True):
         # Usage:
         #   PyTorch:              weights = *.pt
         #   TorchScript:                    *.torchscript
@@ -348,65 +362,68 @@ class DetectMultiBackend(nn.Module):
         fp16 &= pt or jit or onnx or engine or triton  # FP16
         nhwc = coreml or saved_model or pb or tflite or edgetpu  # BHWC formats (vs torch BCWH)
         stride = 32  # default stride
-        cuda = torch.cuda.is_available() and device.type != 'cpu'  # use CUDA
+        cuda = torch.cuda.is_available() and device.type != "cpu"  # use CUDA
         if not (pt or triton):
             w = attempt_download(w)  # download if not local
 
         if pt:  # PyTorch
             model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)
             stride = max(int(model.stride.max()), 32)  # model stride
-            names = model.module.names if hasattr(model, 'module') else model.names  # get class names
+            names = model.module.names if hasattr(model, "module") else model.names  # get class names
             model.half() if fp16 else model.float()
             self.model = model  # explicitly assign for to(), cpu(), cuda(), half()
         elif jit:  # TorchScript
-            LOGGER.info(f'Loading {w} for TorchScript inference...')
+            LOGGER.info(f"Loading {w} for TorchScript inference...")
-            extra_files = {'config.txt': ''}  # model metadata
+            extra_files = {"config.txt": ""}  # model metadata
             model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
             model.half() if fp16 else model.float()
-            if extra_files['config.txt']:  # load metadata dict
+            if extra_files["config.txt"]:  # load metadata dict
-                d = json.loads(extra_files['config.txt'],
+                d = json.loads(
-                               object_hook=lambda d: {
+                    extra_files["config.txt"],
-                                   int(k) if k.isdigit() else k: v
+                    object_hook=lambda d: {int(k) if k.isdigit() else k: v for k, v in d.items()},
-                                   for k, v in d.items()})
+                )
-                stride, names = int(d['stride']), d['names']
+                stride, names = int(d["stride"]), d["names"]
         elif dnn:  # ONNX OpenCV DNN
-            LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
+            LOGGER.info(f"Loading {w} for ONNX OpenCV DNN inference...")
-            check_requirements('opencv-python>=4.5.4')
+            check_requirements("opencv-python>=4.5.4")
             net = cv2.dnn.readNetFromONNX(w)
         elif onnx:  # ONNX Runtime
-            LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
+            LOGGER.info(f"Loading {w} for ONNX Runtime inference...")
-            check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
+            check_requirements(("onnx", "onnxruntime-gpu" if cuda else "onnxruntime"))
             import onnxruntime
-            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
+
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"]
             session = onnxruntime.InferenceSession(w, providers=providers)
             output_names = [x.name for x in session.get_outputs()]
             meta = session.get_modelmeta().custom_metadata_map  # metadata
-            if 'stride' in meta:
+            if "stride" in meta:
-                stride, names = int(meta['stride']), eval(meta['names'])
+                stride, names = int(meta["stride"]), eval(meta["names"])
         elif xml:  # OpenVINO
-            LOGGER.info(f'Loading {w} for OpenVINO inference...')
+            LOGGER.info(f"Loading {w} for OpenVINO inference...")
-            check_requirements('openvino>=2023.0')  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+            check_requirements("openvino>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
             from openvino.runtime import Core, Layout, get_batch
+
             core = Core()
             if not Path(w).is_file():  # if not *.xml
-                w = next(Path(w).glob('*.xml'))  # get *.xml file from *_openvino_model dir
+                w = next(Path(w).glob("*.xml"))  # get *.xml file from *_openvino_model dir
-            ov_model = core.read_model(model=w, weights=Path(w).with_suffix('.bin'))
+            ov_model = core.read_model(model=w, weights=Path(w).with_suffix(".bin"))
             if ov_model.get_parameters()[0].get_layout().empty:
-                ov_model.get_parameters()[0].set_layout(Layout('NCHW'))
+                ov_model.get_parameters()[0].set_layout(Layout("NCHW"))
             batch_dim = get_batch(ov_model)
             if batch_dim.is_static:
                 batch_size = batch_dim.get_length()
-            ov_compiled_model = core.compile_model(ov_model, device_name='AUTO')  # AUTO selects best available device
+            ov_compiled_model = core.compile_model(ov_model, device_name="AUTO")  # AUTO selects best available device
-            stride, names = self._load_metadata(Path(w).with_suffix('.yaml'))  # load metadata
+            stride, names = self._load_metadata(Path(w).with_suffix(".yaml"))  # load metadata
         elif engine:  # TensorRT
-            LOGGER.info(f'Loading {w} for TensorRT inference...')
+            LOGGER.info(f"Loading {w} for TensorRT inference...")
             import tensorrt as trt  # https://developer.nvidia.com/nvidia-tensorrt-download
-            check_version(trt.__version__, '7.0.0', hard=True)  # require tensorrt>=7.0.0
+
-            if device.type == 'cpu':
+            check_version(trt.__version__, "7.0.0", hard=True)  # require tensorrt>=7.0.0
-                device = torch.device('cuda:0')
+            if device.type == "cpu":
-            Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
+                device = torch.device("cuda:0")
+            Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))
             logger = trt.Logger(trt.Logger.INFO)
-            with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
+            with open(w, "rb") as f, trt.Runtime(logger) as runtime:
                 model = runtime.deserialize_cuda_engine(f.read())
             context = model.create_execution_context()
             bindings = OrderedDict()
@@ -428,22 +445,24 @@ class DetectMultiBackend(nn.Module):
                 im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
                 bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
             binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
-            batch_size = bindings['images'].shape[0]  # if dynamic, this is instead max batch size
+            batch_size = bindings["images"].shape[0]  # if dynamic, this is instead max batch size
         elif coreml:  # CoreML
-            LOGGER.info(f'Loading {w} for CoreML inference...')
+            LOGGER.info(f"Loading {w} for CoreML inference...")
             import coremltools as ct
+
             model = ct.models.MLModel(w)
         elif saved_model:  # TF SavedModel
-            LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...')
+            LOGGER.info(f"Loading {w} for TensorFlow SavedModel inference...")
             import tensorflow as tf
+
             keras = False  # assume TF1 saved_model
             model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
         elif pb:  # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
-            LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')
+            LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
             import tensorflow as tf
 
             def wrap_frozen_graph(gd, inputs, outputs):
-                x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=''), [])  # wrapped
+                x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), [])  # wrapped
                 ge = x.graph.as_graph_element
                 return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
 
@@ -452,46 +471,50 @@ class DetectMultiBackend(nn.Module):
                 for node in gd.node:  # tensorflow.core.framework.node_def_pb2.NodeDef
                     name_list.append(node.name)
                     input_list.extend(node.input)
-                return sorted(f'{x}:0' for x in list(set(name_list) - set(input_list)) if not x.startswith('NoOp'))
+                return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
 
             gd = tf.Graph().as_graph_def()  # TF GraphDef
-            with open(w, 'rb') as f:
+            with open(w, "rb") as f:
                 gd.ParseFromString(f.read())
-            frozen_func = wrap_frozen_graph(gd, inputs='x:0', outputs=gd_outputs(gd))
+            frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd))
         elif tflite or edgetpu:  # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
             try:  # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
                 from tflite_runtime.interpreter import Interpreter, load_delegate
             except ImportError:
                 import tensorflow as tf
-                Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,
+
+                Interpreter, load_delegate = (
+                    tf.lite.Interpreter,
+                    tf.lite.experimental.load_delegate,
+                )
             if edgetpu:  # TF Edge TPU https://coral.ai/software/#edgetpu-runtime
-                LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
+                LOGGER.info(f"Loading {w} for TensorFlow Lite Edge TPU inference...")
-                delegate = {
+                delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[
-                    'Linux': 'libedgetpu.so.1',
+                    platform.system()
-                    'Darwin': 'libedgetpu.1.dylib',
+                ]
-                    'Windows': 'edgetpu.dll'}[platform.system()]
                 interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
             else:  # TFLite
-                LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
+                LOGGER.info(f"Loading {w} for TensorFlow Lite inference...")
                 interpreter = Interpreter(model_path=w)  # load TFLite model
             interpreter.allocate_tensors()  # allocate
             input_details = interpreter.get_input_details()  # inputs
             output_details = interpreter.get_output_details()  # outputs
             # load metadata
             with contextlib.suppress(zipfile.BadZipFile):
-                with zipfile.ZipFile(w, 'r') as model:
+                with zipfile.ZipFile(w, "r") as model:
                     meta_file = model.namelist()[0]
-                    meta = ast.literal_eval(model.read(meta_file).decode('utf-8'))
+                    meta = ast.literal_eval(model.read(meta_file).decode("utf-8"))
-                    stride, names = int(meta['stride']), meta['names']
+                    stride, names = int(meta["stride"]), meta["names"]
         elif tfjs:  # TF.js
-            raise NotImplementedError('ERROR: YOLOv5 TF.js inference is not supported')
+            raise NotImplementedError("ERROR: YOLOv5 TF.js inference is not supported")
         elif paddle:  # PaddlePaddle
-            LOGGER.info(f'Loading {w} for PaddlePaddle inference...')
+            LOGGER.info(f"Loading {w} for PaddlePaddle inference...")
-            check_requirements('paddlepaddle-gpu' if cuda else 'paddlepaddle')
+            check_requirements("paddlepaddle-gpu" if cuda else "paddlepaddle")
             import paddle.inference as pdi
+
             if not Path(w).is_file():  # if not *.pdmodel
-                w = next(Path(w).rglob('*.pdmodel'))  # get *.pdmodel file from *_paddle_model dir
+                w = next(Path(w).rglob("*.pdmodel"))  # get *.pdmodel file from *_paddle_model dir
-            weights = Path(w).with_suffix('.pdiparams')
+            weights = Path(w).with_suffix(".pdiparams")
             config = pdi.Config(str(w), str(weights))
             if cuda:
                 config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0)
@@ -499,19 +522,20 @@ class DetectMultiBackend(nn.Module):
             input_handle = predictor.get_input_handle(predictor.get_input_names()[0])
             output_names = predictor.get_output_names()
         elif triton:  # NVIDIA Triton Inference Server
-            LOGGER.info(f'Using {w} as Triton Inference Server...')
+            LOGGER.info(f"Using {w} as Triton Inference Server...")
-            check_requirements('tritonclient[all]')
+            check_requirements("tritonclient[all]")
             from utils.triton import TritonRemoteModel
+
             model = TritonRemoteModel(url=w)
-            nhwc = model.runtime.startswith('tensorflow')
+            nhwc = model.runtime.startswith("tensorflow")
         else:
-            raise NotImplementedError(f'ERROR: {w} is not a supported format')
+            raise NotImplementedError(f"ERROR: {w} is not a supported format")
 
         # class names
-        if 'names' not in locals():
+        if "names" not in locals():
-            names = yaml_load(data)['names'] if data else {i: f'class{i}' for i in range(999)}
+            names = yaml_load(data)["names"] if data else {i: f"class{i}" for i in range(999)}
-        if names[0] == 'n01440764' and len(names) == 1000:  # ImageNet
+        if names[0] == "n01440764" and len(names) == 1000:  # ImageNet
-            names = yaml_load(ROOT / 'data/ImageNet.yaml')['names']  # human-readable names
+            names = yaml_load(ROOT / "data/ImageNet.yaml")["names"]  # human-readable names
 
         self.__dict__.update(locals())  # assign all variables to self
 
@@ -538,26 +562,26 @@ class DetectMultiBackend(nn.Module):
             im = im.cpu().numpy()  # FP32
             y = list(self.ov_compiled_model(im).values())
         elif self.engine:  # TensorRT
-            if self.dynamic and im.shape != self.bindings['images'].shape:
+            if self.dynamic and im.shape != self.bindings["images"].shape:
-                i = self.model.get_binding_index('images')
+                i = self.model.get_binding_index("images")
                 self.context.set_binding_shape(i, im.shape)  # reshape if dynamic
-                self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)
+                self.bindings["images"] = self.bindings["images"]._replace(shape=im.shape)
                 for name in self.output_names:
                     i = self.model.get_binding_index(name)
                     self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))
-            s = self.bindings['images'].shape
+            s = self.bindings["images"].shape
             assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
-            self.binding_addrs['images'] = int(im.data_ptr())
+            self.binding_addrs["images"] = int(im.data_ptr())
             self.context.execute_v2(list(self.binding_addrs.values()))
             y = [self.bindings[x].data for x in sorted(self.output_names)]
         elif self.coreml:  # CoreML
             im = im.cpu().numpy()
-            im = Image.fromarray((im[0] * 255).astype('uint8'))
+            im = Image.fromarray((im[0] * 255).astype("uint8"))
             # im = im.resize((192, 320), Image.BILINEAR)
-            y = self.model.predict({'image': im})  # coordinates are xywh normalized
+            y = self.model.predict({"image": im})  # coordinates are xywh normalized
-            if 'confidence' in y:
+            if "confidence" in y:
-                box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]])  # xyxy pixels
+                box = xywh2xyxy(y["coordinates"] * [[w, h, w, h]])  # xyxy pixels
-                conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
+                conf, cls = y["confidence"].max(1), y["confidence"].argmax(1).astype(np.float)
                 y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
             else:
                 y = list(reversed(y.values()))  # reversed for segmentation models (pred, proto)
@@ -576,17 +600,17 @@ class DetectMultiBackend(nn.Module):
                 y = self.frozen_func(x=self.tf.constant(im))
             else:  # Lite or Edge TPU
                 input = self.input_details[0]
-                int8 = input['dtype'] == np.uint8  # is TFLite quantized uint8 model
+                int8 = input["dtype"] == np.uint8  # is TFLite quantized uint8 model
                 if int8:
-                    scale, zero_point = input['quantization']
+                    scale, zero_point = input["quantization"]
                     im = (im / scale + zero_point).astype(np.uint8)  # de-scale
-                self.interpreter.set_tensor(input['index'], im)
+                self.interpreter.set_tensor(input["index"], im)
                 self.interpreter.invoke()
                 y = []
                 for output in self.output_details:
-                    x = self.interpreter.get_tensor(output['index'])
+                    x = self.interpreter.get_tensor(output["index"])
                     if int8:
-                        scale, zero_point = output['quantization']
+                        scale, zero_point = output["quantization"]
                         x = (x.astype(np.float32) - zero_point) * scale  # re-scale
                     y.append(x)
             y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
@@ -603,32 +627,33 @@ class DetectMultiBackend(nn.Module):
     def warmup(self, imgsz=(1, 3, 640, 640)):
         # Warmup model by running inference once
         warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton
-        if any(warmup_types) and (self.device.type != 'cpu' or self.triton):
+        if any(warmup_types) and (self.device.type != "cpu" or self.triton):
             im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device)  # input
             for _ in range(2 if self.jit else 1):  #
                 self.forward(im)  # warmup
 
     @staticmethod
-    def _model_type(p='path/to/model.pt'):
+    def _model_type(p="path/to/model.pt"):
         # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
         # types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]
         from export import export_formats
         from utils.downloads import is_url
+
         sf = list(export_formats().Suffix)  # export suffixes
         if not is_url(p, check=False):
             check_suffix(p, sf)  # checks
         url = urlparse(p)  # if url may be Triton inference server
         types = [s in Path(p).name for s in sf]
         types[8] &= not types[9]  # tflite &= not edgetpu
-        triton = not any(types) and all([any(s in url.scheme for s in ['http', 'grpc']), url.netloc])
+        triton = not any(types) and all([any(s in url.scheme for s in ["http", "grpc"]), url.netloc])
         return types + [triton]
 
     @staticmethod
-    def _load_metadata(f=Path('path/to/meta.yaml')):
+    def _load_metadata(f=Path("path/to/meta.yaml")):
         # Load metadata from meta.yaml if it exists
         if f.exists():
             d = yaml_load(f)
-            return d['stride'], d['names']  # assign stride, names
+            return d["stride"], d["names"]  # assign stride, names
         return None, None
 
 
@@ -645,8 +670,8 @@ class AutoShape(nn.Module):
     def __init__(self, model, verbose=True):
         super().__init__()
         if verbose:
-            LOGGER.info('Adding AutoShape... ')
+            LOGGER.info("Adding AutoShape... ")
-        copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=())  # copy attributes
+        copy_attr(self, model, include=("yaml", "nc", "hyp", "names", "stride", "abc"), exclude=())  # copy attributes
         self.dmb = isinstance(model, DetectMultiBackend)  # DetectMultiBackend() instance
         self.pt = not self.dmb or model.pt  # PyTorch model
         self.model = model.eval()
@@ -682,7 +707,7 @@ class AutoShape(nn.Module):
             if isinstance(size, int):  # expand
                 size = (size, size)
             p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device)  # param
-            autocast = self.amp and (p.device.type != 'cpu')  # Automatic Mixed Precision (AMP) inference
+            autocast = self.amp and (p.device.type != "cpu")  # Automatic Mixed Precision (AMP) inference
             if isinstance(ims, torch.Tensor):  # torch
                 with amp.autocast(autocast):
                     return self.model(ims.to(p.device).type_as(p), augment=augment)  # inference
@@ -691,13 +716,13 @@ class AutoShape(nn.Module):
             n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims])  # number, list of images
             shape0, shape1, files = [], [], []  # image and inference shapes, filenames
             for i, im in enumerate(ims):
-                f = f'image{i}'  # filename
+                f = f"image{i}"  # filename
                 if isinstance(im, (str, Path)):  # filename or uri
-                    im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
+                    im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im), im
                     im = np.asarray(exif_transpose(im))
                 elif isinstance(im, Image.Image):  # PIL Image
-                    im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f
+                    im, f = np.asarray(exif_transpose(im)), getattr(im, "filename", f) or f
-                files.append(Path(f).with_suffix('.jpg').name)
+                files.append(Path(f).with_suffix(".jpg").name)
                 if im.shape[0] < 5:  # image in CHW
                     im = im.transpose((1, 2, 0))  # reverse dataloader .transpose(2, 0, 1)
                 im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)  # enforce 3ch input
@@ -718,13 +743,15 @@ class AutoShape(nn.Module):
 
             # Post-process
             with dt[2]:
-                y = non_max_suppression(y if self.dmb else y[0],
+                y = non_max_suppression(
+                    y if self.dmb else y[0],
                     self.conf,
                     self.iou,
                     self.classes,
                     self.agnostic,
                     self.multi_label,
-                                        max_det=self.max_det)  # NMS
+                    max_det=self.max_det,
+                )  # NMS
                 for i in range(n):
                     scale_boxes(shape1, y[i][:, :4], shape0[i])
 
@@ -747,40 +774,44 @@ class Detections:
         self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
         self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
         self.n = len(self.pred)  # number of images (batch size)
-        self.t = tuple(x.t / self.n * 1E3 for x in times)  # timestamps (ms)
+        self.t = tuple(x.t / self.n * 1e3 for x in times)  # timestamps (ms)
         self.s = tuple(shape)  # inference BCHW shape
 
-    def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
+    def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path("")):
-        s, crops = '', []
+        s, crops = "", []
         for i, (im, pred) in enumerate(zip(self.ims, self.pred)):
-            s += f'\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '  # string
+            s += f"\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} "  # string
             if pred.shape[0]:
                 for c in pred[:, -1].unique():
                     n = (pred[:, -1] == c).sum()  # detections per class
                     s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string
-                s = s.rstrip(', ')
+                s = s.rstrip(", ")
                 if show or save or render or crop:
                     annotator = Annotator(im, example=str(self.names))
                     for *box, conf, cls in reversed(pred):  # xyxy, confidence, class
-                        label = f'{self.names[int(cls)]} {conf:.2f}'
+                        label = f"{self.names[int(cls)]} {conf:.2f}"
                         if crop:
-                            file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
+                            file = save_dir / "crops" / self.names[int(cls)] / self.files[i] if save else None
-                            crops.append({
+                            crops.append(
-                                'box': box,
+                                {
-                                'conf': conf,
+                                    "box": box,
-                                'cls': cls,
+                                    "conf": conf,
-                                'label': label,
+                                    "cls": cls,
-                                'im': save_one_box(box, im, file=file, save=save)})
+                                    "label": label,
+                                    "im": save_one_box(box, im, file=file, save=save),
+                                }
+                            )
                         else:  # all others
-                            annotator.box_label(box, label if labels else '', color=colors(cls))
+                            annotator.box_label(box, label if labels else "", color=colors(cls))
                     im = annotator.im
             else:
-                s += '(no detections)'
+                s += "(no detections)"
 
             im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im  # from np
             if show:
                 if is_jupyter():
                     from IPython.display import display
+
                     display(im)
                 else:
                     im.show(self.files[i])
@@ -792,22 +823,22 @@ class Detections:
             if render:
                 self.ims[i] = np.asarray(im)
         if pprint:
-            s = s.lstrip('\n')
+            s = s.lstrip("\n")
-            return f'{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t
+            return f"{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}" % self.t
         if crop:
             if save:
-                LOGGER.info(f'Saved results to {save_dir}\n')
+                LOGGER.info(f"Saved results to {save_dir}\n")
             return crops
 
-    @TryExcept('Showing images is not supported in this environment')
+    @TryExcept("Showing images is not supported in this environment")
     def show(self, labels=True):
         self._run(show=True, labels=labels)  # show results
 
-    def save(self, labels=True, save_dir='runs/detect/exp', exist_ok=False):
+    def save(self, labels=True, save_dir="runs/detect/exp", exist_ok=False):
         save_dir = increment_path(save_dir, exist_ok, mkdir=True)  # increment save_dir
         self._run(save=True, labels=labels, save_dir=save_dir)  # save results
 
-    def crop(self, save=True, save_dir='runs/detect/exp', exist_ok=False):
+    def crop(self, save=True, save_dir="runs/detect/exp", exist_ok=False):
         save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None
         return self._run(crop=True, save=save, save_dir=save_dir)  # crop results
 
@@ -818,9 +849,9 @@ class Detections:
     def pandas(self):
         # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
         new = copy(self)  # return copy
-        ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name'  # xyxy columns
+        ca = "xmin", "ymin", "xmax", "ymax", "confidence", "class", "name"  # xyxy columns
-        cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name'  # xywh columns
+        cb = "xcenter", "ycenter", "width", "height", "confidence", "class", "name"  # xywh columns
-        for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
+        for k, c in zip(["xyxy", "xyxyn", "xywh", "xywhn"], [ca, ca, cb, cb]):
             a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)]  # update
             setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
         return new
@@ -844,7 +875,7 @@ class Detections:
         return self._run(pprint=True)  # print results
 
     def __repr__(self):
-        return f'YOLOv5 {self.__class__} instance\n' + self.__str__()
+        return f"YOLOv5 {self.__class__} instance\n" + self.__str__()
 
 
 class Proto(nn.Module):
@@ -852,7 +883,7 @@ class Proto(nn.Module):
     def __init__(self, c1, c_=256, c2=32):  # ch_in, number of protos, number of masks
         super().__init__()
         self.cv1 = Conv(c1, c_, k=3)
-        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
+        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
         self.cv2 = Conv(c_, c_, k=3)
         self.cv3 = Conv(c_, c2)
 
@@ -862,14 +893,9 @@ class Proto(nn.Module):
 
 class Classify(nn.Module):
     # YOLOv5 classification head, i.e. x(b,c1,20,20) to x(b,c2)
-    def __init__(self,
+    def __init__(
-                 c1,
+        self, c1, c2, k=1, s=1, p=None, g=1, dropout_p=0.0
-                 c2,
+    ):  # ch_in, ch_out, kernel, stride, padding, groups, dropout probability
-                 k=1,
-                 s=1,
-                 p=None,
-                 g=1,
-                 dropout_p=0.0):  # ch_in, ch_out, kernel, stride, padding, groups, dropout probability
         super().__init__()
         c_ = 1280  # efficientnet_b0 size
         self.conv = Conv(c1, c_, k, s, autopad(k, p), g)

models/experimental.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Experimental modules."""
-Experimental modules
-"""
 import math
 
 import numpy as np
@@ -38,7 +36,7 @@ class MixConv2d(nn.Module):
         super().__init__()
         n = len(k)  # number of convolutions
         if equal_ch:  # equal c_ per group
-            i = torch.linspace(0, n - 1E-6, c2).floor()  # c2 indices
+            i = torch.linspace(0, n - 1e-6, c2).floor()  # c2 indices
             c_ = [(i == g).sum() for g in range(n)]  # intermediate channels
         else:  # equal weight.numel() per group
             b = [c2] + [0] * n
@@ -48,8 +46,9 @@ class MixConv2d(nn.Module):
             a[0] = 1
             c_ = np.linalg.lstsq(a, b, rcond=None)[0].round()  # solve for equal weight indices, ax = b
 
-        self.m = nn.ModuleList([
+        self.m = nn.ModuleList(
-            nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)])
+            [nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)]
+        )
         self.bn = nn.BatchNorm2d(c2)
         self.act = nn.SiLU()
 
@@ -76,16 +75,16 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
 
     model = Ensemble()
     for w in weights if isinstance(weights, list) else [weights]:
-        ckpt = torch.load(attempt_download(w), map_location='cpu')  # load
+        ckpt = torch.load(attempt_download(w), map_location="cpu")  # load
-        ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model
+        ckpt = (ckpt.get("ema") or ckpt["model"]).to(device).float()  # FP32 model
 
         # Model compatibility updates
-        if not hasattr(ckpt, 'stride'):
+        if not hasattr(ckpt, "stride"):
-            ckpt.stride = torch.tensor([32.])
+            ckpt.stride = torch.tensor([32.0])
-        if hasattr(ckpt, 'names') and isinstance(ckpt.names, (list, tuple)):
+        if hasattr(ckpt, "names") and isinstance(ckpt.names, (list, tuple)):
             ckpt.names = dict(enumerate(ckpt.names))  # convert to dict
 
-        model.append(ckpt.fuse().eval() if fuse and hasattr(ckpt, 'fuse') else ckpt.eval())  # model in eval mode
+        model.append(ckpt.fuse().eval() if fuse and hasattr(ckpt, "fuse") else ckpt.eval())  # model in eval mode
 
     # Module updates
     for m in model.modules():
@@ -93,9 +92,9 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
         if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model):
             m.inplace = inplace
             if t is Detect and not isinstance(m.anchor_grid, list):
-                delattr(m, 'anchor_grid')
+                delattr(m, "anchor_grid")
-                setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl)
+                setattr(m, "anchor_grid", [torch.zeros(1)] * m.nl)
-        elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
+        elif t is nn.Upsample and not hasattr(m, "recompute_scale_factor"):
             m.recompute_scale_factor = None  # torch 1.11.0 compatibility
 
     # Return model
@@ -103,9 +102,9 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
         return model[-1]
 
     # Return detection ensemble
-    print(f'Ensemble created with {weights}\n')
+    print(f"Ensemble created with {weights}\n")
-    for k in 'names', 'nc', 'yaml':
+    for k in "names", "nc", "yaml":
         setattr(model, k, getattr(model[0], k))
     model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride  # max stride
-    assert all(model[0].nc == m.nc for m in model), f'Models have different class counts: {[m.nc for m in model]}'
+    assert all(model[0].nc == m.nc for m in model), f"Models have different class counts: {[m.nc for m in model]}"
     return model

models/tf.py

@@ -27,8 +27,21 @@ import torch
 import torch.nn as nn
 from tensorflow import keras
 
-from models.common import (C3, SPP, SPPF, Bottleneck, BottleneckCSP, C3x, Concat, Conv, CrossConv, DWConv,
+from models.common import (
-                           DWConvTranspose2d, Focus, autopad)
+    C3,
+    SPP,
+    SPPF,
+    Bottleneck,
+    BottleneckCSP,
+    C3x,
+    Concat,
+    Conv,
+    CrossConv,
+    DWConv,
+    DWConvTranspose2d,
+    Focus,
+    autopad,
+)
 from models.experimental import MixConv2d, attempt_load
 from models.yolo import Detect, Segment
 from utils.activations import SiLU
@@ -44,7 +57,8 @@ class TFBN(keras.layers.Layer):
             gamma_initializer=keras.initializers.Constant(w.weight.numpy()),
             moving_mean_initializer=keras.initializers.Constant(w.running_mean.numpy()),
             moving_variance_initializer=keras.initializers.Constant(w.running_var.numpy()),
-            epsilon=w.eps)
+            epsilon=w.eps,
+        )
 
     def call(self, inputs):
         return self.bn(inputs)
@@ -60,7 +74,7 @@ class TFPad(keras.layers.Layer):
             self.pad = tf.constant([[0, 0], [pad[0], pad[0]], [pad[1], pad[1]], [0, 0]])
 
     def call(self, inputs):
-        return tf.pad(inputs, self.pad, mode='constant', constant_values=0)
+        return tf.pad(inputs, self.pad, mode="constant", constant_values=0)
 
 
 class TFConv(keras.layers.Layer):
@@ -75,12 +89,13 @@ class TFConv(keras.layers.Layer):
             filters=c2,
             kernel_size=k,
             strides=s,
-            padding='SAME' if s == 1 else 'VALID',
+            padding="SAME" if s == 1 else "VALID",
-            use_bias=not hasattr(w, 'bn'),
+            use_bias=not hasattr(w, "bn"),
             kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),
-            bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy()))
+            bias_initializer="zeros" if hasattr(w, "bn") else keras.initializers.Constant(w.conv.bias.numpy()),
+        )
         self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])
-        self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity
+        self.bn = TFBN(w.bn) if hasattr(w, "bn") else tf.identity
         self.act = activations(w.act) if act else tf.identity
 
     def call(self, inputs):
@@ -92,17 +107,18 @@ class TFDWConv(keras.layers.Layer):
     def __init__(self, c1, c2, k=1, s=1, p=None, act=True, w=None):
         # ch_in, ch_out, weights, kernel, stride, padding, groups
         super().__init__()
-        assert c2 % c1 == 0, f'TFDWConv() output={c2} must be a multiple of input={c1} channels'
+        assert c2 % c1 == 0, f"TFDWConv() output={c2} must be a multiple of input={c1} channels"
         conv = keras.layers.DepthwiseConv2D(
             kernel_size=k,
             depth_multiplier=c2 // c1,
             strides=s,
-            padding='SAME' if s == 1 else 'VALID',
+            padding="SAME" if s == 1 else "VALID",
-            use_bias=not hasattr(w, 'bn'),
+            use_bias=not hasattr(w, "bn"),
             depthwise_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),
-            bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy()))
+            bias_initializer="zeros" if hasattr(w, "bn") else keras.initializers.Constant(w.conv.bias.numpy()),
+        )
         self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])
-        self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity
+        self.bn = TFBN(w.bn) if hasattr(w, "bn") else tf.identity
         self.act = activations(w.act) if act else tf.identity
 
     def call(self, inputs):
@@ -114,19 +130,23 @@ class TFDWConvTranspose2d(keras.layers.Layer):
     def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0, w=None):
         # ch_in, ch_out, weights, kernel, stride, padding, groups
         super().__init__()
-        assert c1 == c2, f'TFDWConv() output={c2} must be equal to input={c1} channels'
+        assert c1 == c2, f"TFDWConv() output={c2} must be equal to input={c1} channels"
-        assert k == 4 and p1 == 1, 'TFDWConv() only valid for k=4 and p1=1'
+        assert k == 4 and p1 == 1, "TFDWConv() only valid for k=4 and p1=1"
         weight, bias = w.weight.permute(2, 3, 1, 0).numpy(), w.bias.numpy()
         self.c1 = c1
         self.conv = [
-            keras.layers.Conv2DTranspose(filters=1,
+            keras.layers.Conv2DTranspose(
+                filters=1,
                 kernel_size=k,
                 strides=s,
-                                         padding='VALID',
+                padding="VALID",
                 output_padding=p2,
                 use_bias=True,
                 kernel_initializer=keras.initializers.Constant(weight[..., i : i + 1]),
-                                         bias_initializer=keras.initializers.Constant(bias[i])) for i in range(c1)]
+                bias_initializer=keras.initializers.Constant(bias[i]),
+            )
+            for i in range(c1)
+        ]
 
     def call(self, inputs):
         return tf.concat([m(x) for m, x in zip(self.conv, tf.split(inputs, self.c1, 3))], 3)[:, 1:-1, 1:-1]
@@ -176,14 +196,15 @@ class TFConv2d(keras.layers.Layer):
     def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None):
         super().__init__()
         assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
-        self.conv = keras.layers.Conv2D(filters=c2,
+        self.conv = keras.layers.Conv2D(
+            filters=c2,
             kernel_size=k,
             strides=s,
-                                        padding='VALID',
+            padding="VALID",
             use_bias=bias,
-                                        kernel_initializer=keras.initializers.Constant(
+            kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()),
-                                            w.weight.permute(2, 3, 1, 0).numpy()),
+            bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None,
-                                        bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None)
+        )
 
     def call(self, inputs):
         return self.conv(inputs)
@@ -233,8 +254,9 @@ class TFC3x(keras.layers.Layer):
         self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
         self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2)
         self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3)
-        self.m = keras.Sequential([
+        self.m = keras.Sequential(
-            TFCrossConv(c_, c_, k=3, s=1, g=g, e=1.0, shortcut=shortcut, w=w.m[j]) for j in range(n)])
+            [TFCrossConv(c_, c_, k=3, s=1, g=g, e=1.0, shortcut=shortcut, w=w.m[j]) for j in range(n)]
+        )
 
     def call(self, inputs):
         return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3))
@@ -247,7 +269,7 @@ class TFSPP(keras.layers.Layer):
         c_ = c1 // 2  # hidden channels
         self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
         self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2)
-        self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding='SAME') for x in k]
+        self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding="SAME") for x in k]
 
     def call(self, inputs):
         x = self.cv1(inputs)
@@ -261,7 +283,7 @@ class TFSPPF(keras.layers.Layer):
         c_ = c1 // 2  # hidden channels
         self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
         self.cv2 = TFConv(c_ * 4, c2, 1, 1, w=w.cv2)
-        self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding='SAME')
+        self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding="SAME")
 
     def call(self, inputs):
         x = self.cv1(inputs)
@@ -340,11 +362,10 @@ class TFSegment(TFDetect):
 
 
 class TFProto(keras.layers.Layer):
-
     def __init__(self, c1, c_=256, c2=32, w=None):
         super().__init__()
         self.cv1 = TFConv(c1, c_, k=3, w=w.cv1)
-        self.upsample = TFUpsample(None, scale_factor=2, mode='nearest')
+        self.upsample = TFUpsample(None, scale_factor=2, mode="nearest")
         self.cv2 = TFConv(c_, c_, k=3, w=w.cv2)
         self.cv3 = TFConv(c_, c2, w=w.cv3)
 
@@ -356,7 +377,7 @@ class TFUpsample(keras.layers.Layer):
     # TF version of torch.nn.Upsample()
     def __init__(self, size, scale_factor, mode, w=None):  # warning: all arguments needed including 'w'
         super().__init__()
-        assert scale_factor % 2 == 0, 'scale_factor must be multiple of 2'
+        assert scale_factor % 2 == 0, "scale_factor must be multiple of 2"
         self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * scale_factor, x.shape[2] * scale_factor), mode)
         # self.upsample = keras.layers.UpSampling2D(size=scale_factor, interpolation=mode)
         # with default arguments: align_corners=False, half_pixel_centers=False
@@ -371,7 +392,7 @@ class TFConcat(keras.layers.Layer):
     # TF version of torch.concat()
     def __init__(self, dimension=1, w=None):
         super().__init__()
-        assert dimension == 1, 'convert only NCHW to NHWC concat'
+        assert dimension == 1, "convert only NCHW to NHWC concat"
         self.d = 3
 
     def call(self, inputs):
@@ -380,15 +401,20 @@ class TFConcat(keras.layers.Layer):
 
 def parse_model(d, ch, model, imgsz):  # model_dict, input_channels(3)
     LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
-    anchors, nc, gd, gw, ch_mul = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get(
+    anchors, nc, gd, gw, ch_mul = (
-        'channel_multiple')
+        d["anchors"],
+        d["nc"],
+        d["depth_multiple"],
+        d["width_multiple"],
+        d.get("channel_multiple"),
+    )
     na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
     no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
     if not ch_mul:
         ch_mul = 8
 
     layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
-    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
+    for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
         m_str = m
         m = eval(m) if isinstance(m, str) else m  # eval strings
         for j, a in enumerate(args):
@@ -399,8 +425,20 @@ def parse_model(d, ch, model, imgsz):  # model_dict, input_channels(3)
 
         n = max(round(n * gd), 1) if n > 1 else n  # depth gain
         if m in [
-                nn.Conv2d, Conv, DWConv, DWConvTranspose2d, Bottleneck, SPP, SPPF, MixConv2d, Focus, CrossConv,
+            nn.Conv2d,
-                BottleneckCSP, C3, C3x]:
+            Conv,
+            DWConv,
+            DWConvTranspose2d,
+            Bottleneck,
+            SPP,
+            SPPF,
+            MixConv2d,
+            Focus,
+            CrossConv,
+            BottleneckCSP,
+            C3,
+            C3x,
+        ]:
             c1, c2 = ch[f], args[0]
             c2 = make_divisible(c2 * gw, ch_mul) if c2 != no else c2
 
@@ -422,15 +460,18 @@ def parse_model(d, ch, model, imgsz):  # model_dict, input_channels(3)
         else:
             c2 = ch[f]
 
-        tf_m = eval('TF' + m_str.replace('nn.', ''))
+        tf_m = eval("TF" + m_str.replace("nn.", ""))
-        m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) if n > 1 \
+        m_ = (
-            else tf_m(*args, w=model.model[i])  # module
+            keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)])
+            if n > 1
+            else tf_m(*args, w=model.model[i])
+        )  # module
 
         torch_m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
-        t = str(m)[8:-2].replace('__main__.', '')  # module type
+        t = str(m)[8:-2].replace("__main__.", "")  # module type
         np = sum(x.numel() for x in torch_m_.parameters())  # number params
         m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
-        LOGGER.info(f'{i:>3}{str(f):>18}{str(n):>3}{np:>10}  {t:<40}{str(args):<30}')  # print
+        LOGGER.info(f"{i:>3}{str(f):>18}{str(n):>3}{np:>10}  {t:<40}{str(args):<30}")  # print
         save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
         layers.append(m_)
         ch.append(c2)
@@ -439,30 +480,33 @@ def parse_model(d, ch, model, imgsz):  # model_dict, input_channels(3)
 
 class TFModel:
     # TF YOLOv5 model
-    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz=(640, 640)):  # model, channels, classes
+    def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, model=None, imgsz=(640, 640)):  # model, channels, classes
         super().__init__()
         if isinstance(cfg, dict):
             self.yaml = cfg  # model dict
         else:  # is *.yaml
             import yaml  # for torch hub
+
             self.yaml_file = Path(cfg).name
             with open(cfg) as f:
                 self.yaml = yaml.load(f, Loader=yaml.FullLoader)  # model dict
 
         # Define model
-        if nc and nc != self.yaml['nc']:
+        if nc and nc != self.yaml["nc"]:
             LOGGER.info(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}")
-            self.yaml['nc'] = nc  # override yaml value
+            self.yaml["nc"] = nc  # override yaml value
         self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz)
 
-    def predict(self,
+    def predict(
+        self,
         inputs,
         tf_nms=False,
         agnostic_nms=False,
         topk_per_class=100,
         topk_all=100,
         iou_thres=0.45,
-                conf_thres=0.25):
+        conf_thres=0.25,
+    ):
         y = []  # outputs
         x = inputs
         for m in self.model.layers:
@@ -482,13 +526,9 @@ class TFModel:
                 nms = AgnosticNMS()((boxes, classes, scores), topk_all, iou_thres, conf_thres)
             else:
                 boxes = tf.expand_dims(boxes, 2)
-                nms = tf.image.combined_non_max_suppression(boxes,
+                nms = tf.image.combined_non_max_suppression(
-                                                            scores,
+                    boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False
-                                                            topk_per_class,
+                )
-                                                            topk_all,
-                                                            iou_thres,
-                                                            conf_thres,
-                                                            clip_boxes=False)
             return (nms,)
         return x  # output [1,6300,85] = [xywh, conf, class0, class1, ...]
         # x = x[0]  # [x(1,6300,85), ...] to x(6300,85)
@@ -508,36 +548,42 @@ class AgnosticNMS(keras.layers.Layer):
     # TF Agnostic NMS
     def call(self, input, topk_all, iou_thres, conf_thres):
         # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450
-        return tf.map_fn(lambda x: self._nms(x, topk_all, iou_thres, conf_thres),
+        return tf.map_fn(
+            lambda x: self._nms(x, topk_all, iou_thres, conf_thres),
             input,
             fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32),
-                         name='agnostic_nms')
+            name="agnostic_nms",
+        )
 
     @staticmethod
     def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):  # agnostic NMS
         boxes, classes, scores = x
         class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
         scores_inp = tf.reduce_max(scores, -1)
-        selected_inds = tf.image.non_max_suppression(boxes,
+        selected_inds = tf.image.non_max_suppression(
-                                                     scores_inp,
+            boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres
-                                                     max_output_size=topk_all,
+        )
-                                                     iou_threshold=iou_thres,
-                                                     score_threshold=conf_thres)
         selected_boxes = tf.gather(boxes, selected_inds)
-        padded_boxes = tf.pad(selected_boxes,
+        padded_boxes = tf.pad(
+            selected_boxes,
             paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
-                              mode='CONSTANT',
+            mode="CONSTANT",
-                              constant_values=0.0)
+            constant_values=0.0,
+        )
         selected_scores = tf.gather(scores_inp, selected_inds)
-        padded_scores = tf.pad(selected_scores,
+        padded_scores = tf.pad(
+            selected_scores,
             paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
-                               mode='CONSTANT',
+            mode="CONSTANT",
-                               constant_values=-1.0)
+            constant_values=-1.0,
+        )
         selected_classes = tf.gather(class_inds, selected_inds)
-        padded_classes = tf.pad(selected_classes,
+        padded_classes = tf.pad(
+            selected_classes,
             paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
-                                mode='CONSTANT',
+            mode="CONSTANT",
-                                constant_values=-1.0)
+            constant_values=-1.0,
+        )
         valid_detections = tf.shape(selected_inds)[0]
         return padded_boxes, padded_scores, padded_classes, valid_detections
 
@@ -551,7 +597,7 @@ def activations(act=nn.SiLU):
     elif isinstance(act, (nn.SiLU, SiLU)):
         return lambda x: keras.activations.swish(x)
     else:
-        raise Exception(f'no matching TensorFlow activation found for PyTorch activation {act}')
+        raise Exception(f"no matching TensorFlow activation found for PyTorch activation {act}")
 
 
 def representative_dataset_gen(dataset, ncalib=100):
@@ -566,14 +612,14 @@ def representative_dataset_gen(dataset, ncalib=100):
 
 
 def run(
-        weights=ROOT / 'yolov5s.pt',  # weights path
+    weights=ROOT / "yolov5s.pt",  # weights path
     imgsz=(640, 640),  # inference size h,w
     batch_size=1,  # batch size
     dynamic=False,  # dynamic batch size
 ):
     # PyTorch model
     im = torch.zeros((batch_size, 3, *imgsz))  # BCHW image
-    model = attempt_load(weights, device=torch.device('cpu'), inplace=True, fuse=False)
+    model = attempt_load(weights, device=torch.device("cpu"), inplace=True, fuse=False)
     _ = model(im)  # inference
     model.info()
 
@@ -587,15 +633,15 @@ def run(
     keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im))
     keras_model.summary()
 
-    LOGGER.info('PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.')
+    LOGGER.info("PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.")
 
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
+    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path")
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
-    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
+    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
-    parser.add_argument('--dynamic', action='store_true', help='dynamic batch size')
+    parser.add_argument("--dynamic", action="store_true", help="dynamic batch size")
     opt = parser.parse_args()
     opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
     print_args(vars(opt))
@@ -606,6 +652,6 @@ def main(opt):
     run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

models/yolo.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-YOLO-specific modules
+YOLO-specific modules.
 
 Usage:
     $ python models/yolo.py --cfg yolov5s.yaml
@@ -22,18 +22,46 @@ FILE = Path(__file__).resolve()
 ROOT = FILE.parents[1]  # YOLOv5 root directory
 if str(ROOT) not in sys.path:
     sys.path.append(str(ROOT))  # add ROOT to PATH
-if platform.system() != 'Windows':
+if platform.system() != "Windows":
     ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 
-from models.common import (C3, C3SPP, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C3Ghost, C3x, Classify, Concat,
+from models.common import (
-                           Contract, Conv, CrossConv, DetectMultiBackend, DWConv, DWConvTranspose2d, Expand, Focus,
+    C3,
-                           GhostBottleneck, GhostConv, Proto)
+    C3SPP,
+    C3TR,
+    SPP,
+    SPPF,
+    Bottleneck,
+    BottleneckCSP,
+    C3Ghost,
+    C3x,
+    Classify,
+    Concat,
+    Contract,
+    Conv,
+    CrossConv,
+    DetectMultiBackend,
+    DWConv,
+    DWConvTranspose2d,
+    Expand,
+    Focus,
+    GhostBottleneck,
+    GhostConv,
+    Proto,
+)
 from models.experimental import MixConv2d
 from utils.autoanchor import check_anchor_order
 from utils.general import LOGGER, check_version, check_yaml, colorstr, make_divisible, print_args
 from utils.plots import feature_visualization
-from utils.torch_utils import (fuse_conv_and_bn, initialize_weights, model_info, profile, scale_img, select_device,
+from utils.torch_utils import (
-                               time_sync)
+    fuse_conv_and_bn,
+    initialize_weights,
+    model_info,
+    profile,
+    scale_img,
+    select_device,
+    time_sync,
+)
 
 try:
     import thop  # for FLOPs computation
@@ -55,7 +83,7 @@ class Detect(nn.Module):
         self.na = len(anchors[0]) // 2  # number of anchors
         self.grid = [torch.empty(0) for _ in range(self.nl)]  # init grid
         self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]  # init anchor grid
-        self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2))  # shape(nl,na,2)
+        self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))  # shape(nl,na,2)
         self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
         self.inplace = inplace  # use inplace ops (e.g. slice assignment)
 
@@ -84,12 +112,12 @@ class Detect(nn.Module):
 
         return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
 
-    def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, '1.10.0')):
+    def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
         d = self.anchors[i].device
         t = self.anchors[i].dtype
         shape = 1, self.na, ny, nx, 2  # grid shape
         y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
-        yv, xv = torch.meshgrid(y, x, indexing='ij') if torch_1_10 else torch.meshgrid(y, x)  # torch>=0.7 compatibility
+        yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)  # torch>=0.7 compatibility
         grid = torch.stack((xv, yv), 2).expand(shape) - 0.5  # add grid offset, i.e. y = 2.0 * x - 0.5
         anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
         return grid, anchor_grid
@@ -132,23 +160,23 @@ class BaseModel(nn.Module):
 
     def _profile_one_layer(self, m, x, dt):
         c = m == self.model[-1]  # is final layer, copy input as inplace fix
-        o = thop.profile(m, inputs=(x.copy() if c else x, ), verbose=False)[0] / 1E9 * 2 if thop else 0  # FLOPs
+        o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1e9 * 2 if thop else 0  # FLOPs
         t = time_sync()
         for _ in range(10):
             m(x.copy() if c else x)
         dt.append((time_sync() - t) * 100)
         if m == self.model[0]:
             LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s}  module")
-        LOGGER.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}')
+        LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}")
         if c:
             LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")
 
     def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
-        LOGGER.info('Fusing layers... ')
+        LOGGER.info("Fusing layers... ")
         for m in self.model.modules():
-            if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
+            if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"):
                 m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
-                delattr(m, 'bn')  # remove batchnorm
+                delattr(m, "bn")  # remove batchnorm
                 m.forward = m.forward_fuse  # update forward
         self.info()
         return self
@@ -170,27 +198,28 @@ class BaseModel(nn.Module):
 
 class DetectionModel(BaseModel):
     # YOLOv5 detection model
-    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None):  # model, input channels, number of classes
+    def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None):  # model, input channels, number of classes
         super().__init__()
         if isinstance(cfg, dict):
             self.yaml = cfg  # model dict
         else:  # is *.yaml
             import yaml  # for torch hub
+
             self.yaml_file = Path(cfg).name
-            with open(cfg, encoding='ascii', errors='ignore') as f:
+            with open(cfg, encoding="ascii", errors="ignore") as f:
                 self.yaml = yaml.safe_load(f)  # model dict
 
         # Define model
-        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels
+        ch = self.yaml["ch"] = self.yaml.get("ch", ch)  # input channels
-        if nc and nc != self.yaml['nc']:
+        if nc and nc != self.yaml["nc"]:
             LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
-            self.yaml['nc'] = nc  # override yaml value
+            self.yaml["nc"] = nc  # override yaml value
         if anchors:
-            LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
+            LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}")
-            self.yaml['anchors'] = round(anchors)  # override yaml value
+            self.yaml["anchors"] = round(anchors)  # override yaml value
         self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
-        self.names = [str(i) for i in range(self.yaml['nc'])]  # default names
+        self.names = [str(i) for i in range(self.yaml["nc"])]  # default names
-        self.inplace = self.yaml.get('inplace', True)
+        self.inplace = self.yaml.get("inplace", True)
 
         # Build strides, anchors
         m = self.model[-1]  # Detect()
@@ -207,7 +236,7 @@ class DetectionModel(BaseModel):
         # Init weights, biases
         initialize_weights(self)
         self.info()
-        LOGGER.info('')
+        LOGGER.info("")
 
     def forward(self, x, augment=False, profile=False, visualize=False):
         if augment:
@@ -263,7 +292,9 @@ class DetectionModel(BaseModel):
         for mi, s in zip(m.m, m.stride):  # from
             b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
             b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
-            b.data[:, 5:5 + m.nc] += math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())  # cls
+            b.data[:, 5 : 5 + m.nc] += (
+                math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
+            )  # cls
             mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
 
 
@@ -272,7 +303,7 @@ Model = DetectionModel  # retain YOLOv5 'Model' class for backwards compatibilit
 
 class SegmentationModel(DetectionModel):
     # YOLOv5 segmentation model
-    def __init__(self, cfg='yolov5s-seg.yaml', ch=3, nc=None, anchors=None):
+    def __init__(self, cfg="yolov5s-seg.yaml", ch=3, nc=None, anchors=None):
         super().__init__(cfg, ch, nc, anchors)
 
 
@@ -288,9 +319,9 @@ class ClassificationModel(BaseModel):
             model = model.model  # unwrap DetectMultiBackend
         model.model = model.model[:cutoff]  # backbone
         m = model.model[-1]  # last layer
-        ch = m.conv.in_channels if hasattr(m, 'conv') else m.cv1.conv.in_channels  # ch into module
+        ch = m.conv.in_channels if hasattr(m, "conv") else m.cv1.conv.in_channels  # ch into module
         c = Classify(ch, nc)  # Classify()
-        c.i, c.f, c.type = m.i, m.f, 'models.common.Classify'  # index, from, type
+        c.i, c.f, c.type = m.i, m.f, "models.common.Classify"  # index, from, type
         model.model[-1] = c  # replace
         self.model = model.model
         self.stride = model.stride
@@ -305,8 +336,14 @@ class ClassificationModel(BaseModel):
 def parse_model(d, ch):  # model_dict, input_channels(3)
     # Parse a YOLOv5 model.yaml dictionary
     LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
-    anchors, nc, gd, gw, act, ch_mul = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get(
+    anchors, nc, gd, gw, act, ch_mul = (
-        'activation'), d.get('channel_multiple')
+        d["anchors"],
+        d["nc"],
+        d["depth_multiple"],
+        d["width_multiple"],
+        d.get("activation"),
+        d.get("channel_multiple"),
+    )
     if act:
         Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
         LOGGER.info(f"{colorstr('activation:')} {act}")  # print
@@ -316,7 +353,7 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
     no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
 
     layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
-    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
+    for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
         m = eval(m) if isinstance(m, str) else m  # eval strings
         for j, a in enumerate(args):
             with contextlib.suppress(NameError):
@@ -324,8 +361,25 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
 
         n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
         if m in {
-                Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
+            Conv,
-                BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
+            GhostConv,
+            Bottleneck,
+            GhostBottleneck,
+            SPP,
+            SPPF,
+            DWConv,
+            MixConv2d,
+            Focus,
+            CrossConv,
+            BottleneckCSP,
+            C3,
+            C3TR,
+            C3SPP,
+            C3Ghost,
+            nn.ConvTranspose2d,
+            DWConvTranspose2d,
+            C3x,
+        }:
             c1, c2 = ch[f], args[0]
             if c2 != no:  # if not output
                 c2 = make_divisible(c2 * gw, ch_mul)
@@ -353,10 +407,10 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
             c2 = ch[f]
 
         m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
-        t = str(m)[8:-2].replace('__main__.', '')  # module type
+        t = str(m)[8:-2].replace("__main__.", "")  # module type
         np = sum(x.numel() for x in m_.parameters())  # number params
         m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
-        LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}')  # print
+        LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}")  # print
         save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
         layers.append(m_)
         if i == 0:
@@ -365,14 +419,14 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
     return nn.Sequential(*layers), sorted(save)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
+    parser.add_argument("--cfg", type=str, default="yolov5s.yaml", help="model.yaml")
-    parser.add_argument('--batch-size', type=int, default=1, help='total batch size for all GPUs')
+    parser.add_argument("--batch-size", type=int, default=1, help="total batch size for all GPUs")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--profile', action='store_true', help='profile model speed')
+    parser.add_argument("--profile", action="store_true", help="profile model speed")
-    parser.add_argument('--line-profile', action='store_true', help='profile model speed layer by layer')
+    parser.add_argument("--line-profile", action="store_true", help="profile model speed layer by layer")
-    parser.add_argument('--test', action='store_true', help='test all yolo*.yaml')
+    parser.add_argument("--test", action="store_true", help="test all yolo*.yaml")
     opt = parser.parse_args()
     opt.cfg = check_yaml(opt.cfg)  # check YAML
     print_args(vars(opt))
@@ -390,11 +444,11 @@ if __name__ == '__main__':
         results = profile(input=im, ops=[model], n=3)
 
     elif opt.test:  # test all models
-        for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'):
+        for cfg in Path(ROOT / "models").rglob("yolo*.yaml"):
             try:
                 _ = Model(cfg)
             except Exception as e:
-                print(f'Error in {cfg}: {e}')
+                print(f"Error in {cfg}: {e}")
 
     else:  # report fused model summary
         model.fuse()

segment/predict.py

@@ -46,23 +46,36 @@ from ultralytics.utils.plotting import Annotator, colors, save_one_box
 
 from models.common import DetectMultiBackend
 from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
-from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
+from utils.general import (
-                           increment_path, non_max_suppression, print_args, scale_boxes, scale_segments,
+    LOGGER,
-                           strip_optimizer)
+    Profile,
+    check_file,
+    check_img_size,
+    check_imshow,
+    check_requirements,
+    colorstr,
+    cv2,
+    increment_path,
+    non_max_suppression,
+    print_args,
+    scale_boxes,
+    scale_segments,
+    strip_optimizer,
+)
 from utils.segment.general import masks2segments, process_mask, process_mask_native
 from utils.torch_utils import select_device, smart_inference_mode
 
 
 @smart_inference_mode()
 def run(
-    weights=ROOT / 'yolov5s-seg.pt',  # model.pt path(s)
+    weights=ROOT / "yolov5s-seg.pt",  # model.pt path(s)
-    source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
+    source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
-    data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
     imgsz=(640, 640),  # inference size (height, width)
     conf_thres=0.25,  # confidence threshold
     iou_thres=0.45,  # NMS IOU threshold
     max_det=1000,  # maximum detections per image
-    device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     view_img=False,  # show results
     save_txt=False,  # save results to *.txt
     save_conf=False,  # save confidences in --save-txt labels
@@ -73,8 +86,8 @@ def run(
     augment=False,  # augmented inference
     visualize=False,  # visualize features
     update=False,  # update all models
-    project=ROOT / 'runs/predict-seg',  # save results to project/name
+    project=ROOT / "runs/predict-seg",  # save results to project/name
-    name='exp',  # save results to project/name
+    name="exp",  # save results to project/name
     exist_ok=False,  # existing project/name ok, do not increment
     line_thickness=3,  # bounding box thickness (pixels)
     hide_labels=False,  # hide labels
@@ -85,17 +98,17 @@ def run(
     retina_masks=False,
 ):
     source = str(source)
-    save_img = not nosave and not source.endswith('.txt')  # save inference images
+    save_img = not nosave and not source.endswith(".txt")  # save inference images
     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
-    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
+    is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
-    webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
+    webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
-    screenshot = source.lower().startswith('screen')
+    screenshot = source.lower().startswith("screen")
     if is_url and is_file:
         source = check_file(source)  # download
 
     # Directories
     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+    (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
     # Load model
     device = select_device(device)
@@ -143,14 +156,14 @@ def run(
             seen += 1
             if webcam:  # batch_size >= 1
                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
-                s += f'{i}: '
+                s += f"{i}: "
             else:
-                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
+                p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
 
             p = Path(p)  # to Path
             save_path = str(save_dir / p.name)  # im.jpg
-            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
+            txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
-            s += '%gx%g ' % im.shape[2:]  # print string
+            s += "%gx%g " % im.shape[2:]  # print string
             imc = im0.copy() if save_crop else im0  # for save_crop
             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
             if len(det):
@@ -166,7 +179,8 @@ def run(
                 if save_txt:
                     segments = [
                         scale_segments(im0.shape if retina_masks else im.shape[2:], x, im0.shape, normalize=True)
-                        for x in reversed(masks2segments(masks))]
+                        for x in reversed(masks2segments(masks))
+                    ]
 
                 # Print results
                 for c in det[:, 5].unique():
@@ -177,39 +191,42 @@ def run(
                 annotator.masks(
                     masks,
                     colors=[colors(x, True) for x in det[:, 5]],
-                    im_gpu=torch.as_tensor(im0, dtype=torch.float16).to(device).permute(2, 0, 1).flip(0).contiguous() /
+                    im_gpu=torch.as_tensor(im0, dtype=torch.float16).to(device).permute(2, 0, 1).flip(0).contiguous()
-                    255 if retina_masks else im[i])
+                    / 255
+                    if retina_masks
+                    else im[i],
+                )
 
                 # Write results
                 for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
                     if save_txt:  # Write to file
                         seg = segments[j].reshape(-1)  # (n,2) to (n*2)
                         line = (cls, *seg, conf) if save_conf else (cls, *seg)  # label format
-                        with open(f'{txt_path}.txt', 'a') as f:
+                        with open(f"{txt_path}.txt", "a") as f:
-                            f.write(('%g ' * len(line)).rstrip() % line + '\n')
+                            f.write(("%g " * len(line)).rstrip() % line + "\n")
 
                     if save_img or save_crop or view_img:  # Add bbox to image
                         c = int(cls)  # integer class
-                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
+                        label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}")
                         annotator.box_label(xyxy, label, color=colors(c, True))
                         # annotator.draw.polygon(segments[j], outline=colors(c, True), width=3)
                     if save_crop:
-                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
+                        save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True)
 
             # Stream results
             im0 = annotator.result()
             if view_img:
-                if platform.system() == 'Linux' and p not in windows:
+                if platform.system() == "Linux" and p not in windows:
                     windows.append(p)
                     cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                     cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
                 cv2.imshow(str(p), im0)
-                if cv2.waitKey(1) == ord('q'):  # 1 millisecond
+                if cv2.waitKey(1) == ord("q"):  # 1 millisecond
                     exit()
 
             # Save results (image with detections)
             if save_img:
-                if dataset.mode == 'image':
+                if dataset.mode == "image":
                     cv2.imwrite(save_path, im0)
                 else:  # 'video' or 'stream'
                     if vid_path[i] != save_path:  # new video
@@ -222,18 +239,18 @@ def run(
                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                         else:  # stream
                             fps, w, h = 30, im0.shape[1], im0.shape[0]
-                        save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
+                        save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
-                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                     vid_writer[i].write(im0)
 
         # Print time (inference-only)
         LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
 
     # Print results
-    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
-    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
+    LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
     if save_txt or save_img:
-        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     if update:
         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
@@ -241,34 +258,34 @@ def run(
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-seg.pt', help='model path(s)')
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-seg.pt", help="model path(s)")
-    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
+    parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
-    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
+    parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
-    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
+    parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
-    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
+    parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--view-img', action='store_true', help='show results')
+    parser.add_argument("--view-img", action="store_true", help="show results")
-    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
+    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
-    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
+    parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
-    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
+    parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
-    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
+    parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
-    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
+    parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
-    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
+    parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
-    parser.add_argument('--augment', action='store_true', help='augmented inference')
+    parser.add_argument("--augment", action="store_true", help="augmented inference")
-    parser.add_argument('--visualize', action='store_true', help='visualize features')
+    parser.add_argument("--visualize", action="store_true", help="visualize features")
-    parser.add_argument('--update', action='store_true', help='update all models')
+    parser.add_argument("--update", action="store_true", help="update all models")
-    parser.add_argument('--project', default=ROOT / 'runs/predict-seg', help='save results to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/predict-seg", help="save results to project/name")
-    parser.add_argument('--name', default='exp', help='save results to project/name')
+    parser.add_argument("--name", default="exp", help="save results to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
+    parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
-    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
+    parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
-    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
+    parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
-    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
+    parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
-    parser.add_argument('--retina-masks', action='store_true', help='whether to plot masks in native resolution')
+    parser.add_argument("--retina-masks", action="store_true", help="whether to plot masks in native resolution")
     opt = parser.parse_args()
     opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
     print_args(vars(opt))
@@ -276,10 +293,10 @@ def parse_opt():
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

segment/train.py

@@ -1,7 +1,7 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Train a YOLOv5 segment model on a segment dataset
+Train a YOLOv5 segment model on a segment dataset Models and datasets download automatically from the latest YOLOv5
-Models and datasets download automatically from the latest YOLOv5 release.
+release.
 
 Usage - Single-GPU training:
     $ python segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640  # from pretrained (recommended)
@@ -47,47 +47,104 @@ from utils.autoanchor import check_anchors
 from utils.autobatch import check_train_batch_size
 from utils.callbacks import Callbacks
 from utils.downloads import attempt_download, is_url
-from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
+from utils.general import (
-                           check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,
+    LOGGER,
-                           get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
+    TQDM_BAR_FORMAT,
-                           labels_to_image_weights, one_cycle, print_args, print_mutation, strip_optimizer, yaml_save)
+    check_amp,
+    check_dataset,
+    check_file,
+    check_git_info,
+    check_git_status,
+    check_img_size,
+    check_requirements,
+    check_suffix,
+    check_yaml,
+    colorstr,
+    get_latest_run,
+    increment_path,
+    init_seeds,
+    intersect_dicts,
+    labels_to_class_weights,
+    labels_to_image_weights,
+    one_cycle,
+    print_args,
+    print_mutation,
+    strip_optimizer,
+    yaml_save,
+)
 from utils.loggers import GenericLogger
 from utils.plots import plot_evolve, plot_labels
 from utils.segment.dataloaders import create_dataloader
 from utils.segment.loss import ComputeLoss
 from utils.segment.metrics import KEYS, fitness
 from utils.segment.plots import plot_images_and_masks, plot_results_with_masks
-from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
+from utils.torch_utils import (
-                               smart_resume, torch_distributed_zero_first)
+    EarlyStopping,
+    ModelEMA,
+    de_parallel,
+    select_device,
+    smart_DDP,
+    smart_optimizer,
+    smart_resume,
+    torch_distributed_zero_first,
+)
 
-LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
+LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
-WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
+WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
 GIT_INFO = check_git_info()
 
 
 def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictionary
-    save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, mask_ratio = \
+    (
-        Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
+        save_dir,
-        opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze, opt.mask_ratio
+        epochs,
+        batch_size,
+        weights,
+        single_cls,
+        evolve,
+        data,
+        cfg,
+        resume,
+        noval,
+        nosave,
+        workers,
+        freeze,
+        mask_ratio,
+    ) = (
+        Path(opt.save_dir),
+        opt.epochs,
+        opt.batch_size,
+        opt.weights,
+        opt.single_cls,
+        opt.evolve,
+        opt.data,
+        opt.cfg,
+        opt.resume,
+        opt.noval,
+        opt.nosave,
+        opt.workers,
+        opt.freeze,
+        opt.mask_ratio,
+    )
     # callbacks.run('on_pretrain_routine_start')
 
     # Directories
-    w = save_dir / 'weights'  # weights dir
+    w = save_dir / "weights"  # weights dir
     (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
-    last, best = w / 'last.pt', w / 'best.pt'
+    last, best = w / "last.pt", w / "best.pt"
 
     # Hyperparameters
     if isinstance(hyp, str):
-        with open(hyp, errors='ignore') as f:
+        with open(hyp, errors="ignore") as f:
             hyp = yaml.safe_load(f)  # load hyps dict
-    LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
+    LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()))
     opt.hyp = hyp.copy()  # for saving hyps to checkpoints
 
     # Save run settings
     if not evolve:
-        yaml_save(save_dir / 'hyp.yaml', hyp)
+        yaml_save(save_dir / "hyp.yaml", hyp)
-        yaml_save(save_dir / 'opt.yaml', vars(opt))
+        yaml_save(save_dir / "opt.yaml", vars(opt))
 
     # Loggers
     data_dict = None
@@ -97,39 +154,39 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # Config
     plots = not evolve and not opt.noplots  # create plots
     overlap = not opt.no_overlap
-    cuda = device.type != 'cpu'
+    cuda = device.type != "cpu"
     init_seeds(opt.seed + 1 + RANK, deterministic=True)
     with torch_distributed_zero_first(LOCAL_RANK):
         data_dict = data_dict or check_dataset(data)  # check if None
-    train_path, val_path = data_dict['train'], data_dict['val']
+    train_path, val_path = data_dict["train"], data_dict["val"]
-    nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
+    nc = 1 if single_cls else int(data_dict["nc"])  # number of classes
-    names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
+    names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"]  # class names
-    is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')  # COCO dataset
+    is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt")  # COCO dataset
 
     # Model
-    check_suffix(weights, '.pt')  # check weights
+    check_suffix(weights, ".pt")  # check weights
-    pretrained = weights.endswith('.pt')
+    pretrained = weights.endswith(".pt")
     if pretrained:
         with torch_distributed_zero_first(LOCAL_RANK):
             weights = attempt_download(weights)  # download if not found locally
-        ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
+        ckpt = torch.load(weights, map_location="cpu")  # load checkpoint to CPU to avoid CUDA memory leak
-        model = SegmentationModel(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)
+        model = SegmentationModel(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)
-        exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []  # exclude keys
+        exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else []  # exclude keys
-        csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
+        csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
         csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
         model.load_state_dict(csd, strict=False)  # load
-        LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  # report
+        LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}")  # report
     else:
-        model = SegmentationModel(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
+        model = SegmentationModel(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create
     amp = check_amp(model)  # check AMP
 
     # Freeze
-    freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
+    freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
     for k, v in model.named_parameters():
         v.requires_grad = True  # train all layers
         # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
         if any(x in k for x in freeze):
-            LOGGER.info(f'freezing {k}')
+            LOGGER.info(f"freezing {k}")
             v.requires_grad = False
 
     # Image size
@@ -139,20 +196,20 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # Batch size
     if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
         batch_size = check_train_batch_size(model, imgsz, amp)
-        logger.update_params({'batch_size': batch_size})
+        logger.update_params({"batch_size": batch_size})
         # loggers.on_params_update({"batch_size": batch_size})
 
     # Optimizer
     nbs = 64  # nominal batch size
     accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
-    hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
+    hyp["weight_decay"] *= batch_size * accumulate / nbs  # scale weight_decay
-    optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
+    optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"])
 
     # Scheduler
     if opt.cos_lr:
-        lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
+        lf = one_cycle(1, hyp["lrf"], epochs)  # cosine 1->hyp['lrf']
     else:
-        lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
+        lf = lambda x: (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"]  # linear
     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)
 
     # EMA
@@ -168,15 +225,15 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # DP mode
     if cuda and RANK == -1 and torch.cuda.device_count() > 1:
         LOGGER.warning(
-            'WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
+            "WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n"
-            'See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started.'
+            "See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started."
         )
         model = torch.nn.DataParallel(model)
 
     # SyncBatchNorm
     if opt.sync_bn and cuda and RANK != -1:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
-        LOGGER.info('Using SyncBatchNorm()')
+        LOGGER.info("Using SyncBatchNorm()")
 
     # Trainloader
     train_loader, dataset = create_dataloader(
@@ -187,24 +244,25 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
         single_cls,
         hyp=hyp,
         augment=True,
-        cache=None if opt.cache == 'val' else opt.cache,
+        cache=None if opt.cache == "val" else opt.cache,
         rect=opt.rect,
         rank=LOCAL_RANK,
         workers=workers,
         image_weights=opt.image_weights,
         quad=opt.quad,
-        prefix=colorstr('train: '),
+        prefix=colorstr("train: "),
         shuffle=True,
         mask_downsample_ratio=mask_ratio,
         overlap_mask=overlap,
     )
     labels = np.concatenate(dataset.labels, 0)
     mlc = int(labels[:, 0].max())  # max label class
-    assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
+    assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}"
 
     # Process 0
     if RANK in {-1, 0}:
-        val_loader = create_dataloader(val_path,
+        val_loader = create_dataloader(
+            val_path,
             imgsz,
             batch_size // WORLD_SIZE * 2,
             gs,
@@ -217,11 +275,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
             pad=0.5,
             mask_downsample_ratio=mask_ratio,
             overlap_mask=overlap,
-                                       prefix=colorstr('val: '))[0]
+            prefix=colorstr("val: "),
+        )[0]
 
         if not resume:
             if not opt.noautoanchor:
-                check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)  # run AutoAnchor
+                check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz)  # run AutoAnchor
             model.half().float()  # pre-reduce anchor precision
 
             if plots:
@@ -234,10 +293,10 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
 
     # Model attributes
     nl = de_parallel(model).model[-1].nl  # number of detection layers (to scale hyps)
-    hyp['box'] *= 3 / nl  # scale to layers
+    hyp["box"] *= 3 / nl  # scale to layers
-    hyp['cls'] *= nc / 80 * 3 / nl  # scale to classes and layers
+    hyp["cls"] *= nc / 80 * 3 / nl  # scale to classes and layers
-    hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
+    hyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
-    hyp['label_smoothing'] = opt.label_smoothing
+    hyp["label_smoothing"] = opt.label_smoothing
     model.nc = nc  # attach number of classes to model
     model.hyp = hyp  # attach hyperparameters to model
     model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
@@ -246,7 +305,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # Start training
     t0 = time.time()
     nb = len(train_loader)  # number of batches
-    nw = max(round(hyp['warmup_epochs'] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
+    nw = max(round(hyp["warmup_epochs"] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
     # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
     last_opt_step = -1
     maps = np.zeros(nc)  # mAP per class
@@ -256,10 +315,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     stopper, stop = EarlyStopping(patience=opt.patience), False
     compute_loss = ComputeLoss(model, overlap=overlap)  # init loss class
     # callbacks.run('on_train_start')
-    LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
+    LOGGER.info(
+        f'Image sizes {imgsz} train, {imgsz} val\n'
         f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
         f"Logging results to {colorstr('bold', save_dir)}\n"
-                f'Starting training for {epochs} epochs...')
+        f'Starting training for {epochs} epochs...'
+    )
     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
         # callbacks.run('on_train_epoch_start')
         model.train()
@@ -278,8 +339,10 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
         if RANK != -1:
             train_loader.sampler.set_epoch(epoch)
         pbar = enumerate(train_loader)
-        LOGGER.info(('\n' + '%11s' * 8) %
+        LOGGER.info(
-                    ('Epoch', 'GPU_mem', 'box_loss', 'seg_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
+            ("\n" + "%11s" * 8)
+            % ("Epoch", "GPU_mem", "box_loss", "seg_loss", "obj_loss", "cls_loss", "Instances", "Size")
+        )
         if RANK in {-1, 0}:
             pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT)  # progress bar
         optimizer.zero_grad()
@@ -295,9 +358,9 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                 for j, x in enumerate(optimizer.param_groups):
                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
-                    x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
+                    x["lr"] = np.interp(ni, xi, [hyp["warmup_bias_lr"] if j == 0 else 0.0, x["initial_lr"] * lf(epoch)])
-                    if 'momentum' in x:
+                    if "momentum" in x:
-                        x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
+                        x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]])
 
             # Multi-scale
             if opt.multi_scale:
@@ -305,7 +368,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 sf = sz / max(imgs.shape[2:])  # scale factor
                 if sf != 1:
                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
-                    imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
+                    imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
 
             # Forward
             with torch.cuda.amp.autocast(amp):
@@ -314,7 +377,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 if RANK != -1:
                     loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
                 if opt.quad:
-                    loss *= 4.
+                    loss *= 4.0
 
             # Backward
             scaler.scale(loss).backward()
@@ -333,9 +396,11 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
             # Log
             if RANK in {-1, 0}:
                 mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
-                mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'  # (GB)
+                mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G"  # (GB)
-                pbar.set_description(('%11s' * 2 + '%11.4g' * 6) %
+                pbar.set_description(
-                                     (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
+                    ("%11s" * 2 + "%11.4g" * 6)
+                    % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1])
+                )
                 # callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths)
                 # if callbacks.stop_training:
                 #    return
@@ -343,23 +408,24 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 # Mosaic plots
                 if plots:
                     if ni < 3:
-                        plot_images_and_masks(imgs, targets, masks, paths, save_dir / f'train_batch{ni}.jpg')
+                        plot_images_and_masks(imgs, targets, masks, paths, save_dir / f"train_batch{ni}.jpg")
                     if ni == 10:
-                        files = sorted(save_dir.glob('train*.jpg'))
+                        files = sorted(save_dir.glob("train*.jpg"))
-                        logger.log_images(files, 'Mosaics', epoch)
+                        logger.log_images(files, "Mosaics", epoch)
             # end batch ------------------------------------------------------------------------------------------------
 
         # Scheduler
-        lr = [x['lr'] for x in optimizer.param_groups]  # for loggers
+        lr = [x["lr"] for x in optimizer.param_groups]  # for loggers
         scheduler.step()
 
         if RANK in {-1, 0}:
             # mAP
             # callbacks.run('on_train_epoch_end', epoch=epoch)
-            ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
+            ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"])
             final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
             if not noval or final_epoch:  # Calculate mAP
-                results, maps, _ = validate.run(data_dict,
+                results, maps, _ = validate.run(
+                    data_dict,
                     batch_size=batch_size // WORLD_SIZE * 2,
                     imgsz=imgsz,
                     half=amp,
@@ -371,7 +437,8 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                     callbacks=callbacks,
                     compute_loss=compute_loss,
                     mask_downsample_ratio=mask_ratio,
-                                                overlap=overlap)
+                    overlap=overlap,
+                )
 
             # Update best mAP
             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
@@ -387,23 +454,24 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
             # Save model
             if (not nosave) or (final_epoch and not evolve):  # if save
                 ckpt = {
-                    'epoch': epoch,
+                    "epoch": epoch,
-                    'best_fitness': best_fitness,
+                    "best_fitness": best_fitness,
-                    'model': deepcopy(de_parallel(model)).half(),
+                    "model": deepcopy(de_parallel(model)).half(),
-                    'ema': deepcopy(ema.ema).half(),
+                    "ema": deepcopy(ema.ema).half(),
-                    'updates': ema.updates,
+                    "updates": ema.updates,
-                    'optimizer': optimizer.state_dict(),
+                    "optimizer": optimizer.state_dict(),
-                    'opt': vars(opt),
+                    "opt": vars(opt),
-                    'git': GIT_INFO,  # {remote, branch, commit} if a git repo
+                    "git": GIT_INFO,  # {remote, branch, commit} if a git repo
-                    'date': datetime.now().isoformat()}
+                    "date": datetime.now().isoformat(),
+                }
 
                 # Save last, best and delete
                 torch.save(ckpt, last)
                 if best_fitness == fi:
                     torch.save(ckpt, best)
                 if opt.save_period > 0 and epoch % opt.save_period == 0:
-                    torch.save(ckpt, w / f'epoch{epoch}.pt')
+                    torch.save(ckpt, w / f"epoch{epoch}.pt")
-                    logger.log_model(w / f'epoch{epoch}.pt')
+                    logger.log_model(w / f"epoch{epoch}.pt")
                 del ckpt
                 # callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
 
@@ -419,12 +487,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
         # end epoch ----------------------------------------------------------------------------------------------------
     # end training -----------------------------------------------------------------------------------------------------
     if RANK in {-1, 0}:
-        LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
+        LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.")
         for f in last, best:
             if f.exists():
                 strip_optimizer(f)  # strip optimizers
                 if f is best:
-                    LOGGER.info(f'\nValidating {f}...')
+                    LOGGER.info(f"\nValidating {f}...")
                     results, _, _ = validate.run(
                         data_dict,
                         batch_size=batch_size // WORLD_SIZE * 2,
@@ -440,7 +508,8 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                         callbacks=callbacks,
                         compute_loss=compute_loss,
                         mask_downsample_ratio=mask_ratio,
-                        overlap=overlap)  # val best model with plots
+                        overlap=overlap,
+                    )  # val best model with plots
                     if is_coco:
                         # callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
                         metrics_dict = dict(zip(KEYS, list(mloss) + list(results) + lr))
@@ -452,56 +521,56 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
         if not opt.evolve:
             logger.log_model(best, epoch)
         if plots:
-            plot_results_with_masks(file=save_dir / 'results.csv')  # save results.png
+            plot_results_with_masks(file=save_dir / "results.csv")  # save results.png
-            files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
+            files = ["results.png", "confusion_matrix.png", *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R"))]
             files = [(save_dir / f) for f in files if (save_dir / f).exists()]  # filter
             LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
-            logger.log_images(files, 'Results', epoch + 1)
+            logger.log_images(files, "Results", epoch + 1)
-            logger.log_images(sorted(save_dir.glob('val*.jpg')), 'Validation', epoch + 1)
+            logger.log_images(sorted(save_dir.glob("val*.jpg")), "Validation", epoch + 1)
     torch.cuda.empty_cache()
     return results
 
 
 def parse_opt(known=False):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s-seg.pt', help='initial weights path')
+    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s-seg.pt", help="initial weights path")
-    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
+    parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128-seg.yaml', help='dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128-seg.yaml", help="dataset.yaml path")
-    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
+    parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path")
-    parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
+    parser.add_argument("--epochs", type=int, default=100, help="total training epochs")
-    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
+    parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch")
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
+    parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)")
-    parser.add_argument('--rect', action='store_true', help='rectangular training')
+    parser.add_argument("--rect", action="store_true", help="rectangular training")
-    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
+    parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")
-    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
+    parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
-    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
+    parser.add_argument("--noval", action="store_true", help="only validate final epoch")
-    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
+    parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor")
-    parser.add_argument('--noplots', action='store_true', help='save no plot files')
+    parser.add_argument("--noplots", action="store_true", help="save no plot files")
-    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
+    parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations")
-    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
+    parser.add_argument("--bucket", type=str, default="", help="gsutil bucket")
-    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')
+    parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk")
-    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
+    parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
+    parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%")
-    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
+    parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class")
-    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
+    parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer")
-    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
+    parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode")
-    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--project', default=ROOT / 'runs/train-seg', help='save to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/train-seg", help="save to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--quad', action='store_true', help='quad dataloader')
+    parser.add_argument("--quad", action="store_true", help="quad dataloader")
-    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
+    parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler")
-    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
+    parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon")
-    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
+    parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)")
-    parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
+    parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2")
-    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
+    parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)")
-    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
+    parser.add_argument("--seed", type=int, default=0, help="Global training seed")
-    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
+    parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
 
     # Instance Segmentation Args
-    parser.add_argument('--mask-ratio', type=int, default=4, help='Downsample the truth masks to saving memory')
+    parser.add_argument("--mask-ratio", type=int, default=4, help="Downsample the truth masks to saving memory")
-    parser.add_argument('--no-overlap', action='store_true', help='Overlap masks train faster at slightly less mAP')
+    parser.add_argument("--no-overlap", action="store_true", help="Overlap masks train faster at slightly less mAP")
 
     return parser.parse_known_args()[0] if known else parser.parse_args()
 
@@ -511,46 +580,51 @@ def main(opt, callbacks=Callbacks()):
     if RANK in {-1, 0}:
         print_args(vars(opt))
         check_git_status()
-        check_requirements(ROOT / 'requirements.txt')
+        check_requirements(ROOT / "requirements.txt")
 
     # Resume
     if opt.resume and not opt.evolve:  # resume from specified or most recent last.pt
         last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
-        opt_yaml = last.parent.parent / 'opt.yaml'  # train options yaml
+        opt_yaml = last.parent.parent / "opt.yaml"  # train options yaml
         opt_data = opt.data  # original dataset
         if opt_yaml.is_file():
-            with open(opt_yaml, errors='ignore') as f:
+            with open(opt_yaml, errors="ignore") as f:
                 d = yaml.safe_load(f)
         else:
-            d = torch.load(last, map_location='cpu')['opt']
+            d = torch.load(last, map_location="cpu")["opt"]
         opt = argparse.Namespace(**d)  # replace
-        opt.cfg, opt.weights, opt.resume = '', str(last), True  # reinstate
+        opt.cfg, opt.weights, opt.resume = "", str(last), True  # reinstate
         if is_url(opt_data):
             opt.data = check_file(opt_data)  # avoid HUB resume auth timeout
     else:
-        opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
+        opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = (
-            check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project)  # checks
+            check_file(opt.data),
-        assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
+            check_yaml(opt.cfg),
+            check_yaml(opt.hyp),
+            str(opt.weights),
+            str(opt.project),
+        )  # checks
+        assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified"
         if opt.evolve:
-            if opt.project == str(ROOT / 'runs/train-seg'):  # if default project name, rename to runs/evolve-seg
+            if opt.project == str(ROOT / "runs/train-seg"):  # if default project name, rename to runs/evolve-seg
-                opt.project = str(ROOT / 'runs/evolve-seg')
+                opt.project = str(ROOT / "runs/evolve-seg")
             opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resume
-        if opt.name == 'cfg':
+        if opt.name == "cfg":
             opt.name = Path(opt.cfg).stem  # use model.yaml as name
         opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
 
     # DDP mode
     device = select_device(opt.device, batch_size=opt.batch_size)
     if LOCAL_RANK != -1:
-        msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
+        msg = "is not compatible with YOLOv5 Multi-GPU DDP training"
-        assert not opt.image_weights, f'--image-weights {msg}'
+        assert not opt.image_weights, f"--image-weights {msg}"
-        assert not opt.evolve, f'--evolve {msg}'
+        assert not opt.evolve, f"--evolve {msg}"
-        assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
+        assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size"
-        assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
+        assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE"
-        assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
+        assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
         torch.cuda.set_device(LOCAL_RANK)
-        device = torch.device('cuda', LOCAL_RANK)
+        device = torch.device("cuda", LOCAL_RANK)
-        dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo')
+        dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
 
     # Train
     if not opt.evolve:
@@ -560,65 +634,69 @@ def main(opt, callbacks=Callbacks()):
     else:
         # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
         meta = {
-            'lr0': (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
+            "lr0": (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
-            'lrf': (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
+            "lrf": (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
-            'momentum': (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
+            "momentum": (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
-            'weight_decay': (1, 0.0, 0.001),  # optimizer weight decay
+            "weight_decay": (1, 0.0, 0.001),  # optimizer weight decay
-            'warmup_epochs': (1, 0.0, 5.0),  # warmup epochs (fractions ok)
+            "warmup_epochs": (1, 0.0, 5.0),  # warmup epochs (fractions ok)
-            'warmup_momentum': (1, 0.0, 0.95),  # warmup initial momentum
+            "warmup_momentum": (1, 0.0, 0.95),  # warmup initial momentum
-            'warmup_bias_lr': (1, 0.0, 0.2),  # warmup initial bias lr
+            "warmup_bias_lr": (1, 0.0, 0.2),  # warmup initial bias lr
-            'box': (1, 0.02, 0.2),  # box loss gain
+            "box": (1, 0.02, 0.2),  # box loss gain
-            'cls': (1, 0.2, 4.0),  # cls loss gain
+            "cls": (1, 0.2, 4.0),  # cls loss gain
-            'cls_pw': (1, 0.5, 2.0),  # cls BCELoss positive_weight
+            "cls_pw": (1, 0.5, 2.0),  # cls BCELoss positive_weight
-            'obj': (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
+            "obj": (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
-            'obj_pw': (1, 0.5, 2.0),  # obj BCELoss positive_weight
+            "obj_pw": (1, 0.5, 2.0),  # obj BCELoss positive_weight
-            'iou_t': (0, 0.1, 0.7),  # IoU training threshold
+            "iou_t": (0, 0.1, 0.7),  # IoU training threshold
-            'anchor_t': (1, 2.0, 8.0),  # anchor-multiple threshold
+            "anchor_t": (1, 2.0, 8.0),  # anchor-multiple threshold
-            'anchors': (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
+            "anchors": (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
-            'fl_gamma': (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
+            "fl_gamma": (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
-            'hsv_h': (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
+            "hsv_h": (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
-            'hsv_s': (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
+            "hsv_s": (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
-            'hsv_v': (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
+            "hsv_v": (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
-            'degrees': (1, 0.0, 45.0),  # image rotation (+/- deg)
+            "degrees": (1, 0.0, 45.0),  # image rotation (+/- deg)
-            'translate': (1, 0.0, 0.9),  # image translation (+/- fraction)
+            "translate": (1, 0.0, 0.9),  # image translation (+/- fraction)
-            'scale': (1, 0.0, 0.9),  # image scale (+/- gain)
+            "scale": (1, 0.0, 0.9),  # image scale (+/- gain)
-            'shear': (1, 0.0, 10.0),  # image shear (+/- deg)
+            "shear": (1, 0.0, 10.0),  # image shear (+/- deg)
-            'perspective': (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
+            "perspective": (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
-            'flipud': (1, 0.0, 1.0),  # image flip up-down (probability)
+            "flipud": (1, 0.0, 1.0),  # image flip up-down (probability)
-            'fliplr': (0, 0.0, 1.0),  # image flip left-right (probability)
+            "fliplr": (0, 0.0, 1.0),  # image flip left-right (probability)
-            'mosaic': (1, 0.0, 1.0),  # image mixup (probability)
+            "mosaic": (1, 0.0, 1.0),  # image mixup (probability)
-            'mixup': (1, 0.0, 1.0),  # image mixup (probability)
+            "mixup": (1, 0.0, 1.0),  # image mixup (probability)
-            'copy_paste': (1, 0.0, 1.0)}  # segment copy-paste (probability)
+            "copy_paste": (1, 0.0, 1.0),
+        }  # segment copy-paste (probability)
 
-        with open(opt.hyp, errors='ignore') as f:
+        with open(opt.hyp, errors="ignore") as f:
             hyp = yaml.safe_load(f)  # load hyps dict
-            if 'anchors' not in hyp:  # anchors commented in hyp.yaml
+            if "anchors" not in hyp:  # anchors commented in hyp.yaml
-                hyp['anchors'] = 3
+                hyp["anchors"] = 3
         if opt.noautoanchor:
-            del hyp['anchors'], meta['anchors']
+            del hyp["anchors"], meta["anchors"]
         opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch
         # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
-        evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
+        evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv"
         if opt.bucket:
             # download evolve.csv if exists
-            subprocess.run([
+            subprocess.run(
-                'gsutil',
+                [
-                'cp',
+                    "gsutil",
-                f'gs://{opt.bucket}/evolve.csv',
+                    "cp",
-                str(evolve_csv), ])
+                    f"gs://{opt.bucket}/evolve.csv",
+                    str(evolve_csv),
+                ]
+            )
 
         for _ in range(opt.evolve):  # generations to evolve
             if evolve_csv.exists():  # if evolve.csv exists: select best hyps and mutate
                 # Select parent(s)
-                parent = 'single'  # parent selection method: 'single' or 'weighted'
+                parent = "single"  # parent selection method: 'single' or 'weighted'
-                x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
+                x = np.loadtxt(evolve_csv, ndmin=2, delimiter=",", skiprows=1)
                 n = min(5, len(x))  # number of previous results to consider
                 x = x[np.argsort(-fitness(x))][:n]  # top n mutations
-                w = fitness(x) - fitness(x).min() + 1E-6  # weights (sum > 0)
+                w = fitness(x) - fitness(x).min() + 1e-6  # weights (sum > 0)
-                if parent == 'single' or len(x) == 1:
+                if parent == "single" or len(x) == 1:
                     # x = x[random.randint(0, n - 1)]  # random selection
                     x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
-                elif parent == 'weighted':
+                elif parent == "weighted":
                     x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
 
                 # Mutate
@@ -647,9 +725,11 @@ def main(opt, callbacks=Callbacks()):
 
         # Plot results
         plot_evolve(evolve_csv)
-        LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
+        LOGGER.info(
+            f'Hyperparameter evolution finished {opt.evolve} generations\n'
             f"Results saved to {colorstr('bold', save_dir)}\n"
-                    f'Usage example: $ python train.py --hyp {evolve_yaml}')
+            f'Usage example: $ python train.py --hyp {evolve_yaml}'
+        )
 
 
 def run(**kwargs):
@@ -661,6 +741,6 @@ def run(**kwargs):
     return opt
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

segment/val.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Validate a trained YOLOv5 segment model on a segment dataset
+Validate a trained YOLOv5 segment model on a segment dataset.
 
 Usage:
     $ bash data/scripts/get_coco.sh --val --segments  # download COCO-segments val split (1G, 5000 images)
@@ -43,9 +43,24 @@ import torch.nn.functional as F
 from models.common import DetectMultiBackend
 from models.yolo import SegmentationModel
 from utils.callbacks import Callbacks
-from utils.general import (LOGGER, NUM_THREADS, TQDM_BAR_FORMAT, Profile, check_dataset, check_img_size,
+from utils.general import (
-                           check_requirements, check_yaml, coco80_to_coco91_class, colorstr, increment_path,
+    LOGGER,
-                           non_max_suppression, print_args, scale_boxes, xywh2xyxy, xyxy2xywh)
+    NUM_THREADS,
+    TQDM_BAR_FORMAT,
+    Profile,
+    check_dataset,
+    check_img_size,
+    check_requirements,
+    check_yaml,
+    coco80_to_coco91_class,
+    colorstr,
+    increment_path,
+    non_max_suppression,
+    print_args,
+    scale_boxes,
+    xywh2xyxy,
+    xyxy2xywh,
+)
 from utils.metrics import ConfusionMatrix, box_iou
 from utils.plots import output_to_target, plot_val_study
 from utils.segment.dataloaders import create_dataloader
@@ -61,8 +76,8 @@ def save_one_txt(predn, save_conf, shape, file):
     for *xyxy, conf, cls in predn.tolist():
         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
         line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
-        with open(file, 'a') as f:
+        with open(file, "a") as f:
-            f.write(('%g ' * len(line)).rstrip() % line + '\n')
+            f.write(("%g " * len(line)).rstrip() % line + "\n")
 
 
 def save_one_json(predn, jdict, path, class_map, pred_masks):
@@ -70,8 +85,8 @@ def save_one_json(predn, jdict, path, class_map, pred_masks):
     from pycocotools.mask import encode
 
     def single_encode(x):
-        rle = encode(np.asarray(x[:, :, None], order='F', dtype='uint8'))[0]
+        rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
-        rle['counts'] = rle['counts'].decode('utf-8')
+        rle["counts"] = rle["counts"].decode("utf-8")
         return rle
 
     image_id = int(path.stem) if path.stem.isnumeric() else path.stem
@@ -81,12 +96,15 @@ def save_one_json(predn, jdict, path, class_map, pred_masks):
     with ThreadPool(NUM_THREADS) as pool:
         rles = pool.map(single_encode, pred_masks)
     for i, (p, b) in enumerate(zip(predn.tolist(), box.tolist())):
-        jdict.append({
+        jdict.append(
-            'image_id': image_id,
+            {
-            'category_id': class_map[int(p[5])],
+                "image_id": image_id,
-            'bbox': [round(x, 3) for x in b],
+                "category_id": class_map[int(p[5])],
-            'score': round(p[4], 5),
+                "bbox": [round(x, 3) for x in b],
-            'segmentation': rles[i]})
+                "score": round(p[4], 5),
+                "segmentation": rles[i],
+            }
+        )
 
 
 def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, overlap=False, masks=False):
@@ -105,7 +123,7 @@ def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, over
             gt_masks = gt_masks.repeat(nl, 1, 1)  # shape(1,640,640) -> (n,640,640)
             gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
         if gt_masks.shape[1:] != pred_masks.shape[1:]:
-            gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode='bilinear', align_corners=False)[0]
+            gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
             gt_masks = gt_masks.gt_(0.5)
         iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
     else:  # boxes
@@ -135,8 +153,8 @@ def run(
     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
+    task="val",  # train, val, test, speed or study
-        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     workers=8,  # max dataloader workers (per RANK in DDP mode)
     single_cls=False,  # treat as single-class dataset
     augment=False,  # augmented inference
@@ -145,14 +163,14 @@ def run(
     save_hybrid=False,  # save label+prediction hybrid results to *.txt
     save_conf=False,  # save confidences in --save-txt labels
     save_json=False,  # save a COCO-JSON results file
-        project=ROOT / 'runs/val-seg',  # save to project/name
+    project=ROOT / "runs/val-seg",  # save to project/name
-        name='exp',  # save to project/name
+    name="exp",  # save to project/name
     exist_ok=False,  # existing project/name ok, do not increment
     half=True,  # use FP16 half-precision inference
     dnn=False,  # use OpenCV DNN for ONNX inference
     model=None,
     dataloader=None,
-        save_dir=Path(''),
+    save_dir=Path(""),
     plots=True,
     overlap=False,
     mask_downsample_ratio=1,
@@ -160,7 +178,7 @@ def run(
     callbacks=Callbacks(),
 ):
     if save_json:
-        check_requirements('pycocotools>=2.0.6')
+        check_requirements("pycocotools>=2.0.6")
         process = process_mask_native  # more accurate
     else:
         process = process_mask  # faster
@@ -169,7 +187,7 @@ def run(
     training = model is not None
     if training:  # called by train.py
         device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
-        half &= device.type != 'cpu'  # half precision only supported on CUDA
+        half &= device.type != "cpu"  # half precision only supported on CUDA
         model.half() if half else model.float()
         nm = de_parallel(model).model[-1].nm  # number of masks
     else:  # called directly
@@ -177,7 +195,7 @@ def run(
 
         # Directories
         save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+        (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
         # Load model
         model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
@@ -191,16 +209,16 @@ def run(
             device = model.device
             if not (pt or jit):
                 batch_size = 1  # export.py models default to batch-size 1
-                LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
+                LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
 
         # Data
         data = check_dataset(data)  # check
 
     # Configure
     model.eval()
-    cuda = device.type != 'cpu'
+    cuda = device.type != "cpu"
-    is_coco = isinstance(data.get('val'), str) and data['val'].endswith(f'coco{os.sep}val2017.txt')  # COCO dataset
+    is_coco = isinstance(data.get("val"), str) and data["val"].endswith(f"coco{os.sep}val2017.txt")  # COCO dataset
-    nc = 1 if single_cls else int(data['nc'])  # number of classes
+    nc = 1 if single_cls else int(data["nc"])  # number of classes
     iouv = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for mAP@0.5:0.95
     niou = iouv.numel()
 
@@ -208,12 +226,15 @@ def run(
     if not training:
         if pt and not single_cls:  # check --weights are trained on --data
             ncm = model.model.nc
-            assert ncm == nc, f'{weights} ({ncm} classes) trained on different --data than what you passed ({nc} ' \
+            assert ncm == nc, (
-                              f'classes). Pass correct combination of --weights and --data that are trained together.'
+                f"{weights} ({ncm} classes) trained on different --data than what you passed ({nc} "
+                f"classes). Pass correct combination of --weights and --data that are trained together."
+            )
         model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz))  # warmup
-        pad, rect = (0.0, False) if task == 'speed' else (0.5, pt)  # square inference for benchmarks
+        pad, rect = (0.0, False) if task == "speed" else (0.5, pt)  # square inference for benchmarks
-        task = task if task in ('train', 'val', 'test') else 'val'  # path to train/val/test images
+        task = task if task in ("train", "val", "test") else "val"  # path to train/val/test images
-        dataloader = create_dataloader(data[task],
+        dataloader = create_dataloader(
+            data[task],
             imgsz,
             batch_size,
             stride,
@@ -221,18 +242,30 @@ def run(
             pad=pad,
             rect=rect,
             workers=workers,
-                                       prefix=colorstr(f'{task}: '),
+            prefix=colorstr(f"{task}: "),
             overlap_mask=overlap,
-                                       mask_downsample_ratio=mask_downsample_ratio)[0]
+            mask_downsample_ratio=mask_downsample_ratio,
+        )[0]
 
     seen = 0
     confusion_matrix = ConfusionMatrix(nc=nc)
-    names = model.names if hasattr(model, 'names') else model.module.names  # get class names
+    names = model.names if hasattr(model, "names") else model.module.names  # get class names
     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' * 10) % ('Class', 'Images', 'Instances', 'Box(P', 'R', 'mAP50', 'mAP50-95)', 'Mask(P', 'R',
+    s = ("%22s" + "%11s" * 10) % (
-                                  'mAP50', 'mAP50-95)')
+        "Class",
+        "Images",
+        "Instances",
+        "Box(P",
+        "R",
+        "mAP50",
+        "mAP50-95)",
+        "Mask(P",
+        "R",
+        "mAP50",
+        "mAP50-95)",
+    )
     dt = Profile(device=device), Profile(device=device), Profile(device=device)
     metrics = Metrics()
     loss = torch.zeros(4, device=device)
@@ -263,14 +296,9 @@ 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]:
-            preds = non_max_suppression(preds,
+            preds = non_max_suppression(
-                                        conf_thres,
+                preds, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls, max_det=max_det, nm=nm
-                                        iou_thres,
+            )
-                                        labels=lb,
-                                        multi_label=True,
-                                        agnostic=single_cls,
-                                        max_det=max_det,
-                                        nm=nm)
 
         # Metrics
         plot_masks = []  # masks for plotting
@@ -317,10 +345,11 @@ def run(
 
             # Save/log
             if save_txt:
-                save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
+                save_one_txt(predn, save_conf, shape, file=save_dir / "labels" / f"{path.stem}.txt")
             if save_json:
-                pred_masks = scale_image(im[si].shape[1:],
+                pred_masks = scale_image(
-                                         pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(), shape, shapes[si][1])
+                    im[si].shape[1:], pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(), shape, shapes[si][1]
+                )
                 save_one_json(predn, jdict, path, class_map, pred_masks)  # append to COCO-JSON dictionary
             # callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
 
@@ -328,9 +357,15 @@ def run(
         if plots and batch_i < 3:
             if len(plot_masks):
                 plot_masks = torch.cat(plot_masks, dim=0)
-            plot_images_and_masks(im, targets, masks, paths, save_dir / f'val_batch{batch_i}_labels.jpg', names)
+            plot_images_and_masks(im, targets, masks, paths, save_dir / f"val_batch{batch_i}_labels.jpg", names)
-            plot_images_and_masks(im, output_to_target(preds, max_det=15), plot_masks, paths,
+            plot_images_and_masks(
-                                  save_dir / f'val_batch{batch_i}_pred.jpg', names)  # pred
+                im,
+                output_to_target(preds, max_det=15),
+                plot_masks,
+                paths,
+                save_dir / f"val_batch{batch_i}_pred.jpg",
+                names,
+            )  # pred
 
         # callbacks.run('on_val_batch_end')
 
@@ -342,10 +377,10 @@ def run(
     nt = np.bincount(stats[4].astype(int), minlength=nc)  # number of targets per class
 
     # Print results
-    pf = '%22s' + '%11i' * 2 + '%11.3g' * 8  # print format
+    pf = "%22s" + "%11i" * 2 + "%11.3g" * 8  # print format
-    LOGGER.info(pf % ('all', seen, nt.sum(), *metrics.mean_results()))
+    LOGGER.info(pf % ("all", seen, nt.sum(), *metrics.mean_results()))
     if nt.sum() == 0:
-        LOGGER.warning(f'WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels')
+        LOGGER.warning(f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels")
 
     # Print results per class
     if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
@@ -353,10 +388,10 @@ def run(
             LOGGER.info(pf % (names[c], seen, nt[c], *metrics.class_result(i)))
 
     # Print speeds
-    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
     if not training:
         shape = (batch_size, 3, imgsz, imgsz)
-        LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
+        LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}" % t)
 
     # Plots
     if plots:
@@ -367,11 +402,11 @@ def run(
 
     # Save JSON
     if save_json and len(jdict):
-        w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ''  # weights
+        w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ""  # weights
-        anno_json = str(Path('../datasets/coco/annotations/instances_val2017.json'))  # annotations
+        anno_json = str(Path("../datasets/coco/annotations/instances_val2017.json"))  # annotations
-        pred_json = str(save_dir / f'{w}_predictions.json')  # predictions
+        pred_json = str(save_dir / f"{w}_predictions.json")  # predictions
-        LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
+        LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...")
-        with open(pred_json, 'w') as f:
+        with open(pred_json, "w") as f:
             json.dump(jdict, f)
 
         try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
@@ -381,7 +416,7 @@ def run(
             anno = COCO(anno_json)  # init annotations api
             pred = anno.loadRes(pred_json)  # init predictions api
             results = []
-            for eval in COCOeval(anno, pred, 'bbox'), COCOeval(anno, pred, 'segm'):
+            for eval in COCOeval(anno, pred, "bbox"), COCOeval(anno, pred, "segm"):
                 if is_coco:
                     eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files]  # img ID to evaluate
                 eval.evaluate()
@@ -390,12 +425,12 @@ def run(
                 results.extend(eval.stats[:2])  # update results (mAP@0.5:0.95, mAP@0.5)
             map_bbox, map50_bbox, map_mask, map50_mask = results
         except Exception as e:
-            LOGGER.info(f'pycocotools unable to run: {e}')
+            LOGGER.info(f"pycocotools unable to run: {e}")
 
     # Return results
     model.float()  # for training
     if not training:
-        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     final_metric = mp_bbox, mr_bbox, map50_bbox, map_bbox, mp_mask, mr_mask, map50_mask, map_mask
     return (*final_metric, *(loss.cpu() / len(dataloader)).tolist()), metrics.get_maps(nc), t
@@ -403,28 +438,28 @@ def run(
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128-seg.yaml', help='dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128-seg.yaml", help="dataset.yaml path")
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-seg.pt', help='model path(s)')
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-seg.pt", help="model path(s)")
-    parser.add_argument('--batch-size', type=int, default=32, help='batch size')
+    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("--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("--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("--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("--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("--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("--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)')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
+    parser.add_argument("--single-cls", action="store_true", help="treat as single-class dataset")
-    parser.add_argument('--augment', action='store_true', help='augmented inference')
+    parser.add_argument("--augment", action="store_true", help="augmented inference")
-    parser.add_argument('--verbose', action='store_true', help='report mAP by class')
+    parser.add_argument("--verbose", action="store_true", help="report mAP by class")
-    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
+    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
-    parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
+    parser.add_argument("--save-hybrid", action="store_true", help="save label+prediction hybrid results to *.txt")
-    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
+    parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
-    parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
+    parser.add_argument("--save-json", action="store_true", help="save a COCO-JSON results file")
-    parser.add_argument('--project', default=ROOT / 'runs/val-seg', help='save results to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/val-seg", help="save results to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
     opt = parser.parse_args()
     opt.data = check_yaml(opt.data)  # check YAML
     # opt.save_json |= opt.data.endswith('coco.yaml')
@@ -434,40 +469,40 @@ def parse_opt():
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
 
-    if opt.task in ('train', 'val', 'test'):  # run normally
+    if opt.task in ("train", "val", "test"):  # run normally
         if opt.conf_thres > 0.001:  # https://github.com/ultralytics/yolov5/issues/1466
-            LOGGER.warning(f'WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results')
+            LOGGER.warning(f"WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results")
         if opt.save_hybrid:
-            LOGGER.warning('WARNING ⚠️ --save-hybrid returns high mAP from hybrid labels, not from predictions alone')
+            LOGGER.warning("WARNING ⚠️ --save-hybrid returns high mAP from hybrid labels, not from predictions alone")
         run(**vars(opt))
 
     else:
         weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
-        opt.half = torch.cuda.is_available() and opt.device != 'cpu'  # FP16 for fastest results
+        opt.half = torch.cuda.is_available() and opt.device != "cpu"  # FP16 for fastest results
-        if opt.task == 'speed':  # speed benchmarks
+        if opt.task == "speed":  # speed benchmarks
             # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
             opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
             for opt.weights in weights:
                 run(**vars(opt), plots=False)
 
-        elif opt.task == 'study':  # speed vs mAP benchmarks
+        elif opt.task == "study":  # speed vs mAP benchmarks
             # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
             for opt.weights in weights:
-                f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt'  # filename to save to
+                f = f"study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt"  # filename to save to
                 x, y = list(range(256, 1536 + 128, 128)), []  # x axis (image sizes), y axis
                 for opt.imgsz in x:  # img-size
-                    LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
+                    LOGGER.info(f"\nRunning {f} --imgsz {opt.imgsz}...")
                     r, _, t = run(**vars(opt), plots=False)
                     y.append(r + t)  # results and times
-                np.savetxt(f, y, fmt='%10.4g')  # save
+                np.savetxt(f, y, fmt="%10.4g")  # save
-            subprocess.run(['zip', '-r', 'study.zip', 'study_*.txt'])
+            subprocess.run(["zip", "-r", "study.zip", "study_*.txt"])
             plot_val_study(x=x)  # plot
         else:
             raise NotImplementedError(f'--task {opt.task} not in ("train", "val", "test", "speed", "study")')
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

train.py

@@ -1,7 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Train a YOLOv5 model on a custom dataset.
+Train a YOLOv5 model on a custom dataset. Models and datasets download automatically from the latest YOLOv5 release.
-Models and datasets download automatically from the latest YOLOv5 release.
 
 Usage - Single-GPU training:
     $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640  # from pretrained (recommended)
@@ -53,56 +52,97 @@ from utils.autobatch import check_train_batch_size
 from utils.callbacks import Callbacks
 from utils.dataloaders import create_dataloader
 from utils.downloads import attempt_download, is_url
-from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
+from utils.general import (
-                           check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,
+    LOGGER,
-                           get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
+    TQDM_BAR_FORMAT,
-                           labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer,
+    check_amp,
-                           yaml_save)
+    check_dataset,
+    check_file,
+    check_git_info,
+    check_git_status,
+    check_img_size,
+    check_requirements,
+    check_suffix,
+    check_yaml,
+    colorstr,
+    get_latest_run,
+    increment_path,
+    init_seeds,
+    intersect_dicts,
+    labels_to_class_weights,
+    labels_to_image_weights,
+    methods,
+    one_cycle,
+    print_args,
+    print_mutation,
+    strip_optimizer,
+    yaml_save,
+)
 from utils.loggers import LOGGERS, Loggers
 from utils.loggers.comet.comet_utils import check_comet_resume
 from utils.loss import ComputeLoss
 from utils.metrics import fitness
 from utils.plots import plot_evolve
-from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
+from utils.torch_utils import (
-                               smart_resume, torch_distributed_zero_first)
+    EarlyStopping,
+    ModelEMA,
+    de_parallel,
+    select_device,
+    smart_DDP,
+    smart_optimizer,
+    smart_resume,
+    torch_distributed_zero_first,
+)
 
-LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
+LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
-WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
+WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
 GIT_INFO = check_git_info()
 
 
 def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictionary
-    save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \
+    save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = (
-        Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
+        Path(opt.save_dir),
-        opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
+        opt.epochs,
-    callbacks.run('on_pretrain_routine_start')
+        opt.batch_size,
+        opt.weights,
+        opt.single_cls,
+        opt.evolve,
+        opt.data,
+        opt.cfg,
+        opt.resume,
+        opt.noval,
+        opt.nosave,
+        opt.workers,
+        opt.freeze,
+    )
+    callbacks.run("on_pretrain_routine_start")
 
     # Directories
-    w = save_dir / 'weights'  # weights dir
+    w = save_dir / "weights"  # weights dir
     (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
-    last, best = w / 'last.pt', w / 'best.pt'
+    last, best = w / "last.pt", w / "best.pt"
 
     # Hyperparameters
     if isinstance(hyp, str):
-        with open(hyp, errors='ignore') as f:
+        with open(hyp, errors="ignore") as f:
             hyp = yaml.safe_load(f)  # load hyps dict
-    LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
+    LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()))
     opt.hyp = hyp.copy()  # for saving hyps to checkpoints
 
     # Save run settings
     if not evolve:
-        yaml_save(save_dir / 'hyp.yaml', hyp)
+        yaml_save(save_dir / "hyp.yaml", hyp)
-        yaml_save(save_dir / 'opt.yaml', vars(opt))
+        yaml_save(save_dir / "opt.yaml", vars(opt))
 
     # Loggers
     data_dict = None
     if RANK in {-1, 0}:
         include_loggers = list(LOGGERS)
-        if getattr(opt, 'ndjson_console', False):
+        if getattr(opt, "ndjson_console", False):
-            include_loggers.append('ndjson_console')
+            include_loggers.append("ndjson_console")
-        if getattr(opt, 'ndjson_file', False):
+        if getattr(opt, "ndjson_file", False):
-            include_loggers.append('ndjson_file')
+            include_loggers.append("ndjson_file")
 
         loggers = Loggers(
             save_dir=save_dir,
@@ -124,39 +164,39 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
 
     # Config
     plots = not evolve and not opt.noplots  # create plots
-    cuda = device.type != 'cpu'
+    cuda = device.type != "cpu"
     init_seeds(opt.seed + 1 + RANK, deterministic=True)
     with torch_distributed_zero_first(LOCAL_RANK):
         data_dict = data_dict or check_dataset(data)  # check if None
-    train_path, val_path = data_dict['train'], data_dict['val']
+    train_path, val_path = data_dict["train"], data_dict["val"]
-    nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
+    nc = 1 if single_cls else int(data_dict["nc"])  # number of classes
-    names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
+    names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"]  # class names
-    is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')  # COCO dataset
+    is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt")  # COCO dataset
 
     # Model
-    check_suffix(weights, '.pt')  # check weights
+    check_suffix(weights, ".pt")  # check weights
-    pretrained = weights.endswith('.pt')
+    pretrained = weights.endswith(".pt")
     if pretrained:
         with torch_distributed_zero_first(LOCAL_RANK):
             weights = attempt_download(weights)  # download if not found locally
-        ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
+        ckpt = torch.load(weights, map_location="cpu")  # load checkpoint to CPU to avoid CUDA memory leak
-        model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
+        model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create
-        exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []  # exclude keys
+        exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else []  # exclude keys
-        csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
+        csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
         csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
         model.load_state_dict(csd, strict=False)  # load
-        LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  # report
+        LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}")  # report
     else:
-        model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
+        model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create
     amp = check_amp(model)  # check AMP
 
     # Freeze
-    freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
+    freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
     for k, v in model.named_parameters():
         v.requires_grad = True  # train all layers
         # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
         if any(x in k for x in freeze):
-            LOGGER.info(f'freezing {k}')
+            LOGGER.info(f"freezing {k}")
             v.requires_grad = False
 
     # Image size
@@ -166,19 +206,19 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # Batch size
     if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
         batch_size = check_train_batch_size(model, imgsz, amp)
-        loggers.on_params_update({'batch_size': batch_size})
+        loggers.on_params_update({"batch_size": batch_size})
 
     # Optimizer
     nbs = 64  # nominal batch size
     accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
-    hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
+    hyp["weight_decay"] *= batch_size * accumulate / nbs  # scale weight_decay
-    optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
+    optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"])
 
     # Scheduler
     if opt.cos_lr:
-        lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
+        lf = one_cycle(1, hyp["lrf"], epochs)  # cosine 1->hyp['lrf']
     else:
-        lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
+        lf = lambda x: (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"]  # linear
     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)
 
     # EMA
@@ -194,40 +234,43 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # DP mode
     if cuda and RANK == -1 and torch.cuda.device_count() > 1:
         LOGGER.warning(
-            'WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
+            "WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n"
-            'See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started.'
+            "See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started."
         )
         model = torch.nn.DataParallel(model)
 
     # SyncBatchNorm
     if opt.sync_bn and cuda and RANK != -1:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
-        LOGGER.info('Using SyncBatchNorm()')
+        LOGGER.info("Using SyncBatchNorm()")
 
     # Trainloader
-    train_loader, dataset = create_dataloader(train_path,
+    train_loader, dataset = create_dataloader(
+        train_path,
         imgsz,
         batch_size // WORLD_SIZE,
         gs,
         single_cls,
         hyp=hyp,
         augment=True,
-                                              cache=None if opt.cache == 'val' else opt.cache,
+        cache=None if opt.cache == "val" else opt.cache,
         rect=opt.rect,
         rank=LOCAL_RANK,
         workers=workers,
         image_weights=opt.image_weights,
         quad=opt.quad,
-                                              prefix=colorstr('train: '),
+        prefix=colorstr("train: "),
         shuffle=True,
-                                              seed=opt.seed)
+        seed=opt.seed,
+    )
     labels = np.concatenate(dataset.labels, 0)
     mlc = int(labels[:, 0].max())  # max label class
-    assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
+    assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}"
 
     # Process 0
     if RANK in {-1, 0}:
-        val_loader = create_dataloader(val_path,
+        val_loader = create_dataloader(
+            val_path,
             imgsz,
             batch_size // WORLD_SIZE * 2,
             gs,
@@ -238,14 +281,15 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
             rank=-1,
             workers=workers * 2,
             pad=0.5,
-                                       prefix=colorstr('val: '))[0]
+            prefix=colorstr("val: "),
+        )[0]
 
         if not resume:
             if not opt.noautoanchor:
-                check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)  # run AutoAnchor
+                check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz)  # run AutoAnchor
             model.half().float()  # pre-reduce anchor precision
 
-        callbacks.run('on_pretrain_routine_end', labels, names)
+        callbacks.run("on_pretrain_routine_end", labels, names)
 
     # DDP mode
     if cuda and RANK != -1:
@@ -253,10 +297,10 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
 
     # Model attributes
     nl = de_parallel(model).model[-1].nl  # number of detection layers (to scale hyps)
-    hyp['box'] *= 3 / nl  # scale to layers
+    hyp["box"] *= 3 / nl  # scale to layers
-    hyp['cls'] *= nc / 80 * 3 / nl  # scale to classes and layers
+    hyp["cls"] *= nc / 80 * 3 / nl  # scale to classes and layers
-    hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
+    hyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
-    hyp['label_smoothing'] = opt.label_smoothing
+    hyp["label_smoothing"] = opt.label_smoothing
     model.nc = nc  # attach number of classes to model
     model.hyp = hyp  # attach hyperparameters to model
     model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
@@ -265,7 +309,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # Start training
     t0 = time.time()
     nb = len(train_loader)  # number of batches
-    nw = max(round(hyp['warmup_epochs'] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
+    nw = max(round(hyp["warmup_epochs"] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
     # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
     last_opt_step = -1
     maps = np.zeros(nc)  # mAP per class
@@ -274,13 +318,15 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     scaler = torch.cuda.amp.GradScaler(enabled=amp)
     stopper, stop = EarlyStopping(patience=opt.patience), False
     compute_loss = ComputeLoss(model)  # init loss class
-    callbacks.run('on_train_start')
+    callbacks.run("on_train_start")
-    LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
+    LOGGER.info(
+        f'Image sizes {imgsz} train, {imgsz} val\n'
         f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
         f"Logging results to {colorstr('bold', save_dir)}\n"
-                f'Starting training for {epochs} epochs...')
+        f'Starting training for {epochs} epochs...'
+    )
     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
-        callbacks.run('on_train_epoch_start')
+        callbacks.run("on_train_epoch_start")
         model.train()
 
         # Update image weights (optional, single-GPU only)
@@ -297,12 +343,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
         if RANK != -1:
             train_loader.sampler.set_epoch(epoch)
         pbar = enumerate(train_loader)
-        LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
+        LOGGER.info(("\n" + "%11s" * 7) % ("Epoch", "GPU_mem", "box_loss", "obj_loss", "cls_loss", "Instances", "Size"))
         if RANK in {-1, 0}:
             pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT)  # progress bar
         optimizer.zero_grad()
         for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
-            callbacks.run('on_train_batch_start')
+            callbacks.run("on_train_batch_start")
             ni = i + nb * epoch  # number integrated batches (since train start)
             imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0
 
@@ -313,9 +359,9 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                 for j, x in enumerate(optimizer.param_groups):
                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
-                    x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
+                    x["lr"] = np.interp(ni, xi, [hyp["warmup_bias_lr"] if j == 0 else 0.0, x["initial_lr"] * lf(epoch)])
-                    if 'momentum' in x:
+                    if "momentum" in x:
-                        x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
+                        x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]])
 
             # Multi-scale
             if opt.multi_scale:
@@ -323,7 +369,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 sf = sz / max(imgs.shape[2:])  # scale factor
                 if sf != 1:
                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
-                    imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
+                    imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
 
             # Forward
             with torch.cuda.amp.autocast(amp):
@@ -332,7 +378,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                 if RANK != -1:
                     loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
                 if opt.quad:
-                    loss *= 4.
+                    loss *= 4.0
 
             # Backward
             scaler.scale(loss).backward()
@@ -351,25 +397,28 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
             # Log
             if RANK in {-1, 0}:
                 mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
-                mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'  # (GB)
+                mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G"  # (GB)
-                pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
+                pbar.set_description(
-                                     (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
+                    ("%11s" * 2 + "%11.4g" * 5)
-                callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
+                    % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1])
+                )
+                callbacks.run("on_train_batch_end", model, ni, imgs, targets, paths, list(mloss))
                 if callbacks.stop_training:
                     return
             # end batch ------------------------------------------------------------------------------------------------
 
         # Scheduler
-        lr = [x['lr'] for x in optimizer.param_groups]  # for loggers
+        lr = [x["lr"] for x in optimizer.param_groups]  # for loggers
         scheduler.step()
 
         if RANK in {-1, 0}:
             # mAP
-            callbacks.run('on_train_epoch_end', epoch=epoch)
+            callbacks.run("on_train_epoch_end", epoch=epoch)
-            ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
+            ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"])
             final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
             if not noval or final_epoch:  # Calculate mAP
-                results, maps, _ = validate.run(data_dict,
+                results, maps, _ = validate.run(
+                    data_dict,
                     batch_size=batch_size // WORLD_SIZE * 2,
                     imgsz=imgsz,
                     half=amp,
@@ -379,7 +428,8 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                     save_dir=save_dir,
                     plots=False,
                     callbacks=callbacks,
-                                                compute_loss=compute_loss)
+                    compute_loss=compute_loss,
+                )
 
             # Update best mAP
             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
@@ -387,29 +437,30 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
             if fi > best_fitness:
                 best_fitness = fi
             log_vals = list(mloss) + list(results) + lr
-            callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
+            callbacks.run("on_fit_epoch_end", log_vals, epoch, best_fitness, fi)
 
             # Save model
             if (not nosave) or (final_epoch and not evolve):  # if save
                 ckpt = {
-                    'epoch': epoch,
+                    "epoch": epoch,
-                    'best_fitness': best_fitness,
+                    "best_fitness": best_fitness,
-                    'model': deepcopy(de_parallel(model)).half(),
+                    "model": deepcopy(de_parallel(model)).half(),
-                    'ema': deepcopy(ema.ema).half(),
+                    "ema": deepcopy(ema.ema).half(),
-                    'updates': ema.updates,
+                    "updates": ema.updates,
-                    'optimizer': optimizer.state_dict(),
+                    "optimizer": optimizer.state_dict(),
-                    'opt': vars(opt),
+                    "opt": vars(opt),
-                    'git': GIT_INFO,  # {remote, branch, commit} if a git repo
+                    "git": GIT_INFO,  # {remote, branch, commit} if a git repo
-                    'date': datetime.now().isoformat()}
+                    "date": datetime.now().isoformat(),
+                }
 
                 # Save last, best and delete
                 torch.save(ckpt, last)
                 if best_fitness == fi:
                     torch.save(ckpt, best)
                 if opt.save_period > 0 and epoch % opt.save_period == 0:
-                    torch.save(ckpt, w / f'epoch{epoch}.pt')
+                    torch.save(ckpt, w / f"epoch{epoch}.pt")
                 del ckpt
-                callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
+                callbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi)
 
         # EarlyStopping
         if RANK != -1:  # if DDP training
@@ -423,12 +474,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
         # end epoch ----------------------------------------------------------------------------------------------------
     # end training -----------------------------------------------------------------------------------------------------
     if RANK in {-1, 0}:
-        LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
+        LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.")
         for f in last, best:
             if f.exists():
                 strip_optimizer(f)  # strip optimizers
                 if f is best:
-                    LOGGER.info(f'\nValidating {f}...')
+                    LOGGER.info(f"\nValidating {f}...")
                     results, _, _ = validate.run(
                         data_dict,
                         batch_size=batch_size // WORLD_SIZE * 2,
@@ -442,11 +493,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                         verbose=True,
                         plots=plots,
                         callbacks=callbacks,
-                        compute_loss=compute_loss)  # val best model with plots
+                        compute_loss=compute_loss,
+                    )  # val best model with plots
                     if is_coco:
-                        callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
+                        callbacks.run("on_fit_epoch_end", list(mloss) + list(results) + lr, epoch, best_fitness, fi)
 
-        callbacks.run('on_train_end', last, best, epoch, results)
+        callbacks.run("on_train_end", last, best, epoch, results)
 
     torch.cuda.empty_cache()
     return results
@@ -454,55 +506,54 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
 
 def parse_opt(known=False):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
+    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
-    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
+    parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
-    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
+    parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path")
-    parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
+    parser.add_argument("--epochs", type=int, default=100, help="total training epochs")
-    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
+    parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch")
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
+    parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)")
-    parser.add_argument('--rect', action='store_true', help='rectangular training')
+    parser.add_argument("--rect", action="store_true", help="rectangular training")
-    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
+    parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")
-    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
+    parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
-    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
+    parser.add_argument("--noval", action="store_true", help="only validate final epoch")
-    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
+    parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor")
-    parser.add_argument('--noplots', action='store_true', help='save no plot files')
+    parser.add_argument("--noplots", action="store_true", help="save no plot files")
-    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
+    parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations")
-    parser.add_argument('--evolve_population',
+    parser.add_argument(
-                        type=str,
+        "--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population"
-                        default=ROOT / 'data/hyps',
+    )
-                        help='location for loading population')
+    parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation")
-    parser.add_argument('--resume_evolve', type=str, default=None, help='resume evolve from last generation')
+    parser.add_argument("--bucket", type=str, default="", help="gsutil bucket")
-    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
+    parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk")
-    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')
+    parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training")
-    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%")
-    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
+    parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class")
-    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
+    parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer")
-    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
+    parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode")
-    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
+    parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name")
-    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--quad", action="store_true", help="quad dataloader")
-    parser.add_argument('--quad', action='store_true', help='quad dataloader')
+    parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler")
-    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
+    parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon")
-    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
+    parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)")
-    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
+    parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2")
-    parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
+    parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)")
-    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
+    parser.add_argument("--seed", type=int, default=0, help="Global training seed")
-    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
+    parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
-    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
 
     # Logger arguments
-    parser.add_argument('--entity', default=None, help='Entity')
+    parser.add_argument("--entity", default=None, help="Entity")
-    parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
+    parser.add_argument("--upload_dataset", nargs="?", const=True, default=False, help='Upload data, "val" option')
-    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
+    parser.add_argument("--bbox_interval", type=int, default=-1, help="Set bounding-box image logging interval")
-    parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
+    parser.add_argument("--artifact_alias", type=str, default="latest", help="Version of dataset artifact to use")
 
     # NDJSON logging
-    parser.add_argument('--ndjson-console', action='store_true', help='Log ndjson to console')
+    parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console")
-    parser.add_argument('--ndjson-file', action='store_true', help='Log ndjson to file')
+    parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file")
 
     return parser.parse_known_args()[0] if known else parser.parse_args()
 
@@ -512,47 +563,53 @@ def main(opt, callbacks=Callbacks()):
     if RANK in {-1, 0}:
         print_args(vars(opt))
         check_git_status()
-        check_requirements(ROOT / 'requirements.txt')
+        check_requirements(ROOT / "requirements.txt")
 
     # Resume (from specified or most recent last.pt)
     if opt.resume and not check_comet_resume(opt) and not opt.evolve:
         last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
-        opt_yaml = last.parent.parent / 'opt.yaml'  # train options yaml
+        opt_yaml = last.parent.parent / "opt.yaml"  # train options yaml
         opt_data = opt.data  # original dataset
         if opt_yaml.is_file():
-            with open(opt_yaml, errors='ignore') as f:
+            with open(opt_yaml, errors="ignore") as f:
                 d = yaml.safe_load(f)
         else:
-            d = torch.load(last, map_location='cpu')['opt']
+            d = torch.load(last, map_location="cpu")["opt"]
         opt = argparse.Namespace(**d)  # replace
-        opt.cfg, opt.weights, opt.resume = '', str(last), True  # reinstate
+        opt.cfg, opt.weights, opt.resume = "", str(last), True  # reinstate
         if is_url(opt_data):
             opt.data = check_file(opt_data)  # avoid HUB resume auth timeout
     else:
-        opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
+        opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = (
-            check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project)  # checks
+            check_file(opt.data),
-        assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
+            check_yaml(opt.cfg),
+            check_yaml(opt.hyp),
+            str(opt.weights),
+            str(opt.project),
+        )  # checks
+        assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified"
         if opt.evolve:
-            if opt.project == str(ROOT / 'runs/train'):  # if default project name, rename to runs/evolve
+            if opt.project == str(ROOT / "runs/train"):  # if default project name, rename to runs/evolve
-                opt.project = str(ROOT / 'runs/evolve')
+                opt.project = str(ROOT / "runs/evolve")
             opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resume
-        if opt.name == 'cfg':
+        if opt.name == "cfg":
             opt.name = Path(opt.cfg).stem  # use model.yaml as name
         opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
 
     # DDP mode
     device = select_device(opt.device, batch_size=opt.batch_size)
     if LOCAL_RANK != -1:
-        msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
+        msg = "is not compatible with YOLOv5 Multi-GPU DDP training"
-        assert not opt.image_weights, f'--image-weights {msg}'
+        assert not opt.image_weights, f"--image-weights {msg}"
-        assert not opt.evolve, f'--evolve {msg}'
+        assert not opt.evolve, f"--evolve {msg}"
-        assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
+        assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size"
-        assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
+        assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE"
-        assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
+        assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
         torch.cuda.set_device(LOCAL_RANK)
-        device = torch.device('cuda', LOCAL_RANK)
+        device = torch.device("cuda", LOCAL_RANK)
-        dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo',
+        dist.init_process_group(
-                                timeout=timedelta(seconds=10800))
+            backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=10800)
+        )
 
     # Train
     if not opt.evolve:
@@ -562,35 +619,36 @@ def main(opt, callbacks=Callbacks()):
     else:
         # Hyperparameter evolution metadata (including this hyperparameter True-False, lower_limit, upper_limit)
         meta = {
-            'lr0': (False, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
+            "lr0": (False, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
-            'lrf': (False, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
+            "lrf": (False, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
-            'momentum': (False, 0.6, 0.98),  # SGD momentum/Adam beta1
+            "momentum": (False, 0.6, 0.98),  # SGD momentum/Adam beta1
-            'weight_decay': (False, 0.0, 0.001),  # optimizer weight decay
+            "weight_decay": (False, 0.0, 0.001),  # optimizer weight decay
-            'warmup_epochs': (False, 0.0, 5.0),  # warmup epochs (fractions ok)
+            "warmup_epochs": (False, 0.0, 5.0),  # warmup epochs (fractions ok)
-            'warmup_momentum': (False, 0.0, 0.95),  # warmup initial momentum
+            "warmup_momentum": (False, 0.0, 0.95),  # warmup initial momentum
-            'warmup_bias_lr': (False, 0.0, 0.2),  # warmup initial bias lr
+            "warmup_bias_lr": (False, 0.0, 0.2),  # warmup initial bias lr
-            'box': (False, 0.02, 0.2),  # box loss gain
+            "box": (False, 0.02, 0.2),  # box loss gain
-            'cls': (False, 0.2, 4.0),  # cls loss gain
+            "cls": (False, 0.2, 4.0),  # cls loss gain
-            'cls_pw': (False, 0.5, 2.0),  # cls BCELoss positive_weight
+            "cls_pw": (False, 0.5, 2.0),  # cls BCELoss positive_weight
-            'obj': (False, 0.2, 4.0),  # obj loss gain (scale with pixels)
+            "obj": (False, 0.2, 4.0),  # obj loss gain (scale with pixels)
-            'obj_pw': (False, 0.5, 2.0),  # obj BCELoss positive_weight
+            "obj_pw": (False, 0.5, 2.0),  # obj BCELoss positive_weight
-            'iou_t': (False, 0.1, 0.7),  # IoU training threshold
+            "iou_t": (False, 0.1, 0.7),  # IoU training threshold
-            'anchor_t': (False, 2.0, 8.0),  # anchor-multiple threshold
+            "anchor_t": (False, 2.0, 8.0),  # anchor-multiple threshold
-            'anchors': (False, 2.0, 10.0),  # anchors per output grid (0 to ignore)
+            "anchors": (False, 2.0, 10.0),  # anchors per output grid (0 to ignore)
-            'fl_gamma': (False, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
+            "fl_gamma": (False, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
-            'hsv_h': (True, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
+            "hsv_h": (True, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
-            'hsv_s': (True, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
+            "hsv_s": (True, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
-            'hsv_v': (True, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
+            "hsv_v": (True, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
-            'degrees': (True, 0.0, 45.0),  # image rotation (+/- deg)
+            "degrees": (True, 0.0, 45.0),  # image rotation (+/- deg)
-            'translate': (True, 0.0, 0.9),  # image translation (+/- fraction)
+            "translate": (True, 0.0, 0.9),  # image translation (+/- fraction)
-            'scale': (True, 0.0, 0.9),  # image scale (+/- gain)
+            "scale": (True, 0.0, 0.9),  # image scale (+/- gain)
-            'shear': (True, 0.0, 10.0),  # image shear (+/- deg)
+            "shear": (True, 0.0, 10.0),  # image shear (+/- deg)
-            'perspective': (True, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
+            "perspective": (True, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
-            'flipud': (True, 0.0, 1.0),  # image flip up-down (probability)
+            "flipud": (True, 0.0, 1.0),  # image flip up-down (probability)
-            'fliplr': (True, 0.0, 1.0),  # image flip left-right (probability)
+            "fliplr": (True, 0.0, 1.0),  # image flip left-right (probability)
-            'mosaic': (True, 0.0, 1.0),  # image mixup (probability)
+            "mosaic": (True, 0.0, 1.0),  # image mixup (probability)
-            'mixup': (True, 0.0, 1.0),  # image mixup (probability)
+            "mixup": (True, 0.0, 1.0),  # image mixup (probability)
-            'copy_paste': (True, 0.0, 1.0)}  # segment copy-paste (probability)
+            "copy_paste": (True, 0.0, 1.0),
+        }  # segment copy-paste (probability)
 
         # GA configs
         pop_size = 50
@@ -603,22 +661,25 @@ def main(opt, callbacks=Callbacks()):
         tournament_size_min = 2
         tournament_size_max = 10
 
-        with open(opt.hyp, errors='ignore') as f:
+        with open(opt.hyp, errors="ignore") as f:
             hyp = yaml.safe_load(f)  # load hyps dict
-            if 'anchors' not in hyp:  # anchors commented in hyp.yaml
+            if "anchors" not in hyp:  # anchors commented in hyp.yaml
-                hyp['anchors'] = 3
+                hyp["anchors"] = 3
         if opt.noautoanchor:
-            del hyp['anchors'], meta['anchors']
+            del hyp["anchors"], meta["anchors"]
         opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch
         # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
-        evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
+        evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv"
         if opt.bucket:
             # download evolve.csv if exists
-            subprocess.run([
+            subprocess.run(
-                'gsutil',
+                [
-                'cp',
+                    "gsutil",
-                f'gs://{opt.bucket}/evolve.csv',
+                    "cp",
-                str(evolve_csv), ])
+                    f"gs://{opt.bucket}/evolve.csv",
+                    str(evolve_csv),
+                ]
+            )
 
         # Delete the items in meta dictionary whose first value is False
         del_ = []
@@ -644,8 +705,8 @@ def main(opt, callbacks=Callbacks()):
 
         # If resuming evolution from a previous checkpoint
         if opt.resume_evolve is not None:
-            assert os.path.isfile(ROOT / opt.resume_evolve), 'evolve population path is wrong!'
+            assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!"
-            with open(ROOT / opt.resume_evolve, errors='ignore') as f:
+            with open(ROOT / opt.resume_evolve, errors="ignore") as f:
                 evolve_population = yaml.safe_load(f)
                 for value in evolve_population.values():
                     value = np.array([value[k] for k in hyp_GA.keys()])
@@ -653,7 +714,7 @@ def main(opt, callbacks=Callbacks()):
 
         # If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_population
         else:
-            yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith('.yaml')]
+            yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")]
             for file_name in yaml_files:
                 with open(os.path.join(opt.evolve_population, file_name)) as yaml_file:
                     value = yaml.safe_load(yaml_file)
@@ -661,27 +722,28 @@ def main(opt, callbacks=Callbacks()):
                     initial_values.append(list(value))
 
         # Generate random values within the search space for the rest of the population
-        if (initial_values is None):
+        if initial_values is None:
             population = [generate_individual(gene_ranges, len(hyp_GA)) for i in range(pop_size)]
         else:
-            if (pop_size > 1):
+            if pop_size > 1:
                 population = [
-                    generate_individual(gene_ranges, len(hyp_GA)) for i in range(pop_size - len(initial_values))]
+                    generate_individual(gene_ranges, len(hyp_GA)) for i in range(pop_size - len(initial_values))
+                ]
                 for initial_value in initial_values:
                     population = [initial_value] + population
 
         # Run the genetic algorithm for a fixed number of generations
         list_keys = list(hyp_GA.keys())
         for generation in range(opt.evolve):
-            if (generation >= 1):
+            if generation >= 1:
                 save_dict = {}
                 for i in range(len(population)):
                     little_dict = {}
                     for j in range(len(population[i])):
                         little_dict[list_keys[j]] = float(population[i][j])
-                    save_dict['gen' + str(generation) + 'number' + str(i)] = little_dict
+                    save_dict["gen" + str(generation) + "number" + str(i)] = little_dict
 
-                with open(save_dir / 'evolve_population.yaml', 'w') as outfile:
+                with open(save_dir / "evolve_population.yaml", "w") as outfile:
                     yaml.dump(save_dict, outfile, default_flow_style=False)
 
             # Adaptive elite size
@@ -695,8 +757,15 @@ def main(opt, callbacks=Callbacks()):
                 results = train(hyp.copy(), opt, device, callbacks)
                 callbacks = Callbacks()
                 # Write mutation results
-                keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
+                keys = (
-                        'val/box_loss', 'val/obj_loss', 'val/cls_loss')
+                    "metrics/precision",
+                    "metrics/recall",
+                    "metrics/mAP_0.5",
+                    "metrics/mAP_0.5:0.95",
+                    "val/box_loss",
+                    "val/obj_loss",
+                    "val/cls_loss",
+                )
                 print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
                 fitness_scores.append(results[2])
 
@@ -704,8 +773,10 @@ def main(opt, callbacks=Callbacks()):
             selected_indices = []
             for i in range(pop_size - elite_size):
                 # Adaptive tournament size
-                tournament_size = max(max(2, tournament_size_min),
+                tournament_size = max(
-                                      int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))))
+                    max(2, tournament_size_min),
+                    int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))),
+                )
                 # Perform tournament selection to choose the best individual
                 tournament_indices = random.sample(range(pop_size), tournament_size)
                 tournament_fitness = [fitness_scores[j] for j in tournament_indices]
@@ -721,16 +792,18 @@ def main(opt, callbacks=Callbacks()):
                 parent1_index = selected_indices[random.randint(0, pop_size - 1)]
                 parent2_index = selected_indices[random.randint(0, pop_size - 1)]
                 # Adaptive crossover rate
-                crossover_rate = max(crossover_rate_min,
+                crossover_rate = max(
-                                     min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve)))
+                    crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve))
+                )
                 if random.uniform(0, 1) < crossover_rate:
                     crossover_point = random.randint(1, len(hyp_GA) - 1)
                     child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:]
                 else:
                     child = population[parent1_index]
                 # Adaptive mutation rate
-                mutation_rate = max(mutation_rate_min,
+                mutation_rate = max(
-                                    min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)))
+                    mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve))
+                )
                 for j in range(len(hyp_GA)):
                     if random.uniform(0, 1) < mutation_rate:
                         child[j] += random.uniform(-0.1, 0.1)
@@ -741,12 +814,14 @@ def main(opt, callbacks=Callbacks()):
         # Print the best solution found
         best_index = fitness_scores.index(max(fitness_scores))
         best_individual = population[best_index]
-        print('Best solution found:', best_individual)
+        print("Best solution found:", best_individual)
         # Plot results
         plot_evolve(evolve_csv)
-        LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
+        LOGGER.info(
+            f'Hyperparameter evolution finished {opt.evolve} generations\n'
             f"Results saved to {colorstr('bold', save_dir)}\n"
-                    f'Usage example: $ python train.py --hyp {evolve_yaml}')
+            f'Usage example: $ python train.py --hyp {evolve_yaml}'
+        )
 
 
 def generate_individual(input_ranges, individual_length):
@@ -766,6 +841,6 @@ def run(**kwargs):
     return opt
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)

utils/__init__.py

@@ -1,21 +1,19 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""utils/initialization."""
-utils/initialization
-"""
 
 import contextlib
 import platform
 import threading
 
 
-def emojis(str=''):
+def emojis(str=""):
     # Return platform-dependent emoji-safe version of string
-    return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
+    return str.encode().decode("ascii", "ignore") if platform.system() == "Windows" else str
 
 
 class TryExcept(contextlib.ContextDecorator):
     # YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
-    def __init__(self, msg=''):
+    def __init__(self, msg=""):
         self.msg = msg
 
     def __enter__(self):
@@ -43,13 +41,13 @@ def join_threads(verbose=False):
     for t in threading.enumerate():
         if t is not main_thread:
             if verbose:
-                print(f'Joining thread {t.name}')
+                print(f"Joining thread {t.name}")
             t.join()
 
 
 def notebook_init(verbose=True):
     # Check system software and hardware
-    print('Checking setup...')
+    print("Checking setup...")
 
     import os
     import shutil
@@ -63,24 +61,25 @@ def notebook_init(verbose=True):
 
     import psutil
 
-    if check_requirements('wandb', install=False):
+    if check_requirements("wandb", install=False):
-        os.system('pip uninstall -y wandb')  # eliminate unexpected account creation prompt with infinite hang
+        os.system("pip uninstall -y wandb")  # eliminate unexpected account creation prompt with infinite hang
     if is_colab():
-        shutil.rmtree('/content/sample_data', ignore_errors=True)  # remove colab /sample_data directory
+        shutil.rmtree("/content/sample_data", ignore_errors=True)  # remove colab /sample_data directory
 
     # System info
     display = None
     if verbose:
         gb = 1 << 30  # bytes to GiB (1024 ** 3)
         ram = psutil.virtual_memory().total
-        total, used, free = shutil.disk_usage('/')
+        total, used, free = shutil.disk_usage("/")
         with contextlib.suppress(Exception):  # clear display if ipython is installed
             from IPython import display
+
             display.clear_output()
-        s = f'({os.cpu_count()} CPUs, {ram / gb:.1f} GB RAM, {(total - free) / gb:.1f}/{total / gb:.1f} GB disk)'
+        s = f"({os.cpu_count()} CPUs, {ram / gb:.1f} GB RAM, {(total - free) / gb:.1f}/{total / gb:.1f} GB disk)"
     else:
-        s = ''
+        s = ""
 
     select_device(newline=False)
-    print(emojis(f'Setup complete ✅ {s}'))
+    print(emojis(f"Setup complete ✅ {s}"))
     return display

utils/activations.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Activation functions."""
-Activation functions
-"""
 
 import torch
 import torch.nn as nn
@@ -33,7 +31,6 @@ class Mish(nn.Module):
 class MemoryEfficientMish(nn.Module):
     # Mish activation memory-efficient
     class F(torch.autograd.Function):
-
         @staticmethod
         def forward(ctx, x):
             ctx.save_for_backward(x)

utils/augmentations.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Image augmentation functions."""
-Image augmentation functions
-"""
 
 import math
 import random
@@ -23,10 +21,11 @@ class Albumentations:
     # YOLOv5 Albumentations class (optional, only used if package is installed)
     def __init__(self, size=640):
         self.transform = None
-        prefix = colorstr('albumentations: ')
+        prefix = colorstr("albumentations: ")
         try:
             import albumentations as A
-            check_version(A.__version__, '1.0.3', hard=True)  # version requirement
+
+            check_version(A.__version__, "1.0.3", hard=True)  # version requirement
 
             T = [
                 A.RandomResizedCrop(height=size, width=size, scale=(0.8, 1.0), ratio=(0.9, 1.11), p=0.0),
@@ -36,19 +35,20 @@ class Albumentations:
                 A.CLAHE(p=0.01),
                 A.RandomBrightnessContrast(p=0.0),
                 A.RandomGamma(p=0.0),
-                A.ImageCompression(quality_lower=75, p=0.0)]  # transforms
+                A.ImageCompression(quality_lower=75, p=0.0),
-            self.transform = A.Compose(T, bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
+            ]  # transforms
+            self.transform = A.Compose(T, bbox_params=A.BboxParams(format="yolo", label_fields=["class_labels"]))
 
-            LOGGER.info(prefix + ', '.join(f'{x}'.replace('always_apply=False, ', '') for x in T if x.p))
+            LOGGER.info(prefix + ", ".join(f"{x}".replace("always_apply=False, ", "") for x in T if x.p))
         except ImportError:  # package not installed, skip
             pass
         except Exception as e:
-            LOGGER.info(f'{prefix}{e}')
+            LOGGER.info(f"{prefix}{e}")
 
     def __call__(self, im, labels, p=1.0):
         if self.transform and random.random() < p:
             new = self.transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0])  # transformed
-            im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
+            im, labels = new["image"], np.array([[c, *b] for c, b in zip(new["class_labels"], new["bboxes"])])
         return im, labels
 
 
@@ -141,15 +141,9 @@ def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleF
     return im, ratio, (dw, dh)
 
 
-def random_perspective(im,
+def random_perspective(
-                       targets=(),
+    im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.1, shear=10, perspective=0.0, border=(0, 0)
-                       segments=(),
+):
-                       degrees=10,
-                       translate=.1,
-                       scale=.1,
-                       shear=10,
-                       perspective=0.0,
-                       border=(0, 0)):
     # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10))
     # targets = [cls, xyxy]
 
@@ -312,18 +306,20 @@ def classify_albumentations(
     jitter=0.4,
     mean=IMAGENET_MEAN,
     std=IMAGENET_STD,
-        auto_aug=False):
+    auto_aug=False,
+):
     # YOLOv5 classification Albumentations (optional, only used if package is installed)
-    prefix = colorstr('albumentations: ')
+    prefix = colorstr("albumentations: ")
     try:
         import albumentations as A
         from albumentations.pytorch import ToTensorV2
-        check_version(A.__version__, '1.0.3', hard=True)  # version requirement
+
+        check_version(A.__version__, "1.0.3", hard=True)  # version requirement
         if augment:  # Resize and crop
             T = [A.RandomResizedCrop(height=size, width=size, scale=scale, ratio=ratio)]
             if auto_aug:
                 # TODO: implement AugMix, AutoAug & RandAug in albumentation
-                LOGGER.info(f'{prefix}auto augmentations are currently not supported')
+                LOGGER.info(f"{prefix}auto augmentations are currently not supported")
             else:
                 if hflip > 0:
                     T += [A.HorizontalFlip(p=hflip)]
@@ -335,18 +331,18 @@ def classify_albumentations(
         else:  # Use fixed crop for eval set (reproducibility)
             T = [A.SmallestMaxSize(max_size=size), A.CenterCrop(height=size, width=size)]
         T += [A.Normalize(mean=mean, std=std), ToTensorV2()]  # Normalize and convert to Tensor
-        LOGGER.info(prefix + ', '.join(f'{x}'.replace('always_apply=False, ', '') for x in T if x.p))
+        LOGGER.info(prefix + ", ".join(f"{x}".replace("always_apply=False, ", "") for x in T if x.p))
         return A.Compose(T)
 
     except ImportError:  # package not installed, skip
-        LOGGER.warning(f'{prefix}⚠️ not found, install with `pip install albumentations` (recommended)')
+        LOGGER.warning(f"{prefix}⚠️ not found, install with `pip install albumentations` (recommended)")
     except Exception as e:
-        LOGGER.info(f'{prefix}{e}')
+        LOGGER.info(f"{prefix}{e}")
 
 
 def classify_transforms(size=224):
     # Transforms to apply if albumentations not installed
-    assert isinstance(size, int), f'ERROR: classify_transforms size {size} must be integer, not (list, tuple)'
+    assert isinstance(size, int), f"ERROR: classify_transforms size {size} must be integer, not (list, tuple)"
     # T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
     return T.Compose([CenterCrop(size), ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
 

utils/autoanchor.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""AutoAnchor utils."""
-AutoAnchor utils
-"""
 
 import random
 
@@ -13,7 +11,7 @@ from tqdm import tqdm
 from utils import TryExcept
 from utils.general import LOGGER, TQDM_BAR_FORMAT, colorstr
 
-PREFIX = colorstr('AutoAnchor: ')
+PREFIX = colorstr("AutoAnchor: ")
 
 
 def check_anchor_order(m):
@@ -22,14 +20,14 @@ def check_anchor_order(m):
     da = a[-1] - a[0]  # delta a
     ds = m.stride[-1] - m.stride[0]  # delta s
     if da and (da.sign() != ds.sign()):  # same order
-        LOGGER.info(f'{PREFIX}Reversing anchor order')
+        LOGGER.info(f"{PREFIX}Reversing anchor order")
         m.anchors[:] = m.anchors.flip(0)
 
 
-@TryExcept(f'{PREFIX}ERROR')
+@TryExcept(f"{PREFIX}ERROR")
 def check_anchors(dataset, model, thr=4.0, imgsz=640):
     # Check anchor fit to data, recompute if necessary
-    m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1]  # Detect()
+    m = model.module.model[-1] if hasattr(model, "module") else model.model[-1]  # Detect()
     shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
     scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1))  # augment scale
     wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float()  # wh
@@ -45,11 +43,11 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
     stride = m.stride.to(m.anchors.device).view(-1, 1, 1)  # model strides
     anchors = m.anchors.clone() * stride  # current anchors
     bpr, aat = metric(anchors.cpu().view(-1, 2))
-    s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). '
+    s = f"\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). "
     if bpr > 0.98:  # threshold to recompute
-        LOGGER.info(f'{s}Current anchors are a good fit to dataset ✅')
+        LOGGER.info(f"{s}Current anchors are a good fit to dataset ✅")
     else:
-        LOGGER.info(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...')
+        LOGGER.info(f"{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...")
         na = m.anchors.numel() // 2  # number of anchors
         anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
         new_bpr = metric(anchors)[0]
@@ -58,14 +56,15 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
             m.anchors[:] = anchors.clone().view_as(m.anchors)
             check_anchor_order(m)  # must be in pixel-space (not grid-space)
             m.anchors /= stride
-            s = f'{PREFIX}Done ✅ (optional: update model *.yaml to use these anchors in the future)'
+            s = f"{PREFIX}Done ✅ (optional: update model *.yaml to use these anchors in the future)"
         else:
-            s = f'{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)'
+            s = f"{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)"
         LOGGER.info(s)
 
 
-def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
+def kmean_anchors(dataset="./data/coco128.yaml", n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
-    """ Creates kmeans-evolved anchors from training dataset
+    """
+    Creates kmeans-evolved anchors from training dataset.
 
     Arguments:
         dataset: path to data.yaml, or a loaded dataset
@@ -100,20 +99,23 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
         k = k[np.argsort(k.prod(1))]  # sort small to large
         x, best = metric(k, wh0)
         bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n  # best possible recall, anch > thr
-        s = f'{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n' \
+        s = (
-            f'{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' \
+            f"{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n"
-            f'past_thr={x[x > thr].mean():.3f}-mean: '
+            f"{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, "
+            f"past_thr={x[x > thr].mean():.3f}-mean: "
+        )
         for x in k:
-            s += '%i,%i, ' % (round(x[0]), round(x[1]))
+            s += "%i,%i, " % (round(x[0]), round(x[1]))
         if verbose:
             LOGGER.info(s[:-2])
         return k
 
     if isinstance(dataset, str):  # *.yaml file
-        with open(dataset, errors='ignore') as f:
+        with open(dataset, errors="ignore") as f:
             data_dict = yaml.safe_load(f)  # model dict
         from utils.dataloaders import LoadImagesAndLabels
-        dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
+
+        dataset = LoadImagesAndLabels(data_dict["train"], augment=True, rect=True)
 
     # Get label wh
     shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
@@ -122,19 +124,19 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
     # Filter
     i = (wh0 < 3.0).any(1).sum()
     if i:
-        LOGGER.info(f'{PREFIX}WARNING ⚠️ Extremely small objects found: {i} of {len(wh0)} labels are <3 pixels in size')
+        LOGGER.info(f"{PREFIX}WARNING ⚠️ Extremely small objects found: {i} of {len(wh0)} labels are <3 pixels in size")
     wh = wh0[(wh0 >= 2.0).any(1)].astype(np.float32)  # filter > 2 pixels
     # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1)  # multiply by random scale 0-1
 
     # Kmeans init
     try:
-        LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
+        LOGGER.info(f"{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...")
         assert n <= len(wh)  # apply overdetermined constraint
         s = wh.std(0)  # sigmas for whitening
         k = kmeans(wh / s, n, iter=30)[0] * s  # points
         assert n == len(k)  # kmeans may return fewer points than requested if wh is insufficient or too similar
     except Exception:
-        LOGGER.warning(f'{PREFIX}WARNING ⚠️ switching strategies from kmeans to random init')
+        LOGGER.warning(f"{PREFIX}WARNING ⚠️ switching strategies from kmeans to random init")
         k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size  # random init
     wh, wh0 = (torch.tensor(x, dtype=torch.float32) for x in (wh, wh0))
     k = print_results(k, verbose=False)
@@ -162,7 +164,7 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
         fg = anchor_fitness(kg)
         if fg > f:
             f, k = fg, kg.copy()
-            pbar.desc = f'{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
+            pbar.desc = f"{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}"
             if verbose:
                 print_results(k, verbose)
 

utils/autobatch.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Auto-batch utils."""
-Auto-batch utils
-"""
 
 from copy import deepcopy
 
@@ -27,14 +25,14 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
     #     print(autobatch(model))
 
     # Check device
-    prefix = colorstr('AutoBatch: ')
+    prefix = colorstr("AutoBatch: ")
-    LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}')
+    LOGGER.info(f"{prefix}Computing optimal batch size for --imgsz {imgsz}")
     device = next(model.parameters()).device  # get model device
-    if device.type == 'cpu':
+    if device.type == "cpu":
-        LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}')
+        LOGGER.info(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}")
         return batch_size
     if torch.backends.cudnn.benchmark:
-        LOGGER.info(f'{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}')
+        LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}")
         return batch_size
 
     # Inspect CUDA memory
@@ -45,7 +43,7 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
     r = torch.cuda.memory_reserved(device) / gb  # GiB reserved
     a = torch.cuda.memory_allocated(device) / gb  # GiB allocated
     f = t - (r + a)  # GiB free
-    LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free')
+    LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free")
 
     # Profile batch sizes
     batch_sizes = [1, 2, 4, 8, 16]
@@ -53,7 +51,7 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
         img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
         results = profile(img, model, n=3, device=device)
     except Exception as e:
-        LOGGER.warning(f'{prefix}{e}')
+        LOGGER.warning(f"{prefix}{e}")
 
     # Fit a solution
     y = [x[2] for x in results if x]  # memory [2]
@@ -65,8 +63,8 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
             b = batch_sizes[max(i - 1, 0)]  # select prior safe point
     if b < 1 or b > 1024:  # b outside of safe range
         b = batch_size
-        LOGGER.warning(f'{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.')
+        LOGGER.warning(f"{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.")
 
     fraction = (np.polyval(p, b) + r + a) / t  # actual fraction predicted
-    LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅')
+    LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅")
     return b

utils/aws/resume.py

@@ -14,27 +14,27 @@ if str(ROOT) not in sys.path:
     sys.path.append(str(ROOT))  # add ROOT to PATH
 
 port = 0  # --master_port
-path = Path('').resolve()
+path = Path("").resolve()
-for last in path.rglob('*/**/last.pt'):
+for last in path.rglob("*/**/last.pt"):
     ckpt = torch.load(last)
-    if ckpt['optimizer'] is None:
+    if ckpt["optimizer"] is None:
         continue
 
     # Load opt.yaml
-    with open(last.parent.parent / 'opt.yaml', errors='ignore') as f:
+    with open(last.parent.parent / "opt.yaml", errors="ignore") as f:
         opt = yaml.safe_load(f)
 
     # Get device count
-    d = opt['device'].split(',')  # devices
+    d = opt["device"].split(",")  # devices
     nd = len(d)  # number of devices
     ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1)  # distributed data parallel
 
     if ddp:  # multi-GPU
         port += 1
-        cmd = f'python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}'
+        cmd = f"python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}"
     else:  # single-GPU
-        cmd = f'python train.py --resume {last}'
+        cmd = f"python train.py --resume {last}"
 
-    cmd += ' > /dev/null 2>&1 &'  # redirect output to dev/null and run in daemon thread
+    cmd += " > /dev/null 2>&1 &"  # redirect output to dev/null and run in daemon thread
     print(cmd)
     os.system(cmd)

utils/callbacks.py

@@ -1,43 +1,40 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Callback utils."""
-Callback utils
-"""
 
 import threading
 
 
 class Callbacks:
-    """"
+    """" Handles all registered callbacks for YOLOv5 Hooks."""
-    Handles all registered callbacks for YOLOv5 Hooks
-    """
 
     def __init__(self):
         # Define the available callbacks
         self._callbacks = {
-            'on_pretrain_routine_start': [],
+            "on_pretrain_routine_start": [],
-            'on_pretrain_routine_end': [],
+            "on_pretrain_routine_end": [],
-            'on_train_start': [],
+            "on_train_start": [],
-            'on_train_epoch_start': [],
+            "on_train_epoch_start": [],
-            'on_train_batch_start': [],
+            "on_train_batch_start": [],
-            'optimizer_step': [],
+            "optimizer_step": [],
-            'on_before_zero_grad': [],
+            "on_before_zero_grad": [],
-            'on_train_batch_end': [],
+            "on_train_batch_end": [],
-            'on_train_epoch_end': [],
+            "on_train_epoch_end": [],
-            'on_val_start': [],
+            "on_val_start": [],
-            'on_val_batch_start': [],
+            "on_val_batch_start": [],
-            'on_val_image_end': [],
+            "on_val_image_end": [],
-            'on_val_batch_end': [],
+            "on_val_batch_end": [],
-            'on_val_end': [],
+            "on_val_end": [],
-            'on_fit_epoch_end': [],  # fit = train + val
+            "on_fit_epoch_end": [],  # fit = train + val
-            'on_model_save': [],
+            "on_model_save": [],
-            'on_train_end': [],
+            "on_train_end": [],
-            'on_params_update': [],
+            "on_params_update": [],
-            'teardown': [], }
+            "teardown": [],
+        }
         self.stop_training = False  # set True to interrupt training
 
-    def register_action(self, hook, name='', callback=None):
+    def register_action(self, hook, name="", callback=None):
         """
-        Register a new action to a callback hook
+        Register a new action to a callback hook.
 
         Args:
             hook: The callback hook name to register the action to
@@ -46,11 +43,11 @@ class Callbacks:
         """
         assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
         assert callable(callback), f"callback '{callback}' is not callable"
-        self._callbacks[hook].append({'name': name, 'callback': callback})
+        self._callbacks[hook].append({"name": name, "callback": callback})
 
     def get_registered_actions(self, hook=None):
-        """"
+        """
-        Returns all the registered actions by callback hook
+        " Returns all the registered actions by callback hook.
 
         Args:
             hook: The name of the hook to check, defaults to all
@@ -59,7 +56,7 @@ class Callbacks:
 
     def run(self, hook, *args, thread=False, **kwargs):
         """
-        Loop through the registered actions and fire all callbacks on main thread
+        Loop through the registered actions and fire all callbacks on main thread.
 
         Args:
             hook: The name of the hook to check, defaults to all
@@ -71,6 +68,6 @@ class Callbacks:
         assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
         for logger in self._callbacks[hook]:
             if thread:
-                threading.Thread(target=logger['callback'], args=args, kwargs=kwargs, daemon=True).start()
+                threading.Thread(target=logger["callback"], args=args, kwargs=kwargs, daemon=True).start()
             else:
-                logger['callback'](*args, **kwargs)
+                logger["callback"](*args, **kwargs)

utils/downloads.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Download utils."""
-Download utils
-"""
 
 import logging
 import subprocess
@@ -23,89 +21,90 @@ def is_url(url, check=True):
         return False
 
 
-def gsutil_getsize(url=''):
+def gsutil_getsize(url=""):
     # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
-    output = subprocess.check_output(['gsutil', 'du', url], shell=True, encoding='utf-8')
+    output = subprocess.check_output(["gsutil", "du", url], shell=True, encoding="utf-8")
     if output:
         return int(output.split()[0])
     return 0
 
 
-def url_getsize(url='https://ultralytics.com/images/bus.jpg'):
+def url_getsize(url="https://ultralytics.com/images/bus.jpg"):
     # Return downloadable file size in bytes
     response = requests.head(url, allow_redirects=True)
-    return int(response.headers.get('content-length', -1))
+    return int(response.headers.get("content-length", -1))
 
 
 def curl_download(url, filename, *, silent: bool = False) -> bool:
-    """
+    """Download a file from a url to a filename using curl."""
-    Download a file from a url to a filename using curl.
+    silent_option = "sS" if silent else ""  # silent
-    """
+    proc = subprocess.run(
-    silent_option = 'sS' if silent else ''  # silent
+        [
-    proc = subprocess.run([
+            "curl",
-        'curl',
+            "-#",
-        '-#',
+            f"-{silent_option}L",
-        f'-{silent_option}L',
             url,
-        '--output',
+            "--output",
             filename,
-        '--retry',
+            "--retry",
-        '9',
+            "9",
-        '-C',
+            "-C",
-        '-', ])
+            "-",
+        ]
+    )
     return proc.returncode == 0
 
 
-def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
+def safe_download(file, url, url2=None, min_bytes=1e0, error_msg=""):
     # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
     from utils.general import LOGGER
 
     file = Path(file)
     assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
     try:  # url1
-        LOGGER.info(f'Downloading {url} to {file}...')
+        LOGGER.info(f"Downloading {url} to {file}...")
         torch.hub.download_url_to_file(url, str(file), progress=LOGGER.level <= logging.INFO)
         assert file.exists() and file.stat().st_size > min_bytes, assert_msg  # check
     except Exception as e:  # url2
         if file.exists():
             file.unlink()  # remove partial downloads
-        LOGGER.info(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
+        LOGGER.info(f"ERROR: {e}\nRe-attempting {url2 or url} to {file}...")
         # curl download, retry and resume on fail
         curl_download(url2 or url, file)
     finally:
         if not file.exists() or file.stat().st_size < min_bytes:  # check
             if file.exists():
                 file.unlink()  # remove partial downloads
-            LOGGER.info(f'ERROR: {assert_msg}\n{error_msg}')
+            LOGGER.info(f"ERROR: {assert_msg}\n{error_msg}")
-        LOGGER.info('')
+        LOGGER.info("")
 
 
-def attempt_download(file, repo='ultralytics/yolov5', release='v7.0'):
+def attempt_download(file, repo="ultralytics/yolov5", release="v7.0"):
     # Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v7.0', etc.
     from utils.general import LOGGER
 
-    def github_assets(repository, version='latest'):
+    def github_assets(repository, version="latest"):
         # Return GitHub repo tag (i.e. 'v7.0') and assets (i.e. ['yolov5s.pt', 'yolov5m.pt', ...])
-        if version != 'latest':
+        if version != "latest":
-            version = f'tags/{version}'  # i.e. tags/v7.0
+            version = f"tags/{version}"  # i.e. tags/v7.0
-        response = requests.get(f'https://api.github.com/repos/{repository}/releases/{version}').json()  # github api
+        response = requests.get(f"https://api.github.com/repos/{repository}/releases/{version}").json()  # github api
-        return response['tag_name'], [x['name'] for x in response['assets']]  # tag, assets
+        return response["tag_name"], [x["name"] for x in response["assets"]]  # tag, assets
 
-    file = Path(str(file).strip().replace("'", ''))
+    file = Path(str(file).strip().replace("'", ""))
     if not file.exists():
         # URL specified
         name = Path(urllib.parse.unquote(str(file))).name  # decode '%2F' to '/' etc.
-        if str(file).startswith(('http:/', 'https:/')):  # download
+        if str(file).startswith(("http:/", "https:/")):  # download
-            url = str(file).replace(':/', '://')  # Pathlib turns :// -> :/
+            url = str(file).replace(":/", "://")  # Pathlib turns :// -> :/
-            file = name.split('?')[0]  # parse authentication https://url.com/file.txt?auth...
+            file = name.split("?")[0]  # parse authentication https://url.com/file.txt?auth...
             if Path(file).is_file():
-                LOGGER.info(f'Found {url} locally at {file}')  # file already exists
+                LOGGER.info(f"Found {url} locally at {file}")  # file already exists
             else:
-                safe_download(file=file, url=url, min_bytes=1E5)
+                safe_download(file=file, url=url, min_bytes=1e5)
             return file
 
         # GitHub assets
-        assets = [f'yolov5{size}{suffix}.pt' for size in 'nsmlx' for suffix in ('', '6', '-cls', '-seg')]  # default
+        assets = [f"yolov5{size}{suffix}.pt" for size in "nsmlx" for suffix in ("", "6", "-cls", "-seg")]  # default
         try:
             tag, assets = github_assets(repo, release)
         except Exception:
@@ -113,15 +112,17 @@ def attempt_download(file, repo='ultralytics/yolov5', release='v7.0'):
                 tag, assets = github_assets(repo)  # latest release
             except Exception:
                 try:
-                    tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
+                    tag = subprocess.check_output("git tag", shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
                 except Exception:
                     tag = release
 
         if name in assets:
             file.parent.mkdir(parents=True, exist_ok=True)  # make parent dir (if required)
-            safe_download(file,
+            safe_download(
-                          url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
+                file,
-                          min_bytes=1E5,
+                url=f"https://github.com/{repo}/releases/download/{tag}/{name}",
-                          error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/{tag}')
+                min_bytes=1e5,
+                error_msg=f"{file} missing, try downloading from https://github.com/{repo}/releases/{tag}",
+            )
 
     return str(file)

utils/flask_rest_api/example_request.py

@@ -1,19 +1,17 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Perform test request."""
-Perform test request
-"""
 
 import pprint
 
 import requests
 
-DETECTION_URL = 'http://localhost:5000/v1/object-detection/yolov5s'
+DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
-IMAGE = 'zidane.jpg'
+IMAGE = "zidane.jpg"
 
 # Read image
-with open(IMAGE, 'rb') as f:
+with open(IMAGE, "rb") as f:
     image_data = f.read()
 
-response = requests.post(DETECTION_URL, files={'image': image_data}).json()
+response = requests.post(DETECTION_URL, files={"image": image_data}).json()
 
 pprint.pprint(response)

utils/flask_rest_api/restapi.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Run a Flask REST API exposing one or more YOLOv5s models."""
-Run a Flask REST API exposing one or more YOLOv5s models
-"""
 
 import argparse
 import io
@@ -13,36 +11,36 @@ from PIL import Image
 app = Flask(__name__)
 models = {}
 
-DETECTION_URL = '/v1/object-detection/<model>'
+DETECTION_URL = "/v1/object-detection/<model>"
 
 
-@app.route(DETECTION_URL, methods=['POST'])
+@app.route(DETECTION_URL, methods=["POST"])
 def predict(model):
-    if request.method != 'POST':
+    if request.method != "POST":
         return
 
-    if request.files.get('image'):
+    if request.files.get("image"):
         # Method 1
         # with request.files["image"] as f:
         #     im = Image.open(io.BytesIO(f.read()))
 
         # Method 2
-        im_file = request.files['image']
+        im_file = request.files["image"]
         im_bytes = im_file.read()
         im = Image.open(io.BytesIO(im_bytes))
 
         if model in models:
             results = models[model](im, size=640)  # reduce size=320 for faster inference
-            return results.pandas().xyxy[0].to_json(orient='records')
+            return results.pandas().xyxy[0].to_json(orient="records")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Flask API exposing YOLOv5 model')
+    parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
-    parser.add_argument('--port', default=5000, type=int, help='port number')
+    parser.add_argument("--port", default=5000, type=int, help="port number")
-    parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s')
+    parser.add_argument("--model", nargs="+", default=["yolov5s"], help="model(s) to run, i.e. --model yolov5n yolov5s")
     opt = parser.parse_args()
 
     for m in opt.model:
-        models[m] = torch.hub.load('ultralytics/yolov5', m, force_reload=True, skip_validation=True)
+        models[m] = torch.hub.load("ultralytics/yolov5", m, force_reload=True, skip_validation=True)
 
-    app.run(host='0.0.0.0', port=opt.port)  # debug=True causes Restarting with stat
+    app.run(host="0.0.0.0", port=opt.port)  # debug=True causes Restarting with stat

utils/general.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""General utils."""
-General utils
-"""
 
 import contextlib
 import glob
@@ -40,9 +38,9 @@ import yaml
 try:
     import ultralytics
 
-    assert hasattr(ultralytics, '__version__')  # verify package is not directory
+    assert hasattr(ultralytics, "__version__")  # verify package is not directory
 except (ImportError, AssertionError):
-    os.system('pip install -U ultralytics')
+    os.system("pip install -U ultralytics")
     import ultralytics
 
 from ultralytics.utils.checks import check_requirements
@@ -53,67 +51,67 @@ from utils.metrics import box_iou, fitness
 
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[1]  # YOLOv5 root directory
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
 
 # Settings
 NUM_THREADS = min(8, max(1, os.cpu_count() - 1))  # number of YOLOv5 multiprocessing threads
-DATASETS_DIR = Path(os.getenv('YOLOv5_DATASETS_DIR', ROOT.parent / 'datasets'))  # global datasets directory
+DATASETS_DIR = Path(os.getenv("YOLOv5_DATASETS_DIR", ROOT.parent / "datasets"))  # global datasets directory
-AUTOINSTALL = str(os.getenv('YOLOv5_AUTOINSTALL', True)).lower() == 'true'  # global auto-install mode
+AUTOINSTALL = str(os.getenv("YOLOv5_AUTOINSTALL", True)).lower() == "true"  # global auto-install mode
-VERBOSE = str(os.getenv('YOLOv5_VERBOSE', True)).lower() == 'true'  # global verbose mode
+VERBOSE = str(os.getenv("YOLOv5_VERBOSE", True)).lower() == "true"  # global verbose mode
-TQDM_BAR_FORMAT = '{l_bar}{bar:10}{r_bar}'  # tqdm bar format
+TQDM_BAR_FORMAT = "{l_bar}{bar:10}{r_bar}"  # tqdm bar format
-FONT = 'Arial.ttf'  # https://ultralytics.com/assets/Arial.ttf
+FONT = "Arial.ttf"  # https://ultralytics.com/assets/Arial.ttf
 
-torch.set_printoptions(linewidth=320, precision=5, profile='long')
+torch.set_printoptions(linewidth=320, precision=5, profile="long")
-np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format})  # format short g, %precision=5
+np.set_printoptions(linewidth=320, formatter={"float_kind": "{:11.5g}".format})  # format short g, %precision=5
 pd.options.display.max_columns = 10
 cv2.setNumThreads(0)  # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
-os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS)  # NumExpr max threads
+os.environ["NUMEXPR_MAX_THREADS"] = str(NUM_THREADS)  # NumExpr max threads
-os.environ['OMP_NUM_THREADS'] = '1' if platform.system() == 'darwin' else str(NUM_THREADS)  # OpenMP (PyTorch and SciPy)
+os.environ["OMP_NUM_THREADS"] = "1" if platform.system() == "darwin" else str(NUM_THREADS)  # OpenMP (PyTorch and SciPy)
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # suppress verbose TF compiler warnings in Colab
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # suppress verbose TF compiler warnings in Colab
 
 
-def is_ascii(s=''):
+def is_ascii(s=""):
     # Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
     s = str(s)  # convert list, tuple, None, etc. to str
-    return len(s.encode().decode('ascii', 'ignore')) == len(s)
+    return len(s.encode().decode("ascii", "ignore")) == len(s)
 
 
-def is_chinese(s='人工智能'):
+def is_chinese(s="人工智能"):
     # Is string composed of any Chinese characters?
-    return bool(re.search('[\u4e00-\u9fff]', str(s)))
+    return bool(re.search("[\u4e00-\u9fff]", str(s)))
 
 
 def is_colab():
     # Is environment a Google Colab instance?
-    return 'google.colab' in sys.modules
+    return "google.colab" in sys.modules
 
 
 def is_jupyter():
     """
-    Check if the current script is running inside a Jupyter Notebook.
+    Check if the current script is running inside a Jupyter Notebook. Verified on Colab, Jupyterlab, Kaggle, Paperspace.
-    Verified on Colab, Jupyterlab, Kaggle, Paperspace.
 
     Returns:
         bool: True if running inside a Jupyter Notebook, False otherwise.
     """
     with contextlib.suppress(Exception):
         from IPython import get_ipython
+
         return get_ipython() is not None
     return False
 
 
 def is_kaggle():
     # Is environment a Kaggle Notebook?
-    return os.environ.get('PWD') == '/kaggle/working' and os.environ.get('KAGGLE_URL_BASE') == 'https://www.kaggle.com'
+    return os.environ.get("PWD") == "/kaggle/working" and os.environ.get("KAGGLE_URL_BASE") == "https://www.kaggle.com"
 
 
 def is_docker() -> bool:
     """Check if the process runs inside a docker container."""
-    if Path('/.dockerenv').exists():
+    if Path("/.dockerenv").exists():
         return True
     try:  # check if docker is in control groups
-        with open('/proc/self/cgroup') as file:
+        with open("/proc/self/cgroup") as file:
-            return any('docker' in line for line in file)
+            return any("docker" in line for line in file)
     except OSError:
         return False
 
@@ -122,9 +120,9 @@ def is_writeable(dir, test=False):
     # Return True if directory has write permissions, test opening a file with write permissions if test=True
     if not test:
         return os.access(dir, os.W_OK)  # possible issues on Windows
-    file = Path(dir) / 'tmp.txt'
+    file = Path(dir) / "tmp.txt"
     try:
-        with open(file, 'w'):  # open file with write permissions
+        with open(file, "w"):  # open file with write permissions
             pass
         file.unlink()  # remove file
         return True
@@ -132,47 +130,52 @@ def is_writeable(dir, test=False):
         return False
 
 
-LOGGING_NAME = 'yolov5'
+LOGGING_NAME = "yolov5"
 
 
 def set_logging(name=LOGGING_NAME, verbose=True):
     # sets up logging for the given name
-    rank = int(os.getenv('RANK', -1))  # rank in world for Multi-GPU trainings
+    rank = int(os.getenv("RANK", -1))  # rank in world for Multi-GPU trainings
     level = logging.INFO if verbose and rank in {-1, 0} else logging.ERROR
-    logging.config.dictConfig({
+    logging.config.dictConfig(
-        'version': 1,
+        {
-        'disable_existing_loggers': False,
+            "version": 1,
-        'formatters': {
+            "disable_existing_loggers": False,
+            "formatters": {name: {"format": "%(message)s"}},
+            "handlers": {
                 name: {
-                'format': '%(message)s'}},
+                    "class": "logging.StreamHandler",
-        'handlers': {
+                    "formatter": name,
+                    "level": level,
+                }
+            },
+            "loggers": {
                 name: {
-                'class': 'logging.StreamHandler',
+                    "level": level,
-                'formatter': name,
+                    "handlers": [name],
-                'level': level, }},
+                    "propagate": False,
-        'loggers': {
+                }
-            name: {
+            },
-                'level': level,
+        }
-                'handlers': [name],
+    )
-                'propagate': False, }}})
 
 
 set_logging(LOGGING_NAME)  # run before defining LOGGER
 LOGGER = logging.getLogger(LOGGING_NAME)  # define globally (used in train.py, val.py, detect.py, etc.)
-if platform.system() == 'Windows':
+if platform.system() == "Windows":
     for fn in LOGGER.info, LOGGER.warning:
         setattr(LOGGER, fn.__name__, lambda x: fn(emojis(x)))  # emoji safe logging
 
 
-def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
+def user_config_dir(dir="Ultralytics", env_var="YOLOV5_CONFIG_DIR"):
     # Return path of user configuration directory. Prefer environment variable if exists. Make dir if required.
     env = os.getenv(env_var)
     if env:
         path = Path(env)  # use environment variable
     else:
-        cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'}  # 3 OS dirs
+        cfg = {"Windows": "AppData/Roaming", "Linux": ".config", "Darwin": "Library/Application Support"}  # 3 OS dirs
-        path = Path.home() / cfg.get(platform.system(), '')  # OS-specific config dir
+        path = Path.home() / cfg.get(platform.system(), "")  # OS-specific config dir
-        path = (path if is_writeable(path) else Path('/tmp')) / dir  # GCP and AWS lambda fix, only /tmp is writeable
+        path = (path if is_writeable(path) else Path("/tmp")) / dir  # GCP and AWS lambda fix, only /tmp is writeable
     path.mkdir(exist_ok=True)  # make if required
     return path
 
@@ -185,7 +188,7 @@ class Profile(contextlib.ContextDecorator):
     def __init__(self, t=0.0, device: torch.device = None):
         self.t = t
         self.device = device
-        self.cuda = True if (device and str(device)[:4] == 'cuda') else False
+        self.cuda = True if (device and str(device)[:4] == "cuda") else False
 
     def __enter__(self):
         self.start = self.time()
@@ -203,7 +206,7 @@ class Profile(contextlib.ContextDecorator):
 
 class Timeout(contextlib.ContextDecorator):
     # YOLOv5 Timeout class. Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager
-    def __init__(self, seconds, *, timeout_msg='', suppress_timeout_errors=True):
+    def __init__(self, seconds, *, timeout_msg="", suppress_timeout_errors=True):
         self.seconds = int(seconds)
         self.timeout_message = timeout_msg
         self.suppress = bool(suppress_timeout_errors)
@@ -212,12 +215,12 @@ class Timeout(contextlib.ContextDecorator):
         raise TimeoutError(self.timeout_message)
 
     def __enter__(self):
-        if platform.system() != 'Windows':  # not supported on Windows
+        if platform.system() != "Windows":  # not supported on Windows
             signal.signal(signal.SIGALRM, self._timeout_handler)  # Set handler for SIGALRM
             signal.alarm(self.seconds)  # start countdown for SIGALRM to be raised
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if platform.system() != 'Windows':
+        if platform.system() != "Windows":
             signal.alarm(0)  # Cancel SIGALRM if it's scheduled
             if self.suppress and exc_type is TimeoutError:  # Suppress TimeoutError
                 return True
@@ -238,7 +241,7 @@ class WorkingDirectory(contextlib.ContextDecorator):
 
 def methods(instance):
     # Get class/instance methods
-    return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith('__')]
+    return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith("__")]
 
 
 def print_args(args: Optional[dict] = None, show_file=True, show_func=False):
@@ -249,11 +252,11 @@ def print_args(args: Optional[dict] = None, show_file=True, show_func=False):
         args, _, _, frm = inspect.getargvalues(x)
         args = {k: v for k, v in frm.items() if k in args}
     try:
-        file = Path(file).resolve().relative_to(ROOT).with_suffix('')
+        file = Path(file).resolve().relative_to(ROOT).with_suffix("")
     except ValueError:
         file = Path(file).stem
-    s = (f'{file}: ' if show_file else '') + (f'{func}: ' if show_func else '')
+    s = (f"{file}: " if show_file else "") + (f"{func}: " if show_func else "")
-    LOGGER.info(colorstr(s) + ', '.join(f'{k}={v}' for k, v in args.items()))
+    LOGGER.info(colorstr(s) + ", ".join(f"{k}={v}" for k, v in args.items()))
 
 
 def init_seeds(seed=0, deterministic=False):
@@ -264,11 +267,11 @@ def init_seeds(seed=0, deterministic=False):
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)  # for Multi-GPU, exception safe
     # torch.backends.cudnn.benchmark = True  # AutoBatch problem https://github.com/ultralytics/yolov5/issues/9287
-    if deterministic and check_version(torch.__version__, '1.12.0'):  # https://github.com/ultralytics/yolov5/pull/8213
+    if deterministic and check_version(torch.__version__, "1.12.0"):  # https://github.com/ultralytics/yolov5/pull/8213
         torch.use_deterministic_algorithms(True)
         torch.backends.cudnn.deterministic = True
-        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-        os.environ['PYTHONHASHSEED'] = str(seed)
+        os.environ["PYTHONHASHSEED"] = str(seed)
 
 
 def intersect_dicts(da, db, exclude=()):
@@ -282,22 +285,22 @@ def get_default_args(func):
     return {k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty}
 
 
-def get_latest_run(search_dir='.'):
+def get_latest_run(search_dir="."):
     # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
-    last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
+    last_list = glob.glob(f"{search_dir}/**/last*.pt", recursive=True)
-    return max(last_list, key=os.path.getctime) if last_list else ''
+    return max(last_list, key=os.path.getctime) if last_list else ""
 
 
 def file_age(path=__file__):
     # Return days since last file update
-    dt = (datetime.now() - datetime.fromtimestamp(Path(path).stat().st_mtime))  # delta
+    dt = datetime.now() - datetime.fromtimestamp(Path(path).stat().st_mtime)  # delta
     return dt.days  # + dt.seconds / 86400  # fractional days
 
 
 def file_date(path=__file__):
     # Return human-readable file modification date, i.e. '2021-3-26'
     t = datetime.fromtimestamp(Path(path).stat().st_mtime)
-    return f'{t.year}-{t.month}-{t.day}'
+    return f"{t.year}-{t.month}-{t.day}"
 
 
 def file_size(path):
@@ -307,7 +310,7 @@ def file_size(path):
     if path.is_file():
         return path.stat().st_size / mb
     elif path.is_dir():
-        return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / mb
+        return sum(f.stat().st_size for f in path.glob("**/*") if f.is_file()) / mb
     else:
         return 0.0
 
@@ -319,7 +322,7 @@ def check_online():
     def run_once():
         # Check once
         try:
-            socket.create_connection(('1.1.1.1', 443), 5)  # check host accessibility
+            socket.create_connection(("1.1.1.1", 443), 5)  # check host accessibility
             return True
         except OSError:
             return False
@@ -330,68 +333,69 @@ def check_online():
 def git_describe(path=ROOT):  # path must be a directory
     # Return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
     try:
-        assert (Path(path) / '.git').is_dir()
+        assert (Path(path) / ".git").is_dir()
-        return check_output(f'git -C {path} describe --tags --long --always', shell=True).decode()[:-1]
+        return check_output(f"git -C {path} describe --tags --long --always", shell=True).decode()[:-1]
     except Exception:
-        return ''
+        return ""
 
 
 @TryExcept()
 @WorkingDirectory(ROOT)
-def check_git_status(repo='ultralytics/yolov5', branch='master'):
+def check_git_status(repo="ultralytics/yolov5", branch="master"):
     # YOLOv5 status check, recommend 'git pull' if code is out of date
-    url = f'https://github.com/{repo}'
+    url = f"https://github.com/{repo}"
-    msg = f', for updates see {url}'
+    msg = f", for updates see {url}"
-    s = colorstr('github: ')  # string
+    s = colorstr("github: ")  # string
-    assert Path('.git').exists(), s + 'skipping check (not a git repository)' + msg
+    assert Path(".git").exists(), s + "skipping check (not a git repository)" + msg
-    assert check_online(), s + 'skipping check (offline)' + msg
+    assert check_online(), s + "skipping check (offline)" + msg
 
-    splits = re.split(pattern=r'\s', string=check_output('git remote -v', shell=True).decode())
+    splits = re.split(pattern=r"\s", string=check_output("git remote -v", shell=True).decode())
     matches = [repo in s for s in splits]
     if any(matches):
         remote = splits[matches.index(True) - 1]
     else:
-        remote = 'ultralytics'
+        remote = "ultralytics"
-        check_output(f'git remote add {remote} {url}', shell=True)
+        check_output(f"git remote add {remote} {url}", shell=True)
-    check_output(f'git fetch {remote}', shell=True, timeout=5)  # git fetch
+    check_output(f"git fetch {remote}", shell=True, timeout=5)  # git fetch
-    local_branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip()  # checked out
+    local_branch = check_output("git rev-parse --abbrev-ref HEAD", shell=True).decode().strip()  # checked out
-    n = int(check_output(f'git rev-list {local_branch}..{remote}/{branch} --count', shell=True))  # commits behind
+    n = int(check_output(f"git rev-list {local_branch}..{remote}/{branch} --count", shell=True))  # commits behind
     if n > 0:
-        pull = 'git pull' if remote == 'origin' else f'git pull {remote} {branch}'
+        pull = "git pull" if remote == "origin" else f"git pull {remote} {branch}"
         s += f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use '{pull}' or 'git clone {url}' to update."
     else:
-        s += f'up to date with {url} ✅'
+        s += f"up to date with {url} ✅"
     LOGGER.info(s)
 
 
 @WorkingDirectory(ROOT)
-def check_git_info(path='.'):
+def check_git_info(path="."):
     # YOLOv5 git info check, return {remote, branch, commit}
-    check_requirements('gitpython')
+    check_requirements("gitpython")
     import git
+
     try:
         repo = git.Repo(path)
-        remote = repo.remotes.origin.url.replace('.git', '')  # i.e. 'https://github.com/ultralytics/yolov5'
+        remote = repo.remotes.origin.url.replace(".git", "")  # i.e. 'https://github.com/ultralytics/yolov5'
         commit = repo.head.commit.hexsha  # i.e. '3134699c73af83aac2a481435550b968d5792c0d'
         try:
             branch = repo.active_branch.name  # i.e. 'main'
         except TypeError:  # not on any branch
             branch = None  # i.e. 'detached HEAD' state
-        return {'remote': remote, 'branch': branch, 'commit': commit}
+        return {"remote": remote, "branch": branch, "commit": commit}
     except git.exc.InvalidGitRepositoryError:  # path is not a git dir
-        return {'remote': None, 'branch': None, 'commit': None}
+        return {"remote": None, "branch": None, "commit": None}
 
 
-def check_python(minimum='3.8.0'):
+def check_python(minimum="3.8.0"):
     # Check current python version vs. required python version
-    check_version(platform.python_version(), minimum, name='Python ', hard=True)
+    check_version(platform.python_version(), minimum, name="Python ", hard=True)
 
 
-def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False, verbose=False):
+def check_version(current="0.0.0", minimum="0.0.0", name="version ", pinned=False, hard=False, verbose=False):
     # Check version vs. required version
     current, minimum = (pkg.parse_version(x) for x in (current, minimum))
     result = (current == minimum) if pinned else (current >= minimum)  # bool
-    s = f'WARNING ⚠️ {name}{minimum} is required by YOLOv5, but {name}{current} is currently installed'  # string
+    s = f"WARNING ⚠️ {name}{minimum} is required by YOLOv5, but {name}{current} is currently installed"  # string
     if hard:
         assert result, emojis(s)  # assert min requirements met
     if verbose and not result:
@@ -407,7 +411,7 @@ def check_img_size(imgsz, s=32, floor=0):
         imgsz = list(imgsz)  # convert to list if tuple
         new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
     if new_size != imgsz:
-        LOGGER.warning(f'WARNING ⚠️ --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
+        LOGGER.warning(f"WARNING ⚠️ --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}")
     return new_size
 
 
@@ -416,18 +420,18 @@ def check_imshow(warn=False):
     try:
         assert not is_jupyter()
         assert not is_docker()
-        cv2.imshow('test', np.zeros((1, 1, 3)))
+        cv2.imshow("test", np.zeros((1, 1, 3)))
         cv2.waitKey(1)
         cv2.destroyAllWindows()
         cv2.waitKey(1)
         return True
     except Exception as e:
         if warn:
-            LOGGER.warning(f'WARNING ⚠️ Environment does not support cv2.imshow() or PIL Image.show()\n{e}')
+            LOGGER.warning(f"WARNING ⚠️ Environment does not support cv2.imshow() or PIL Image.show()\n{e}")
         return False
 
 
-def check_suffix(file='yolov5s.pt', suffix=('.pt', ), msg=''):
+def check_suffix(file="yolov5s.pt", suffix=(".pt",), msg=""):
     # Check file(s) for acceptable suffix
     if file and suffix:
         if isinstance(suffix, str):
@@ -435,38 +439,40 @@ def check_suffix(file='yolov5s.pt', suffix=('.pt', ), msg=''):
         for f in file if isinstance(file, (list, tuple)) else [file]:
             s = Path(f).suffix.lower()  # file suffix
             if len(s):
-                assert s in suffix, f'{msg}{f} acceptable suffix is {suffix}'
+                assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}"
 
 
-def check_yaml(file, suffix=('.yaml', '.yml')):
+def check_yaml(file, suffix=(".yaml", ".yml")):
     # Search/download YAML file (if necessary) and return path, checking suffix
     return check_file(file, suffix)
 
 
-def check_file(file, suffix=''):
+def check_file(file, suffix=""):
     # Search/download file (if necessary) and return path
     check_suffix(file, suffix)  # optional
     file = str(file)  # convert to str()
     if os.path.isfile(file) or not file:  # exists
         return file
-    elif file.startswith(('http:/', 'https:/')):  # download
+    elif file.startswith(("http:/", "https:/")):  # download
         url = file  # warning: Pathlib turns :// -> :/
-        file = Path(urllib.parse.unquote(file).split('?')[0]).name  # '%2F' to '/', split https://url.com/file.txt?auth
+        file = Path(urllib.parse.unquote(file).split("?")[0]).name  # '%2F' to '/', split https://url.com/file.txt?auth
         if os.path.isfile(file):
-            LOGGER.info(f'Found {url} locally at {file}')  # file already exists
+            LOGGER.info(f"Found {url} locally at {file}")  # file already exists
         else:
-            LOGGER.info(f'Downloading {url} to {file}...')
+            LOGGER.info(f"Downloading {url} to {file}...")
             torch.hub.download_url_to_file(url, file)
-            assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}'  # check
+            assert Path(file).exists() and Path(file).stat().st_size > 0, f"File download failed: {url}"  # check
         return file
-    elif file.startswith('clearml://'):  # ClearML Dataset ID
+    elif file.startswith("clearml://"):  # ClearML Dataset ID
-        assert 'clearml' in sys.modules, "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
+        assert (
+            "clearml" in sys.modules
+        ), "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
         return file
     else:  # search
         files = []
-        for d in 'data', 'models', 'utils':  # search directories
+        for d in "data", "models", "utils":  # search directories
-            files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True))  # find file
+            files.extend(glob.glob(str(ROOT / d / "**" / file), recursive=True))  # find file
-        assert len(files), f'File not found: {file}'  # assert file was found
+        assert len(files), f"File not found: {file}"  # assert file was found
         assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}"  # assert unique
         return files[0]  # return file
 
@@ -476,8 +482,8 @@ def check_font(font=FONT, progress=False):
     font = Path(font)
     file = CONFIG_DIR / font.name
     if not font.exists() and not file.exists():
-        url = f'https://ultralytics.com/assets/{font.name}'
+        url = f"https://ultralytics.com/assets/{font.name}"
-        LOGGER.info(f'Downloading {url} to {file}...')
+        LOGGER.info(f"Downloading {url} to {file}...")
         torch.hub.download_url_to_file(url, str(file), progress=progress)
 
 
@@ -485,10 +491,10 @@ def check_dataset(data, autodownload=True):
     # Download, check and/or unzip dataset if not found locally
 
     # Download (optional)
-    extract_dir = ''
+    extract_dir = ""
     if isinstance(data, (str, Path)) and (is_zipfile(data) or is_tarfile(data)):
-        download(data, dir=f'{DATASETS_DIR}/{Path(data).stem}', unzip=True, delete=False, curl=False, threads=1)
+        download(data, dir=f"{DATASETS_DIR}/{Path(data).stem}", unzip=True, delete=False, curl=False, threads=1)
-        data = next((DATASETS_DIR / Path(data).stem).rglob('*.yaml'))
+        data = next((DATASETS_DIR / Path(data).stem).rglob("*.yaml"))
         extract_dir, autodownload = data.parent, False
 
     # Read yaml (optional)
@@ -496,54 +502,54 @@ def check_dataset(data, autodownload=True):
         data = yaml_load(data)  # dictionary
 
     # Checks
-    for k in 'train', 'val', 'names':
+    for k in "train", "val", "names":
         assert k in data, emojis(f"data.yaml '{k}:' field missing ❌")
-    if isinstance(data['names'], (list, tuple)):  # old array format
+    if isinstance(data["names"], (list, tuple)):  # old array format
-        data['names'] = dict(enumerate(data['names']))  # convert to dict
+        data["names"] = dict(enumerate(data["names"]))  # convert to dict
-    assert all(isinstance(k, int) for k in data['names'].keys()), 'data.yaml names keys must be integers, i.e. 2: car'
+    assert all(isinstance(k, int) for k in data["names"].keys()), "data.yaml names keys must be integers, i.e. 2: car"
-    data['nc'] = len(data['names'])
+    data["nc"] = len(data["names"])
 
     # Resolve paths
-    path = Path(extract_dir or data.get('path') or '')  # optional 'path' default to '.'
+    path = Path(extract_dir or data.get("path") or "")  # optional 'path' default to '.'
     if not path.is_absolute():
         path = (ROOT / path).resolve()
-        data['path'] = path  # download scripts
+        data["path"] = path  # download scripts
-    for k in 'train', 'val', 'test':
+    for k in "train", "val", "test":
         if data.get(k):  # prepend path
             if isinstance(data[k], str):
                 x = (path / data[k]).resolve()
-                if not x.exists() and data[k].startswith('../'):
+                if not x.exists() and data[k].startswith("../"):
                     x = (path / data[k][3:]).resolve()
                 data[k] = str(x)
             else:
                 data[k] = [str((path / x).resolve()) for x in data[k]]
 
     # Parse yaml
-    train, val, test, s = (data.get(x) for x in ('train', 'val', 'test', 'download'))
+    train, val, test, s = (data.get(x) for x in ("train", "val", "test", "download"))
     if val:
         val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])]  # val path
         if not all(x.exists() for x in val):
-            LOGGER.info('\nDataset not found ⚠️, missing paths %s' % [str(x) for x in val if not x.exists()])
+            LOGGER.info("\nDataset not found ⚠️, missing paths %s" % [str(x) for x in val if not x.exists()])
             if not s or not autodownload:
-                raise Exception('Dataset not found ❌')
+                raise Exception("Dataset not found ❌")
             t = time.time()
-            if s.startswith('http') and s.endswith('.zip'):  # URL
+            if s.startswith("http") and s.endswith(".zip"):  # URL
                 f = Path(s).name  # filename
-                LOGGER.info(f'Downloading {s} to {f}...')
+                LOGGER.info(f"Downloading {s} to {f}...")
                 torch.hub.download_url_to_file(s, f)
                 Path(DATASETS_DIR).mkdir(parents=True, exist_ok=True)  # create root
                 unzip_file(f, path=DATASETS_DIR)  # unzip
                 Path(f).unlink()  # remove zip
                 r = None  # success
-            elif s.startswith('bash '):  # bash script
+            elif s.startswith("bash "):  # bash script
-                LOGGER.info(f'Running {s} ...')
+                LOGGER.info(f"Running {s} ...")
                 r = subprocess.run(s, shell=True)
             else:  # python script
-                r = exec(s, {'yaml': data})  # return None
+                r = exec(s, {"yaml": data})  # return None
-            dt = f'({round(time.time() - t, 1)}s)'
+            dt = f"({round(time.time() - t, 1)}s)"
-            s = f"success ✅ {dt}, saved to {colorstr('bold', DATASETS_DIR)}" if r in (0, None) else f'failure {dt} ❌'
+            s = f"success ✅ {dt}, saved to {colorstr('bold', DATASETS_DIR)}" if r in (0, None) else f"failure {dt} ❌"
-            LOGGER.info(f'Dataset download {s}')
+            LOGGER.info(f"Dataset download {s}")
-    check_font('Arial.ttf' if is_ascii(data['names']) else 'Arial.Unicode.ttf', progress=True)  # download fonts
+    check_font("Arial.ttf" if is_ascii(data["names"]) else "Arial.Unicode.ttf", progress=True)  # download fonts
     return data  # dictionary
 
 
@@ -559,35 +565,35 @@ def check_amp(model):
         b = m(im).xywhn[0]  # AMP inference
         return a.shape == b.shape and torch.allclose(a, b, atol=0.1)  # close to 10% absolute tolerance
 
-    prefix = colorstr('AMP: ')
+    prefix = colorstr("AMP: ")
     device = next(model.parameters()).device  # get model device
-    if device.type in ('cpu', 'mps'):
+    if device.type in ("cpu", "mps"):
         return False  # AMP only used on CUDA devices
-    f = ROOT / 'data' / 'images' / 'bus.jpg'  # image to check
+    f = ROOT / "data" / "images" / "bus.jpg"  # image to check
-    im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if check_online() else np.ones((640, 640, 3))
+    im = f if f.exists() else "https://ultralytics.com/images/bus.jpg" if check_online() else np.ones((640, 640, 3))
     try:
-        assert amp_allclose(deepcopy(model), im) or amp_allclose(DetectMultiBackend('yolov5n.pt', device), im)
+        assert amp_allclose(deepcopy(model), im) or amp_allclose(DetectMultiBackend("yolov5n.pt", device), im)
-        LOGGER.info(f'{prefix}checks passed ✅')
+        LOGGER.info(f"{prefix}checks passed ✅")
         return True
     except Exception:
-        help_url = 'https://github.com/ultralytics/yolov5/issues/7908'
+        help_url = "https://github.com/ultralytics/yolov5/issues/7908"
-        LOGGER.warning(f'{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}')
+        LOGGER.warning(f"{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}")
         return False
 
 
-def yaml_load(file='data.yaml'):
+def yaml_load(file="data.yaml"):
     # Single-line safe yaml loading
-    with open(file, errors='ignore') as f:
+    with open(file, errors="ignore") as f:
         return yaml.safe_load(f)
 
 
-def yaml_save(file='data.yaml', data={}):
+def yaml_save(file="data.yaml", data={}):
     # Single-line safe yaml saving
-    with open(file, 'w') as f:
+    with open(file, "w") as f:
         yaml.safe_dump({k: str(v) if isinstance(v, Path) else v for k, v in data.items()}, f, sort_keys=False)
 
 
-def unzip_file(file, path=None, exclude=('.DS_Store', '__MACOSX')):
+def unzip_file(file, path=None, exclude=(".DS_Store", "__MACOSX")):
     # Unzip a *.zip file to path/, excluding files containing strings in exclude list
     if path is None:
         path = Path(file).parent  # default path
@@ -599,11 +605,11 @@ def unzip_file(file, path=None, exclude=('.DS_Store', '__MACOSX')):
 
 def url2file(url):
     # Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt
-    url = str(Path(url)).replace(':/', '://')  # Pathlib turns :// -> :/
+    url = str(Path(url)).replace(":/", "://")  # Pathlib turns :// -> :/
-    return Path(urllib.parse.unquote(url)).name.split('?')[0]  # '%2F' to '/', split https://url.com/file.txt?auth
+    return Path(urllib.parse.unquote(url)).name.split("?")[0]  # '%2F' to '/', split https://url.com/file.txt?auth
 
 
-def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1, retry=3):
+def download(url, dir=".", unzip=True, delete=True, curl=False, threads=1, retry=3):
     # Multithreaded file download and unzip function, used in data.yaml for autodownload
     def download_one(url, dir):
         # Download 1 file
@@ -612,7 +618,7 @@ def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1, retry
             f = Path(url)  # filename
         else:  # does not exist
             f = dir / Path(url).name
-            LOGGER.info(f'Downloading {url} to {f}...')
+            LOGGER.info(f"Downloading {url} to {f}...")
             for i in range(retry + 1):
                 if curl:
                     success = curl_download(url, f, silent=(threads > 1))
@@ -622,18 +628,18 @@ def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1, retry
                 if success:
                     break
                 elif i < retry:
-                    LOGGER.warning(f'⚠️ Download failure, retrying {i + 1}/{retry} {url}...')
+                    LOGGER.warning(f"⚠️ Download failure, retrying {i + 1}/{retry} {url}...")
                 else:
-                    LOGGER.warning(f'❌ Failed to download {url}...')
+                    LOGGER.warning(f"❌ Failed to download {url}...")
 
-        if unzip and success and (f.suffix == '.gz' or is_zipfile(f) or is_tarfile(f)):
+        if unzip and success and (f.suffix == ".gz" or is_zipfile(f) or is_tarfile(f)):
-            LOGGER.info(f'Unzipping {f}...')
+            LOGGER.info(f"Unzipping {f}...")
             if is_zipfile(f):
                 unzip_file(f, dir)  # unzip
             elif is_tarfile(f):
-                subprocess.run(['tar', 'xf', f, '--directory', f.parent], check=True)  # unzip
+                subprocess.run(["tar", "xf", f, "--directory", f.parent], check=True)  # unzip
-            elif f.suffix == '.gz':
+            elif f.suffix == ".gz":
-                subprocess.run(['tar', 'xfz', f, '--directory', f.parent], check=True)  # unzip
+                subprocess.run(["tar", "xfz", f, "--directory", f.parent], check=True)  # unzip
             if delete:
                 f.unlink()  # remove zip
 
@@ -658,7 +664,7 @@ def make_divisible(x, divisor):
 
 def clean_str(s):
     # Cleans a string by replacing special characters with underscore _
-    return re.sub(pattern='[|@#!¡·$€%&()=?¿^*;:,¨´><+]', repl='_', string=s)
+    return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
 
 
 def one_cycle(y1=0.0, y2=1.0, steps=100):
@@ -668,28 +674,29 @@ def one_cycle(y1=0.0, y2=1.0, steps=100):
 
 def colorstr(*input):
     # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e.  colorstr('blue', 'hello world')
-    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])  # color arguments, string
+    *args, string = input if len(input) > 1 else ("blue", "bold", input[0])  # color arguments, string
     colors = {
-        'black': '\033[30m',  # basic colors
+        "black": "\033[30m",  # basic colors
-        'red': '\033[31m',
+        "red": "\033[31m",
-        'green': '\033[32m',
+        "green": "\033[32m",
-        'yellow': '\033[33m',
+        "yellow": "\033[33m",
-        'blue': '\033[34m',
+        "blue": "\033[34m",
-        'magenta': '\033[35m',
+        "magenta": "\033[35m",
-        'cyan': '\033[36m',
+        "cyan": "\033[36m",
-        'white': '\033[37m',
+        "white": "\033[37m",
-        'bright_black': '\033[90m',  # bright colors
+        "bright_black": "\033[90m",  # bright colors
-        'bright_red': '\033[91m',
+        "bright_red": "\033[91m",
-        'bright_green': '\033[92m',
+        "bright_green": "\033[92m",
-        'bright_yellow': '\033[93m',
+        "bright_yellow": "\033[93m",
-        'bright_blue': '\033[94m',
+        "bright_blue": "\033[94m",
-        'bright_magenta': '\033[95m',
+        "bright_magenta": "\033[95m",
-        'bright_cyan': '\033[96m',
+        "bright_cyan": "\033[96m",
-        'bright_white': '\033[97m',
+        "bright_white": "\033[97m",
-        'end': '\033[0m',  # misc
+        "end": "\033[0m",  # misc
-        'bold': '\033[1m',
+        "bold": "\033[1m",
-        'underline': '\033[4m'}
+        "underline": "\033[4m",
-    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
+    }
+    return "".join(colors[x] for x in args) + f"{string}" + colors["end"]
 
 
 def labels_to_class_weights(labels, nc=80):
@@ -725,9 +732,87 @@ def coco80_to_coco91_class():  # converts 80-index (val2014) to 91-index (paper)
     # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)]  # darknet to coco
     # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)]  # coco to darknet
     return [
-        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
+        1,
-        35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
+        2,
-        64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        24,
+        25,
+        27,
+        28,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        40,
+        41,
+        42,
+        43,
+        44,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        53,
+        54,
+        55,
+        56,
+        57,
+        58,
+        59,
+        60,
+        61,
+        62,
+        63,
+        64,
+        65,
+        67,
+        70,
+        72,
+        73,
+        74,
+        75,
+        76,
+        77,
+        78,
+        79,
+        80,
+        81,
+        82,
+        84,
+        85,
+        86,
+        87,
+        88,
+        89,
+        90,
+    ]
 
 
 def xyxy2xywh(x):
@@ -784,7 +869,10 @@ def segment2box(segment, width=640, height=640):
     # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
     x, y = segment.T  # segment xy
     inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
-    x, y, = x[inside], y[inside]
+    (
+        x,
+        y,
+    ) = x[inside], y[inside]
     return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4))  # xyxy
 
 
@@ -875,20 +963,21 @@ def non_max_suppression(
     max_det=300,
     nm=0,  # number of masks
 ):
-    """Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
+    """
+    Non-Maximum Suppression (NMS) on inference results to reject overlapping detections.
 
     Returns:
          list of detections, on (n,6) tensor per image [xyxy, conf, cls]
     """
 
     # Checks
-    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
-    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
     if isinstance(prediction, (list, tuple)):  # YOLOv5 model in validation model, output = (inference_out, loss_out)
         prediction = prediction[0]  # select only inference output
 
     device = prediction.device
-    mps = 'mps' in device.type  # Apple MPS
+    mps = "mps" in device.type  # Apple MPS
     if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
         prediction = prediction.cpu()
     bs = prediction.shape[0]  # batch size
@@ -959,7 +1048,7 @@ def non_max_suppression(
         boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
         i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
         i = i[:max_det]  # limit detections
-        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
+        if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
             # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
             iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
             weights = iou * scores[None]  # box weights
@@ -971,31 +1060,31 @@ def non_max_suppression(
         if mps:
             output[xi] = output[xi].to(device)
         if (time.time() - t) > time_limit:
-            LOGGER.warning(f'WARNING ⚠️ NMS time limit {time_limit:.3f}s exceeded')
+            LOGGER.warning(f"WARNING ⚠️ NMS time limit {time_limit:.3f}s exceeded")
             break  # time limit exceeded
 
     return output
 
 
-def strip_optimizer(f='best.pt', s=''):  # from utils.general import *; strip_optimizer()
+def strip_optimizer(f="best.pt", s=""):  # from utils.general import *; strip_optimizer()
     # Strip optimizer from 'f' to finalize training, optionally save as 's'
-    x = torch.load(f, map_location=torch.device('cpu'))
+    x = torch.load(f, map_location=torch.device("cpu"))
-    if x.get('ema'):
+    if x.get("ema"):
-        x['model'] = x['ema']  # replace model with ema
+        x["model"] = x["ema"]  # replace model with ema
-    for k in 'optimizer', 'best_fitness', 'ema', 'updates':  # keys
+    for k in "optimizer", "best_fitness", "ema", "updates":  # keys
         x[k] = None
-    x['epoch'] = -1
+    x["epoch"] = -1
-    x['model'].half()  # to FP16
+    x["model"].half()  # to FP16
-    for p in x['model'].parameters():
+    for p in x["model"].parameters():
         p.requires_grad = False
     torch.save(x, s or f)
-    mb = os.path.getsize(s or f) / 1E6  # filesize
+    mb = os.path.getsize(s or f) / 1e6  # filesize
     LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
 
 
-def print_mutation(keys, results, hyp, save_dir, bucket, prefix=colorstr('evolve: ')):
+def print_mutation(keys, results, hyp, save_dir, bucket, prefix=colorstr("evolve: ")):
-    evolve_csv = save_dir / 'evolve.csv'
+    evolve_csv = save_dir / "evolve.csv"
-    evolve_yaml = save_dir / 'hyp_evolve.yaml'
+    evolve_yaml = save_dir / "hyp_evolve.yaml"
     keys = tuple(keys) + tuple(hyp.keys())  # [results + hyps]
     keys = tuple(x.strip() for x in keys)
     vals = results + tuple(hyp.values())
@@ -1003,33 +1092,48 @@ def print_mutation(keys, results, hyp, save_dir, bucket, prefix=colorstr('evolve
 
     # Download (optional)
     if bucket:
-        url = f'gs://{bucket}/evolve.csv'
+        url = f"gs://{bucket}/evolve.csv"
         if gsutil_getsize(url) > (evolve_csv.stat().st_size if evolve_csv.exists() else 0):
-            subprocess.run(['gsutil', 'cp', f'{url}', f'{save_dir}'])  # download evolve.csv if larger than local
+            subprocess.run(["gsutil", "cp", f"{url}", f"{save_dir}"])  # download evolve.csv if larger than local
 
     # Log to evolve.csv
-    s = '' if evolve_csv.exists() else (('%20s,' * n % keys).rstrip(',') + '\n')  # add header
+    s = "" if evolve_csv.exists() else (("%20s," * n % keys).rstrip(",") + "\n")  # add header
-    with open(evolve_csv, 'a') as f:
+    with open(evolve_csv, "a") as f:
-        f.write(s + ('%20.5g,' * n % vals).rstrip(',') + '\n')
+        f.write(s + ("%20.5g," * n % vals).rstrip(",") + "\n")
 
     # Save yaml
-    with open(evolve_yaml, 'w') as f:
+    with open(evolve_yaml, "w") as f:
         data = pd.read_csv(evolve_csv, skipinitialspace=True)
         data = data.rename(columns=lambda x: x.strip())  # strip keys
         i = np.argmax(fitness(data.values[:, :4]))  #
         generations = len(data)
-        f.write('# YOLOv5 Hyperparameter Evolution Results\n' + f'# Best generation: {i}\n' +
+        f.write(
-                f'# Last generation: {generations - 1}\n' + '# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) +
+            "# YOLOv5 Hyperparameter Evolution Results\n"
-                '\n' + '# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n')
+            + f"# Best generation: {i}\n"
+            + f"# Last generation: {generations - 1}\n"
+            + "# "
+            + ", ".join(f"{x.strip():>20s}" for x in keys[:7])
+            + "\n"
+            + "# "
+            + ", ".join(f"{x:>20.5g}" for x in data.values[i, :7])
+            + "\n\n"
+        )
         yaml.safe_dump(data.loc[i][7:].to_dict(), f, sort_keys=False)
 
     # Print to screen
-    LOGGER.info(prefix + f'{generations} generations finished, current result:\n' + prefix +
+    LOGGER.info(
-                ', '.join(f'{x.strip():>20s}' for x in keys) + '\n' + prefix + ', '.join(f'{x:20.5g}'
+        prefix
-                                                                                         for x in vals) + '\n\n')
+        + f"{generations} generations finished, current result:\n"
+        + prefix
+        + ", ".join(f"{x.strip():>20s}" for x in keys)
+        + "\n"
+        + prefix
+        + ", ".join(f"{x:20.5g}" for x in vals)
+        + "\n\n"
+    )
 
     if bucket:
-        subprocess.run(['gsutil', 'cp', f'{evolve_csv}', f'{evolve_yaml}', f'gs://{bucket}'])  # upload
+        subprocess.run(["gsutil", "cp", f"{evolve_csv}", f"{evolve_yaml}", f"gs://{bucket}"])  # upload
 
 
 def apply_classifier(x, model, img, im0):
@@ -1067,15 +1171,15 @@ def apply_classifier(x, model, img, im0):
     return x
 
 
-def increment_path(path, exist_ok=False, sep='', mkdir=False):
+def increment_path(path, exist_ok=False, sep="", mkdir=False):
     # Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
     path = Path(path)  # os-agnostic
     if path.exists() and not exist_ok:
-        path, suffix = (path.with_suffix(''), path.suffix) if path.is_file() else (path, '')
+        path, suffix = (path.with_suffix(""), path.suffix) if path.is_file() else (path, "")
 
         # Method 1
         for n in range(2, 9999):
-            p = f'{path}{sep}{n}{suffix}'  # increment path
+            p = f"{path}{sep}{n}{suffix}"  # increment path
             if not os.path.exists(p):  #
                 break
         path = Path(p)
@@ -1110,7 +1214,7 @@ def imwrite(filename, img):
 
 
 def imshow(path, im):
-    imshow_(path.encode('unicode_escape').decode(), im)
+    imshow_(path.encode("unicode_escape").decode(), im)
 
 
 if Path(inspect.stack()[0].filename).parent.parent.as_posix() in inspect.stack()[-1].filename:

utils/loggers/__init__.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Logging utils."""
-Logging utils
-"""
 import json
 import os
 import warnings
@@ -16,8 +14,8 @@ from utils.loggers.wandb.wandb_utils import WandbLogger
 from utils.plots import plot_images, plot_labels, plot_results
 from utils.torch_utils import de_parallel
 
-LOGGERS = ('csv', 'tb', 'wandb', 'clearml', 'comet')  # *.csv, TensorBoard, Weights & Biases, ClearML
+LOGGERS = ("csv", "tb", "wandb", "clearml", "comet")  # *.csv, TensorBoard, Weights & Biases, ClearML
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
 
 try:
     from torch.utils.tensorboard import SummaryWriter
@@ -27,8 +25,8 @@ except ImportError:
 try:
     import wandb
 
-    assert hasattr(wandb, '__version__')  # verify package import not local dir
+    assert hasattr(wandb, "__version__")  # verify package import not local dir
-    if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in {0, -1}:
+    if pkg.parse_version(wandb.__version__) >= pkg.parse_version("0.12.2") and RANK in {0, -1}:
         try:
             wandb_login_success = wandb.login(timeout=30)
         except wandb.errors.UsageError:  # known non-TTY terminal issue
@@ -41,7 +39,7 @@ except (ImportError, AssertionError):
 try:
     import clearml
 
-    assert hasattr(clearml, '__version__')  # verify package import not local dir
+    assert hasattr(clearml, "__version__")  # verify package import not local dir
 except (ImportError, AssertionError):
     clearml = None
 
@@ -49,7 +47,7 @@ try:
     if RANK in {0, -1}:
         import comet_ml
 
-        assert hasattr(comet_ml, '__version__')  # verify package import not local dir
+        assert hasattr(comet_ml, "__version__")  # verify package import not local dir
         from utils.loggers.comet import CometLogger
 
     else:
@@ -59,7 +57,11 @@ except (ImportError, AssertionError):
 
 
 def _json_default(value):
-    """Format `value` for JSON serialization (e.g. unwrap tensors). Fall back to strings."""
+    """
+    Format `value` for JSON serialization (e.g. unwrap tensors).
+
+    Fall back to strings.
+    """
     if isinstance(value, torch.Tensor):
         try:
             value = value.item()
@@ -70,7 +72,7 @@ def _json_default(value):
     return str(value)
 
 
-class Loggers():
+class Loggers:
     # YOLOv5 Loggers class
     def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
         self.save_dir = save_dir
@@ -81,62 +83,65 @@ class Loggers():
         self.logger = logger  # for printing results to console
         self.include = include
         self.keys = [
-            'train/box_loss',
+            "train/box_loss",
-            'train/obj_loss',
+            "train/obj_loss",
-            'train/cls_loss',  # train loss
+            "train/cls_loss",  # train loss
-            'metrics/precision',
+            "metrics/precision",
-            'metrics/recall',
+            "metrics/recall",
-            'metrics/mAP_0.5',
+            "metrics/mAP_0.5",
-            'metrics/mAP_0.5:0.95',  # metrics
+            "metrics/mAP_0.5:0.95",  # metrics
-            'val/box_loss',
+            "val/box_loss",
-            'val/obj_loss',
+            "val/obj_loss",
-            'val/cls_loss',  # val loss
+            "val/cls_loss",  # val loss
-            'x/lr0',
+            "x/lr0",
-            'x/lr1',
+            "x/lr1",
-            'x/lr2']  # params
+            "x/lr2",
-        self.best_keys = ['best/epoch', 'best/precision', 'best/recall', 'best/mAP_0.5', 'best/mAP_0.5:0.95']
+        ]  # params
+        self.best_keys = ["best/epoch", "best/precision", "best/recall", "best/mAP_0.5", "best/mAP_0.5:0.95"]
         for k in LOGGERS:
             setattr(self, k, None)  # init empty logger dictionary
         self.csv = True  # always log to csv
-        self.ndjson_console = ('ndjson_console' in self.include)  # log ndjson to console
+        self.ndjson_console = "ndjson_console" in self.include  # log ndjson to console
-        self.ndjson_file = ('ndjson_file' in self.include)  # log ndjson to file
+        self.ndjson_file = "ndjson_file" in self.include  # log ndjson to file
 
         # Messages
         if not comet_ml:
-            prefix = colorstr('Comet: ')
+            prefix = colorstr("Comet: ")
             s = f"{prefix}run 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet"
             self.logger.info(s)
         # TensorBoard
         s = self.save_dir
-        if 'tb' in self.include and not self.opt.evolve:
+        if "tb" in self.include and not self.opt.evolve:
-            prefix = colorstr('TensorBoard: ')
+            prefix = colorstr("TensorBoard: ")
             self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/")
             self.tb = SummaryWriter(str(s))
 
         # W&B
-        if wandb and 'wandb' in self.include:
+        if wandb and "wandb" in self.include:
             self.opt.hyp = self.hyp  # add hyperparameters
             self.wandb = WandbLogger(self.opt)
         else:
             self.wandb = None
 
         # ClearML
-        if clearml and 'clearml' in self.include:
+        if clearml and "clearml" in self.include:
             try:
                 self.clearml = ClearmlLogger(self.opt, self.hyp)
             except Exception:
                 self.clearml = None
-                prefix = colorstr('ClearML: ')
+                prefix = colorstr("ClearML: ")
-                LOGGER.warning(f'{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging.'
+                LOGGER.warning(
-                               f' See https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration#readme')
+                    f"{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging."
+                    f" See https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration#readme"
+                )
 
         else:
             self.clearml = None
 
         # Comet
-        if comet_ml and 'comet' in self.include:
+        if comet_ml and "comet" in self.include:
-            if isinstance(self.opt.resume, str) and self.opt.resume.startswith('comet://'):
+            if isinstance(self.opt.resume, str) and self.opt.resume.startswith("comet://"):
-                run_id = self.opt.resume.split('/')[-1]
+                run_id = self.opt.resume.split("/")[-1]
                 self.comet_logger = CometLogger(self.opt, self.hyp, run_id=run_id)
 
             else:
@@ -170,9 +175,9 @@ class Loggers():
         # Callback runs on pre-train routine end
         if self.plots:
             plot_labels(labels, names, self.save_dir)
-            paths = self.save_dir.glob('*labels*.jpg')  # training labels
+            paths = self.save_dir.glob("*labels*.jpg")  # training labels
             if self.wandb:
-                self.wandb.log({'Labels': [wandb.Image(str(x), caption=x.name) for x in paths]})
+                self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
             if self.comet_logger:
                 self.comet_logger.on_pretrain_routine_end(paths)
             if self.clearml:
@@ -185,16 +190,16 @@ class Loggers():
         # ni: number integrated batches (since train start)
         if self.plots:
             if ni < 3:
-                f = self.save_dir / f'train_batch{ni}.jpg'  # filename
+                f = self.save_dir / f"train_batch{ni}.jpg"  # filename
                 plot_images(imgs, targets, paths, f)
                 if ni == 0 and self.tb and not self.opt.sync_bn:
                     log_tensorboard_graph(self.tb, model, imgsz=(self.opt.imgsz, self.opt.imgsz))
             if ni == 10 and (self.wandb or self.clearml):
-                files = sorted(self.save_dir.glob('train*.jpg'))
+                files = sorted(self.save_dir.glob("train*.jpg"))
                 if self.wandb:
-                    self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
+                    self.wandb.log({"Mosaics": [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
                 if self.clearml:
-                    self.clearml.log_debug_samples(files, title='Mosaics')
+                    self.clearml.log_debug_samples(files, title="Mosaics")
 
         if self.comet_logger:
             self.comet_logger.on_train_batch_end(log_dict, step=ni)
@@ -225,11 +230,11 @@ class Loggers():
     def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
         # Callback runs on val end
         if self.wandb or self.clearml:
-            files = sorted(self.save_dir.glob('val*.jpg'))
+            files = sorted(self.save_dir.glob("val*.jpg"))
         if self.wandb:
-            self.wandb.log({'Validation': [wandb.Image(str(f), caption=f.name) for f in files]})
+            self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
         if self.clearml:
-            self.clearml.log_debug_samples(files, title='Validation')
+            self.clearml.log_debug_samples(files, title="Validation")
 
         if self.comet_logger:
             self.comet_logger.on_val_end(nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
@@ -238,18 +243,18 @@ class Loggers():
         # Callback runs at the end of each fit (train+val) epoch
         x = dict(zip(self.keys, vals))
         if self.csv:
-            file = self.save_dir / 'results.csv'
+            file = self.save_dir / "results.csv"
             n = len(x) + 1  # number of cols
-            s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n')  # add header
+            s = "" if file.exists() else (("%20s," * n % tuple(["epoch"] + self.keys)).rstrip(",") + "\n")  # add header
-            with open(file, 'a') as f:
+            with open(file, "a") as f:
-                f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
+                f.write(s + ("%20.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n")
         if self.ndjson_console or self.ndjson_file:
             json_data = json.dumps(dict(epoch=epoch, **x), default=_json_default)
             if self.ndjson_console:
                 print(json_data)
             if self.ndjson_file:
-                file = self.save_dir / 'results.ndjson'
+                file = self.save_dir / "results.ndjson"
-                with open(file, 'a') as f:
+                with open(file, "a") as f:
                     print(json_data, file=f)
 
         if self.tb:
@@ -279,9 +284,9 @@ class Loggers():
             if self.wandb:
                 self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
             if self.clearml:
-                self.clearml.task.update_output_model(model_path=str(last),
+                self.clearml.task.update_output_model(
-                                                      model_name='Latest Model',
+                    model_path=str(last), model_name="Latest Model", auto_delete_file=False
-                                                      auto_delete_file=False)
+                )
 
         if self.comet_logger:
             self.comet_logger.on_model_save(last, epoch, final_epoch, best_fitness, fi)
@@ -289,31 +294,34 @@ class Loggers():
     def on_train_end(self, last, best, epoch, results):
         # Callback runs on training end, i.e. saving best model
         if self.plots:
-            plot_results(file=self.save_dir / 'results.csv')  # save results.png
+            plot_results(file=self.save_dir / "results.csv")  # save results.png
-        files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
+        files = ["results.png", "confusion_matrix.png", *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R"))]
         files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()]  # filter
         self.logger.info(f"Results saved to {colorstr('bold', self.save_dir)}")
 
         if self.tb and not self.clearml:  # These images are already captured by ClearML by now, we don't want doubles
             for f in files:
-                self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
+                self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats="HWC")
 
         if self.wandb:
             self.wandb.log(dict(zip(self.keys[3:10], results)))
-            self.wandb.log({'Results': [wandb.Image(str(f), caption=f.name) for f in files]})
+            self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
             # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
             if not self.opt.evolve:
-                wandb.log_artifact(str(best if best.exists() else last),
+                wandb.log_artifact(
-                                   type='model',
+                    str(best if best.exists() else last),
-                                   name=f'run_{self.wandb.wandb_run.id}_model',
+                    type="model",
-                                   aliases=['latest', 'best', 'stripped'])
+                    name=f"run_{self.wandb.wandb_run.id}_model",
+                    aliases=["latest", "best", "stripped"],
+                )
             self.wandb.finish_run()
 
         if self.clearml and not self.opt.evolve:
             self.clearml.log_summary(dict(zip(self.keys[3:10], results)))
             [self.clearml.log_plot(title=f.stem, plot_path=f) for f in files]
-            self.clearml.log_model(str(best if best.exists() else last),
+            self.clearml.log_model(
-                                   "Best Model" if best.exists() else "Last Model", epoch)
+                str(best if best.exists() else last), "Best Model" if best.exists() else "Last Model", epoch
+            )
 
         if self.comet_logger:
             final_results = dict(zip(self.keys[3:10], results))
@@ -339,38 +347,41 @@ class GenericLogger:
         include:         loggers to include
     """
 
-    def __init__(self, opt, console_logger, include=('tb', 'wandb', 'clearml')):
+    def __init__(self, opt, console_logger, include=("tb", "wandb", "clearml")):
         # init default loggers
         self.save_dir = Path(opt.save_dir)
         self.include = include
         self.console_logger = console_logger
-        self.csv = self.save_dir / 'results.csv'  # CSV logger
+        self.csv = self.save_dir / "results.csv"  # CSV logger
-        if 'tb' in self.include:
+        if "tb" in self.include:
-            prefix = colorstr('TensorBoard: ')
+            prefix = colorstr("TensorBoard: ")
             self.console_logger.info(
-                f"{prefix}Start with 'tensorboard --logdir {self.save_dir.parent}', view at http://localhost:6006/")
+                f"{prefix}Start with 'tensorboard --logdir {self.save_dir.parent}', view at http://localhost:6006/"
+            )
             self.tb = SummaryWriter(str(self.save_dir))
 
-        if wandb and 'wandb' in self.include:
+        if wandb and "wandb" in self.include:
-            self.wandb = wandb.init(project=web_project_name(str(opt.project)),
+            self.wandb = wandb.init(
-                                    name=None if opt.name == 'exp' else opt.name,
+                project=web_project_name(str(opt.project)), name=None if opt.name == "exp" else opt.name, config=opt
-                                    config=opt)
+            )
         else:
             self.wandb = None
 
-        if clearml and 'clearml' in self.include:
+        if clearml and "clearml" in self.include:
             try:
                 # Hyp is not available in classification mode
-                if 'hyp' not in opt:
+                if "hyp" not in opt:
                     hyp = {}
                 else:
                     hyp = opt.hyp
                 self.clearml = ClearmlLogger(opt, hyp)
             except Exception:
                 self.clearml = None
-                prefix = colorstr('ClearML: ')
+                prefix = colorstr("ClearML: ")
-                LOGGER.warning(f'{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging.'
+                LOGGER.warning(
-                               f' See https://github.com/ultralytics/yolov5/tree/master/utils/loggers/clearml#readme')
+                    f"{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging."
+                    f" See https://github.com/ultralytics/yolov5/tree/master/utils/loggers/clearml#readme"
+                )
         else:
             self.clearml = None
 
@@ -379,9 +390,9 @@ class GenericLogger:
         if self.csv:
             keys, vals = list(metrics.keys()), list(metrics.values())
             n = len(metrics) + 1  # number of cols
-            s = '' if self.csv.exists() else (('%23s,' * n % tuple(['epoch'] + keys)).rstrip(',') + '\n')  # header
+            s = "" if self.csv.exists() else (("%23s," * n % tuple(["epoch"] + keys)).rstrip(",") + "\n")  # header
-            with open(self.csv, 'a') as f:
+            with open(self.csv, "a") as f:
-                f.write(s + ('%23.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
+                f.write(s + ("%23.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n")
 
         if self.tb:
             for k, v in metrics.items():
@@ -393,20 +404,20 @@ class GenericLogger:
         if self.clearml:
             self.clearml.log_scalars(metrics, epoch)
 
-    def log_images(self, files, name='Images', epoch=0):
+    def log_images(self, files, name="Images", epoch=0):
         # Log images to all loggers
         files = [Path(f) for f in (files if isinstance(files, (tuple, list)) else [files])]  # to Path
         files = [f for f in files if f.exists()]  # filter by exists
 
         if self.tb:
             for f in files:
-                self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
+                self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats="HWC")
 
         if self.wandb:
             self.wandb.log({name: [wandb.Image(str(f), caption=f.name) for f in files]}, step=epoch)
 
         if self.clearml:
-            if name == 'Results':
+            if name == "Results":
                 [self.clearml.log_plot(f.stem, f) for f in files]
             else:
                 self.clearml.log_debug_samples(files, title=name)
@@ -419,7 +430,7 @@ class GenericLogger:
     def log_model(self, model_path, epoch=0, metadata={}):
         # Log model to all loggers
         if self.wandb:
-            art = wandb.Artifact(name=f'run_{wandb.run.id}_model', type='model', metadata=metadata)
+            art = wandb.Artifact(name=f"run_{wandb.run.id}_model", type="model", metadata=metadata)
             art.add_file(str(model_path))
             wandb.log_artifact(art)
         if self.clearml:
@@ -440,15 +451,15 @@ def log_tensorboard_graph(tb, model, imgsz=(640, 640)):
         imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz  # expand
         im = torch.zeros((1, 3, *imgsz)).to(p.device).type_as(p)  # input image (WARNING: must be zeros, not empty)
         with warnings.catch_warnings():
-            warnings.simplefilter('ignore')  # suppress jit trace warning
+            warnings.simplefilter("ignore")  # suppress jit trace warning
             tb.add_graph(torch.jit.trace(de_parallel(model), im, strict=False), [])
     except Exception as e:
-        LOGGER.warning(f'WARNING ⚠️ TensorBoard graph visualization failure {e}')
+        LOGGER.warning(f"WARNING ⚠️ TensorBoard graph visualization failure {e}")
 
 
 def web_project_name(project):
     # Convert local project name to web project name
-    if not project.startswith('runs/train'):
+    if not project.startswith("runs/train"):
         return project
-    suffix = '-Classify' if project.endswith('-cls') else '-Segment' if project.endswith('-seg') else ''
+    suffix = "-Classify" if project.endswith("-cls") else "-Segment" if project.endswith("-seg") else ""
-    return f'YOLOv5{suffix}'
+    return f"YOLOv5{suffix}"

utils/loggers/clearml/clearml_utils.py

@@ -13,55 +13,63 @@ try:
     import clearml
     from clearml import Dataset, Task
 
-    assert hasattr(clearml, '__version__')  # verify package import not local dir
+    assert hasattr(clearml, "__version__")  # verify package import not local dir
 except (ImportError, AssertionError):
     clearml = None
 
 
 def construct_dataset(clearml_info_string):
-    """Load in a clearml dataset and fill the internal data_dict with its contents.
+    """Load in a clearml dataset and fill the internal data_dict with its contents."""
-    """
+    dataset_id = clearml_info_string.replace("clearml://", "")
-    dataset_id = clearml_info_string.replace('clearml://', '')
     dataset = Dataset.get(dataset_id=dataset_id)
     dataset_root_path = Path(dataset.get_local_copy())
 
     # We'll search for the yaml file definition in the dataset
-    yaml_filenames = list(glob.glob(str(dataset_root_path / '*.yaml')) + glob.glob(str(dataset_root_path / '*.yml')))
+    yaml_filenames = list(glob.glob(str(dataset_root_path / "*.yaml")) + glob.glob(str(dataset_root_path / "*.yml")))
     if len(yaml_filenames) > 1:
-        raise ValueError('More than one yaml file was found in the dataset root, cannot determine which one contains '
+        raise ValueError(
-                         'the dataset definition this way.')
+            "More than one yaml file was found in the dataset root, cannot determine which one contains "
+            "the dataset definition this way."
+        )
     elif len(yaml_filenames) == 0:
-        raise ValueError('No yaml definition found in dataset root path, check that there is a correct yaml file '
+        raise ValueError(
-                         'inside the dataset root path.')
+            "No yaml definition found in dataset root path, check that there is a correct yaml file "
+            "inside the dataset root path."
+        )
     with open(yaml_filenames[0]) as f:
         dataset_definition = yaml.safe_load(f)
 
-    assert set(dataset_definition.keys()).issuperset(
+    assert set(
-        {'train', 'test', 'val', 'nc', 'names'}
+        dataset_definition.keys()
+    ).issuperset(
+        {"train", "test", "val", "nc", "names"}
     ), "The right keys were not found in the yaml file, make sure it at least has the following keys: ('train', 'test', 'val', 'nc', 'names')"
 
     data_dict = dict()
-    data_dict['train'] = str(
+    data_dict["train"] = (
-        (dataset_root_path / dataset_definition['train']).resolve()) if dataset_definition['train'] else None
+        str((dataset_root_path / dataset_definition["train"]).resolve()) if dataset_definition["train"] else None
-    data_dict['test'] = str(
+    )
-        (dataset_root_path / dataset_definition['test']).resolve()) if dataset_definition['test'] else None
+    data_dict["test"] = (
-    data_dict['val'] = str(
+        str((dataset_root_path / dataset_definition["test"]).resolve()) if dataset_definition["test"] else None
-        (dataset_root_path / dataset_definition['val']).resolve()) if dataset_definition['val'] else None
+    )
-    data_dict['nc'] = dataset_definition['nc']
+    data_dict["val"] = (
-    data_dict['names'] = dataset_definition['names']
+        str((dataset_root_path / dataset_definition["val"]).resolve()) if dataset_definition["val"] else None
+    )
+    data_dict["nc"] = dataset_definition["nc"]
+    data_dict["names"] = dataset_definition["names"]
 
     return data_dict
 
 
 class ClearmlLogger:
-    """Log training runs, datasets, models, and predictions to ClearML.
+    """
+    Log training runs, datasets, models, and predictions to ClearML.
 
-    This logger sends information to ClearML at app.clear.ml or to your own hosted server. By default,
+    This logger sends information to ClearML at app.clear.ml or to your own hosted server. By default, this information
-    this information includes hyperparameters, system configuration and metrics, model metrics, code information and
+    includes hyperparameters, system configuration and metrics, model metrics, code information and basic data metrics
-    basic data metrics and analyses.
+    and analyses.
 
-    By providing additional command line arguments to train.py, datasets,
+    By providing additional command line arguments to train.py, datasets, models and predictions can also be logged.
-    models and predictions can also be logged.
     """
 
     def __init__(self, opt, hyp):
@@ -81,36 +89,36 @@ class ClearmlLogger:
         self.max_imgs_to_log_per_epoch = 16
         # Get the interval of epochs when bounding box images should be logged
         # Only for detection task though!
-        if 'bbox_interval' in opt:
+        if "bbox_interval" in opt:
             self.bbox_interval = opt.bbox_interval
         self.clearml = clearml
         self.task = None
         self.data_dict = None
         if self.clearml:
             self.task = Task.init(
-                project_name=opt.project if not str(opt.project).startswith('runs/') else 'YOLOv5',
+                project_name=opt.project if not str(opt.project).startswith("runs/") else "YOLOv5",
-                task_name=opt.name if opt.name != 'exp' else 'Training',
+                task_name=opt.name if opt.name != "exp" else "Training",
-                tags=['YOLOv5'],
+                tags=["YOLOv5"],
                 output_uri=True,
                 reuse_last_task_id=opt.exist_ok,
-                auto_connect_frameworks={
+                auto_connect_frameworks={"pytorch": False, "matplotlib": False},
-                    'pytorch': False,
-                    'matplotlib': False}
                 # We disconnect pytorch auto-detection, because we added manual model save points in the code
             )
             # ClearML's hooks will already grab all general parameters
             # Only the hyperparameters coming from the yaml config file
             # will have to be added manually!
-            self.task.connect(hyp, name='Hyperparameters')
+            self.task.connect(hyp, name="Hyperparameters")
-            self.task.connect(opt, name='Args')
+            self.task.connect(opt, name="Args")
 
             # Make sure the code is easily remotely runnable by setting the docker image to use by the remote agent
-            self.task.set_base_docker('ultralytics/yolov5:latest',
+            self.task.set_base_docker(
+                "ultralytics/yolov5:latest",
                 docker_arguments='--ipc=host -e="CLEARML_AGENT_SKIP_PYTHON_ENV_INSTALL=1"',
-                                      docker_setup_bash_script='pip install clearml')
+                docker_setup_bash_script="pip install clearml",
+            )
 
             # Get ClearML Dataset Version if requested
-            if opt.data.startswith('clearml://'):
+            if opt.data.startswith("clearml://"):
                 # data_dict should have the following keys:
                 # names, nc (number of classes), test, train, val (all three relative paths to ../datasets)
                 self.data_dict = construct_dataset(opt.data)
@@ -120,33 +128,32 @@ class ClearmlLogger:
 
     def log_scalars(self, metrics, epoch):
         """
-        Log scalars/metrics to ClearML
+        Log scalars/metrics to ClearML.
 
         arguments:
         metrics (dict) Metrics in dict format: {"metrics/mAP": 0.8, ...}
         epoch (int) iteration number for the current set of metrics
         """
         for k, v in metrics.items():
-            title, series = k.split('/')
+            title, series = k.split("/")
             self.task.get_logger().report_scalar(title, series, v, epoch)
 
     def log_model(self, model_path, model_name, epoch=0):
         """
-        Log model weights to ClearML
+        Log model weights to ClearML.
 
         arguments:
         model_path (PosixPath or str) Path to the model weights
         model_name (str) Name of the model visible in ClearML
         epoch (int) Iteration / epoch of the model weights
         """
-        self.task.update_output_model(model_path=str(model_path),
+        self.task.update_output_model(
-                                      name=model_name,
+            model_path=str(model_path), name=model_name, iteration=epoch, auto_delete_file=False
-                                      iteration=epoch,
+        )
-                                      auto_delete_file=False)
 
     def log_summary(self, metrics):
         """
-        Log final metrics to a summary table
+        Log final metrics to a summary table.
 
         arguments:
         metrics (dict) Metrics in dict format: {"metrics/mAP": 0.8, ...}
@@ -156,7 +163,7 @@ class ClearmlLogger:
 
     def log_plot(self, title, plot_path):
         """
-        Log image as plot in the plot section of ClearML
+        Log image as plot in the plot section of ClearML.
 
         arguments:
         title (str) Title of the plot
@@ -164,12 +171,12 @@ class ClearmlLogger:
         """
         img = mpimg.imread(plot_path)
         fig = plt.figure()
-        ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect='auto', xticks=[], yticks=[])  # no ticks
+        ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect="auto", xticks=[], yticks=[])  # no ticks
         ax.imshow(img)
 
         self.task.get_logger().report_matplotlib_figure(title, "", figure=fig, report_interactive=False)
 
-    def log_debug_samples(self, files, title='Debug Samples'):
+    def log_debug_samples(self, files, title="Debug Samples"):
         """
         Log files (images) as debug samples in the ClearML task.
 
@@ -179,12 +186,11 @@ class ClearmlLogger:
         """
         for f in files:
             if f.exists():
-                it = re.search(r'_batch(\d+)', f.name)
+                it = re.search(r"_batch(\d+)", f.name)
                 iteration = int(it.groups()[0]) if it else 0
-                self.task.get_logger().report_image(title=title,
+                self.task.get_logger().report_image(
-                                                    series=f.name.replace(f"_batch{iteration}", ''),
+                    title=title, series=f.name.replace(f"_batch{iteration}", ""), local_path=str(f), iteration=iteration
-                                                    local_path=str(f),
+                )
-                                                    iteration=iteration)
 
     def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_threshold=0.25):
         """
@@ -206,15 +212,14 @@ class ClearmlLogger:
 
                     class_name = class_names[int(class_nr)]
                     confidence_percentage = round(float(conf) * 100, 2)
-                    label = f'{class_name}: {confidence_percentage}%'
+                    label = f"{class_name}: {confidence_percentage}%"
 
                     if conf > conf_threshold:
                         annotator.rectangle(box.cpu().numpy(), outline=color)
                         annotator.box_label(box.cpu().numpy(), label=label, color=color)
 
                 annotated_image = annotator.result()
-                self.task.get_logger().report_image(title='Bounding Boxes',
+                self.task.get_logger().report_image(
-                                                    series=image_path.name,
+                    title="Bounding Boxes", series=image_path.name, iteration=self.current_epoch, image=annotated_image
-                                                    iteration=self.current_epoch,
+                )
-                                                    image=annotated_image)
                 self.current_epoch_logged_images.add(image_path)

utils/loggers/clearml/hpo.py

@@ -1,18 +1,21 @@
 from clearml import Task
+
 # Connecting ClearML with the current process,
 # from here on everything is logged automatically
 from clearml.automation import HyperParameterOptimizer, UniformParameterRange
 from clearml.automation.optuna import OptimizerOptuna
 
-task = Task.init(project_name='Hyper-Parameter Optimization',
+task = Task.init(
-                 task_name='YOLOv5',
+    project_name="Hyper-Parameter Optimization",
+    task_name="YOLOv5",
     task_type=Task.TaskTypes.optimizer,
-                 reuse_last_task_id=False)
+    reuse_last_task_id=False,
+)
 
 # Example use case:
 optimizer = HyperParameterOptimizer(
     # This is the experiment we want to optimize
-    base_task_id='<your_template_task_id>',
+    base_task_id="<your_template_task_id>",
     # here we define the hyper-parameters to optimize
     # Notice: The parameter name should exactly match what you see in the UI: <section_name>/<parameter>
     # For Example, here we see in the base experiment a section Named: "General"
@@ -20,39 +23,40 @@ optimizer = HyperParameterOptimizer(
     # If you have `argparse` for example, then arguments will appear under the "Args" section,
     # and you should instead pass "Args/batch_size"
     hyper_parameters=[
-        UniformParameterRange('Hyperparameters/lr0', min_value=1e-5, max_value=1e-1),
+        UniformParameterRange("Hyperparameters/lr0", min_value=1e-5, max_value=1e-1),
-        UniformParameterRange('Hyperparameters/lrf', min_value=0.01, max_value=1.0),
+        UniformParameterRange("Hyperparameters/lrf", min_value=0.01, max_value=1.0),
-        UniformParameterRange('Hyperparameters/momentum', min_value=0.6, max_value=0.98),
+        UniformParameterRange("Hyperparameters/momentum", min_value=0.6, max_value=0.98),
-        UniformParameterRange('Hyperparameters/weight_decay', min_value=0.0, max_value=0.001),
+        UniformParameterRange("Hyperparameters/weight_decay", min_value=0.0, max_value=0.001),
-        UniformParameterRange('Hyperparameters/warmup_epochs', min_value=0.0, max_value=5.0),
+        UniformParameterRange("Hyperparameters/warmup_epochs", min_value=0.0, max_value=5.0),
-        UniformParameterRange('Hyperparameters/warmup_momentum', min_value=0.0, max_value=0.95),
+        UniformParameterRange("Hyperparameters/warmup_momentum", min_value=0.0, max_value=0.95),
-        UniformParameterRange('Hyperparameters/warmup_bias_lr', min_value=0.0, max_value=0.2),
+        UniformParameterRange("Hyperparameters/warmup_bias_lr", min_value=0.0, max_value=0.2),
-        UniformParameterRange('Hyperparameters/box', min_value=0.02, max_value=0.2),
+        UniformParameterRange("Hyperparameters/box", min_value=0.02, max_value=0.2),
-        UniformParameterRange('Hyperparameters/cls', min_value=0.2, max_value=4.0),
+        UniformParameterRange("Hyperparameters/cls", min_value=0.2, max_value=4.0),
-        UniformParameterRange('Hyperparameters/cls_pw', min_value=0.5, max_value=2.0),
+        UniformParameterRange("Hyperparameters/cls_pw", min_value=0.5, max_value=2.0),
-        UniformParameterRange('Hyperparameters/obj', min_value=0.2, max_value=4.0),
+        UniformParameterRange("Hyperparameters/obj", min_value=0.2, max_value=4.0),
-        UniformParameterRange('Hyperparameters/obj_pw', min_value=0.5, max_value=2.0),
+        UniformParameterRange("Hyperparameters/obj_pw", min_value=0.5, max_value=2.0),
-        UniformParameterRange('Hyperparameters/iou_t', min_value=0.1, max_value=0.7),
+        UniformParameterRange("Hyperparameters/iou_t", min_value=0.1, max_value=0.7),
-        UniformParameterRange('Hyperparameters/anchor_t', min_value=2.0, max_value=8.0),
+        UniformParameterRange("Hyperparameters/anchor_t", min_value=2.0, max_value=8.0),
-        UniformParameterRange('Hyperparameters/fl_gamma', min_value=0.0, max_value=4.0),
+        UniformParameterRange("Hyperparameters/fl_gamma", min_value=0.0, max_value=4.0),
-        UniformParameterRange('Hyperparameters/hsv_h', min_value=0.0, max_value=0.1),
+        UniformParameterRange("Hyperparameters/hsv_h", min_value=0.0, max_value=0.1),
-        UniformParameterRange('Hyperparameters/hsv_s', min_value=0.0, max_value=0.9),
+        UniformParameterRange("Hyperparameters/hsv_s", min_value=0.0, max_value=0.9),
-        UniformParameterRange('Hyperparameters/hsv_v', min_value=0.0, max_value=0.9),
+        UniformParameterRange("Hyperparameters/hsv_v", min_value=0.0, max_value=0.9),
-        UniformParameterRange('Hyperparameters/degrees', min_value=0.0, max_value=45.0),
+        UniformParameterRange("Hyperparameters/degrees", min_value=0.0, max_value=45.0),
-        UniformParameterRange('Hyperparameters/translate', min_value=0.0, max_value=0.9),
+        UniformParameterRange("Hyperparameters/translate", min_value=0.0, max_value=0.9),
-        UniformParameterRange('Hyperparameters/scale', min_value=0.0, max_value=0.9),
+        UniformParameterRange("Hyperparameters/scale", min_value=0.0, max_value=0.9),
-        UniformParameterRange('Hyperparameters/shear', min_value=0.0, max_value=10.0),
+        UniformParameterRange("Hyperparameters/shear", min_value=0.0, max_value=10.0),
-        UniformParameterRange('Hyperparameters/perspective', min_value=0.0, max_value=0.001),
+        UniformParameterRange("Hyperparameters/perspective", min_value=0.0, max_value=0.001),
-        UniformParameterRange('Hyperparameters/flipud', min_value=0.0, max_value=1.0),
+        UniformParameterRange("Hyperparameters/flipud", min_value=0.0, max_value=1.0),
-        UniformParameterRange('Hyperparameters/fliplr', min_value=0.0, max_value=1.0),
+        UniformParameterRange("Hyperparameters/fliplr", min_value=0.0, max_value=1.0),
-        UniformParameterRange('Hyperparameters/mosaic', min_value=0.0, max_value=1.0),
+        UniformParameterRange("Hyperparameters/mosaic", min_value=0.0, max_value=1.0),
-        UniformParameterRange('Hyperparameters/mixup', min_value=0.0, max_value=1.0),
+        UniformParameterRange("Hyperparameters/mixup", min_value=0.0, max_value=1.0),
-        UniformParameterRange('Hyperparameters/copy_paste', min_value=0.0, max_value=1.0)],
+        UniformParameterRange("Hyperparameters/copy_paste", min_value=0.0, max_value=1.0),
+    ],
     # this is the objective metric we want to maximize/minimize
-    objective_metric_title='metrics',
+    objective_metric_title="metrics",
-    objective_metric_series='mAP_0.5',
+    objective_metric_series="mAP_0.5",
     # now we decide if we want to maximize it or minimize it (accuracy we maximize)
-    objective_metric_sign='max',
+    objective_metric_sign="max",
     # let us limit the number of concurrent experiments,
     # this in turn will make sure we do dont bombard the scheduler with experiments.
     # if we have an auto-scaler connected, this, by proxy, will limit the number of machine
@@ -81,4 +85,4 @@ optimizer.wait()
 # make sure background optimization stopped
 optimizer.stop()
 
-print('We are done, good bye')
+print("We are done, good bye")

utils/loggers/comet/__init__.py

@@ -17,7 +17,7 @@ try:
 
     # Project Configuration
     config = comet_ml.config.get_config()
-    COMET_PROJECT_NAME = config.get_string(os.getenv('COMET_PROJECT_NAME'), 'comet.project_name', default='yolov5')
+    COMET_PROJECT_NAME = config.get_string(os.getenv("COMET_PROJECT_NAME"), "comet.project_name", default="yolov5")
 except ImportError:
     comet_ml = None
     COMET_PROJECT_NAME = None
@@ -31,42 +31,40 @@ from utils.dataloaders import img2label_paths
 from utils.general import check_dataset, scale_boxes, xywh2xyxy
 from utils.metrics import box_iou
 
-COMET_PREFIX = 'comet://'
+COMET_PREFIX = "comet://"
 
-COMET_MODE = os.getenv('COMET_MODE', 'online')
+COMET_MODE = os.getenv("COMET_MODE", "online")
 
 # Model Saving Settings
-COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5')
+COMET_MODEL_NAME = os.getenv("COMET_MODEL_NAME", "yolov5")
 
 # Dataset Artifact Settings
-COMET_UPLOAD_DATASET = os.getenv('COMET_UPLOAD_DATASET', 'false').lower() == 'true'
+COMET_UPLOAD_DATASET = os.getenv("COMET_UPLOAD_DATASET", "false").lower() == "true"
 
 # Evaluation Settings
-COMET_LOG_CONFUSION_MATRIX = (os.getenv('COMET_LOG_CONFUSION_MATRIX', 'true').lower() == 'true')
+COMET_LOG_CONFUSION_MATRIX = os.getenv("COMET_LOG_CONFUSION_MATRIX", "true").lower() == "true"
-COMET_LOG_PREDICTIONS = os.getenv('COMET_LOG_PREDICTIONS', 'true').lower() == 'true'
+COMET_LOG_PREDICTIONS = os.getenv("COMET_LOG_PREDICTIONS", "true").lower() == "true"
-COMET_MAX_IMAGE_UPLOADS = int(os.getenv('COMET_MAX_IMAGE_UPLOADS', 100))
+COMET_MAX_IMAGE_UPLOADS = int(os.getenv("COMET_MAX_IMAGE_UPLOADS", 100))
 
 # Confusion Matrix Settings
-CONF_THRES = float(os.getenv('CONF_THRES', 0.001))
+CONF_THRES = float(os.getenv("CONF_THRES", 0.001))
-IOU_THRES = float(os.getenv('IOU_THRES', 0.6))
+IOU_THRES = float(os.getenv("IOU_THRES", 0.6))
 
 # Batch Logging Settings
-COMET_LOG_BATCH_METRICS = (os.getenv('COMET_LOG_BATCH_METRICS', 'false').lower() == 'true')
+COMET_LOG_BATCH_METRICS = os.getenv("COMET_LOG_BATCH_METRICS", "false").lower() == "true"
-COMET_BATCH_LOGGING_INTERVAL = os.getenv('COMET_BATCH_LOGGING_INTERVAL', 1)
+COMET_BATCH_LOGGING_INTERVAL = os.getenv("COMET_BATCH_LOGGING_INTERVAL", 1)
-COMET_PREDICTION_LOGGING_INTERVAL = os.getenv('COMET_PREDICTION_LOGGING_INTERVAL', 1)
+COMET_PREDICTION_LOGGING_INTERVAL = os.getenv("COMET_PREDICTION_LOGGING_INTERVAL", 1)
-COMET_LOG_PER_CLASS_METRICS = (os.getenv('COMET_LOG_PER_CLASS_METRICS', 'false').lower() == 'true')
+COMET_LOG_PER_CLASS_METRICS = os.getenv("COMET_LOG_PER_CLASS_METRICS", "false").lower() == "true"
 
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
 
 to_pil = T.ToPILImage()
 
 
 class CometLogger:
-    """Log metrics, parameters, source code, models and much more
+    """Log metrics, parameters, source code, models and much more with Comet."""
-    with Comet
-    """
 
-    def __init__(self, opt, hyp, run_id=None, job_type='Training', **experiment_kwargs) -> None:
+    def __init__(self, opt, hyp, run_id=None, job_type="Training", **experiment_kwargs) -> None:
         self.job_type = job_type
         self.opt = opt
         self.hyp = hyp
@@ -87,57 +85,58 @@ class CometLogger:
 
         # Default parameters to pass to Experiment objects
         self.default_experiment_kwargs = {
-            'log_code': False,
+            "log_code": False,
-            'log_env_gpu': True,
+            "log_env_gpu": True,
-            'log_env_cpu': True,
+            "log_env_cpu": True,
-            'project_name': COMET_PROJECT_NAME, }
+            "project_name": COMET_PROJECT_NAME,
+        }
         self.default_experiment_kwargs.update(experiment_kwargs)
         self.experiment = self._get_experiment(self.comet_mode, run_id)
         self.experiment.set_name(self.opt.name)
 
         self.data_dict = self.check_dataset(self.opt.data)
-        self.class_names = self.data_dict['names']
+        self.class_names = self.data_dict["names"]
-        self.num_classes = self.data_dict['nc']
+        self.num_classes = self.data_dict["nc"]
 
         self.logged_images_count = 0
         self.max_images = COMET_MAX_IMAGE_UPLOADS
 
         if run_id is None:
-            self.experiment.log_other('Created from', 'YOLOv5')
+            self.experiment.log_other("Created from", "YOLOv5")
             if not isinstance(self.experiment, comet_ml.OfflineExperiment):
-                workspace, project_name, experiment_id = self.experiment.url.split('/')[-3:]
+                workspace, project_name, experiment_id = self.experiment.url.split("/")[-3:]
                 self.experiment.log_other(
-                    'Run Path',
+                    "Run Path",
-                    f'{workspace}/{project_name}/{experiment_id}',
+                    f"{workspace}/{project_name}/{experiment_id}",
                 )
             self.log_parameters(vars(opt))
             self.log_parameters(self.opt.hyp)
             self.log_asset_data(
                 self.opt.hyp,
-                name='hyperparameters.json',
+                name="hyperparameters.json",
-                metadata={'type': 'hyp-config-file'},
+                metadata={"type": "hyp-config-file"},
             )
             self.log_asset(
-                f'{self.opt.save_dir}/opt.yaml',
+                f"{self.opt.save_dir}/opt.yaml",
-                metadata={'type': 'opt-config-file'},
+                metadata={"type": "opt-config-file"},
             )
 
         self.comet_log_confusion_matrix = COMET_LOG_CONFUSION_MATRIX
 
-        if hasattr(self.opt, 'conf_thres'):
+        if hasattr(self.opt, "conf_thres"):
             self.conf_thres = self.opt.conf_thres
         else:
             self.conf_thres = CONF_THRES
-        if hasattr(self.opt, 'iou_thres'):
+        if hasattr(self.opt, "iou_thres"):
             self.iou_thres = self.opt.iou_thres
         else:
             self.iou_thres = IOU_THRES
 
-        self.log_parameters({'val_iou_threshold': self.iou_thres, 'val_conf_threshold': self.conf_thres})
+        self.log_parameters({"val_iou_threshold": self.iou_thres, "val_conf_threshold": self.conf_thres})
 
         self.comet_log_predictions = COMET_LOG_PREDICTIONS
         if self.opt.bbox_interval == -1:
-            self.comet_log_prediction_interval = (1 if self.opt.epochs < 10 else self.opt.epochs // 10)
+            self.comet_log_prediction_interval = 1 if self.opt.epochs < 10 else self.opt.epochs // 10
         else:
             self.comet_log_prediction_interval = self.opt.bbox_interval
 
@@ -147,30 +146,35 @@ class CometLogger:
 
         self.comet_log_per_class_metrics = COMET_LOG_PER_CLASS_METRICS
 
-        self.experiment.log_others({
+        self.experiment.log_others(
-            'comet_mode': COMET_MODE,
+            {
-            'comet_max_image_uploads': COMET_MAX_IMAGE_UPLOADS,
+                "comet_mode": COMET_MODE,
-            'comet_log_per_class_metrics': COMET_LOG_PER_CLASS_METRICS,
+                "comet_max_image_uploads": COMET_MAX_IMAGE_UPLOADS,
-            'comet_log_batch_metrics': COMET_LOG_BATCH_METRICS,
+                "comet_log_per_class_metrics": COMET_LOG_PER_CLASS_METRICS,
-            'comet_log_confusion_matrix': COMET_LOG_CONFUSION_MATRIX,
+                "comet_log_batch_metrics": COMET_LOG_BATCH_METRICS,
-            'comet_model_name': COMET_MODEL_NAME, })
+                "comet_log_confusion_matrix": COMET_LOG_CONFUSION_MATRIX,
+                "comet_model_name": COMET_MODEL_NAME,
+            }
+        )
 
         # Check if running the Experiment with the Comet Optimizer
-        if hasattr(self.opt, 'comet_optimizer_id'):
+        if hasattr(self.opt, "comet_optimizer_id"):
-            self.experiment.log_other('optimizer_id', self.opt.comet_optimizer_id)
+            self.experiment.log_other("optimizer_id", self.opt.comet_optimizer_id)
-            self.experiment.log_other('optimizer_objective', self.opt.comet_optimizer_objective)
+            self.experiment.log_other("optimizer_objective", self.opt.comet_optimizer_objective)
-            self.experiment.log_other('optimizer_metric', self.opt.comet_optimizer_metric)
+            self.experiment.log_other("optimizer_metric", self.opt.comet_optimizer_metric)
-            self.experiment.log_other('optimizer_parameters', json.dumps(self.hyp))
+            self.experiment.log_other("optimizer_parameters", json.dumps(self.hyp))
 
     def _get_experiment(self, mode, experiment_id=None):
-        if mode == 'offline':
+        if mode == "offline":
             if experiment_id is not None:
                 return comet_ml.ExistingOfflineExperiment(
                     previous_experiment=experiment_id,
                     **self.default_experiment_kwargs,
                 )
 
-            return comet_ml.OfflineExperiment(**self.default_experiment_kwargs, )
+            return comet_ml.OfflineExperiment(
+                **self.default_experiment_kwargs,
+            )
 
         else:
             try:
@@ -183,11 +187,13 @@ class CometLogger:
                 return comet_ml.Experiment(**self.default_experiment_kwargs)
 
             except ValueError:
-                logger.warning('COMET WARNING: '
+                logger.warning(
-                               'Comet credentials have not been set. '
+                    "COMET WARNING: "
-                               'Comet will default to offline logging. '
+                    "Comet credentials have not been set. "
-                               'Please set your credentials to enable online logging.')
+                    "Comet will default to offline logging. "
-                return self._get_experiment('offline', experiment_id)
+                    "Please set your credentials to enable online logging."
+                )
+                return self._get_experiment("offline", experiment_id)
 
         return
 
@@ -211,12 +217,13 @@ class CometLogger:
             return
 
         model_metadata = {
-            'fitness_score': fitness_score[-1],
+            "fitness_score": fitness_score[-1],
-            'epochs_trained': epoch + 1,
+            "epochs_trained": epoch + 1,
-            'save_period': opt.save_period,
+            "save_period": opt.save_period,
-            'total_epochs': opt.epochs, }
+            "total_epochs": opt.epochs,
+        }
 
-        model_files = glob.glob(f'{path}/*.pt')
+        model_files = glob.glob(f"{path}/*.pt")
         for model_path in model_files:
             name = Path(model_path).name
 
@@ -232,14 +239,14 @@ class CometLogger:
         with open(data_file) as f:
             data_config = yaml.safe_load(f)
 
-        path = data_config.get('path')
+        path = data_config.get("path")
         if path and path.startswith(COMET_PREFIX):
-            path = data_config['path'].replace(COMET_PREFIX, '')
+            path = data_config["path"].replace(COMET_PREFIX, "")
             data_dict = self.download_dataset_artifact(path)
 
             return data_dict
 
-        self.log_asset(self.opt.data, metadata={'type': 'data-config-file'})
+        self.log_asset(self.opt.data, metadata={"type": "data-config-file"})
 
         return check_dataset(data_file)
 
@@ -255,8 +262,8 @@ class CometLogger:
         filtered_detections = detections[mask]
         filtered_labels = labelsn[mask]
 
-        image_id = path.split('/')[-1].split('.')[0]
+        image_id = path.split("/")[-1].split(".")[0]
-        image_name = f'{image_id}_curr_epoch_{self.experiment.curr_epoch}'
+        image_name = f"{image_id}_curr_epoch_{self.experiment.curr_epoch}"
         if image_name not in self.logged_image_names:
             native_scale_image = PIL.Image.open(path)
             self.log_image(native_scale_image, name=image_name)
@@ -264,23 +271,21 @@ class CometLogger:
 
         metadata = []
         for cls, *xyxy in filtered_labels.tolist():
-            metadata.append({
+            metadata.append(
-                'label': f'{self.class_names[int(cls)]}-gt',
+                {
-                'score': 100,
+                    "label": f"{self.class_names[int(cls)]}-gt",
-                'box': {
+                    "score": 100,
-                    'x': xyxy[0],
+                    "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]},
-                    'y': xyxy[1],
+                }
-                    'x2': xyxy[2],
+            )
-                    'y2': xyxy[3]}, })
         for *xyxy, conf, cls in filtered_detections.tolist():
-            metadata.append({
+            metadata.append(
-                'label': f'{self.class_names[int(cls)]}',
+                {
-                'score': conf * 100,
+                    "label": f"{self.class_names[int(cls)]}",
-                'box': {
+                    "score": conf * 100,
-                    'x': xyxy[0],
+                    "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]},
-                    'y': xyxy[1],
+                }
-                    'x2': xyxy[2],
+            )
-                    'y2': xyxy[3]}, })
 
         self.metadata_dict[image_name] = metadata
         self.logged_images_count += 1
@@ -307,7 +312,7 @@ class CometLogger:
         return predn, labelsn
 
     def add_assets_to_artifact(self, artifact, path, asset_path, split):
-        img_paths = sorted(glob.glob(f'{asset_path}/*'))
+        img_paths = sorted(glob.glob(f"{asset_path}/*"))
         label_paths = img2label_paths(img_paths)
 
         for image_file, label_file in zip(img_paths, label_paths):
@@ -317,33 +322,33 @@ class CometLogger:
                 artifact.add(
                     image_file,
                     logical_path=image_logical_path,
-                    metadata={'split': split},
+                    metadata={"split": split},
                 )
                 artifact.add(
                     label_file,
                     logical_path=label_logical_path,
-                    metadata={'split': split},
+                    metadata={"split": split},
                 )
             except ValueError as e:
-                logger.error('COMET ERROR: Error adding file to Artifact. Skipping file.')
+                logger.error("COMET ERROR: Error adding file to Artifact. Skipping file.")
-                logger.error(f'COMET ERROR: {e}')
+                logger.error(f"COMET ERROR: {e}")
                 continue
 
         return artifact
 
     def upload_dataset_artifact(self):
-        dataset_name = self.data_dict.get('dataset_name', 'yolov5-dataset')
+        dataset_name = self.data_dict.get("dataset_name", "yolov5-dataset")
-        path = str((ROOT / Path(self.data_dict['path'])).resolve())
+        path = str((ROOT / Path(self.data_dict["path"])).resolve())
 
         metadata = self.data_dict.copy()
-        for key in ['train', 'val', 'test']:
+        for key in ["train", "val", "test"]:
             split_path = metadata.get(key)
             if split_path is not None:
-                metadata[key] = split_path.replace(path, '')
+                metadata[key] = split_path.replace(path, "")
 
-        artifact = comet_ml.Artifact(name=dataset_name, artifact_type='dataset', metadata=metadata)
+        artifact = comet_ml.Artifact(name=dataset_name, artifact_type="dataset", metadata=metadata)
         for key in metadata.keys():
-            if key in ['train', 'val', 'test']:
+            if key in ["train", "val", "test"]:
                 if isinstance(self.upload_dataset, str) and (key != self.upload_dataset):
                     continue
 
@@ -362,26 +367,27 @@ class CometLogger:
 
         metadata = logged_artifact.metadata
         data_dict = metadata.copy()
-        data_dict['path'] = artifact_save_dir
+        data_dict["path"] = artifact_save_dir
 
-        metadata_names = metadata.get('names')
+        metadata_names = metadata.get("names")
         if isinstance(metadata_names, dict):
-            data_dict['names'] = {int(k): v for k, v in metadata.get('names').items()}
+            data_dict["names"] = {int(k): v for k, v in metadata.get("names").items()}
         elif isinstance(metadata_names, list):
-            data_dict['names'] = {int(k): v for k, v in zip(range(len(metadata_names)), metadata_names)}
+            data_dict["names"] = {int(k): v for k, v in zip(range(len(metadata_names)), metadata_names)}
         else:
             raise "Invalid 'names' field in dataset yaml file. Please use a list or dictionary"
 
         return self.update_data_paths(data_dict)
 
     def update_data_paths(self, data_dict):
-        path = data_dict.get('path', '')
+        path = data_dict.get("path", "")
 
-        for split in ['train', 'val', 'test']:
+        for split in ["train", "val", "test"]:
             if data_dict.get(split):
                 split_path = data_dict.get(split)
-                data_dict[split] = (f'{path}/{split_path}' if isinstance(split, str) else [
+                data_dict[split] = (
-                    f'{path}/{x}' for x in split_path])
+                    f"{path}/{split_path}" if isinstance(split, str) else [f"{path}/{x}" for x in split_path]
+                )
 
         return data_dict
 
@@ -422,11 +428,11 @@ class CometLogger:
     def on_train_end(self, files, save_dir, last, best, epoch, results):
         if self.comet_log_predictions:
             curr_epoch = self.experiment.curr_epoch
-            self.experiment.log_asset_data(self.metadata_dict, 'image-metadata.json', epoch=curr_epoch)
+            self.experiment.log_asset_data(self.metadata_dict, "image-metadata.json", epoch=curr_epoch)
 
         for f in files:
-            self.log_asset(f, metadata={'epoch': epoch})
+            self.log_asset(f, metadata={"epoch": epoch})
-        self.log_asset(f'{save_dir}/results.csv', metadata={'epoch': epoch})
+        self.log_asset(f"{save_dir}/results.csv", metadata={"epoch": epoch})
 
         if not self.opt.evolve:
             model_path = str(best if best.exists() else last)
@@ -440,9 +446,9 @@ class CometLogger:
                 )
 
         # Check if running Experiment with Comet Optimizer
-        if hasattr(self.opt, 'comet_optimizer_id'):
+        if hasattr(self.opt, "comet_optimizer_id"):
             metric = results.get(self.opt.comet_optimizer_metric)
-            self.experiment.log_other('optimizer_metric_value', metric)
+            self.experiment.log_other("optimizer_metric_value", metric)
 
         self.finish_run()
 
@@ -477,21 +483,22 @@ class CometLogger:
                     class_name = self.class_names[c]
                     self.experiment.log_metrics(
                         {
-                            'mAP@.5': ap50[i],
+                            "mAP@.5": ap50[i],
-                            'mAP@.5:.95': ap[i],
+                            "mAP@.5:.95": ap[i],
-                            'precision': p[i],
+                            "precision": p[i],
-                            'recall': r[i],
+                            "recall": r[i],
-                            'f1': f1[i],
+                            "f1": f1[i],
-                            'true_positives': tp[i],
+                            "true_positives": tp[i],
-                            'false_positives': fp[i],
+                            "false_positives": fp[i],
-                            'support': nt[c], },
+                            "support": nt[c],
+                        },
                         prefix=class_name,
                     )
 
         if self.comet_log_confusion_matrix:
             epoch = self.experiment.curr_epoch
             class_names = list(self.class_names.values())
-            class_names.append('background')
+            class_names.append("background")
             num_classes = len(class_names)
 
             self.experiment.log_confusion_matrix(
@@ -499,9 +506,9 @@ class CometLogger:
                 max_categories=num_classes,
                 labels=class_names,
                 epoch=epoch,
-                column_label='Actual Category',
+                column_label="Actual Category",
-                row_label='Predicted Category',
+                row_label="Predicted Category",
-                file_name=f'confusion-matrix-epoch-{epoch}.json',
+                file_name=f"confusion-matrix-epoch-{epoch}.json",
             )
 
     def on_fit_epoch_end(self, result, epoch):

utils/loggers/comet/comet_utils.py

@@ -11,28 +11,28 @@ import yaml
 
 logger = logging.getLogger(__name__)
 
-COMET_PREFIX = 'comet://'
+COMET_PREFIX = "comet://"
-COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5')
+COMET_MODEL_NAME = os.getenv("COMET_MODEL_NAME", "yolov5")
-COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv('COMET_DEFAULT_CHECKPOINT_FILENAME', 'last.pt')
+COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv("COMET_DEFAULT_CHECKPOINT_FILENAME", "last.pt")
 
 
 def download_model_checkpoint(opt, experiment):
-    model_dir = f'{opt.project}/{experiment.name}'
+    model_dir = f"{opt.project}/{experiment.name}"
     os.makedirs(model_dir, exist_ok=True)
 
     model_name = COMET_MODEL_NAME
     model_asset_list = experiment.get_model_asset_list(model_name)
 
     if len(model_asset_list) == 0:
-        logger.error(f'COMET ERROR: No checkpoints found for model name : {model_name}')
+        logger.error(f"COMET ERROR: No checkpoints found for model name : {model_name}")
         return
 
     model_asset_list = sorted(
         model_asset_list,
-        key=lambda x: x['step'],
+        key=lambda x: x["step"],
         reverse=True,
     )
-    logged_checkpoint_map = {asset['fileName']: asset['assetId'] for asset in model_asset_list}
+    logged_checkpoint_map = {asset["fileName"]: asset["assetId"] for asset in model_asset_list}
 
     resource_url = urlparse(opt.weights)
     checkpoint_filename = resource_url.query
@@ -44,28 +44,28 @@ def download_model_checkpoint(opt, experiment):
         checkpoint_filename = COMET_DEFAULT_CHECKPOINT_FILENAME
 
     if asset_id is None:
-        logger.error(f'COMET ERROR: Checkpoint {checkpoint_filename} not found in the given Experiment')
+        logger.error(f"COMET ERROR: Checkpoint {checkpoint_filename} not found in the given Experiment")
         return
 
     try:
-        logger.info(f'COMET INFO: Downloading checkpoint {checkpoint_filename}')
+        logger.info(f"COMET INFO: Downloading checkpoint {checkpoint_filename}")
         asset_filename = checkpoint_filename
 
-        model_binary = experiment.get_asset(asset_id, return_type='binary', stream=False)
+        model_binary = experiment.get_asset(asset_id, return_type="binary", stream=False)
-        model_download_path = f'{model_dir}/{asset_filename}'
+        model_download_path = f"{model_dir}/{asset_filename}"
-        with open(model_download_path, 'wb') as f:
+        with open(model_download_path, "wb") as f:
             f.write(model_binary)
 
         opt.weights = model_download_path
 
     except Exception as e:
-        logger.warning('COMET WARNING: Unable to download checkpoint from Comet')
+        logger.warning("COMET WARNING: Unable to download checkpoint from Comet")
         logger.exception(e)
 
 
 def set_opt_parameters(opt, experiment):
-    """Update the opts Namespace with parameters
+    """
-    from Comet's ExistingExperiment when resuming a run
+    Update the opts Namespace with parameters from Comet's ExistingExperiment when resuming a run.
 
     Args:
         opt (argparse.Namespace): Namespace of command line options
@@ -75,9 +75,9 @@ def set_opt_parameters(opt, experiment):
     resume_string = opt.resume
 
     for asset in asset_list:
-        if asset['fileName'] == 'opt.yaml':
+        if asset["fileName"] == "opt.yaml":
-            asset_id = asset['assetId']
+            asset_id = asset["assetId"]
-            asset_binary = experiment.get_asset(asset_id, return_type='binary', stream=False)
+            asset_binary = experiment.get_asset(asset_id, return_type="binary", stream=False)
             opt_dict = yaml.safe_load(asset_binary)
             for key, value in opt_dict.items():
                 setattr(opt, key, value)
@@ -85,18 +85,18 @@ def set_opt_parameters(opt, experiment):
 
     # Save hyperparameters to YAML file
     # Necessary to pass checks in training script
-    save_dir = f'{opt.project}/{experiment.name}'
+    save_dir = f"{opt.project}/{experiment.name}"
     os.makedirs(save_dir, exist_ok=True)
 
-    hyp_yaml_path = f'{save_dir}/hyp.yaml'
+    hyp_yaml_path = f"{save_dir}/hyp.yaml"
-    with open(hyp_yaml_path, 'w') as f:
+    with open(hyp_yaml_path, "w") as f:
         yaml.dump(opt.hyp, f)
     opt.hyp = hyp_yaml_path
 
 
 def check_comet_weights(opt):
-    """Downloads model weights from Comet and updates the
+    """
-    weights path to point to saved weights location
+    Downloads model weights from Comet and updates the weights path to point to saved weights location.
 
     Args:
         opt (argparse.Namespace): Command Line arguments passed
@@ -113,7 +113,7 @@ def check_comet_weights(opt):
         if opt.weights.startswith(COMET_PREFIX):
             api = comet_ml.API()
             resource = urlparse(opt.weights)
-            experiment_path = f'{resource.netloc}{resource.path}'
+            experiment_path = f"{resource.netloc}{resource.path}"
             experiment = api.get(experiment_path)
             download_model_checkpoint(opt, experiment)
             return True
@@ -122,8 +122,8 @@ def check_comet_weights(opt):
 
 
 def check_comet_resume(opt):
-    """Restores run parameters to its original state based on the model checkpoint
+    """
-    and logged Experiment parameters.
+    Restores run parameters to its original state based on the model checkpoint and logged Experiment parameters.
 
     Args:
         opt (argparse.Namespace): Command Line arguments passed
@@ -140,7 +140,7 @@ def check_comet_resume(opt):
         if opt.resume.startswith(COMET_PREFIX):
             api = comet_ml.API()
             resource = urlparse(opt.resume)
-            experiment_path = f'{resource.netloc}{resource.path}'
+            experiment_path = f"{resource.netloc}{resource.path}"
             experiment = api.get(experiment_path)
             set_opt_parameters(opt, experiment)
             download_model_checkpoint(opt, experiment)

utils/loggers/comet/hpo.py

@@ -21,77 +21,79 @@ from utils.torch_utils import select_device
 
 # Project Configuration
 config = comet_ml.config.get_config()
-COMET_PROJECT_NAME = config.get_string(os.getenv('COMET_PROJECT_NAME'), 'comet.project_name', default='yolov5')
+COMET_PROJECT_NAME = config.get_string(os.getenv("COMET_PROJECT_NAME"), "comet.project_name", default="yolov5")
 
 
 def get_args(known=False):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
+    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
-    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
+    parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
-    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
-    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
+    parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path")
-    parser.add_argument('--epochs', type=int, default=300, help='total training epochs')
+    parser.add_argument("--epochs", type=int, default=300, help="total training epochs")
-    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
+    parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch")
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
+    parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)")
-    parser.add_argument('--rect', action='store_true', help='rectangular training')
+    parser.add_argument("--rect", action="store_true", help="rectangular training")
-    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
+    parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")
-    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
+    parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")
-    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
+    parser.add_argument("--noval", action="store_true", help="only validate final epoch")
-    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
+    parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor")
-    parser.add_argument('--noplots', action='store_true', help='save no plot files')
+    parser.add_argument("--noplots", action="store_true", help="save no plot files")
-    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
+    parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations")
-    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
+    parser.add_argument("--bucket", type=str, default="", help="gsutil bucket")
-    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
+    parser.add_argument("--cache", type=str, nargs="?", const="ram", help='--cache images in "ram" (default) or "disk"')
-    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
+    parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training")
-    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
+    parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%")
-    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
+    parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class")
-    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
+    parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer")
-    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
+    parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode")
-    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--quad', action='store_true', help='quad dataloader')
+    parser.add_argument("--quad", action="store_true", help="quad dataloader")
-    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
+    parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler")
-    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
+    parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon")
-    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
+    parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)")
-    parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
+    parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2")
-    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
+    parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)")
-    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
+    parser.add_argument("--seed", type=int, default=0, help="Global training seed")
-    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
+    parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify")
 
     # Weights & Biases arguments
-    parser.add_argument('--entity', default=None, help='W&B: Entity')
+    parser.add_argument("--entity", default=None, help="W&B: Entity")
-    parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option')
+    parser.add_argument("--upload_dataset", nargs="?", const=True, default=False, help='W&B: Upload data, "val" option')
-    parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
+    parser.add_argument("--bbox_interval", type=int, default=-1, help="W&B: Set bounding-box image logging interval")
-    parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')
+    parser.add_argument("--artifact_alias", type=str, default="latest", help="W&B: Version of dataset artifact to use")
 
     # Comet Arguments
-    parser.add_argument('--comet_optimizer_config', type=str, help='Comet: Path to a Comet Optimizer Config File.')
+    parser.add_argument("--comet_optimizer_config", type=str, help="Comet: Path to a Comet Optimizer Config File.")
-    parser.add_argument('--comet_optimizer_id', type=str, help='Comet: ID of the Comet Optimizer sweep.')
+    parser.add_argument("--comet_optimizer_id", type=str, help="Comet: ID of the Comet Optimizer sweep.")
-    parser.add_argument('--comet_optimizer_objective', type=str, help="Comet: Set to 'minimize' or 'maximize'.")
+    parser.add_argument("--comet_optimizer_objective", type=str, help="Comet: Set to 'minimize' or 'maximize'.")
-    parser.add_argument('--comet_optimizer_metric', type=str, help='Comet: Metric to Optimize.')
+    parser.add_argument("--comet_optimizer_metric", type=str, help="Comet: Metric to Optimize.")
-    parser.add_argument('--comet_optimizer_workers',
+    parser.add_argument(
+        "--comet_optimizer_workers",
         type=int,
         default=1,
-                        help='Comet: Number of Parallel Workers to use with the Comet Optimizer.')
+        help="Comet: Number of Parallel Workers to use with the Comet Optimizer.",
+    )
 
     return parser.parse_known_args()[0] if known else parser.parse_args()
 
 
 def run(parameters, opt):
-    hyp_dict = {k: v for k, v in parameters.items() if k not in ['epochs', 'batch_size']}
+    hyp_dict = {k: v for k, v in parameters.items() if k not in ["epochs", "batch_size"]}
 
     opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve))
-    opt.batch_size = parameters.get('batch_size')
+    opt.batch_size = parameters.get("batch_size")
-    opt.epochs = parameters.get('epochs')
+    opt.epochs = parameters.get("epochs")
 
     device = select_device(opt.device, batch_size=opt.batch_size)
     train(hyp_dict, opt, device, callbacks=Callbacks())
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = get_args(known=True)
 
     opt.weights = str(opt.weights)
@@ -99,7 +101,7 @@ if __name__ == '__main__':
     opt.data = str(opt.data)
     opt.project = str(opt.project)
 
-    optimizer_id = os.getenv('COMET_OPTIMIZER_ID')
+    optimizer_id = os.getenv("COMET_OPTIMIZER_ID")
     if optimizer_id is None:
         with open(opt.comet_optimizer_config) as f:
             optimizer_config = json.load(f)
@@ -110,9 +112,9 @@ if __name__ == '__main__':
     opt.comet_optimizer_id = optimizer.id
     status = optimizer.status()
 
-    opt.comet_optimizer_objective = status['spec']['objective']
+    opt.comet_optimizer_objective = status["spec"]["objective"]
-    opt.comet_optimizer_metric = status['spec']['metric']
+    opt.comet_optimizer_metric = status["spec"]["metric"]
 
-    logger.info('COMET INFO: Starting Hyperparameter Sweep')
+    logger.info("COMET INFO: Starting Hyperparameter Sweep")
     for parameter in optimizer.get_parameters():
-        run(parameter['parameters'], opt)
+        run(parameter["parameters"], opt)

utils/loggers/wandb/wandb_utils.py

@@ -15,34 +15,35 @@ FILE = Path(__file__).resolve()
 ROOT = FILE.parents[3]  # YOLOv5 root directory
 if str(ROOT) not in sys.path:
     sys.path.append(str(ROOT))  # add ROOT to PATH
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
-DEPRECATION_WARNING = f"{colorstr('wandb')}: WARNING ⚠️ wandb is deprecated and will be removed in a future release. " \
+DEPRECATION_WARNING = (
+    f"{colorstr('wandb')}: WARNING ⚠️ wandb is deprecated and will be removed in a future release. "
     f'See supported integrations at https://github.com/ultralytics/yolov5#integrations.'
+)
 
 try:
     import wandb
 
-    assert hasattr(wandb, '__version__')  # verify package import not local dir
+    assert hasattr(wandb, "__version__")  # verify package import not local dir
     LOGGER.warning(DEPRECATION_WARNING)
 except (ImportError, AssertionError):
     wandb = None
 
 
-class WandbLogger():
+class WandbLogger:
-    """Log training runs, datasets, models, and predictions to Weights & Biases.
+    """
+    Log training runs, datasets, models, and predictions to Weights & Biases.
 
-    This logger sends information to W&B at wandb.ai. By default, this information
+    This logger sends information to W&B at wandb.ai. By default, this information includes hyperparameters, system
-    includes hyperparameters, system configuration and metrics, model metrics,
+    configuration and metrics, model metrics, and basic data metrics and analyses.
-    and basic data metrics and analyses.
 
-    By providing additional command line arguments to train.py, datasets,
+    By providing additional command line arguments to train.py, datasets, models and predictions can also be logged.
-    models and predictions can also be logged.
 
     For more on how this logger is used, see the Weights & Biases documentation:
     https://docs.wandb.com/guides/integrations/yolov5
     """
 
-    def __init__(self, opt, run_id=None, job_type='Training'):
+    def __init__(self, opt, run_id=None, job_type="Training"):
         """
         - Initialize WandbLogger instance
         - Upload dataset if opt.upload_dataset is True
@@ -64,17 +65,23 @@ class WandbLogger():
         self.max_imgs_to_log = 16
         self.data_dict = None
         if self.wandb:
-            self.wandb_run = wandb.init(config=opt,
+            self.wandb_run = (
-                                        resume='allow',
+                wandb.init(
-                                        project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
+                    config=opt,
+                    resume="allow",
+                    project="YOLOv5" if opt.project == "runs/train" else Path(opt.project).stem,
                     entity=opt.entity,
-                                        name=opt.name if opt.name != 'exp' else None,
+                    name=opt.name if opt.name != "exp" else None,
                     job_type=job_type,
                     id=run_id,
-                                        allow_val_change=True) if not wandb.run else wandb.run
+                    allow_val_change=True,
+                )
+                if not wandb.run
+                else wandb.run
+            )
 
         if self.wandb_run:
-            if self.job_type == 'Training':
+            if self.job_type == "Training":
                 if isinstance(opt.data, dict):
                     # This means another dataset manager has already processed the dataset info (e.g. ClearML)
                     # and they will have stored the already processed dict in opt.data
@@ -97,11 +104,17 @@ class WandbLogger():
         if isinstance(opt.resume, str):
             model_dir, _ = self.download_model_artifact(opt)
             if model_dir:
-                self.weights = Path(model_dir) / 'last.pt'
+                self.weights = Path(model_dir) / "last.pt"
                 config = self.wandb_run.config
-                opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp, opt.imgsz = str(
+                opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp, opt.imgsz = (
-                    self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \
+                    str(self.weights),
-                    config.hyp, config.imgsz
+                    config.save_period,
+                    config.batch_size,
+                    config.bbox_interval,
+                    config.epochs,
+                    config.hyp,
+                    config.imgsz,
+                )
 
         if opt.bbox_interval == -1:
             self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
@@ -110,7 +123,7 @@ class WandbLogger():
 
     def log_model(self, path, opt, epoch, fitness_score, best_model=False):
         """
-        Log the model checkpoint as W&B artifact
+        Log the model checkpoint as W&B artifact.
 
         arguments:
         path (Path)   -- Path of directory containing the checkpoints
@@ -119,26 +132,30 @@ class WandbLogger():
         fitness_score (float) -- fitness score for current epoch
         best_model (boolean) -- Boolean representing if the current checkpoint is the best yet.
         """
-        model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model',
+        model_artifact = wandb.Artifact(
-                                        type='model',
+            "run_" + wandb.run.id + "_model",
+            type="model",
             metadata={
-                                            'original_url': str(path),
+                "original_url": str(path),
-                                            'epochs_trained': epoch + 1,
+                "epochs_trained": epoch + 1,
-                                            'save period': opt.save_period,
+                "save period": opt.save_period,
-                                            'project': opt.project,
+                "project": opt.project,
-                                            'total_epochs': opt.epochs,
+                "total_epochs": opt.epochs,
-                                            'fitness_score': fitness_score})
+                "fitness_score": fitness_score,
-        model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
+            },
-        wandb.log_artifact(model_artifact,
+        )
-                           aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
+        model_artifact.add_file(str(path / "last.pt"), name="last.pt")
-        LOGGER.info(f'Saving model artifact on epoch {epoch + 1}')
+        wandb.log_artifact(
+            model_artifact, aliases=["latest", "last", "epoch " + str(self.current_epoch), "best" if best_model else ""]
+        )
+        LOGGER.info(f"Saving model artifact on epoch {epoch + 1}")
 
     def val_one_image(self, pred, predn, path, names, im):
         pass
 
     def log(self, log_dict):
         """
-        save the metrics to the logging dictionary
+        Save the metrics to the logging dictionary.
 
         arguments:
         log_dict (Dict) -- metrics/media to be logged in current step
@@ -149,7 +166,7 @@ class WandbLogger():
 
     def end_epoch(self):
         """
-        commit the log_dict, model artifacts and Tables to W&B and flush the log_dict.
+        Commit the log_dict, model artifacts and Tables to W&B and flush the log_dict.
 
         arguments:
         best_result (boolean): Boolean representing if the result of this evaluation is best or not
@@ -160,16 +177,14 @@ class WandbLogger():
                     wandb.log(self.log_dict)
                 except BaseException as e:
                     LOGGER.info(
-                        f'An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}'
+                        f"An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}"
                     )
                     self.wandb_run.finish()
                     self.wandb_run = None
                 self.log_dict = {}
 
     def finish_run(self):
-        """
+        """Log metrics if any and finish the current W&B run."""
-        Log metrics if any and finish the current W&B run
-        """
         if self.wandb_run:
             if self.log_dict:
                 with all_logging_disabled():

utils/loss.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Loss functions."""
-Loss functions
-"""
 
 import torch
 import torch.nn as nn
@@ -19,7 +17,7 @@ class BCEBlurWithLogitsLoss(nn.Module):
     # BCEwithLogitLoss() with reduced missing label effects.
     def __init__(self, alpha=0.05):
         super().__init__()
-        self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none')  # must be nn.BCEWithLogitsLoss()
+        self.loss_fcn = nn.BCEWithLogitsLoss(reduction="none")  # must be nn.BCEWithLogitsLoss()
         self.alpha = alpha
 
     def forward(self, pred, true):
@@ -40,7 +38,7 @@ class FocalLoss(nn.Module):
         self.gamma = gamma
         self.alpha = alpha
         self.reduction = loss_fcn.reduction
-        self.loss_fcn.reduction = 'none'  # required to apply FL to each element
+        self.loss_fcn.reduction = "none"  # required to apply FL to each element
 
     def forward(self, pred, true):
         loss = self.loss_fcn(pred, true)
@@ -54,9 +52,9 @@ class FocalLoss(nn.Module):
         modulating_factor = (1.0 - p_t) ** self.gamma
         loss *= alpha_factor * modulating_factor
 
-        if self.reduction == 'mean':
+        if self.reduction == "mean":
             return loss.mean()
-        elif self.reduction == 'sum':
+        elif self.reduction == "sum":
             return loss.sum()
         else:  # 'none'
             return loss
@@ -70,7 +68,7 @@ class QFocalLoss(nn.Module):
         self.gamma = gamma
         self.alpha = alpha
         self.reduction = loss_fcn.reduction
-        self.loss_fcn.reduction = 'none'  # required to apply FL to each element
+        self.loss_fcn.reduction = "none"  # required to apply FL to each element
 
     def forward(self, pred, true):
         loss = self.loss_fcn(pred, true)
@@ -80,9 +78,9 @@ class QFocalLoss(nn.Module):
         modulating_factor = torch.abs(true - pred_prob) ** self.gamma
         loss *= alpha_factor * modulating_factor
 
-        if self.reduction == 'mean':
+        if self.reduction == "mean":
             return loss.mean()
-        elif self.reduction == 'sum':
+        elif self.reduction == "sum":
             return loss.sum()
         else:  # 'none'
             return loss
@@ -97,14 +95,14 @@ class ComputeLoss:
         h = model.hyp  # hyperparameters
 
         # Define criteria
-        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
+        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device))
-        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
+        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device))
 
         # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
-        self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0))  # positive, negative BCE targets
+        self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0))  # positive, negative BCE targets
 
         # Focal loss
-        g = h['fl_gamma']  # focal loss gamma
+        g = h["fl_gamma"]  # focal loss gamma
         if g > 0:
             BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
 
@@ -167,9 +165,9 @@ class ComputeLoss:
 
         if self.autobalance:
             self.balance = [x / self.balance[self.ssi] for x in self.balance]
-        lbox *= self.hyp['box']
+        lbox *= self.hyp["box"]
-        lobj *= self.hyp['obj']
+        lobj *= self.hyp["obj"]
-        lcls *= self.hyp['cls']
+        lcls *= self.hyp["cls"]
         bs = tobj.shape[0]  # batch size
 
         return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()
@@ -183,7 +181,8 @@ class ComputeLoss:
         targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2)  # append anchor indices
 
         g = 0.5  # bias
-        off = torch.tensor(
+        off = (
+            torch.tensor(
                 [
                     [0, 0],
                     [1, 0],
@@ -192,7 +191,10 @@ class ComputeLoss:
                     [0, -1],  # j,k,l,m
                     # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                 ],
-            device=self.device).float() * g  # offsets
+                device=self.device,
+            ).float()
+            * g
+        )  # offsets
 
         for i in range(self.nl):
             anchors, shape = self.anchors[i], p[i].shape
@@ -203,7 +205,7 @@ class ComputeLoss:
             if nt:
                 # Matches
                 r = t[..., 4:6] / anchors[:, None]  # wh ratio
-                j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare
+                j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"]  # compare
                 # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                 t = t[j]  # filter
 

utils/metrics.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Model validation metrics."""
-Model validation metrics
-"""
 
 import math
 import warnings
@@ -25,11 +23,13 @@ def smooth(y, f=0.05):
     nf = round(len(y) * f * 2) // 2 + 1  # number of filter elements (must be odd)
     p = np.ones(nf // 2)  # ones padding
     yp = np.concatenate((p * y[0], y, p * y[-1]), 0)  # y padded
-    return np.convolve(yp, np.ones(nf) / nf, mode='valid')  # y-smoothed
+    return np.convolve(yp, np.ones(nf) / nf, mode="valid")  # y-smoothed
 
 
-def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16, prefix=''):
+def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=".", names=(), eps=1e-16, prefix=""):
-    """ Compute the average precision, given the recall and precision curves.
+    """
+    Compute the average precision, given the recall and precision curves.
+
     Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
     # Arguments
         tp:  True positives (nparray, nx1 or nx10).
@@ -83,10 +83,10 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
     names = [v for k, v in names.items() if k in unique_classes]  # list: only classes that have data
     names = dict(enumerate(names))  # to dict
     if plot:
-        plot_pr_curve(px, py, ap, Path(save_dir) / f'{prefix}PR_curve.png', names)
+        plot_pr_curve(px, py, ap, Path(save_dir) / f"{prefix}PR_curve.png", names)
-        plot_mc_curve(px, f1, Path(save_dir) / f'{prefix}F1_curve.png', names, ylabel='F1')
+        plot_mc_curve(px, f1, Path(save_dir) / f"{prefix}F1_curve.png", names, ylabel="F1")
-        plot_mc_curve(px, p, Path(save_dir) / f'{prefix}P_curve.png', names, ylabel='Precision')
+        plot_mc_curve(px, p, Path(save_dir) / f"{prefix}P_curve.png", names, ylabel="Precision")
-        plot_mc_curve(px, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall')
+        plot_mc_curve(px, r, Path(save_dir) / f"{prefix}R_curve.png", names, ylabel="Recall")
 
     i = smooth(f1.mean(0), 0.1).argmax()  # max F1 index
     p, r, f1 = p[:, i], r[:, i], f1[:, i]
@@ -112,8 +112,8 @@ def compute_ap(recall, precision):
     mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
 
     # Integrate area under curve
-    method = 'interp'  # methods: 'continuous', 'interp'
+    method = "interp"  # methods: 'continuous', 'interp'
-    if method == 'interp':
+    if method == "interp":
         x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
         ap = np.trapz(np.interp(x, mrec, mpre), x)  # integrate
     else:  # 'continuous'
@@ -134,6 +134,7 @@ class ConfusionMatrix:
     def process_batch(self, detections, labels):
         """
         Return intersection-over-union (Jaccard index) of boxes.
+
         Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
         Arguments:
             detections (Array[N, 6]), x1, y1, x2, y2, conf, class
@@ -183,40 +184,41 @@ class ConfusionMatrix:
         # fn = self.matrix.sum(0) - tp  # false negatives (missed detections)
         return tp[:-1], fp[:-1]  # remove background class
 
-    @TryExcept('WARNING ⚠️ ConfusionMatrix plot failure')
+    @TryExcept("WARNING ⚠️ ConfusionMatrix plot failure")
-    def plot(self, normalize=True, save_dir='', names=()):
+    def plot(self, normalize=True, save_dir="", names=()):
         import seaborn as sn
 
-        array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-9) if normalize else 1)  # normalize columns
+        array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1e-9) if normalize else 1)  # normalize columns
         array[array < 0.005] = np.nan  # don't annotate (would appear as 0.00)
 
         fig, ax = plt.subplots(1, 1, figsize=(12, 9), tight_layout=True)
         nc, nn = self.nc, len(names)  # number of classes, names
         sn.set(font_scale=1.0 if nc < 50 else 0.8)  # for label size
         labels = (0 < nn < 99) and (nn == nc)  # apply names to ticklabels
-        ticklabels = (names + ['background']) if labels else 'auto'
+        ticklabels = (names + ["background"]) if labels else "auto"
         with warnings.catch_warnings():
-            warnings.simplefilter('ignore')  # suppress empty matrix RuntimeWarning: All-NaN slice encountered
+            warnings.simplefilter("ignore")  # suppress empty matrix RuntimeWarning: All-NaN slice encountered
-            sn.heatmap(array,
+            sn.heatmap(
+                array,
                 ax=ax,
                 annot=nc < 30,
-                       annot_kws={
+                annot_kws={"size": 8},
-                           'size': 8},
+                cmap="Blues",
-                       cmap='Blues',
+                fmt=".2f",
-                       fmt='.2f',
                 square=True,
                 vmin=0.0,
                 xticklabels=ticklabels,
-                       yticklabels=ticklabels).set_facecolor((1, 1, 1))
+                yticklabels=ticklabels,
-        ax.set_xlabel('True')
+            ).set_facecolor((1, 1, 1))
-        ax.set_ylabel('Predicted')
+        ax.set_xlabel("True")
-        ax.set_title('Confusion Matrix')
+        ax.set_ylabel("Predicted")
-        fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
+        ax.set_title("Confusion Matrix")
+        fig.savefig(Path(save_dir) / "confusion_matrix.png", dpi=250)
         plt.close(fig)
 
     def print(self):
         for i in range(self.nc + 1):
-            print(' '.join(map(str, self.matrix[i])))
+            print(" ".join(map(str, self.matrix[i])))
 
 
 def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
@@ -235,8 +237,9 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7
         w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)
 
     # Intersection area
-    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * \
+    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * (
-            (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp(0)
+        b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)
+    ).clamp(0)
 
     # Union Area
     union = w1 * h1 + w2 * h2 - inter + eps
@@ -264,6 +267,7 @@ def box_iou(box1, box2, eps=1e-7):
     # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
     """
     Return intersection-over-union (Jaccard index) of boxes.
+
     Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
     Arguments:
         box1 (Tensor[N, 4])
@@ -282,7 +286,10 @@ def box_iou(box1, box2, eps=1e-7):
 
 
 def bbox_ioa(box1, box2, eps=1e-7):
-    """ Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2
+    """
+    Returns the intersection over box2 area given box1, box2.
+
+    Boxes are x1y1x2y2
     box1:       np.array of shape(4)
     box2:       np.array of shape(nx4)
     returns:    np.array of shape(n)
@@ -293,8 +300,9 @@ def bbox_ioa(box1, box2, eps=1e-7):
     b2_x1, b2_y1, b2_x2, b2_y2 = box2.T
 
     # Intersection area
-    inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \
+    inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * (
-                 (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)
+        np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)
+    ).clip(0)
 
     # box2 area
     box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + eps
@@ -315,46 +323,46 @@ def wh_iou(wh1, wh2, eps=1e-7):
 
 
 @threaded
-def plot_pr_curve(px, py, ap, save_dir=Path('pr_curve.png'), names=()):
+def plot_pr_curve(px, py, ap, save_dir=Path("pr_curve.png"), names=()):
     # Precision-recall curve
     fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
     py = np.stack(py, axis=1)
 
     if 0 < len(names) < 21:  # display per-class legend if < 21 classes
         for i, y in enumerate(py.T):
-            ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}')  # plot(recall, precision)
+            ax.plot(px, y, linewidth=1, label=f"{names[i]} {ap[i, 0]:.3f}")  # plot(recall, precision)
     else:
-        ax.plot(px, py, linewidth=1, color='grey')  # plot(recall, precision)
+        ax.plot(px, py, linewidth=1, color="grey")  # plot(recall, precision)
 
-    ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean())
+    ax.plot(px, py.mean(1), linewidth=3, color="blue", label="all classes %.3f mAP@0.5" % ap[:, 0].mean())
-    ax.set_xlabel('Recall')
+    ax.set_xlabel("Recall")
-    ax.set_ylabel('Precision')
+    ax.set_ylabel("Precision")
     ax.set_xlim(0, 1)
     ax.set_ylim(0, 1)
-    ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left')
+    ax.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
-    ax.set_title('Precision-Recall Curve')
+    ax.set_title("Precision-Recall Curve")
     fig.savefig(save_dir, dpi=250)
     plt.close(fig)
 
 
 @threaded
-def plot_mc_curve(px, py, save_dir=Path('mc_curve.png'), names=(), xlabel='Confidence', ylabel='Metric'):
+def plot_mc_curve(px, py, save_dir=Path("mc_curve.png"), names=(), xlabel="Confidence", ylabel="Metric"):
     # Metric-confidence curve
     fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
 
     if 0 < len(names) < 21:  # display per-class legend if < 21 classes
         for i, y in enumerate(py):
-            ax.plot(px, y, linewidth=1, label=f'{names[i]}')  # plot(confidence, metric)
+            ax.plot(px, y, linewidth=1, label=f"{names[i]}")  # plot(confidence, metric)
     else:
-        ax.plot(px, py.T, linewidth=1, color='grey')  # plot(confidence, metric)
+        ax.plot(px, py.T, linewidth=1, color="grey")  # plot(confidence, metric)
 
     y = smooth(py.mean(0), 0.05)
-    ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}')
+    ax.plot(px, y, linewidth=3, color="blue", label=f"all classes {y.max():.2f} at {px[y.argmax()]:.3f}")
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
     ax.set_xlim(0, 1)
     ax.set_ylim(0, 1)
-    ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left')
+    ax.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
-    ax.set_title(f'{ylabel}-Confidence Curve')
+    ax.set_title(f"{ylabel}-Confidence Curve")
     fig.savefig(save_dir, dpi=250)
     plt.close(fig)

utils/plots.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Plotting utils."""
-Plotting utils
-"""
 
 import contextlib
 import math
@@ -25,18 +23,38 @@ from utils.general import LOGGER, clip_boxes, increment_path, xywh2xyxy, xyxy2xy
 from utils.metrics import fitness
 
 # Settings
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
-matplotlib.rc('font', **{'size': 11})
+matplotlib.rc("font", **{"size": 11})
-matplotlib.use('Agg')  # for writing to files only
+matplotlib.use("Agg")  # for writing to files only
 
 
 class Colors:
     # Ultralytics color palette https://ultralytics.com/
     def __init__(self):
         # hex = matplotlib.colors.TABLEAU_COLORS.values()
-        hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
+        hexs = (
-                '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
+            "FF3838",
-        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
+            "FF9D97",
+            "FF701F",
+            "FFB21D",
+            "CFD231",
+            "48F90A",
+            "92CC17",
+            "3DDB86",
+            "1A9334",
+            "00D4BB",
+            "2C99A8",
+            "00C2FF",
+            "344593",
+            "6473FF",
+            "0018EC",
+            "8438FF",
+            "520085",
+            "CB38FF",
+            "FF95C8",
+            "FF37C7",
+        )
+        self.palette = [self.hex2rgb(f"#{c}") for c in hexs]
         self.n = len(self.palette)
 
     def __call__(self, i, bgr=False):
@@ -51,7 +69,7 @@ class Colors:
 colors = Colors()  # create instance for 'from utils.plots import colors'
 
 
-def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')):
+def feature_visualization(x, module_type, stage, n=32, save_dir=Path("runs/detect/exp")):
     """
     x:              Features to be visualized
     module_type:    Module type
@@ -59,9 +77,9 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
     n:              Maximum number of feature maps to plot
     save_dir:       Directory to save results
     """
-    if ('Detect'
+    if ("Detect" not in module_type) and (
-            not in module_type) and ('Segment'
+        "Segment" not in module_type
-                                     not in module_type):  # 'Detect' for Object Detect task,'Segment' for Segment task
+    ):  # 'Detect' for Object Detect task,'Segment' for Segment task
         batch, channels, height, width = x.shape  # batch, channels, height, width
         if height > 1 and width > 1:
             f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png"  # filename
@@ -73,12 +91,12 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
             plt.subplots_adjust(wspace=0.05, hspace=0.05)
             for i in range(n):
                 ax[i].imshow(blocks[i].squeeze())  # cmap='gray'
-                ax[i].axis('off')
+                ax[i].axis("off")
 
-            LOGGER.info(f'Saving {f}... ({n}/{channels})')
+            LOGGER.info(f"Saving {f}... ({n}/{channels})")
-            plt.savefig(f, dpi=300, bbox_inches='tight')
+            plt.savefig(f, dpi=300, bbox_inches="tight")
             plt.close()
-            np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy())  # npy save
+            np.save(str(f.with_suffix(".npy")), x[0].cpu().numpy())  # npy save
 
 
 def hist2d(x, y, n=100):
@@ -97,7 +115,7 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
     def butter_lowpass(cutoff, fs, order):
         nyq = 0.5 * fs
         normal_cutoff = cutoff / nyq
-        return butter(order, normal_cutoff, btype='low', analog=False)
+        return butter(order, normal_cutoff, btype="low", analog=False)
 
     b, a = butter_lowpass(cutoff, fs, order=order)
     return filtfilt(b, a, data)  # forward-backward filter
@@ -114,7 +132,7 @@ def output_to_target(output, max_det=300):
 
 
 @threaded
-def plot_images(images, targets, paths=None, fname='images.jpg', names=None):
+def plot_images(images, targets, paths=None, fname="images.jpg", names=None):
     # Plot image grid with labels
     if isinstance(images, torch.Tensor):
         images = images.cpu().float().numpy()
@@ -156,7 +174,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None):
         if len(targets) > 0:
             ti = targets[targets[:, 0] == i]  # image targets
             boxes = xywh2xyxy(ti[:, 2:6]).T
-            classes = ti[:, 1].astype('int')
+            classes = ti[:, 1].astype("int")
             labels = ti.shape[1] == 6  # labels if no conf column
             conf = None if labels else ti[:, 6]  # check for confidence presence (label vs pred)
 
@@ -173,59 +191,59 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None):
                 color = colors(cls)
                 cls = names[cls] if names else cls
                 if labels or conf[j] > 0.25:  # 0.25 conf thresh
-                    label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}'
+                    label = f"{cls}" if labels else f"{cls} {conf[j]:.1f}"
                     annotator.box_label(box, label, color=color)
     annotator.im.save(fname)  # save
 
 
-def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
+def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=""):
     # Plot LR simulating training for full epochs
     optimizer, scheduler = copy(optimizer), copy(scheduler)  # do not modify originals
     y = []
     for _ in range(epochs):
         scheduler.step()
-        y.append(optimizer.param_groups[0]['lr'])
+        y.append(optimizer.param_groups[0]["lr"])
-    plt.plot(y, '.-', label='LR')
+    plt.plot(y, ".-", label="LR")
-    plt.xlabel('epoch')
+    plt.xlabel("epoch")
-    plt.ylabel('LR')
+    plt.ylabel("LR")
     plt.grid()
     plt.xlim(0, epochs)
     plt.ylim(0)
-    plt.savefig(Path(save_dir) / 'LR.png', dpi=200)
+    plt.savefig(Path(save_dir) / "LR.png", dpi=200)
     plt.close()
 
 
 def plot_val_txt():  # from utils.plots import *; plot_val()
     # Plot val.txt histograms
-    x = np.loadtxt('val.txt', dtype=np.float32)
+    x = np.loadtxt("val.txt", dtype=np.float32)
     box = xyxy2xywh(x[:, :4])
     cx, cy = box[:, 0], box[:, 1]
 
     fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True)
     ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0)
-    ax.set_aspect('equal')
+    ax.set_aspect("equal")
-    plt.savefig('hist2d.png', dpi=300)
+    plt.savefig("hist2d.png", dpi=300)
 
     fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True)
     ax[0].hist(cx, bins=600)
     ax[1].hist(cy, bins=600)
-    plt.savefig('hist1d.png', dpi=200)
+    plt.savefig("hist1d.png", dpi=200)
 
 
 def plot_targets_txt():  # from utils.plots import *; plot_targets_txt()
     # Plot targets.txt histograms
-    x = np.loadtxt('targets.txt', dtype=np.float32).T
+    x = np.loadtxt("targets.txt", dtype=np.float32).T
-    s = ['x targets', 'y targets', 'width targets', 'height targets']
+    s = ["x targets", "y targets", "width targets", "height targets"]
     fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
     ax = ax.ravel()
     for i in range(4):
-        ax[i].hist(x[i], bins=100, label=f'{x[i].mean():.3g} +/- {x[i].std():.3g}')
+        ax[i].hist(x[i], bins=100, label=f"{x[i].mean():.3g} +/- {x[i].std():.3g}")
         ax[i].legend()
         ax[i].set_title(s[i])
-    plt.savefig('targets.jpg', dpi=200)
+    plt.savefig("targets.jpg", dpi=200)
 
 
-def plot_val_study(file='', dir='', x=None):  # from utils.plots import *; plot_val_study()
+def plot_val_study(file="", dir="", x=None):  # from utils.plots import *; plot_val_study()
     # Plot file=study.txt generated by val.py (or plot all study*.txt in dir)
     save_dir = Path(file).parent if file else Path(dir)
     plot2 = False  # plot additional results
@@ -234,69 +252,74 @@ def plot_val_study(file='', dir='', x=None):  # from utils.plots import *; plot_
 
     fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
     # for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov5n6', 'yolov5s6', 'yolov5m6', 'yolov5l6', 'yolov5x6']]:
-    for f in sorted(save_dir.glob('study*.txt')):
+    for f in sorted(save_dir.glob("study*.txt")):
         y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
         x = np.arange(y.shape[1]) if x is None else np.array(x)
         if plot2:
-            s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_preprocess (ms/img)', 't_inference (ms/img)', 't_NMS (ms/img)']
+            s = ["P", "R", "mAP@.5", "mAP@.5:.95", "t_preprocess (ms/img)", "t_inference (ms/img)", "t_NMS (ms/img)"]
             for i in range(7):
-                ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
+                ax[i].plot(x, y[i], ".-", linewidth=2, markersize=8)
                 ax[i].set_title(s[i])
 
         j = y[3].argmax() + 1
-        ax2.plot(y[5, 1:j],
+        ax2.plot(
-                 y[3, 1:j] * 1E2,
+            y[5, 1:j],
-                 '.-',
+            y[3, 1:j] * 1e2,
+            ".-",
             linewidth=2,
             markersize=8,
-                 label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
+            label=f.stem.replace("study_coco_", "").replace("yolo", "YOLO"),
+        )
 
-    ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
+    ax2.plot(
-             'k.-',
+        1e3 / np.array([209, 140, 97, 58, 35, 18]),
+        [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
+        "k.-",
         linewidth=2,
         markersize=8,
-             alpha=.25,
+        alpha=0.25,
-             label='EfficientDet')
+        label="EfficientDet",
+    )
 
     ax2.grid(alpha=0.2)
     ax2.set_yticks(np.arange(20, 60, 5))
     ax2.set_xlim(0, 57)
     ax2.set_ylim(25, 55)
-    ax2.set_xlabel('GPU Speed (ms/img)')
+    ax2.set_xlabel("GPU Speed (ms/img)")
-    ax2.set_ylabel('COCO AP val')
+    ax2.set_ylabel("COCO AP val")
-    ax2.legend(loc='lower right')
+    ax2.legend(loc="lower right")
-    f = save_dir / 'study.png'
+    f = save_dir / "study.png"
-    print(f'Saving {f}...')
+    print(f"Saving {f}...")
     plt.savefig(f, dpi=300)
 
 
 @TryExcept()  # known issue https://github.com/ultralytics/yolov5/issues/5395
-def plot_labels(labels, names=(), save_dir=Path('')):
+def plot_labels(labels, names=(), save_dir=Path("")):
     # plot dataset labels
     LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
     c, b = labels[:, 0], labels[:, 1:].transpose()  # classes, boxes
     nc = int(c.max() + 1)  # number of classes
-    x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
+    x = pd.DataFrame(b.transpose(), columns=["x", "y", "width", "height"])
 
     # seaborn correlogram
-    sn.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
+    sn.pairplot(x, corner=True, diag_kind="auto", kind="hist", diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
-    plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200)
+    plt.savefig(save_dir / "labels_correlogram.jpg", dpi=200)
     plt.close()
 
     # matplotlib labels
-    matplotlib.use('svg')  # faster
+    matplotlib.use("svg")  # faster
     ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
     y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
     with contextlib.suppress(Exception):  # color histogram bars by class
         [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)]  # known issue #3195
-    ax[0].set_ylabel('instances')
+    ax[0].set_ylabel("instances")
     if 0 < len(names) < 30:
         ax[0].set_xticks(range(len(names)))
         ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10)
     else:
-        ax[0].set_xlabel('classes')
+        ax[0].set_xlabel("classes")
-    sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
+    sn.histplot(x, x="x", y="y", ax=ax[2], bins=50, pmax=0.9)
-    sn.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
+    sn.histplot(x, x="width", y="height", ax=ax[3], bins=50, pmax=0.9)
 
     # rectangles
     labels[:, 1:3] = 0.5  # center
@@ -305,24 +328,25 @@ def plot_labels(labels, names=(), save_dir=Path('')):
     for cls, *box in labels[:1000]:
         ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls))  # plot
     ax[1].imshow(img)
-    ax[1].axis('off')
+    ax[1].axis("off")
 
     for a in [0, 1, 2, 3]:
-        for s in ['top', 'right', 'left', 'bottom']:
+        for s in ["top", "right", "left", "bottom"]:
             ax[a].spines[s].set_visible(False)
 
-    plt.savefig(save_dir / 'labels.jpg', dpi=200)
+    plt.savefig(save_dir / "labels.jpg", dpi=200)
-    matplotlib.use('Agg')
+    matplotlib.use("Agg")
     plt.close()
 
 
-def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f=Path('images.jpg')):
+def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f=Path("images.jpg")):
     # Show classification image grid with labels (optional) and predictions (optional)
     from utils.augmentations import denormalize
 
-    names = names or [f'class{i}' for i in range(1000)]
+    names = names or [f"class{i}" for i in range(1000)]
-    blocks = torch.chunk(denormalize(im.clone()).cpu().float(), len(im),
+    blocks = torch.chunk(
-                         dim=0)  # select batch index 0, block by channels
+        denormalize(im.clone()).cpu().float(), len(im), dim=0
+    )  # select batch index 0, block by channels
     n = min(len(blocks), nmax)  # number of plots
     m = min(8, round(n**0.5))  # 8 x 8 default
     fig, ax = plt.subplots(math.ceil(n / m), m)  # 8 rows x n/8 cols
@@ -330,22 +354,22 @@ def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f
     # plt.subplots_adjust(wspace=0.05, hspace=0.05)
     for i in range(n):
         ax[i].imshow(blocks[i].squeeze().permute((1, 2, 0)).numpy().clip(0.0, 1.0))
-        ax[i].axis('off')
+        ax[i].axis("off")
         if labels is not None:
-            s = names[labels[i]] + (f'—{names[pred[i]]}' if pred is not None else '')
+            s = names[labels[i]] + (f"—{names[pred[i]]}" if pred is not None else "")
-            ax[i].set_title(s, fontsize=8, verticalalignment='top')
+            ax[i].set_title(s, fontsize=8, verticalalignment="top")
-    plt.savefig(f, dpi=300, bbox_inches='tight')
+    plt.savefig(f, dpi=300, bbox_inches="tight")
     plt.close()
     if verbose:
-        LOGGER.info(f'Saving {f}')
+        LOGGER.info(f"Saving {f}")
         if labels is not None:
-            LOGGER.info('True:     ' + ' '.join(f'{names[i]:3s}' for i in labels[:nmax]))
+            LOGGER.info("True:     " + " ".join(f"{names[i]:3s}" for i in labels[:nmax]))
         if pred is not None:
-            LOGGER.info('Predicted:' + ' '.join(f'{names[i]:3s}' for i in pred[:nmax]))
+            LOGGER.info("Predicted:" + " ".join(f"{names[i]:3s}" for i in pred[:nmax]))
     return f
 
 
-def plot_evolve(evolve_csv='path/to/evolve.csv'):  # from utils.plots import *; plot_evolve()
+def plot_evolve(evolve_csv="path/to/evolve.csv"):  # from utils.plots import *; plot_evolve()
     # Plot evolve.csv hyp evolution results
     evolve_csv = Path(evolve_csv)
     data = pd.read_csv(evolve_csv)
@@ -354,83 +378,83 @@ def plot_evolve(evolve_csv='path/to/evolve.csv'):  # from utils.plots import *;
     f = fitness(x)
     j = np.argmax(f)  # max fitness index
     plt.figure(figsize=(10, 12), tight_layout=True)
-    matplotlib.rc('font', **{'size': 8})
+    matplotlib.rc("font", **{"size": 8})
-    print(f'Best results from row {j} of {evolve_csv}:')
+    print(f"Best results from row {j} of {evolve_csv}:")
     for i, k in enumerate(keys[7:]):
         v = x[:, 7 + i]
         mu = v[j]  # best single result
         plt.subplot(6, 5, i + 1)
-        plt.scatter(v, f, c=hist2d(v, f, 20), cmap='viridis', alpha=.8, edgecolors='none')
+        plt.scatter(v, f, c=hist2d(v, f, 20), cmap="viridis", alpha=0.8, edgecolors="none")
-        plt.plot(mu, f.max(), 'k+', markersize=15)
+        plt.plot(mu, f.max(), "k+", markersize=15)
-        plt.title(f'{k} = {mu:.3g}', fontdict={'size': 9})  # limit to 40 characters
+        plt.title(f"{k} = {mu:.3g}", fontdict={"size": 9})  # limit to 40 characters
         if i % 5 != 0:
             plt.yticks([])
-        print(f'{k:>15}: {mu:.3g}')
+        print(f"{k:>15}: {mu:.3g}")
-    f = evolve_csv.with_suffix('.png')  # filename
+    f = evolve_csv.with_suffix(".png")  # filename
     plt.savefig(f, dpi=200)
     plt.close()
-    print(f'Saved {f}')
+    print(f"Saved {f}")
 
 
-def plot_results(file='path/to/results.csv', dir=''):
+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'))
+    files = list(save_dir.glob("results*.csv"))
-    assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.'
+    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 = 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)  # actual results
+                ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8)  # actual results
-                ax[i].plot(x, gaussian_filter1d(y, sigma=3), ':', label='smooth', linewidth=2)  # smoothing line
+                ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2)  # smoothing line
                 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:
-            LOGGER.info(f'Warning: Plotting error for {f}: {e}')
+            LOGGER.info(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()
 
 
-def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
+def profile_idetection(start=0, stop=0, labels=(), save_dir=""):
     # Plot iDetection '*.txt' per-image logs. from utils.plots import *; profile_idetection()
     ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel()
-    s = ['Images', 'Free Storage (GB)', 'RAM Usage (GB)', 'Battery', 'dt_raw (ms)', 'dt_smooth (ms)', 'real-world FPS']
+    s = ["Images", "Free Storage (GB)", "RAM Usage (GB)", "Battery", "dt_raw (ms)", "dt_smooth (ms)", "real-world FPS"]
-    files = list(Path(save_dir).glob('frames*.txt'))
+    files = list(Path(save_dir).glob("frames*.txt"))
     for fi, f in enumerate(files):
         try:
             results = np.loadtxt(f, ndmin=2).T[:, 90:-30]  # clip first and last rows
             n = results.shape[1]  # number of rows
             x = np.arange(start, min(stop, n) if stop else n)
             results = results[:, x]
-            t = (results[0] - results[0].min())  # set t0=0s
+            t = results[0] - results[0].min()  # set t0=0s
             results[0] = x
             for i, a in enumerate(ax):
                 if i < len(results):
-                    label = labels[fi] if len(labels) else f.stem.replace('frames_', '')
+                    label = labels[fi] if len(labels) else f.stem.replace("frames_", "")
-                    a.plot(t, results[i], marker='.', label=label, linewidth=1, markersize=5)
+                    a.plot(t, results[i], marker=".", label=label, linewidth=1, markersize=5)
                     a.set_title(s[i])
-                    a.set_xlabel('time (s)')
+                    a.set_xlabel("time (s)")
                     # if fi == len(files) - 1:
                     #     a.set_ylim(bottom=0)
-                    for side in ['top', 'right']:
+                    for side in ["top", "right"]:
                         a.spines[side].set_visible(False)
                 else:
                     a.remove()
         except Exception as e:
-            print(f'Warning: Plotting error for {f}; {e}')
+            print(f"Warning: Plotting error for {f}; {e}")
     ax[1].legend()
-    plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
+    plt.savefig(Path(save_dir) / "idetection_profile.png", dpi=200)
 
 
-def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
+def save_one_box(xyxy, im, file=Path("im.jpg"), gain=1.02, pad=10, square=False, BGR=False, save=True):
     # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
     xyxy = torch.tensor(xyxy).view(-1, 4)
     b = xyxy2xywh(xyxy)  # boxes
@@ -442,7 +466,7 @@ def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False,
     crop = im[int(xyxy[0, 1]) : int(xyxy[0, 3]), int(xyxy[0, 0]) : int(xyxy[0, 2]), :: (1 if BGR else -1)]
     if save:
         file.parent.mkdir(parents=True, exist_ok=True)  # make directory
-        f = str(increment_path(file).with_suffix('.jpg'))
+        f = str(increment_path(file).with_suffix(".jpg"))
         # cv2.imwrite(f, crop)  # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
         Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0)  # save RGB
     return crop

utils/segment/augmentations.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Image augmentation functions."""
-Image augmentation functions
-"""
 
 import math
 import random
@@ -22,15 +20,9 @@ def mixup(im, labels, segments, im2, labels2, segments2):
     return im, labels, segments
 
 
-def random_perspective(im,
+def random_perspective(
-                       targets=(),
+    im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.1, shear=10, perspective=0.0, border=(0, 0)
-                       segments=(),
+):
-                       degrees=10,
-                       translate=.1,
-                       scale=.1,
-                       shear=10,
-                       perspective=0.0,
-                       border=(0, 0)):
     # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
     # targets = [cls, xyxy]
 
@@ -62,8 +54,8 @@ def random_perspective(im,
 
     # Translation
     T = np.eye(3)
-    T[0, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * width)  # x translation (pixels)
+    T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width  # x translation (pixels)
-    T[1, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * height)  # y translation (pixels)
+    T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height  # y translation (pixels)
 
     # Combined rotation matrix
     M = T @ S @ R @ P @ C  # order of operations (right to left) is IMPORTANT
@@ -89,7 +81,7 @@ def random_perspective(im,
             xy = np.ones((len(segment), 3))
             xy[:, :2] = segment
             xy = xy @ M.T  # transform
-            xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2])  # perspective rescale or affine
+            xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]  # perspective rescale or affine
 
             # clip
             new[i] = segment2box(xy, width, height)

utils/segment/dataloaders.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Dataloaders."""
-Dataloaders
-"""
 
 import os
 import random
@@ -17,10 +15,11 @@ from ..general import LOGGER, xyn2xy, xywhn2xyxy, xyxy2xywhn
 from ..torch_utils import torch_distributed_zero_first
 from .augmentations import mixup, random_perspective
 
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
 
 
-def create_dataloader(path,
+def create_dataloader(
+    path,
     imgsz,
     batch_size,
     stride,
@@ -34,13 +33,14 @@ def create_dataloader(path,
     workers=8,
     image_weights=False,
     quad=False,
-                      prefix='',
+    prefix="",
     shuffle=False,
     mask_downsample_ratio=1,
     overlap_mask=False,
-                      seed=0):
+    seed=0,
+):
     if rect and shuffle:
-        LOGGER.warning('WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False')
+        LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False")
         shuffle = False
     with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
         dataset = LoadImagesAndLabelsAndMasks(
@@ -58,7 +58,8 @@ def create_dataloader(path,
             prefix=prefix,
             downsample_ratio=mask_downsample_ratio,
             overlap=overlap_mask,
-            rank=rank)
+            rank=rank,
+        )
 
     batch_size = min(batch_size, len(dataset))
     nd = torch.cuda.device_count()  # number of CUDA devices
@@ -81,7 +82,6 @@ def create_dataloader(path,
 
 
 class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
-
     def __init__(
         self,
         path,
@@ -96,14 +96,29 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
         stride=32,
         pad=0,
         min_items=0,
-        prefix='',
+        prefix="",
         downsample_ratio=1,
         overlap=False,
         rank=-1,
         seed=0,
     ):
-        super().__init__(path, img_size, batch_size, augment, hyp, rect, image_weights, cache_images, single_cls,
+        super().__init__(
-                         stride, pad, min_items, prefix, rank, seed)
+            path,
+            img_size,
+            batch_size,
+            augment,
+            hyp,
+            rect,
+            image_weights,
+            cache_images,
+            single_cls,
+            stride,
+            pad,
+            min_items,
+            prefix,
+            rank,
+            seed,
+        )
         self.downsample_ratio = downsample_ratio
         self.overlap = overlap
 
@@ -111,7 +126,7 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
         index = self.indices[index]  # linear, shuffled, or image_weights
 
         hyp = self.hyp
-        mosaic = self.mosaic and random.random() < hyp['mosaic']
+        mosaic = self.mosaic and random.random() < hyp["mosaic"]
         masks = []
         if mosaic:
             # Load mosaic
@@ -119,7 +134,7 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
             shapes = None
 
             # MixUp augmentation
-            if random.random() < hyp['mixup']:
+            if random.random() < hyp["mixup"]:
                 img, labels, segments = mixup(img, labels, segments, *self.load_mosaic(random.randint(0, self.n - 1)))
 
         else:
@@ -147,30 +162,36 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
                 labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
 
             if self.augment:
-                img, labels, segments = random_perspective(img,
+                img, labels, segments = random_perspective(
+                    img,
                     labels,
                     segments=segments,
-                                                           degrees=hyp['degrees'],
+                    degrees=hyp["degrees"],
-                                                           translate=hyp['translate'],
+                    translate=hyp["translate"],
-                                                           scale=hyp['scale'],
+                    scale=hyp["scale"],
-                                                           shear=hyp['shear'],
+                    shear=hyp["shear"],
-                                                           perspective=hyp['perspective'])
+                    perspective=hyp["perspective"],
+                )
 
         nl = len(labels)  # number of labels
         if nl:
             labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
             if self.overlap:
-                masks, sorted_idx = polygons2masks_overlap(img.shape[:2],
+                masks, sorted_idx = polygons2masks_overlap(
-                                                           segments,
+                    img.shape[:2], segments, downsample_ratio=self.downsample_ratio
-                                                           downsample_ratio=self.downsample_ratio)
+                )
                 masks = masks[None]  # (640, 640) -> (1, 640, 640)
                 labels = labels[sorted_idx]
             else:
                 masks = polygons2masks(img.shape[:2], segments, color=1, downsample_ratio=self.downsample_ratio)
 
-        masks = (torch.from_numpy(masks) if len(masks) else torch.zeros(1 if self.overlap else nl, img.shape[0] //
+        masks = (
-                                                                        self.downsample_ratio, img.shape[1] //
+            torch.from_numpy(masks)
-                                                                        self.downsample_ratio))
+            if len(masks)
+            else torch.zeros(
+                1 if self.overlap else nl, img.shape[0] // self.downsample_ratio, img.shape[1] // self.downsample_ratio
+            )
+        )
         # TODO: albumentations support
         if self.augment:
             # Albumentations
@@ -180,17 +201,17 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
             nl = len(labels)  # update after albumentations
 
             # HSV color-space
-            augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
+            augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
 
             # Flip up-down
-            if random.random() < hyp['flipud']:
+            if random.random() < hyp["flipud"]:
                 img = np.flipud(img)
                 if nl:
                     labels[:, 2] = 1 - labels[:, 2]
                     masks = torch.flip(masks, dims=[1])
 
             # Flip left-right
-            if random.random() < hyp['fliplr']:
+            if random.random() < hyp["fliplr"]:
                 img = np.fliplr(img)
                 if nl:
                     labels[:, 1] = 1 - labels[:, 1]
@@ -254,16 +275,18 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
         # img4, labels4 = replicate(img4, labels4)  # replicate
 
         # Augment
-        img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])
+        img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"])
-        img4, labels4, segments4 = random_perspective(img4,
+        img4, labels4, segments4 = random_perspective(
+            img4,
             labels4,
             segments4,
-                                                      degrees=self.hyp['degrees'],
+            degrees=self.hyp["degrees"],
-                                                      translate=self.hyp['translate'],
+            translate=self.hyp["translate"],
-                                                      scale=self.hyp['scale'],
+            scale=self.hyp["scale"],
-                                                      shear=self.hyp['shear'],
+            shear=self.hyp["shear"],
-                                                      perspective=self.hyp['perspective'],
+            perspective=self.hyp["perspective"],
-                                                      border=self.mosaic_border)  # border to remove
+            border=self.mosaic_border,
+        )  # border to remove
         return img4, labels4, segments4
 
     @staticmethod
@@ -312,8 +335,10 @@ def polygons2masks(img_size, polygons, color, downsample_ratio=1):
 
 def polygons2masks_overlap(img_size, segments, downsample_ratio=1):
     """Return a (640, 640) overlap mask."""
-    masks = np.zeros((img_size[0] // downsample_ratio, img_size[1] // downsample_ratio),
+    masks = np.zeros(
-                     dtype=np.int32 if len(segments) > 255 else np.uint8)
+        (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio),
+        dtype=np.int32 if len(segments) > 255 else np.uint8,
+    )
     areas = []
     ms = []
     for si in range(len(segments)):

utils/segment/general.py

@@ -6,8 +6,7 @@ import torch.nn.functional as F
 
 def crop_mask(masks, boxes):
     """
-    "Crop" predicted masks by zeroing out everything not in the predicted bbox.
+    "Crop" predicted masks by zeroing out everything not in the predicted bbox. Vectorized by Chong (thanks Chong).
-    Vectorized by Chong (thanks Chong).
 
     Args:
         - masks should be a size [n, h, w] tensor of masks
@@ -35,7 +34,7 @@ def process_mask_upsample(protos, masks_in, bboxes, shape):
 
     c, mh, mw = protos.shape  # CHW
     masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
-    masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]  # CHW
+    masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
     masks = crop_mask(masks, bboxes)  # CHW
     return masks.gt_(0.5)
 
@@ -63,7 +62,7 @@ def process_mask(protos, masks_in, bboxes, shape, upsample=False):
 
     masks = crop_mask(masks, downsampled_bboxes)  # CHW
     if upsample:
-        masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]  # CHW
+        masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
     return masks.gt_(0.5)
 
 
@@ -85,7 +84,7 @@ def process_mask_native(protos, masks_in, bboxes, shape):
     bottom, right = int(mh - pad[1]), int(mw - pad[0])
     masks = masks[:, top:bottom, left:right]
 
-    masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]  # CHW
+    masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
     masks = crop_mask(masks, bboxes)  # CHW
     return masks.gt_(0.5)
 
@@ -144,17 +143,17 @@ def masks_iou(mask1, mask2, eps=1e-7):
     return intersection / (union + eps)
 
 
-def masks2segments(masks, strategy='largest'):
+def masks2segments(masks, strategy="largest"):
     # Convert masks(n,160,160) into segments(n,xy)
     segments = []
-    for x in masks.int().cpu().numpy().astype('uint8'):
+    for x in masks.int().cpu().numpy().astype("uint8"):
         c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
         if c:
-            if strategy == 'concat':  # concatenate all segments
+            if strategy == "concat":  # concatenate all segments
                 c = np.concatenate([x.reshape(-1, 2) for x in c])
-            elif strategy == 'largest':  # select largest segment
+            elif strategy == "largest":  # select largest segment
                 c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
         else:
             c = np.zeros((0, 2))  # no segments found
-        segments.append(c.astype('float32'))
+        segments.append(c.astype("float32"))
     return segments

utils/segment/loss.py

@@ -18,14 +18,14 @@ class ComputeLoss:
         h = model.hyp  # hyperparameters
 
         # Define criteria
-        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
+        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device))
-        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
+        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device))
 
         # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
-        self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0))  # positive, negative BCE targets
+        self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0))  # positive, negative BCE targets
 
         # Focal loss
-        g = h['fl_gamma']  # focal loss gamma
+        g = h["fl_gamma"]  # focal loss gamma
         if g > 0:
             BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
 
@@ -82,7 +82,7 @@ class ComputeLoss:
 
                 # Mask regression
                 if tuple(masks.shape[-2:]) != (mask_h, mask_w):  # downsample
-                    masks = F.interpolate(masks[None], (mask_h, mask_w), mode='nearest')[0]
+                    masks = F.interpolate(masks[None], (mask_h, mask_w), mode="nearest")[0]
                 marea = xywhn[i][:, 2:].prod(1)  # mask width, height normalized
                 mxyxy = xywh2xyxy(xywhn[i] * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device))
                 for bi in b.unique():
@@ -100,10 +100,10 @@ class ComputeLoss:
 
         if self.autobalance:
             self.balance = [x / self.balance[self.ssi] for x in self.balance]
-        lbox *= self.hyp['box']
+        lbox *= self.hyp["box"]
-        lobj *= self.hyp['obj']
+        lobj *= self.hyp["obj"]
-        lcls *= self.hyp['cls']
+        lcls *= self.hyp["cls"]
-        lseg *= self.hyp['box'] / bs
+        lseg *= self.hyp["box"] / bs
 
         loss = lbox + lobj + lcls + lseg
         return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach()
@@ -111,7 +111,7 @@ class ComputeLoss:
     def single_mask_loss(self, gt_mask, pred, proto, xyxy, area):
         # Mask loss for one image
         pred_mask = (pred @ proto.view(self.nm, -1)).view(-1, *proto.shape[1:])  # (n,32) @ (32,80,80) -> (n,80,80)
-        loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction='none')
+        loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
         return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).mean()
 
     def build_targets(self, p, targets):
@@ -132,7 +132,8 @@ class ComputeLoss:
         targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None], ti[..., None]), 2)  # append anchor indices
 
         g = 0.5  # bias
-        off = torch.tensor(
+        off = (
+            torch.tensor(
                 [
                     [0, 0],
                     [1, 0],
@@ -141,7 +142,10 @@ class ComputeLoss:
                     [0, -1],  # j,k,l,m
                     # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                 ],
-            device=self.device).float() * g  # offsets
+                device=self.device,
+            ).float()
+            * g
+        )  # offsets
 
         for i in range(self.nl):
             anchors, shape = self.anchors[i], p[i].shape
@@ -152,7 +156,7 @@ class ComputeLoss:
             if nt:
                 # Matches
                 r = t[..., 4:6] / anchors[:, None]  # wh ratio
-                j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare
+                j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"]  # compare
                 # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                 t = t[j]  # filter
 

utils/segment/metrics.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""Model validation metrics."""
-Model validation metrics
-"""
 
 import numpy as np
 
@@ -21,7 +19,7 @@ def ap_per_class_box_and_mask(
     pred_cls,
     target_cls,
     plot=False,
-        save_dir='.',
+    save_dir=".",
     names=(),
 ):
     """
@@ -30,41 +28,33 @@ def ap_per_class_box_and_mask(
         tp_m: tp of masks.
         other arguments see `func: ap_per_class`.
     """
-    results_boxes = ap_per_class(tp_b,
+    results_boxes = ap_per_class(
-                                 conf,
+        tp_b, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Box"
-                                 pred_cls,
+    )[2:]
-                                 target_cls,
+    results_masks = ap_per_class(
-                                 plot=plot,
+        tp_m, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Mask"
-                                 save_dir=save_dir,
+    )[2:]
-                                 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': {
+        "boxes": {
-            'p': results_boxes[0],
+            "p": results_boxes[0],
-            'r': results_boxes[1],
+            "r": results_boxes[1],
-            'ap': results_boxes[3],
+            "ap": results_boxes[3],
-            'f1': results_boxes[2],
+            "f1": results_boxes[2],
-            'ap_class': results_boxes[4]},
+            "ap_class": results_boxes[4],
-        'masks': {
+        },
-            'p': results_masks[0],
+        "masks": {
-            'r': results_masks[1],
+            "p": results_masks[0],
-            'ap': results_masks[3],
+            "r": results_masks[1],
-            'f1': results_masks[2],
+            "ap": results_masks[3],
-            'ap_class': results_masks[4]}}
+            "f1": results_masks[2],
+            "ap_class": results_masks[4],
+        },
+    }
     return results
 
 
 class Metric:
-
     def __init__(self) -> None:
         self.p = []  # (nc, )
         self.r = []  # (nc, )
@@ -74,7 +64,9 @@ class Metric:
 
     @property
     def ap50(self):
-        """AP@0.5 of all classes.
+        """
+        AP@0.5 of all classes.
+
         Return:
             (nc, ) or [].
         """
@@ -90,7 +82,9 @@ class Metric:
 
     @property
     def mp(self):
-        """mean precision of all classes.
+        """
+        Mean precision of all classes.
+
         Return:
             float.
         """
@@ -98,7 +92,9 @@ class Metric:
 
     @property
     def mr(self):
-        """mean recall of all classes.
+        """
+        Mean recall of all classes.
+
         Return:
             float.
         """
@@ -106,7 +102,9 @@ class Metric:
 
     @property
     def map50(self):
-        """Mean AP@0.5 of all classes.
+        """
+        Mean AP@0.5 of all classes.
+
         Return:
             float.
         """
@@ -114,18 +112,20 @@ class Metric:
 
     @property
     def map(self):
-        """Mean AP@0.5:0.95 of all classes.
+        """
+        Mean AP@0.5:0.95 of all classes.
+
         Return:
             float.
         """
         return self.all_ap.mean() if len(self.all_ap) else 0.0
 
     def mean_results(self):
-        """Mean of results, return mp, mr, map50, map"""
+        """Mean of results, return mp, mr, map50, map."""
         return (self.mp, self.mr, self.map50, self.map)
 
     def class_result(self, i):
-        """class-aware result, return p[i], r[i], ap50[i], ap[i]"""
+        """Class-aware result, return p[i], r[i], ap50[i], ap[i]"""
         return (self.p[i], self.r[i], self.ap50[i], self.ap[i])
 
     def get_maps(self, nc):
@@ -159,8 +159,8 @@ class Metrics:
         Args:
             results: Dict{'boxes': Dict{}, 'masks': Dict{}}
         """
-        self.metric_box.update(list(results['boxes'].values()))
+        self.metric_box.update(list(results["boxes"].values()))
-        self.metric_mask.update(list(results['masks'].values()))
+        self.metric_mask.update(list(results["masks"].values()))
 
     def mean_results(self):
         return self.metric_box.mean_results() + self.metric_mask.mean_results()
@@ -178,33 +178,35 @@ class Metrics:
 
 
 KEYS = [
-    'train/box_loss',
+    "train/box_loss",
-    'train/seg_loss',  # train loss
+    "train/seg_loss",  # train loss
-    'train/obj_loss',
+    "train/obj_loss",
-    'train/cls_loss',
+    "train/cls_loss",
-    'metrics/precision(B)',
+    "metrics/precision(B)",
-    'metrics/recall(B)',
+    "metrics/recall(B)",
-    'metrics/mAP_0.5(B)',
+    "metrics/mAP_0.5(B)",
-    'metrics/mAP_0.5:0.95(B)',  # metrics
+    "metrics/mAP_0.5:0.95(B)",  # metrics
-    'metrics/precision(M)',
+    "metrics/precision(M)",
-    'metrics/recall(M)',
+    "metrics/recall(M)",
-    'metrics/mAP_0.5(M)',
+    "metrics/mAP_0.5(M)",
-    'metrics/mAP_0.5:0.95(M)',  # metrics
+    "metrics/mAP_0.5:0.95(M)",  # metrics
-    'val/box_loss',
+    "val/box_loss",
-    'val/seg_loss',  # val loss
+    "val/seg_loss",  # val loss
-    'val/obj_loss',
+    "val/obj_loss",
-    'val/cls_loss',
+    "val/cls_loss",
-    'x/lr0',
+    "x/lr0",
-    'x/lr1',
+    "x/lr1",
-    'x/lr2', ]
+    "x/lr2",
+]
 
 BEST_KEYS = [
-    'best/epoch',
+    "best/epoch",
-    'best/precision(B)',
+    "best/precision(B)",
-    'best/recall(B)',
+    "best/recall(B)",
-    'best/mAP_0.5(B)',
+    "best/mAP_0.5(B)",
-    'best/mAP_0.5:0.95(B)',
+    "best/mAP_0.5:0.95(B)",
-    'best/precision(M)',
+    "best/precision(M)",
-    'best/recall(M)',
+    "best/recall(M)",
-    'best/mAP_0.5(M)',
+    "best/mAP_0.5(M)",
-    'best/mAP_0.5:0.95(M)', ]
+    "best/mAP_0.5:0.95(M)",
+]

utils/segment/plots.py

@@ -14,7 +14,7 @@ from ..plots import Annotator, colors
 
 
 @threaded
-def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg', names=None):
+def plot_images_and_masks(images, targets, masks, paths=None, fname="images.jpg", names=None):
     # Plot image grid with labels
     if isinstance(images, torch.Tensor):
         images = images.cpu().float().numpy()
@@ -60,7 +60,7 @@ def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg'
             ti = targets[idx]  # image targets
 
             boxes = xywh2xyxy(ti[:, 2:6]).T
-            classes = ti[:, 1].astype('int')
+            classes = ti[:, 1].astype("int")
             labels = ti.shape[1] == 6  # labels if no conf column
             conf = None if labels else ti[:, 6]  # check for confidence presence (label vs pred)
 
@@ -77,7 +77,7 @@ def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg'
                 color = colors(cls)
                 cls = names[cls] if names else cls
                 if labels or conf[j] > 0.25:  # 0.25 conf thresh
-                    label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}'
+                    label = f"{cls}" if labels else f"{cls} {conf[j]:.1f}"
                     annotator.box_label(box, label, color=color)
 
             # Plot masks
@@ -103,41 +103,44 @@ def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg'
                         else:
                             mask = image_masks[j].astype(bool)
                         with contextlib.suppress(Exception):
-                            im[y:y + h, x:x + w, :][mask] = im[y:y + h, x:x + w, :][mask] * 0.4 + np.array(color) * 0.6
+                            im[y : y + h, x : x + w, :][mask] = (
+                                im[y : y + h, x : x + w, :][mask] * 0.4 + np.array(color) * 0.6
+                            )
                 annotator.fromarray(im)
     annotator.im.save(fname)  # save
 
 
-def plot_results_with_masks(file='path/to/results.csv', dir='', best=True):
+def plot_results_with_masks(file="path/to/results.csv", dir="", best=True):
     # 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'))
+    files = list(save_dir.glob("results*.csv"))
-    assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.'
+    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] +
+            index = np.argmax(
-                              0.1 * data.values[:, 11])
+                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)
+                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].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)}')
+                    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].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].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}')
+            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()

utils/torch_utils.py

@@ -1,7 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-"""
+"""PyTorch utils."""
-PyTorch utils
-"""
 
 import math
 import os
@@ -21,9 +19,9 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 
 from utils.general import LOGGER, check_version, colorstr, file_date, git_describe
 
-LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
+LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
-RANK = int(os.getenv('RANK', -1))
+RANK = int(os.getenv("RANK", -1))
-WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
+WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
 
 try:
     import thop  # for FLOPs computation
@@ -31,11 +29,11 @@ except ImportError:
     thop = None
 
 # Suppress PyTorch warnings
-warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')
+warnings.filterwarnings("ignore", message="User provided device_type of 'cuda', but CUDA is not available. Disabling")
-warnings.filterwarnings('ignore', category=UserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
 
 
-def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
+def smart_inference_mode(torch_1_9=check_version(torch.__version__, "1.9.0")):
     # Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator
     def decorate(fn):
         return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
@@ -45,19 +43,20 @@ def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
 
 def smartCrossEntropyLoss(label_smoothing=0.0):
     # Returns nn.CrossEntropyLoss with label smoothing enabled for torch>=1.10.0
-    if check_version(torch.__version__, '1.10.0'):
+    if check_version(torch.__version__, "1.10.0"):
         return nn.CrossEntropyLoss(label_smoothing=label_smoothing)
     if label_smoothing > 0:
-        LOGGER.warning(f'WARNING ⚠️ label smoothing {label_smoothing} requires torch>=1.10.0')
+        LOGGER.warning(f"WARNING ⚠️ label smoothing {label_smoothing} requires torch>=1.10.0")
     return nn.CrossEntropyLoss()
 
 
 def smart_DDP(model):
     # Model DDP creation with checks
-    assert not check_version(torch.__version__, '1.12.0', pinned=True), \
+    assert not check_version(torch.__version__, "1.12.0", pinned=True), (
-        'torch==1.12.0 torchvision==0.13.0 DDP training is not supported due to a known issue. ' \
+        "torch==1.12.0 torchvision==0.13.0 DDP training is not supported due to a known issue. "
-        'Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395'
+        "Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395"
-    if check_version(torch.__version__, '1.11.0'):
+    )
+    if check_version(torch.__version__, "1.11.0"):
         return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK, static_graph=True)
     else:
         return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
@@ -66,7 +65,8 @@ def smart_DDP(model):
 def reshape_classifier_output(model, n=1000):
     # Update a TorchVision classification model to class count 'n' if required
     from models.common import Classify
-    name, m = list((model.model if hasattr(model, 'model') else model).named_children())[-1]  # last module
+
+    name, m = list((model.model if hasattr(model, "model") else model).named_children())[-1]  # last module
     if isinstance(m, Classify):  # YOLOv5 Classify() head
         if m.linear.out_features != n:
             m.linear = nn.Linear(m.linear.in_features, n)
@@ -97,43 +97,44 @@ def torch_distributed_zero_first(local_rank: int):
 
 def device_count():
     # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Supports Linux and Windows
-    assert platform.system() in ('Linux', 'Windows'), 'device_count() only supported on Linux or Windows'
+    assert platform.system() in ("Linux", "Windows"), "device_count() only supported on Linux or Windows"
     try:
-        cmd = 'nvidia-smi -L | wc -l' if platform.system() == 'Linux' else 'nvidia-smi -L | find /c /v ""'  # Windows
+        cmd = "nvidia-smi -L | wc -l" if platform.system() == "Linux" else 'nvidia-smi -L | find /c /v ""'  # Windows
         return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])
     except Exception:
         return 0
 
 
-def select_device(device='', batch_size=0, newline=True):
+def select_device(device="", batch_size=0, newline=True):
     # device = None or 'cpu' or 0 or '0' or '0,1,2,3'
-    s = f'YOLOv5 🚀 {git_describe() or file_date()} Python-{platform.python_version()} torch-{torch.__version__} '
+    s = f"YOLOv5 🚀 {git_describe() or file_date()} Python-{platform.python_version()} torch-{torch.__version__} "
-    device = str(device).strip().lower().replace('cuda:', '').replace('none', '')  # to string, 'cuda:0' to '0'
+    device = str(device).strip().lower().replace("cuda:", "").replace("none", "")  # to string, 'cuda:0' to '0'
-    cpu = device == 'cpu'
+    cpu = device == "cpu"
-    mps = device == 'mps'  # Apple Metal Performance Shaders (MPS)
+    mps = device == "mps"  # Apple Metal Performance Shaders (MPS)
     if cpu or mps:
-        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # force torch.cuda.is_available() = False
     elif device:  # non-cpu device requested
-        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable - must be before assert is_available()
+        os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable - must be before assert is_available()
-        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
+        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(
-            f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
+            device.replace(",", "")
+        ), f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
 
     if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available
-        devices = device.split(',') if device else '0'  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
+        devices = device.split(",") if device else "0"  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
         n = len(devices)  # device count
         if n > 1 and batch_size > 0:  # check batch_size is divisible by device_count
-            assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
+            assert batch_size % n == 0, f"batch-size {batch_size} not multiple of GPU count {n}"
-        space = ' ' * (len(s) + 1)
+        space = " " * (len(s) + 1)
         for i, d in enumerate(devices):
             p = torch.cuda.get_device_properties(i)
             s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n"  # bytes to MB
-        arg = 'cuda:0'
+        arg = "cuda:0"
-    elif mps and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available():  # prefer MPS if available
+    elif mps and getattr(torch, "has_mps", False) and torch.backends.mps.is_available():  # prefer MPS if available
-        s += 'MPS\n'
+        s += "MPS\n"
-        arg = 'mps'
+        arg = "mps"
     else:  # revert to CPU
-        s += 'CPU\n'
+        s += "CPU\n"
-        arg = 'cpu'
+        arg = "cpu"
 
     if not newline:
         s = s.rstrip()
@@ -159,18 +160,20 @@ def profile(input, ops, n=10, device=None):
     results = []
     if not isinstance(device, torch.device):
         device = select_device(device)
-    print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
+    print(
-          f"{'input':>24s}{'output':>24s}")
+        f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
+        f"{'input':>24s}{'output':>24s}"
+    )
 
     for x in input if isinstance(input, list) else [input]:
         x = x.to(device)
         x.requires_grad = True
         for m in ops if isinstance(ops, list) else [ops]:
-            m = m.to(device) if hasattr(m, 'to') else m  # device
+            m = m.to(device) if hasattr(m, "to") else m  # device
-            m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
+            m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
             tf, tb, t = 0, 0, [0, 0, 0]  # dt forward, backward
             try:
-                flops = thop.profile(m, inputs=(x, ), verbose=False)[0] / 1E9 * 2  # GFLOPs
+                flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1e9 * 2  # GFLOPs
             except Exception:
                 flops = 0
 
@@ -184,13 +187,13 @@ def profile(input, ops, n=10, device=None):
                         t[2] = time_sync()
                     except Exception:  # no backward method
                         # print(e)  # for debug
-                        t[2] = float('nan')
+                        t[2] = float("nan")
                     tf += (t[1] - t[0]) * 1000 / n  # ms per op forward
                     tb += (t[2] - t[1]) * 1000 / n  # ms per op backward
-                mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0  # (GB)
+                mem = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0  # (GB)
-                s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' for x in (x, y))  # shapes
+                s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else "list" for x in (x, y))  # shapes
                 p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0  # parameters
-                print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}')
+                print(f"{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}")
                 results.append([p, flops, mem, tf, tb, s_in, s_out])
             except Exception as e:
                 print(e)
@@ -238,23 +241,30 @@ def sparsity(model):
 def prune(model, amount=0.3):
     # Prune model to requested global sparsity
     import torch.nn.utils.prune as prune
+
     for name, m in model.named_modules():
         if isinstance(m, nn.Conv2d):
-            prune.l1_unstructured(m, name='weight', amount=amount)  # prune
+            prune.l1_unstructured(m, name="weight", amount=amount)  # prune
-            prune.remove(m, 'weight')  # make permanent
+            prune.remove(m, "weight")  # make permanent
-    LOGGER.info(f'Model pruned to {sparsity(model):.3g} global sparsity')
+    LOGGER.info(f"Model pruned to {sparsity(model):.3g} global sparsity")
 
 
 def fuse_conv_and_bn(conv, bn):
     # Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
-    fusedconv = nn.Conv2d(conv.in_channels,
+    fusedconv = (
+        nn.Conv2d(
+            conv.in_channels,
             conv.out_channels,
             kernel_size=conv.kernel_size,
             stride=conv.stride,
             padding=conv.padding,
             dilation=conv.dilation,
             groups=conv.groups,
-                          bias=True).requires_grad_(False).to(conv.weight.device)
+            bias=True,
+        )
+        .requires_grad_(False)
+        .to(conv.weight.device)
+    )
 
     # Prepare filters
     w_conv = conv.weight.clone().view(conv.out_channels, -1)
@@ -276,22 +286,24 @@ def model_info(model, verbose=False, imgsz=640):
     if verbose:
         print(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}")
         for i, (name, p) in enumerate(model.named_parameters()):
-            name = name.replace('module_list.', '')
+            name = name.replace("module_list.", "")
-            print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
+            print(
-                  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
+                "%5g %40s %9s %12g %20s %10.3g %10.3g"
+                % (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())
+            )
 
     try:  # FLOPs
         p = next(model.parameters())
-        stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32  # max stride
+        stride = max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32  # max stride
         im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
-        flops = thop.profile(deepcopy(model), inputs=(im, ), verbose=False)[0] / 1E9 * 2  # stride GFLOPs
+        flops = thop.profile(deepcopy(model), inputs=(im,), verbose=False)[0] / 1e9 * 2  # stride GFLOPs
         imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz]  # expand if int/float
-        fs = f', {flops * imgsz[0] / stride * imgsz[1] / stride:.1f} GFLOPs'  # 640x640 GFLOPs
+        fs = f", {flops * imgsz[0] / stride * imgsz[1] / stride:.1f} GFLOPs"  # 640x640 GFLOPs
     except Exception:
-        fs = ''
+        fs = ""
 
-    name = Path(model.yaml_file).stem.replace('yolov5', 'YOLOv5') if hasattr(model, 'yaml_file') else 'Model'
+    name = Path(model.yaml_file).stem.replace("yolov5", "YOLOv5") if hasattr(model, "yaml_file") else "Model"
-    LOGGER.info(f'{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}')
+    LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
 
 
 def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
@@ -300,7 +312,7 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
         return img
     h, w = img.shape[2:]
     s = (int(h * ratio), int(w * ratio))  # new size
-    img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
+    img = F.interpolate(img, size=s, mode="bilinear", align_corners=False)  # resize
     if not same_shape:  # pad/crop img
         h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
     return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
@@ -309,72 +321,76 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
 def copy_attr(a, b, include=(), exclude=()):
     # Copy attributes from b to a, options to only include [...] and to exclude [...]
     for k, v in b.__dict__.items():
-        if (len(include) and k not in include) or k.startswith('_') or k in exclude:
+        if (len(include) and k not in include) or k.startswith("_") or k in exclude:
             continue
         else:
             setattr(a, k, v)
 
 
-def smart_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
+def smart_optimizer(model, name="Adam", lr=0.001, momentum=0.9, decay=1e-5):
     # YOLOv5 3-param group optimizer: 0) weights with decay, 1) weights no decay, 2) biases no decay
     g = [], [], []  # optimizer parameter groups
-    bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k)  # normalization layers, i.e. BatchNorm2d()
+    bn = tuple(v for k, v in nn.__dict__.items() if "Norm" in k)  # normalization layers, i.e. BatchNorm2d()
     for v in model.modules():
         for p_name, p in v.named_parameters(recurse=0):
-            if p_name == 'bias':  # bias (no decay)
+            if p_name == "bias":  # bias (no decay)
                 g[2].append(p)
-            elif p_name == 'weight' and isinstance(v, bn):  # weight (no decay)
+            elif p_name == "weight" and isinstance(v, bn):  # weight (no decay)
                 g[1].append(p)
             else:
                 g[0].append(p)  # weight (with decay)
 
-    if name == 'Adam':
+    if name == "Adam":
         optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999))  # adjust beta1 to momentum
-    elif name == 'AdamW':
+    elif name == "AdamW":
         optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
-    elif name == 'RMSProp':
+    elif name == "RMSProp":
         optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum)
-    elif name == 'SGD':
+    elif name == "SGD":
         optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
     else:
-        raise NotImplementedError(f'Optimizer {name} not implemented.')
+        raise NotImplementedError(f"Optimizer {name} not implemented.")
 
-    optimizer.add_param_group({'params': g[0], 'weight_decay': decay})  # add g0 with weight_decay
+    optimizer.add_param_group({"params": g[0], "weight_decay": decay})  # add g0 with weight_decay
-    optimizer.add_param_group({'params': g[1], 'weight_decay': 0.0})  # add g1 (BatchNorm2d weights)
+    optimizer.add_param_group({"params": g[1], "weight_decay": 0.0})  # add g1 (BatchNorm2d weights)
-    LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "
+    LOGGER.info(
-                f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias')
+        f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "
+        f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias'
+    )
     return optimizer
 
 
-def smart_hub_load(repo='ultralytics/yolov5', model='yolov5s', **kwargs):
+def smart_hub_load(repo="ultralytics/yolov5", model="yolov5s", **kwargs):
     # YOLOv5 torch.hub.load() wrapper with smart error/issue handling
-    if check_version(torch.__version__, '1.9.1'):
+    if check_version(torch.__version__, "1.9.1"):
-        kwargs['skip_validation'] = True  # validation causes GitHub API rate limit errors
+        kwargs["skip_validation"] = True  # validation causes GitHub API rate limit errors
-    if check_version(torch.__version__, '1.12.0'):
+    if check_version(torch.__version__, "1.12.0"):
-        kwargs['trust_repo'] = True  # argument required starting in torch 0.12
+        kwargs["trust_repo"] = True  # argument required starting in torch 0.12
     try:
         return torch.hub.load(repo, model, **kwargs)
     except Exception:
         return torch.hub.load(repo, model, force_reload=True, **kwargs)
 
 
-def smart_resume(ckpt, optimizer, ema=None, weights='yolov5s.pt', epochs=300, resume=True):
+def smart_resume(ckpt, optimizer, ema=None, weights="yolov5s.pt", epochs=300, resume=True):
     # Resume training from a partially trained checkpoint
     best_fitness = 0.0
-    start_epoch = ckpt['epoch'] + 1
+    start_epoch = ckpt["epoch"] + 1
-    if ckpt['optimizer'] is not None:
+    if ckpt["optimizer"] is not None:
-        optimizer.load_state_dict(ckpt['optimizer'])  # optimizer
+        optimizer.load_state_dict(ckpt["optimizer"])  # optimizer
-        best_fitness = ckpt['best_fitness']
+        best_fitness = ckpt["best_fitness"]
-    if ema and ckpt.get('ema'):
+    if ema and ckpt.get("ema"):
-        ema.ema.load_state_dict(ckpt['ema'].float().state_dict())  # EMA
+        ema.ema.load_state_dict(ckpt["ema"].float().state_dict())  # EMA
-        ema.updates = ckpt['updates']
+        ema.updates = ckpt["updates"]
     if resume:
-        assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.\n' \
+        assert start_epoch > 0, (
+            f"{weights} training to {epochs} epochs is finished, nothing to resume.\n"
             f"Start a new training without --resume, i.e. 'python train.py --weights {weights}'"
-        LOGGER.info(f'Resuming training from {weights} from epoch {start_epoch} to {epochs} total epochs')
+        )
+        LOGGER.info(f"Resuming training from {weights} from epoch {start_epoch} to {epochs} total epochs")
     if epochs < start_epoch:
         LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")
-        epochs += ckpt['epoch']  # finetune additional epochs
+        epochs += ckpt["epoch"]  # finetune additional epochs
     return best_fitness, start_epoch, epochs
 
 
@@ -383,7 +399,7 @@ class EarlyStopping:
     def __init__(self, patience=30):
         self.best_fitness = 0.0  # i.e. mAP
         self.best_epoch = 0
-        self.patience = patience or float('inf')  # epochs to wait after fitness stops improving to stop
+        self.patience = patience or float("inf")  # epochs to wait after fitness stops improving to stop
         self.possible_stop = False  # possible stop may occur next epoch
 
     def __call__(self, epoch, fitness):
@@ -394,10 +410,12 @@ class EarlyStopping:
         self.possible_stop = delta >= (self.patience - 1)  # possible stop may occur next epoch
         stop = delta >= self.patience  # stop training if patience exceeded
         if stop:
-            LOGGER.info(f'Stopping training early as no improvement observed in last {self.patience} epochs. '
+            LOGGER.info(
-                        f'Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n'
+                f"Stopping training early as no improvement observed in last {self.patience} epochs. "
-                        f'To update EarlyStopping(patience={self.patience}) pass a new patience value, '
+                f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n"
-                        f'i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping.')
+                f"To update EarlyStopping(patience={self.patience}) pass a new patience value, "
+                f"i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping."
+            )
         return stop
 
 
@@ -427,6 +445,6 @@ class ModelEMA:
                 v += (1 - d) * msd[k].detach()
         # assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype} and model {msd[k].dtype} must be FP32'
 
-    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
+    def update_attr(self, model, include=(), exclude=("process_group", "reducer")):
         # Update EMA attributes
         copy_attr(self.ema, model, include, exclude)

utils/triton.py

@@ -1,6 +1,5 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
-""" Utils to interact with the Triton Inference Server
+"""Utils to interact with the Triton Inference Server."""
-"""
 
 import typing
 from urllib.parse import urlparse
@@ -9,9 +8,11 @@ import torch
 
 
 class TritonRemoteModel:
-    """ A wrapper over a model served by the Triton Inference Server. It can
+    """
-    be configured to communicate over GRPC or HTTP. It accepts Torch Tensors
+    A wrapper over a model served by the Triton Inference Server.
-    as input and returns them as outputs.
+
+    It can be configured to communicate over GRPC or HTTP. It accepts Torch Tensors as input and returns them as
+    outputs.
     """
 
     def __init__(self, url: str):
@@ -21,7 +22,7 @@ class TritonRemoteModel:
         """
 
         parsed_url = urlparse(url)
-        if parsed_url.scheme == 'grpc':
+        if parsed_url.scheme == "grpc":
             from tritonclient.grpc import InferenceServerClient, InferInput
 
             self.client = InferenceServerClient(parsed_url.netloc)  # Triton GRPC client
@@ -31,51 +32,55 @@ class TritonRemoteModel:
 
             def create_input_placeholders() -> typing.List[InferInput]:
                 return [
-                    InferInput(i['name'], [int(s) for s in i['shape']], i['datatype']) for i in self.metadata['inputs']]
+                    InferInput(i["name"], [int(s) for s in i["shape"]], i["datatype"]) for i in self.metadata["inputs"]
+                ]
 
         else:
             from tritonclient.http import InferenceServerClient, InferInput
 
             self.client = InferenceServerClient(parsed_url.netloc)  # Triton HTTP client
             model_repository = self.client.get_model_repository_index()
-            self.model_name = model_repository[0]['name']
+            self.model_name = model_repository[0]["name"]
             self.metadata = self.client.get_model_metadata(self.model_name)
 
             def create_input_placeholders() -> typing.List[InferInput]:
                 return [
-                    InferInput(i['name'], [int(s) for s in i['shape']], i['datatype']) for i in self.metadata['inputs']]
+                    InferInput(i["name"], [int(s) for s in i["shape"]], i["datatype"]) for i in self.metadata["inputs"]
+                ]
 
         self._create_input_placeholders_fn = create_input_placeholders
 
     @property
     def runtime(self):
-        """Returns the model runtime"""
+        """Returns the model runtime."""
-        return self.metadata.get('backend', self.metadata.get('platform'))
+        return self.metadata.get("backend", self.metadata.get("platform"))
 
     def __call__(self, *args, **kwargs) -> typing.Union[torch.Tensor, typing.Tuple[torch.Tensor, ...]]:
-        """ Invokes the model. Parameters can be provided via args or kwargs.
+        """
-        args, if provided, are assumed to match the order of inputs of the model.
+        Invokes the model.
-        kwargs are matched with the model input names.
+
+        Parameters can be provided via args or kwargs. args, if provided, are assumed to match the order of inputs of
+        the model. kwargs are matched with the model input names.
         """
         inputs = self._create_inputs(*args, **kwargs)
         response = self.client.infer(model_name=self.model_name, inputs=inputs)
         result = []
-        for output in self.metadata['outputs']:
+        for output in self.metadata["outputs"]:
-            tensor = torch.as_tensor(response.as_numpy(output['name']))
+            tensor = torch.as_tensor(response.as_numpy(output["name"]))
             result.append(tensor)
         return result[0] if len(result) == 1 else result
 
     def _create_inputs(self, *args, **kwargs):
         args_len, kwargs_len = len(args), len(kwargs)
         if not args_len and not kwargs_len:
-            raise RuntimeError('No inputs provided.')
+            raise RuntimeError("No inputs provided.")
         if args_len and kwargs_len:
-            raise RuntimeError('Cannot specify args and kwargs at the same time')
+            raise RuntimeError("Cannot specify args and kwargs at the same time")
 
         placeholders = self._create_input_placeholders_fn()
         if args_len:
             if args_len != len(placeholders):
-                raise RuntimeError(f'Expected {len(placeholders)} inputs, got {args_len}.')
+                raise RuntimeError(f"Expected {len(placeholders)} inputs, got {args_len}.")
             for input, value in zip(placeholders, args):
                 input.set_data_from_numpy(value.cpu().numpy())
         else:

val.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 """
-Validate a trained YOLOv5 detection model on a detection dataset
+Validate a trained YOLOv5 detection model on a detection dataset.
 
 Usage:
     $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
@@ -39,9 +39,23 @@ ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 from models.common import DetectMultiBackend
 from utils.callbacks import Callbacks
 from utils.dataloaders import create_dataloader
-from utils.general import (LOGGER, TQDM_BAR_FORMAT, Profile, check_dataset, check_img_size, check_requirements,
+from utils.general import (
-                           check_yaml, coco80_to_coco91_class, colorstr, increment_path, non_max_suppression,
+    LOGGER,
-                           print_args, scale_boxes, xywh2xyxy, xyxy2xywh)
+    TQDM_BAR_FORMAT,
+    Profile,
+    check_dataset,
+    check_img_size,
+    check_requirements,
+    check_yaml,
+    coco80_to_coco91_class,
+    colorstr,
+    increment_path,
+    non_max_suppression,
+    print_args,
+    scale_boxes,
+    xywh2xyxy,
+    xyxy2xywh,
+)
 from utils.metrics import ConfusionMatrix, ap_per_class, box_iou
 from utils.plots import output_to_target, plot_images, plot_val_study
 from utils.torch_utils import select_device, smart_inference_mode
@@ -53,8 +67,8 @@ def save_one_txt(predn, save_conf, shape, file):
     for *xyxy, conf, cls in predn.tolist():
         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
         line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
-        with open(file, 'a') as f:
+        with open(file, "a") as f:
-            f.write(('%g ' * len(line)).rstrip() % line + '\n')
+            f.write(("%g " * len(line)).rstrip() % line + "\n")
 
 
 def save_one_json(predn, jdict, path, class_map):
@@ -63,11 +77,14 @@ def save_one_json(predn, jdict, path, class_map):
     box = xyxy2xywh(predn[:, :4])  # xywh
     box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
     for p, b in zip(predn.tolist(), box.tolist()):
-        jdict.append({
+        jdict.append(
-            'image_id': image_id,
+            {
-            'category_id': class_map[int(p[5])],
+                "image_id": image_id,
-            'bbox': [round(x, 3) for x in b],
+                "category_id": class_map[int(p[5])],
-            'score': round(p[4], 5)})
+                "bbox": [round(x, 3) for x in b],
+                "score": round(p[4], 5),
+            }
+        )
 
 
 def process_batch(detections, labels, iouv):
@@ -105,8 +122,8 @@ def run(
     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
+    task="val",  # train, val, test, speed or study
-        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
     workers=8,  # max dataloader workers (per RANK in DDP mode)
     single_cls=False,  # treat as single-class dataset
     augment=False,  # augmented inference
@@ -115,14 +132,14 @@ def run(
     save_hybrid=False,  # save label+prediction hybrid results to *.txt
     save_conf=False,  # save confidences in --save-txt labels
     save_json=False,  # save a COCO-JSON results file
-        project=ROOT / 'runs/val',  # save to project/name
+    project=ROOT / "runs/val",  # save to project/name
-        name='exp',  # save to project/name
+    name="exp",  # save to project/name
     exist_ok=False,  # existing project/name ok, do not increment
     half=True,  # use FP16 half-precision inference
     dnn=False,  # use OpenCV DNN for ONNX inference
     model=None,
     dataloader=None,
-        save_dir=Path(''),
+    save_dir=Path(""),
     plots=True,
     callbacks=Callbacks(),
     compute_loss=None,
@@ -131,14 +148,14 @@ def run(
     training = model is not None
     if training:  # called by train.py
         device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
-        half &= device.type != 'cpu'  # half precision only supported on CUDA
+        half &= device.type != "cpu"  # half precision only supported on CUDA
         model.half() if half else model.float()
     else:  # called directly
         device = select_device(device, batch_size=batch_size)
 
         # Directories
         save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+        (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
         # Load model
         model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
@@ -151,16 +168,16 @@ def run(
             device = model.device
             if not (pt or jit):
                 batch_size = 1  # export.py models default to batch-size 1
-                LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
+                LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
 
         # Data
         data = check_dataset(data)  # check
 
     # Configure
     model.eval()
-    cuda = device.type != 'cpu'
+    cuda = device.type != "cpu"
-    is_coco = isinstance(data.get('val'), str) and data['val'].endswith(f'coco{os.sep}val2017.txt')  # COCO dataset
+    is_coco = isinstance(data.get("val"), str) and data["val"].endswith(f"coco{os.sep}val2017.txt")  # COCO dataset
-    nc = 1 if single_cls else int(data['nc'])  # number of classes
+    nc = 1 if single_cls else int(data["nc"])  # number of classes
     iouv = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for mAP@0.5:0.95
     niou = iouv.numel()
 
@@ -168,12 +185,15 @@ def run(
     if not training:
         if pt and not single_cls:  # check --weights are trained on --data
             ncm = model.model.nc
-            assert ncm == nc, f'{weights} ({ncm} classes) trained on different --data than what you passed ({nc} ' \
+            assert ncm == nc, (
-                              f'classes). Pass correct combination of --weights and --data that are trained together.'
+                f"{weights} ({ncm} classes) trained on different --data than what you passed ({nc} "
+                f"classes). Pass correct combination of --weights and --data that are trained together."
+            )
         model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz))  # warmup
-        pad, rect = (0.0, False) if task == 'speed' else (0.5, pt)  # square inference for benchmarks
+        pad, rect = (0.0, False) if task == "speed" else (0.5, pt)  # square inference for benchmarks
-        task = task if task in ('train', 'val', 'test') else 'val'  # path to train/val/test images
+        task = task if task in ("train", "val", "test") else "val"  # path to train/val/test images
-        dataloader = create_dataloader(data[task],
+        dataloader = create_dataloader(
+            data[task],
             imgsz,
             batch_size,
             stride,
@@ -181,23 +201,24 @@ def run(
             pad=pad,
             rect=rect,
             workers=workers,
-                                       prefix=colorstr(f'{task}: '))[0]
+            prefix=colorstr(f"{task}: "),
+        )[0]
 
     seen = 0
     confusion_matrix = ConfusionMatrix(nc=nc)
-    names = model.names if hasattr(model, 'names') else model.module.names  # get class names
+    names = model.names if hasattr(model, "names") else model.module.names  # get class names
     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', 'mAP50', 'mAP50-95')
+    s = ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "P", "R", "mAP50", "mAP50-95")
     tp, fp, p, r, f1, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
     dt = Profile(device=device), Profile(device=device), Profile(device=device)  # profiling times
     loss = torch.zeros(3, device=device)
     jdict, stats, ap, ap_class = [], [], [], []
-    callbacks.run('on_val_start')
+    callbacks.run("on_val_start")
     pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT)  # progress bar
     for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
-        callbacks.run('on_val_batch_start')
+        callbacks.run("on_val_batch_start")
         with dt[0]:
             if cuda:
                 im = im.to(device, non_blocking=True)
@@ -218,13 +239,9 @@ 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]:
-            preds = non_max_suppression(preds,
+            preds = non_max_suppression(
-                                        conf_thres,
+                preds, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls, max_det=max_det
-                                        iou_thres,
+            )
-                                        labels=lb,
-                                        multi_label=True,
-                                        agnostic=single_cls,
-                                        max_det=max_det)
 
         # Metrics
         for si, pred in enumerate(preds):
@@ -259,18 +276,18 @@ def run(
 
             # Save/log
             if save_txt:
-                (save_dir / 'labels').mkdir(parents=True, exist_ok=True)
+                (save_dir / "labels").mkdir(parents=True, exist_ok=True)
-                save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
+                save_one_txt(predn, save_conf, shape, file=save_dir / "labels" / f"{path.stem}.txt")
             if save_json:
                 save_one_json(predn, jdict, path, class_map)  # append to COCO-JSON dictionary
-            callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
+            callbacks.run("on_val_image_end", pred, predn, path, names, im[si])
 
         # 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, targets, paths, save_dir / f"val_batch{batch_i}_labels.jpg", names)  # labels
-            plot_images(im, output_to_target(preds), 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, preds)
+        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
@@ -281,10 +298,10 @@ def run(
     nt = np.bincount(stats[3].astype(int), minlength=nc)  # number of targets per class
 
     # Print results
-    pf = '%22s' + '%11i' * 2 + '%11.3g' * 4  # print format
+    pf = "%22s" + "%11i" * 2 + "%11.3g" * 4  # print format
-    LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
+    LOGGER.info(pf % ("all", seen, nt.sum(), mp, mr, map50, map))
     if nt.sum() == 0:
-        LOGGER.warning(f'WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels')
+        LOGGER.warning(f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels")
 
     # Print results per class
     if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
@@ -292,35 +309,35 @@ def run(
             LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
 
     # Print speeds
-    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
     if not training:
         shape = (batch_size, 3, imgsz, imgsz)
-        LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
+        LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}" % t)
 
     # Plots
     if plots:
         confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
-        callbacks.run('on_val_end', nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
+        callbacks.run("on_val_end", nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
 
     # Save JSON
     if save_json and len(jdict):
-        w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ''  # weights
+        w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ""  # weights
-        anno_json = str(Path('../datasets/coco/annotations/instances_val2017.json'))  # annotations
+        anno_json = str(Path("../datasets/coco/annotations/instances_val2017.json"))  # annotations
         if not os.path.exists(anno_json):
-            anno_json = os.path.join(data['path'], 'annotations', 'instances_val2017.json')
+            anno_json = os.path.join(data["path"], "annotations", "instances_val2017.json")
-        pred_json = str(save_dir / f'{w}_predictions.json')  # predictions
+        pred_json = str(save_dir / f"{w}_predictions.json")  # predictions
-        LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
+        LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...")
-        with open(pred_json, 'w') as f:
+        with open(pred_json, "w") as f:
             json.dump(jdict, f)
 
         try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
-            check_requirements('pycocotools>=2.0.6')
+            check_requirements("pycocotools>=2.0.6")
             from pycocotools.coco import COCO
             from pycocotools.cocoeval import COCOeval
 
             anno = COCO(anno_json)  # init annotations api
             pred = anno.loadRes(pred_json)  # init predictions api
-            eval = COCOeval(anno, pred, 'bbox')
+            eval = COCOeval(anno, pred, "bbox")
             if is_coco:
                 eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files]  # image IDs to evaluate
             eval.evaluate()
@@ -328,12 +345,12 @@ def run(
             eval.summarize()
             map, map50 = eval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)
         except Exception as e:
-            LOGGER.info(f'pycocotools unable to run: {e}')
+            LOGGER.info(f"pycocotools unable to run: {e}")
 
     # Return results
     model.float()  # for training
     if not training:
-        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     maps = np.zeros(nc) + map
     for i, c in enumerate(ap_class):
@@ -343,71 +360,71 @@ 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("--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 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("--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("--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("--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("--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("--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("--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("--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)')
+    parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)")
-    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
+    parser.add_argument("--single-cls", action="store_true", help="treat as single-class dataset")
-    parser.add_argument('--augment', action='store_true', help='augmented inference')
+    parser.add_argument("--augment", action="store_true", help="augmented inference")
-    parser.add_argument('--verbose', action='store_true', help='report mAP by class')
+    parser.add_argument("--verbose", action="store_true", help="report mAP by class")
-    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
+    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
-    parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
+    parser.add_argument("--save-hybrid", action="store_true", help="save label+prediction hybrid results to *.txt")
-    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
+    parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
-    parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
+    parser.add_argument("--save-json", action="store_true", help="save a COCO-JSON results file")
-    parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name')
+    parser.add_argument("--project", default=ROOT / "runs/val", help="save to project/name")
-    parser.add_argument('--name', default='exp', help='save to project/name')
+    parser.add_argument("--name", default="exp", help="save to project/name")
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
-    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
-    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
     opt = parser.parse_args()
     opt.data = check_yaml(opt.data)  # check YAML
-    opt.save_json |= opt.data.endswith('coco.yaml')
+    opt.save_json |= opt.data.endswith("coco.yaml")
     opt.save_txt |= opt.save_hybrid
     print_args(vars(opt))
     return opt
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
 
-    if opt.task in ('train', 'val', 'test'):  # run normally
+    if opt.task in ("train", "val", "test"):  # run normally
         if opt.conf_thres > 0.001:  # https://github.com/ultralytics/yolov5/issues/1466
-            LOGGER.info(f'WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results')
+            LOGGER.info(f"WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results")
         if opt.save_hybrid:
-            LOGGER.info('WARNING ⚠️ --save-hybrid will return high mAP from hybrid labels, not from predictions alone')
+            LOGGER.info("WARNING ⚠️ --save-hybrid will return high mAP from hybrid labels, not from predictions alone")
         run(**vars(opt))
 
     else:
         weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
-        opt.half = torch.cuda.is_available() and opt.device != 'cpu'  # FP16 for fastest results
+        opt.half = torch.cuda.is_available() and opt.device != "cpu"  # FP16 for fastest results
-        if opt.task == 'speed':  # speed benchmarks
+        if opt.task == "speed":  # speed benchmarks
             # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
             opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
             for opt.weights in weights:
                 run(**vars(opt), plots=False)
 
-        elif opt.task == 'study':  # speed vs mAP benchmarks
+        elif opt.task == "study":  # speed vs mAP benchmarks
             # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
             for opt.weights in weights:
-                f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt'  # filename to save to
+                f = f"study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt"  # filename to save to
                 x, y = list(range(256, 1536 + 128, 128)), []  # x axis (image sizes), y axis
                 for opt.imgsz in x:  # img-size
-                    LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
+                    LOGGER.info(f"\nRunning {f} --imgsz {opt.imgsz}...")
                     r, _, t = run(**vars(opt), plots=False)
                     y.append(r + t)  # results and times
-                np.savetxt(f, y, fmt='%10.4g')  # save
+                np.savetxt(f, y, fmt="%10.4g")  # save
-            subprocess.run(['zip', '-r', 'study.zip', 'study_*.txt'])
+            subprocess.run(["zip", "-r", "study.zip", "study_*.txt"])
             plot_val_study(x=x)  # plot
         else:
             raise NotImplementedError(f'--task {opt.task} not in ("train", "val", "test", "speed", "study")')
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     opt = parse_opt()
     main(opt)