From ad97dc6a6d8b4384eb7224e98c8584ecd8bcfd2a Mon Sep 17 00:00:00 2001 From: Paula Derrenger <107626595+pderrenger@users.noreply.github.com> Date: Mon, 8 Jan 2024 01:29:14 +0100 Subject: [PATCH] Update Actions with Lychee and GitHub Token (#12592) * Add Ultralytics Actions * Auto-format by Ultralytics actions * Update format.yml Signed-off-by: Glenn Jocher --------- Signed-off-by: Glenn Jocher Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- .github/workflows/format.yml | 25 + benchmarks.py | 116 ++--- classify/predict.py | 127 ++--- classify/train.py | 255 +++++----- classify/val.py | 81 ++-- detect.py | 184 +++---- export.py | 616 +++++++++++++----------- hubconf.py | 72 +-- models/common.py | 316 ++++++------ models/experimental.py | 35 +- models/tf.py | 252 ++++++---- models/yolo.py | 150 ++++-- segment/predict.py | 143 +++--- segment/train.py | 522 +++++++++++--------- segment/val.py | 297 +++++++----- train.py | 595 +++++++++++++---------- utils/__init__.py | 29 +- utils/activations.py | 9 +- utils/augmentations.py | 76 ++- utils/autoanchor.py | 72 +-- utils/autobatch.py | 24 +- utils/aws/resume.py | 16 +- utils/callbacks.py | 63 ++- utils/dataloaders.py | 602 ++++++++++++----------- utils/downloads.py | 91 ++-- utils/flask_rest_api/example_request.py | 12 +- utils/flask_rest_api/restapi.py | 28 +- utils/general.py | 574 +++++++++++++--------- utils/loggers/__init__.py | 195 ++++---- utils/loggers/clearml/clearml_utils.py | 123 ++--- utils/loggers/clearml/hpo.py | 78 +-- utils/loggers/comet/__init__.py | 243 +++++----- utils/loggers/comet/comet_utils.py | 54 +-- utils/loggers/comet/hpo.py | 114 ++--- utils/loggers/wandb/wandb_utils.py | 109 +++-- utils/loss.py | 58 +-- utils/metrics.py | 118 ++--- utils/plots.py | 254 +++++----- utils/segment/augmentations.py | 22 +- utils/segment/dataloaders.py | 149 +++--- utils/segment/general.py | 19 +- utils/segment/loss.py | 46 +- utils/segment/metrics.py | 156 +++--- utils/segment/plots.py | 39 +- utils/torch_utils.py | 230 +++++---- utils/triton.py | 43 +- val.py | 269 ++++++----- 47 files changed, 4228 insertions(+), 3473 deletions(-) create mode 100644 .github/workflows/format.yml diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml new file mode 100644 index 000000000000..8d623f4f6e7e --- /dev/null +++ b/.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 diff --git a/benchmarks.py b/benchmarks.py index b590ff63cb01..09e82e588a2a 100644 --- a/benchmarks.py +++ b/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,115 +50,115 @@ def run( - weights=ROOT / 'yolov5s.pt', # weights path - imgsz=640, # inference size (pixels) - batch_size=1, # batch size - data=ROOT / 'data/coco128.yaml', # dataset.yaml path - 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 - hard_fail=False, # throw error on benchmark failure + weights=ROOT / "yolov5s.pt", # weights path + imgsz=640, # inference size (pixels) + batch_size=1, # batch size + data=ROOT / "data/coco128.yaml", # dataset.yaml path + 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 + hard_fail=False, # throw error on benchmark failure ): y, t = [], time.time() device = select_device(device) 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 != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13' # CoreML - if 'cpu' in device.type: - assert cpu, 'inference not supported on CPU' - if 'cuda' in device.type: - assert gpu, 'inference not supported on GPU' + 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 + if "cpu" in device.type: + assert cpu, "inference not supported on CPU" + if "cuda" in device.type: + assert gpu, "inference not supported on GPU" # Export - if f == '-': + if f == "-": w = weights # PyTorch format else: - w = export.run(weights=weights, - imgsz=[imgsz], - include=[f], - batch_size=batch_size, - device=device, - half=half)[-1] # all others - assert suffix in str(w), 'export failed' + w = export.run( + weights=weights, imgsz=[imgsz], include=[f], batch_size=batch_size, 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}' - LOGGER.warning(f'WARNING ⚠️ Benchmark failure for {name}: {e}') + assert type(e) is AssertionError, f"Benchmark --hard-fail 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 - imgsz=640, # inference size (pixels) - batch_size=1, # batch size - data=ROOT / 'data/coco128.yaml', # dataset.yaml path - 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 - hard_fail=False, # throw error on benchmark failure + weights=ROOT / "yolov5s.pt", # weights path + imgsz=640, # inference size (pixels) + batch_size=1, # batch size + data=ROOT / "data/coco128.yaml", # dataset.yaml path + 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 + hard_fail=False, # throw error on benchmark failure ): y, t = [], time.time() 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 \ - export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # weights - assert suffix in str(w), 'export failed' + w = ( + weights + 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']) - LOGGER.info(f'\nExports complete ({time.time() - t:.2f}s)') + py = pd.DataFrame(y, columns=["Format", "Export"]) + 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('--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('--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('--half', action='store_true', help='use FP16 half-precision inference') - 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('--hard-fail', nargs='?', const=True, default=False, help='Exception on error or < min metric') + 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("--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("--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("--test", action="store_true", help="test exports 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") 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) diff --git a/classify/predict.py b/classify/predict.py index b056a0cd707b..b7d2f05d7bce 100644 --- a/classify/predict.py +++ b/classify/predict.py @@ -48,43 +48,54 @@ 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, - increment_path, print_args, strip_optimizer) +from utils.general import ( + 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) - source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam) - 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 - 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 - 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 + weights=ROOT / "yolov5s-cls.pt", # model.pt path(s) + source=ROOT / "data/images", # file/dir/URL/glob/screen/0(webcam) + 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 + 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 + 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://')) - webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file) - screenshot = source.lower().startswith('screen') + is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://")) + webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file) + 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: - f.write(text + '\n') + with open(f"{txt_path}.txt", "a") as f: + 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 - vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) + 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].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 - LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t) + 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) 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('--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('--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('--view-img', action='store_true', help='show results') - 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('--augment', action='store_true', help='augmented inference') - parser.add_argument('--visualize', action='store_true', help='visualize features') - 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('--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('--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('--vid-stride', type=int, default=1, help='video frame-rate stride') + 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("--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("--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("--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("--augment", action="store_true", help="augmented inference") + parser.add_argument("--visualize", action="store_true", help="visualize features") + 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("--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("--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("--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) diff --git a/classify/train.py b/classify/train.py index ecbea1d8c0de..63befed0f780 100644 --- a/classify/train.py +++ b/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 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, - check_requirements, colorstr, download, increment_path, init_seeds, print_args, yaml_save) +from utils.general import ( + 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, - smart_optimizer, smartCrossEntropyLoss, torch_distributed_zero_first) - -LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html -RANK = int(os.getenv('RANK', -1)) -WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) +from utils.torch_utils import ( + 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 +RANK = int(os.getenv("RANK", -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 = \ - opt.save_dir, Path(opt.data), opt.batch_size, opt.epochs, min(os.cpu_count() - 1, opt.workers), \ - opt.imgsz, str(opt.pretrained).lower() == 'true' - cuda = device.type != 'cpu' + 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.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': - subprocess.run(['bash', str(ROOT / 'data/scripts/get_imagenet.sh')], shell=True, check=True) + if str(data) == "imagenet": + 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 - 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) - - test_dir = data_dir / 'test' if (data_dir / 'test').exists() else data_dir / 'val' # data/test or data/val + 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", + imgsz=imgsz, + batch_size=bs // WORLD_SIZE, + augment=True, + cache=opt.cache, + rank=LOCAL_RANK, + workers=nw, + ) + + 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, - imgsz=imgsz, - batch_size=bs // WORLD_SIZE * 2, - augment=False, - cache=opt.cache, - rank=-1, - workers=nw) + testloader = create_classification_dataloader( + path=test_dir, + imgsz=imgsz, + batch_size=bs // WORLD_SIZE * 2, + augment=False, + cache=opt.cache, + rank=-1, + workers=nw, + ) # Model with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT): - if Path(opt.model).is_file() or opt.model.endswith('.pt'): - model = attempt_load(opt.model, device='cpu', fuse=False) + if Path(opt.model).is_file() or opt.model.endswith(".pt"): + 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 - raise ModuleNotFoundError(f'--model {opt.model} not found. Available models are: \n' + '\n'.join(m)) + 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)) 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') - logger.log_images(file, name='Train Examples') + 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_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' - 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}") + 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}" + ) 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) - pbar.desc = f"{f'{epoch + 1}/{epochs}':>10}{mem:>10}{tloss:>12.3g}" + ' ' * 36 + 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 # Test if i == len(pbar) - 1: # last batch - top1, top5, vloss = validate.run(model=ema.ema, - dataloader=testloader, - criterion=criterion, - pbar=pbar) # test accuracy, loss + top1, top5, vloss = validate.run( + model=ema.ema, dataloader=testloader, criterion=criterion, 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, - f'{val}/loss': vloss, - 'metrics/accuracy_top1': top1, - 'metrics/accuracy_top5': top5, - 'lr/0': optimizer.param_groups[0]['lr']} # learning rate + "train/loss": tloss, + f"{val}/loss": vloss, + "metrics/accuracy_top1": top1, + "metrics/accuracy_top5": top5, + "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, - 'best_fitness': best_fitness, - 'model': deepcopy(ema.ema).half(), # deepcopy(de_parallel(model)).half(), - 'ema': None, # deepcopy(ema.ema).half(), - 'updates': ema.updates, - 'optimizer': None, # optimizer.state_dict(), - 'opt': vars(opt), - 'git': GIT_INFO, # {remote, branch, commit} if a git repo - 'date': datetime.now().isoformat()} + "epoch": epoch, + "best_fitness": best_fitness, + "model": deepcopy(ema.ema).half(), # deepcopy(de_parallel(model)).half(), + "ema": None, # deepcopy(ema.ema).half(), + "updates": ema.updates, + "optimizer": None, # optimizer.state_dict(), + "opt": vars(opt), + "git": GIT_INFO, # {remote, branch, commit} if a git repo + "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)' - 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') + 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' + ) # 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()} - logger.log_images(file, name='Test Examples (true-predicted)', epoch=epoch) + meta = {"epochs": epochs, "top1_acc": best_fitness, "date": datetime.now().isoformat()} + 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('--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('--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('--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('--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('--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('--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('--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('--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('--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('--verbose', action='store_true', help='Verbose mode') - 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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--dropout", type=float, default=None, help="Dropout (fraction)") + 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("--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 % 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 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 torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command" torch.cuda.set_device(LOCAL_RANK) - device = torch.device('cuda', LOCAL_RANK) - dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo') + device = torch.device("cuda", LOCAL_RANK) + 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) diff --git a/classify/val.py b/classify/val.py index 6814c4d780e1..b170253d6e0c 100644 --- a/classify/val.py +++ b/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 @@ 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, - increment_path, print_args) +from utils.general import ( + 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 - weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s) + data=ROOT / "../datasets/mnist", # dataset dir + 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 - name='exp', # save to project/name + project=ROOT / "runs/val-cls", # 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 - dataloader = create_classification_dataloader(path=test_dir, - imgsz=imgsz, - batch_size=batch_size, - augment=False, - rank=-1, - workers=workers) + test_dir = data / "test" if (data / "test").exists() else data / "val" # data/test or data/val + dataloader = create_classification_dataloader( + path=test_dir, imgsz=imgsz, 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' - desc = f'{pbar.desc[:-36]}{action:>36}' if pbar else f'{action}' + action = "validating" if dataloader.dataset.root.stem == "val" else "testing" + 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('--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('--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('--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('--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('--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('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference') + 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("--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("--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("--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("--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("--half", action="store_true", help="use FP16 half-precision 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) diff --git a/detect.py b/detect.py index 03bc29de999e..b7d77ef431d4 100644 --- a/detect.py +++ b/detect.py @@ -47,54 +47,68 @@ 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, - increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh) +from utils.general import ( + 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 - source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam) - 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 - view_img=False, # show results - save_txt=False, # save results to *.txt - save_csv=False, # save results in CSV format - save_conf=False, # save confidences in --save-txt labels - save_crop=False, # save cropped prediction boxes - nosave=False, # do not save images/videos - classes=None, # filter by class: --class 0, or --class 0 2 3 - agnostic_nms=False, # class-agnostic NMS - augment=False, # augmented inference - visualize=False, # visualize features - update=False, # update all models - project=ROOT / 'runs/detect', # 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 - hide_conf=False, # hide confidences - half=False, # use FP16 half-precision inference - dnn=False, # use OpenCV DNN for ONNX inference - vid_stride=1, # video frame-rate stride + weights=ROOT / "yolov5s.pt", # model path or triton URL + source=ROOT / "data/images", # file/dir/URL/glob/screen/0(webcam) + 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 + view_img=False, # show results + save_txt=False, # save results to *.txt + save_csv=False, # save results in CSV format + save_conf=False, # save confidences in --save-txt labels + save_crop=False, # save cropped prediction boxes + nosave=False, # do not save images/videos + classes=None, # filter by class: --class 0, or --class 0 2 3 + agnostic_nms=False, # class-agnostic NMS + augment=False, # augmented inference + visualize=False, # visualize features + update=False, # update all models + project=ROOT / "runs/detect", # 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 + hide_conf=False, # hide confidences + 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://')) - webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file) - screenshot = source.lower().startswith('screen') + is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://")) + webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file) + 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} - with open(csv_path, mode='a', newline='') as f: + data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence} + 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 write_to_csv(image_name, prediction, confidence): 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 - s += '%gx%g ' % im.shape[2:] # print string + 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 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 write_to_csv(image_name, prediction, confidence): # 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 write_to_csv(image_name, prediction, confidence): 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: - f.write(('%g ' * len(line)).rstrip() % line + '\n') + with open(f"{txt_path}.txt", "a") as f: + 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 write_to_csv(image_name, prediction, confidence): # 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 write_to_csv(image_name, prediction, confidence): 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 - vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) + 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].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 - LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t) + 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) 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 write_to_csv(image_name, prediction, confidence): 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('--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('--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('--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('--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('--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-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('--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('--agnostic-nms', action='store_true', help='class-agnostic NMS') - parser.add_argument('--augment', action='store_true', help='augmented inference') - parser.add_argument('--visualize', action='store_true', help='visualize features') - 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('--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('--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-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('--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("--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("--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("--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("--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("--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-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-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("--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("--augment", action="store_true", help="augmented inference") + parser.add_argument("--visualize", action="store_true", help="visualize features") + 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("--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("--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-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("--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") 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) diff --git a/export.py b/export.py index d550a85fd99f..74701c37a947 100644 --- a/export.py +++ b/export.py @@ -64,30 +64,42 @@ 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, - check_yaml, colorstr, file_size, get_default_args, print_args, url2file, yaml_save) +from utils.general import ( + 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 @@ def forward(self, x): def export_formats(): # YOLOv5 export formats x = [ - ['PyTorch', '-', '.pt', True, True], - ['TorchScript', 'torchscript', '.torchscript', True, True], - ['ONNX', 'onnx', '.onnx', True, True], - ['OpenVINO', 'openvino', '_openvino_model', True, False], - ['TensorRT', 'engine', '.engine', False, True], - ['CoreML', 'coreml', '.mlmodel', True, False], - ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True], - ['TensorFlow GraphDef', 'pb', '.pb', True, True], - ['TensorFlow Lite', 'tflite', '.tflite', True, False], - ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False], - ['TensorFlow.js', 'tfjs', '_web_model', False, False], - ['PaddlePaddle', 'paddle', '_paddle_model', True, True], ] - return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU']) + ["PyTorch", "-", ".pt", True, True], + ["TorchScript", "torchscript", ".torchscript", True, True], + ["ONNX", "onnx", ".onnx", True, True], + ["OpenVINO", "openvino", "_openvino_model", True, False], + ["TensorRT", "engine", ".engine", False, True], + ["CoreML", "coreml", ".mlmodel", True, False], + ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True], + ["TensorFlow GraphDef", "pb", ".pb", True, True], + ["TensorFlow Lite", "tflite", ".tflite", True, False], + ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False], + ["TensorFlow.js", "tfjs", "_web_model", False, False], + ["PaddlePaddle", "paddle", "_paddle_model", True, True], + ] + 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__}...') - f = file.with_suffix('.torchscript') + LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...") + 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} - extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap() + d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names} + 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__}...') - f = str(file.with_suffix('.onnx')) + LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") + 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['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) + 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) 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__}...') - f = str(file).replace(file.suffix, f'_openvino_model{os.sep}') - f_onnx = file.with_suffix('.onnx') - f_ov = str(Path(f) / file.with_suffix('.xml').name) + LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...") + f = str(file).replace(file.suffix, f"_openvino_model{os.sep}") + f_onnx = file.with_suffix(".onnx") + 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 prepare_input_tensor(image: np.ndarray): 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], - imgsz=imgsz, - batch_size=1, - stride=32, - pad=0.5, - single_cls=False, - rect=False, - workers=workers)[0] + dataloader = create_dataloader( + data[task], imgsz=imgsz, batch_size=1, 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 transform_fn(data_item): 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__}...') - f = str(file).replace('.pt', f'_paddle_model{os.sep}') + LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...") + f = str(file).replace(".pt", f"_paddle_model{os.sep}") - 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 + 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 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__}...') - f = file.with_suffix('.mlmodel') + LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...") + 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])]) - bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None) + 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) 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': - check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com') + if platform.system() == "Linux": + 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__}...') - assert onnx.exists(), f'failed to export ONNX file: {onnx}' - f = file.with_suffix('.engine') # TensorRT engine file + LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...") + assert onnx.exists(), f"failed to export ONNX file: {onnx}" + 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,33 +371,35 @@ 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, - im, - file, - dynamic, - tf_nms=False, - agnostic_nms=False, - topk_per_class=100, - topk_all=100, - iou_thres=0.45, - conf_thres=0.25, - keras=False, - prefix=colorstr('TensorFlow SavedModel:')): +def export_saved_model( + model, + im, + file, + dynamic, + tf_nms=False, + agnostic_nms=False, + topk_per_class=100, + topk_all=100, + iou_thres=0.45, + conf_thres=0.25, + keras=False, + 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__}...') - if tf.__version__ > '2.13.1': - helper_url = 'https://github.com/ultralytics/yolov5/issues/12489' + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + if tf.__version__ > "2.13.1": + 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, - f, - options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version( - tf.__version__, '2.6') else tf.saved_model.SaveOptions()) + tf.saved_model.save( + tfm, + f, + options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) + 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__}...') - f = file.with_suffix('.pb') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + 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, - prefix=colorstr('TensorFlow Lite:')): +def export_tflite( + 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' - help_url = 'https://coral.ai/docs/edgetpu/compiler/' - 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: - 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 + cmd = "edgetpu_compiler --version" + help_url = "https://coral.ai/docs/edgetpu/compiler/" + 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: + 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 for c in ( - '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'): - subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True) + "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", + ): + 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}...') - f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model - f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model - - subprocess.run([ - 'edgetpu_compiler', - '-s', - '-d', - '-k', - '10', - '--out_dir', - str(file.parent), - f_tfl, ], check=True) + LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...") + f = str(file).replace(".pt", "-int8_edgetpu.tflite") # Edge TPU model + f_tfl = str(file).replace(".pt", "-int8.tflite") # TFLite model + + subprocess.run( + [ + "edgetpu_compiler", + "-s", + "-d", + "-k", + "10", + "--out_dir", + str(file.parent), + 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__}...') - f = str(file).replace('.pt', '_web_model') # js dir - f_pb = file.with_suffix('.pb') # *.pb path - f_json = f'{f}/model.json' # *.json path + LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...") + f = str(file).replace(".pt", "_web_model") # js dir + f_pb = file.with_suffix(".pb") # *.pb path + f_json = f"{f}/model.json" # *.json path args = [ - 'tensorflowjs_converter', - '--input_format=tf_frozen_model', - '--quantize_uint8' if int8 else '', - '--output_node_names=Identity,Identity_1,Identity_2,Identity_3', + "tensorflowjs_converter", + "--input_format=tf_frozen_model", + "--quantize_uint8" if int8 else "", + "--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') - with open(tmp_file, 'w') as meta_f: + tmp_file = Path("/tmp/meta.txt") + 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': - img = Image.new('RGB', (w, h)) # img(192 width, 320 height) + if platform.system() == "Darwin": + 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[1].name = 'coordinates' + nms_spec.description.output[0].name = "confidence" + 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.coordinatesOutputFeatureName = 'coordinates' - nms.iouThresholdInputFeatureName = 'iouThreshold' - nms.confidenceThresholdInputFeatureName = 'confidenceThreshold' + nms.confidenceOutputFeatureName = "confidence" + nms.coordinatesOutputFeatureName = "coordinates" + nms.iouThresholdInputFeatureName = "iouThreshold" + 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)), - ('iouThreshold', ct.models.datatypes.Double()), - ('confidenceThreshold', ct.models.datatypes.Double())], - output_features=['confidence', 'coordinates']) + pipeline = ct.models.pipeline.Pipeline( + input_features=[ + ("image", ct.models.datatypes.Array(3, ny, nx)), + ("iouThreshold", ct.models.datatypes.Double()), + ("confidenceThreshold", ct.models.datatypes.Double()), + ], + output_features=["confidence", "coordinates"], + ) pipeline.add_model(model) pipeline.add_model(nms_model) @@ -681,73 +714,77 @@ 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.shortDescription = 'https://github.com/ultralytics/yolov5' - 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.userDefined.update({ - 'classes': ','.join(names.values()), - 'iou_threshold': str(nms.iouThreshold), - 'confidence_threshold': str(nms.confidenceThreshold)}) + pipeline.spec.description.metadata.versionString = "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.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE" + pipeline.spec.description.metadata.userDefined.update( + { + "classes": ",".join(names.values()), + "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['iouThreshold'] = f'(optional) IOU Threshold override (default: {nms.iouThreshold})' - model.input_description['confidenceThreshold'] = \ - f'(optional) Confidence Threshold override (default: {nms.confidenceThreshold})' - 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.input_description["image"] = "Input image" + model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})" + model.input_description[ + "confidenceThreshold" + ] = f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})" + 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' - 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 - include=('torchscript', 'onnx'), # include formats - half=False, # FP16 half-precision export - inplace=False, # set YOLOv5 Detect() inplace=True - keras=False, # use Keras - optimize=False, # TorchScript: optimize for mobile - int8=False, # CoreML/TF INT8 quantization - per_tensor=False, # TF per tensor quantization - dynamic=False, # ONNX/TF/TensorRT: dynamic axes - simplify=False, # ONNX: simplify model - opset=12, # ONNX: opset version - verbose=False, # TensorRT: verbose log - workspace=4, # TensorRT: workspace size (GB) - nms=False, # TF: add NMS to model - agnostic_nms=False, # TF: add agnostic NMS to model - topk_per_class=100, # TF.js NMS: topk per class to keep - topk_all=100, # TF.js NMS: topk for all classes to keep - iou_thres=0.45, # TF.js NMS: IoU threshold - conf_thres=0.25, # TF.js NMS: confidence threshold + data=ROOT / "data/coco128.yaml", # 'dataset.yaml 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 + include=("torchscript", "onnx"), # include formats + half=False, # FP16 half-precision export + inplace=False, # set YOLOv5 Detect() inplace=True + keras=False, # use Keras + optimize=False, # TorchScript: optimize for mobile + int8=False, # CoreML/TF INT8 quantization + per_tensor=False, # TF per tensor quantization + dynamic=False, # ONNX/TF/TensorRT: dynamic axes + simplify=False, # ONNX: simplify model + opset=12, # ONNX: opset version + verbose=False, # TensorRT: verbose log + workspace=4, # TensorRT: workspace size (GB) + nms=False, # TF: add NMS to model + agnostic_nms=False, # TF: add agnostic NMS to model + topk_per_class=100, # TF.js NMS: topk per class to keep + topk_all=100, # TF.js NMS: topk for all classes to keep + iou_thres=0.45, # TF.js NMS: IoU threshold + conf_thres=0.25, # TF.js NMS: confidence threshold ): 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 not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both' + 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" 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 - warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning + f = [""] * len(fmts) # exported filenames + 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,30 +823,27 @@ 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 isinstance(model, ClassificationModel), 'ClassificationModel export to TF formats not yet supported.' - f[5], s_model = export_saved_model(model.cpu(), - im, - file, - dynamic, - tf_nms=nms or agnostic_nms or tfjs, - agnostic_nms=agnostic_nms or tfjs, - topk_per_class=topk_per_class, - topk_all=topk_all, - iou_thres=iou_thres, - conf_thres=conf_thres, - keras=keras) + 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." + f[5], s_model = export_saved_model( + model.cpu(), + im, + file, + dynamic, + tf_nms=nms or agnostic_nms or tfjs, + agnostic_nms=agnostic_nms or tfjs, + topk_per_class=topk_per_class, + topk_all=topk_all, + iou_thres=iou_thres, + conf_thres=conf_thres, + 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, - im, - file, - int8 or edgetpu, - per_tensor, - data=data, - nms=nms, - agnostic_nms=agnostic_nms) + f[7], _ = export_tflite( + s_model, im, 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 '') - h = '--half' if half else '' # --half FP16 inference arg - s = '# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference' 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') + dir = Path("segment" if seg else "classify" if cls else "") + h = "--half" if half else "" # --half FP16 inference arg + s = ( + "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference" + 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' + ) 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('--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('--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('--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('--keras', action='store_true', help='TF: use Keras') - 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('--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('--simplify', action='store_true', help='ONNX: simplify model') - 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('--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('--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-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('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold') + 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("--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("--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("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True") + 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("--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("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes") + 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("--verbose", action="store_true", help="TensorRT: verbose log") + 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("--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-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("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold") parser.add_argument( - '--include', - nargs='+', - default=['torchscript'], - help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle') + "--include", + nargs="+", + default=["torchscript"], + 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) diff --git a/hubconf.py b/hubconf.py index f0192698fbe3..691d8eb64749 100644 --- a/hubconf.py +++ b/hubconf.py @@ -14,7 +14,8 @@ 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. ' - 'You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224).') + LOGGER.warning( + "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. ' - 'You will not be able to run inference with this model.') + LOGGER.warning( + "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 = intersect_dicts(csd, model.state_dict(), exclude=['anchors']) # intersect + csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 + csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"]) # intersect model.load_state_dict(csd, strict=False) # load - if len(ckpt['model'].names) == classes: - model.names = ckpt['model'].names # set class names attribute + if len(ckpt["model"].names) == classes: + 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' - s = f'{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help.' + 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." 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 @@ def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T # 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 @@ def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T # Images imgs = [ - 'data/images/zidane.jpg', # filename - Path('data/images/zidane.jpg'), # Path - 'https://ultralytics.com/images/zidane.jpg', # URI - cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV - Image.open('data/images/bus.jpg'), # PIL - np.zeros((320, 640, 3))] # numpy + "data/images/zidane.jpg", # filename + Path("data/images/zidane.jpg"), # Path + "https://ultralytics.com/images/zidane.jpg", # URI + cv2.imread("data/images/bus.jpg")[:, :, ::-1], # OpenCV + Image.open("data/images/bus.jpg"), # PIL + np.zeros((320, 640, 3)), + ] # numpy # Inference results = model(imgs, size=320) # batched inference diff --git a/models/common.py b/models/common.py index 75cc4e97bbc7..09e7560f4d84 100644 --- a/models/common.py +++ b/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 @@ 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, - increment_path, is_jupyter, make_divisible, non_max_suppression, scale_boxes, xywh2xyxy, - xyxy2xywh, yaml_load) +from utils.general import ( + LOGGER, + 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 @@ def __init__(self, c1, c2, k=(5, 9, 13)): 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 @@ def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13)) 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 @@ def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride 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 - self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, - act=False)) if s == 2 else nn.Identity() + GhostConv(c_, c2, 1, 1, act=False), + ) # pw-linear + 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) @@ -309,9 +323,9 @@ def __init__(self, gain=2): def forward(self, x): b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain' s = self.gain - x = x.view(b, s, s, c // s ** 2, h, w) # x(1,2,2,16,80,80) + x = x.view(b, s, s, c // s**2, h, w) # x(1,2,2,16,80,80) x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2) - return x.view(b, c // s ** 2, h * s, w * s) # x(1,16,160,160) + return x.view(b, c // s**2, h * s, w * s) # x(1,16,160,160) class Concat(nn.Module): @@ -326,7 +340,7 @@ def forward(self, x): 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 @@ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, 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...') - extra_files = {'config.txt': ''} # model metadata + LOGGER.info(f"Loading {w} for TorchScript inference...") + 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 - d = json.loads(extra_files['config.txt'], - object_hook=lambda d: { - int(k) if k.isdigit() else k: v - for k, v in d.items()}) - stride, names = int(d['stride']), d['names'] + if extra_files["config.txt"]: # load metadata dict + d = json.loads( + extra_files["config.txt"], + object_hook=lambda d: {int(k) if k.isdigit() else k: v for k, v in d.items()}, + ) + stride, names = int(d["stride"]), d["names"] elif dnn: # ONNX OpenCV DNN - LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...') - check_requirements('opencv-python>=4.5.4') + LOGGER.info(f"Loading {w} for ONNX OpenCV DNN inference...") + 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...') - check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime')) + LOGGER.info(f"Loading {w} for ONNX Runtime inference...") + 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: - stride, names = int(meta['stride']), eval(meta['names']) + if "stride" in meta: + stride, names = int(meta["stride"]), eval(meta["names"]) elif xml: # OpenVINO - LOGGER.info(f'Loading {w} for OpenVINO inference...') - check_requirements('openvino>=2023.0') # requires openvino-dev: https://pypi.org/project/openvino-dev/ + LOGGER.info(f"Loading {w} for OpenVINO inference...") + 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 - ov_model = core.read_model(model=w, weights=Path(w).with_suffix('.bin')) + 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")) 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 - stride, names = self._load_metadata(Path(w).with_suffix('.yaml')) # load metadata + 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 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': - device = torch.device('cuda:0') - Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr')) + + check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0 + if device.type == "cpu": + 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 @@ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, 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 @@ def gd_outputs(gd): 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...') - delegate = { - 'Linux': 'libedgetpu.so.1', - 'Darwin': 'libedgetpu.1.dylib', - 'Windows': 'edgetpu.dll'}[platform.system()] + LOGGER.info(f"Loading {w} for TensorFlow Lite Edge TPU inference...") + delegate = {"Linux": "libedgetpu.so.1", "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')) - stride, names = int(meta['stride']), meta['names'] + meta = ast.literal_eval(model.read(meta_file).decode("utf-8")) + 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...') - check_requirements('paddlepaddle-gpu' if cuda else 'paddlepaddle') + LOGGER.info(f"Loading {w} for PaddlePaddle inference...") + 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 - weights = Path(w).with_suffix('.pdiparams') + w = next(Path(w).rglob("*.pdmodel")) # get *.pdmodel file from *_paddle_model dir + 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 @@ def gd_outputs(gd): 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...') - check_requirements('tritonclient[all]') + LOGGER.info(f"Using {w} as Triton Inference Server...") + 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(): - 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 - names = yaml_load(ROOT / 'data/ImageNet.yaml')['names'] # human-readable names + if "names" not in locals(): + 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 + names = yaml_load(ROOT / "data/ImageNet.yaml")["names"] # human-readable names self.__dict__.update(locals()) # assign all variables to self @@ -538,26 +562,26 @@ def forward(self, im, augment=False, visualize=False): 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: - i = self.model.get_binding_index('images') + if self.dynamic and im.shape != self.bindings["images"].shape: + 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 - if 'confidence' in y: - box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels - conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float) + y = self.model.predict({"image": im}) # coordinates are xywh normalized + if "confidence" in y: + box = xywh2xyxy(y["coordinates"] * [[w, h, w, h]]) # xyxy pixels + 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 @@ def forward(self, im, augment=False, visualize=False): 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 @@ def from_numpy(self, x): 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... ') - copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes + LOGGER.info("Adding AutoShape... ") + 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 @@ def forward(self, ims, size=640, augment=False, profile=False): 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 @@ def forward(self, ims, size=640, augment=False, profile=False): 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 - files.append(Path(f).with_suffix('.jpg').name) + im, f = np.asarray(exif_transpose(im)), getattr(im, "filename", f) or f + 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 @@ def forward(self, ims, size=640, augment=False, profile=False): # Post-process with dt[2]: - 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 + 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 for i in range(n): scale_boxes(shape1, y[i][:, :4], shape0[i]) @@ -747,40 +774,44 @@ def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None): 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('')): - s, crops = '', [] + def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path("")): + 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 - crops.append({ - 'box': box, - 'conf': conf, - 'cls': cls, - 'label': label, - 'im': save_one_box(box, im, file=file, save=save)}) + file = save_dir / "crops" / self.names[int(cls)] / self.files[i] if save else None + crops.append( + { + "box": box, + "conf": conf, + "cls": cls, + "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 @@ def _run(self, pprint=False, show=False, save=False, crop=False, render=False, l if render: self.ims[i] = np.asarray(im) if pprint: - s = s.lstrip('\n') - return f'{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t + s = s.lstrip("\n") + 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 @@ def render(self, labels=True): 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 - cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns - for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]): + ca = "xmin", "ymin", "xmax", "ymax", "confidence", "class", "name" # xyxy columns + cb = "xcenter", "ycenter", "width", "height", "confidence", "class", "name" # xywh columns + 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 @@ def __str__(self): # override print(results) 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 @@ def forward(self, x): class Classify(nn.Module): # YOLOv5 classification head, i.e. x(b,c1,20,20) to x(b,c2) - def __init__(self, - c1, - c2, - k=1, - s=1, - p=None, - g=1, - dropout_p=0.0): # ch_in, ch_out, kernel, stride, padding, groups, dropout probability + def __init__( + self, c1, c2, 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) diff --git a/models/experimental.py b/models/experimental.py index 11f75e2254b3..c242364bdec5 100644 --- a/models/experimental.py +++ b/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 @@ def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kern 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 @@ def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kern a[0] = 1 c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b - 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_)]) + 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_)] + ) 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 = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model + ckpt = torch.load(attempt_download(w), map_location="cpu") # load + ckpt = (ckpt.get("ema") or ckpt["model"]).to(device).float() # FP32 model # Model compatibility updates - if not hasattr(ckpt, 'stride'): - ckpt.stride = torch.tensor([32.]) - if hasattr(ckpt, 'names') and isinstance(ckpt.names, (list, tuple)): + if not hasattr(ckpt, "stride"): + ckpt.stride = torch.tensor([32.0]) + 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') - setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl) - elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'): + delattr(m, "anchor_grid") + setattr(m, "anchor_grid", [torch.zeros(1)] * m.nl) + 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') - for k in 'names', 'nc', 'yaml': + print(f"Ensemble created with {weights}\n") + 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 diff --git a/models/tf.py b/models/tf.py index 17cca1e54fcf..53520b52c086 100644 --- a/models/tf.py +++ b/models/tf.py @@ -27,8 +27,21 @@ import torch.nn as nn from tensorflow import keras -from models.common import (C3, SPP, SPPF, Bottleneck, BottleneckCSP, C3x, Concat, Conv, CrossConv, DWConv, - DWConvTranspose2d, Focus, autopad) +from models.common import ( + 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 @@ def __init__(self, w=None): 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 @@ def __init__(self, pad): 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 @@ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): filters=c2, kernel_size=k, strides=s, - padding='SAME' if s == 1 else 'VALID', - use_bias=not hasattr(w, 'bn'), + padding="SAME" if s == 1 else "VALID", + 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', - use_bias=not hasattr(w, 'bn'), + padding="SAME" if s == 1 else "VALID", + 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 k == 4 and p1 == 1, 'TFDWConv() only valid for k=4 and p1=1' + 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" weight, bias = w.weight.permute(2, 3, 1, 0).numpy(), w.bias.numpy() self.c1 = c1 self.conv = [ - keras.layers.Conv2DTranspose(filters=1, - kernel_size=k, - strides=s, - 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)] + keras.layers.Conv2DTranspose( + filters=1, + kernel_size=k, + strides=s, + 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) + ] 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, - kernel_size=k, - strides=s, - padding='VALID', - use_bias=bias, - kernel_initializer=keras.initializers.Constant( - w.weight.permute(2, 3, 1, 0).numpy()), - bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None) + self.conv = keras.layers.Conv2D( + filters=c2, + kernel_size=k, + strides=s, + padding="VALID", + use_bias=bias, + kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()), + bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None, + ) def call(self, inputs): return self.conv(inputs) @@ -233,8 +254,9 @@ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): 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([ - TFCrossConv(c_, c_, k=3, s=1, g=g, e=1.0, shortcut=shortcut, w=w.m[j]) for j in range(n)]) + 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)] + ) def call(self, inputs): return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3)) @@ -247,7 +269,7 @@ def __init__(self, c1, c2, k=(5, 9, 13), w=None): 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 @@ def __init__(self, c1, c2, k=5, w=None): 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) @@ -307,10 +329,10 @@ def call(self, inputs): # Normalize xywh to 0-1 to reduce calibration error xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32) wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32) - y = tf.concat([xy, wh, tf.sigmoid(y[..., 4:5 + self.nc]), y[..., 5 + self.nc:]], -1) + y = tf.concat([xy, wh, tf.sigmoid(y[..., 4 : 5 + self.nc]), y[..., 5 + self.nc :]], -1) z.append(tf.reshape(y, [-1, self.na * ny * nx, self.no])) - return tf.transpose(x, [0, 2, 1, 3]) if self.training else (tf.concat(z, 1), ) + return tf.transpose(x, [0, 2, 1, 3]) if self.training else (tf.concat(z, 1),) @staticmethod def _make_grid(nx=20, ny=20): @@ -340,11 +362,10 @@ def call(self, x): 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 @@ def call(self, inputs): 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( - 'channel_multiple') + anchors, nc, gd, gw, ch_mul = ( + 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, - BottleneckCSP, C3, C3x]: + nn.Conv2d, + 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.', '')) - m_ = 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 + 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 + 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, - inputs, - tf_nms=False, - agnostic_nms=False, - topk_per_class=100, - topk_all=100, - iou_thres=0.45, - conf_thres=0.25): + 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, + ): y = [] # outputs x = inputs for m in self.model.layers: @@ -482,14 +526,10 @@ def predict(self, 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, - scores, - topk_per_class, - topk_all, - iou_thres, - conf_thres, - clip_boxes=False) - return (nms, ) + nms = tf.image.combined_non_max_suppression( + boxes, scores, 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) # xywh = x[..., :4] # x(6300,4) boxes @@ -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), - input, - fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32), - name='agnostic_nms') + 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", + ) @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, - scores_inp, - max_output_size=topk_all, - iou_threshold=iou_thres, - score_threshold=conf_thres) + selected_inds = tf.image.non_max_suppression( + boxes, scores_inp, 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, - paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], - mode='CONSTANT', - constant_values=0.0) + padded_boxes = tf.pad( + selected_boxes, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], + mode="CONSTANT", + constant_values=0.0, + ) selected_scores = tf.gather(scores_inp, selected_inds) - padded_scores = tf.pad(selected_scores, - paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], - mode='CONSTANT', - constant_values=-1.0) + padded_scores = tf.pad( + selected_scores, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], + mode="CONSTANT", + constant_values=-1.0, + ) selected_classes = tf.gather(class_inds, selected_inds) - padded_classes = tf.pad(selected_classes, - paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], - mode='CONSTANT', - constant_values=-1.0) + padded_classes = tf.pad( + selected_classes, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], + mode="CONSTANT", + 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 - imgsz=(640, 640), # inference size h,w - batch_size=1, # batch size - dynamic=False, # dynamic batch size + 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('--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('--dynamic', action='store_true', help='dynamic batch size') + 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("--batch-size", type=int, default=1, help="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) diff --git a/models/yolo.py b/models/yolo.py index f6cdbcb5d2d8..e98351b98691 100644 --- a/models/yolo.py +++ b/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 @@ 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, - Contract, Conv, CrossConv, DetectMultiBackend, DWConv, DWConvTranspose2d, Expand, Focus, - GhostBottleneck, GhostConv, Proto) +from models.common import ( + C3, + 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, - time_sync) +from utils.torch_utils import ( + 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 @@ def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer 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) @@ -82,14 +110,14 @@ def forward(self, x): y = torch.cat((xy, wh, conf), 4) z.append(y.view(bs, self.na * nx * ny, self.no)) - return x if self.training else (torch.cat(z, 1), ) if self.export else (torch.cat(z, 1), x) + 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 @@ def _forward_once(self, x, profile=False, visualize=False): 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 @@ def _apply(self, fn): 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 - if nc and nc != self.yaml['nc']: + ch = self.yaml["ch"] = self.yaml.get("ch", ch) # input channels + 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}') - self.yaml['anchors'] = round(anchors) # override yaml value + LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}") + 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.inplace = self.yaml.get('inplace', True) + self.names = [str(i) for i in range(self.yaml["nc"])] # default names + self.inplace = self.yaml.get("inplace", True) # Build strides, anchors m = self.model[-1] # Detect() @@ -207,7 +236,7 @@ def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, i # Init weights, biases initialize_weights(self) self.info() - LOGGER.info('') + LOGGER.info("") def forward(self, x, augment=False, profile=False, visualize=False): if augment: @@ -248,9 +277,9 @@ def _descale_pred(self, p, flips, scale, img_size): def _clip_augmented(self, y): # Clip YOLOv5 augmented inference tails nl = self.model[-1].nl # number of detection layers (P3-P5) - g = sum(4 ** x for x in range(nl)) # grid points + g = sum(4**x for x in range(nl)) # grid points e = 1 # exclude layer count - i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e)) # indices + i = (y[0].shape[1] // g) * sum(4**x for x in range(e)) # indices y[0] = y[0][:, :-i] # large i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # indices y[-1] = y[-1][:, i:] # small @@ -263,7 +292,9 @@ def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is 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 @@ def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is 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 @@ def _from_detection_model(self, model, nc=1000, cutoff=10): 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 @@ def _from_yaml(self, cfg): 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( - 'activation'), d.get('channel_multiple') + anchors, nc, gd, gw, act, ch_mul = ( + 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, - BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}: + Conv, + 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('--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('--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('--test', action='store_true', help='test all yolo*.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("--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("--line-profile", action="store_true", help="profile model speed layer by layer") + 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 @@ def parse_model(d, ch): # model_dict, input_channels(3) 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() diff --git a/segment/predict.py b/segment/predict.py index 8e3d97dfeb92..23a4e3538509 100644 --- a/segment/predict.py +++ b/segment/predict.py @@ -46,23 +46,36 @@ 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, - increment_path, non_max_suppression, print_args, scale_boxes, scale_segments, - strip_optimizer) +from utils.general import ( + LOGGER, + 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) - source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam) - data=ROOT / 'data/coco128.yaml', # dataset.yaml path + weights=ROOT / "yolov5s-seg.pt", # model.pt path(s) + source=ROOT / "data/images", # file/dir/URL/glob/screen/0(webcam) + 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 - name='exp', # save results to project/name + project=ROOT / "runs/predict-seg", # 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://')) - webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file) - screenshot = source.lower().startswith('screen') + is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://")) + webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file) + 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 - s += '%gx%g ' % im.shape[2:] # print string + 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 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() / - 255 if retina_masks else im[i]) + 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], + ) # 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: - f.write(('%g ' * len(line)).rstrip() % line + '\n') + with open(f"{txt_path}.txt", "a") as f: + 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 - vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) + 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].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 - LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t) + 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) 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('--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('--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('--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('--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('--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-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('--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('--augment', action='store_true', help='augmented inference') - parser.add_argument('--visualize', action='store_true', help='visualize features') - 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('--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('--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-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('--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('--retina-masks', action='store_true', help='whether to plot masks in native resolution') + 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("--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("--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("--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("--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-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("--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("--agnostic-nms", action="store_true", help="class-agnostic NMS") + parser.add_argument("--augment", action="store_true", help="augmented inference") + parser.add_argument("--visualize", action="store_true", help="visualize features") + 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("--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("--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-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("--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("--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) diff --git a/segment/train.py b/segment/train.py index 2ae09c1cbf66..fe262348fae4 100644 --- a/segment/train.py +++ b/segment/train.py @@ -1,7 +1,7 @@ # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license """ -Train a YOLOv5 segment model on a segment dataset -Models and datasets download automatically from the latest YOLOv5 release. +Train a YOLOv5 segment model on a segment dataset Models and datasets download automatically from the latest YOLOv5 +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.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, - 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.general import ( + LOGGER, + TQDM_BAR_FORMAT, + 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, - smart_resume, torch_distributed_zero_first) - -LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html -RANK = int(os.getenv('RANK', -1)) -WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) +from utils.torch_utils import ( + 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 +RANK = int(os.getenv("RANK", -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, \ - opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze, opt.mask_ratio + ( + 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, + 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 / 'opt.yaml', vars(opt)) + yaml_save(save_dir / "hyp.yaml", hyp) + 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'] - 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 - is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset + train_path, val_path = data_dict["train"], data_dict["val"] + 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 + is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset # Model - check_suffix(weights, '.pt') # check weights - pretrained = weights.endswith('.pt') + check_suffix(weights, ".pt") # check weights + 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 - 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 - csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 + 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) + 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 = 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 - optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['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"]) # 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' - 'See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started.' + "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." ) 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,41 +244,43 @@ 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, - imgsz, - batch_size // WORLD_SIZE * 2, - gs, - single_cls, - hyp=hyp, - cache=None if noval else opt.cache, - rect=True, - rank=-1, - workers=workers * 2, - pad=0.5, - mask_downsample_ratio=mask_ratio, - overlap_mask=overlap, - prefix=colorstr('val: '))[0] + val_loader = create_dataloader( + val_path, + imgsz, + batch_size // WORLD_SIZE * 2, + gs, + single_cls, + hyp=hyp, + cache=None if noval else opt.cache, + rect=True, + rank=-1, + workers=workers * 2, + pad=0.5, + mask_downsample_ratio=mask_ratio, + overlap_mask=overlap, + 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['cls'] *= nc / 80 * 3 / nl # scale to classes and layers - hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers - hyp['label_smoothing'] = opt.label_smoothing + hyp["box"] *= 3 / nl # scale to 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["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' - 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...') + 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...' + ) 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) % - ('Epoch', 'GPU_mem', 'box_loss', 'seg_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size')) + LOGGER.info( + ("\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)]) - if 'momentum' in x: - x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) + 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: + 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) - pbar.set_description(('%11s' * 2 + '%11.4g' * 6) % - (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1])) + 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) + % (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,35 +408,37 @@ 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')) - logger.log_images(files, 'Mosaics', epoch) + files = sorted(save_dir.glob("train*.jpg")) + 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, - batch_size=batch_size // WORLD_SIZE * 2, - imgsz=imgsz, - half=amp, - model=ema.ema, - single_cls=single_cls, - dataloader=val_loader, - save_dir=save_dir, - plots=False, - callbacks=callbacks, - compute_loss=compute_loss, - mask_downsample_ratio=mask_ratio, - overlap=overlap) + results, maps, _ = validate.run( + data_dict, + batch_size=batch_size // WORLD_SIZE * 2, + imgsz=imgsz, + half=amp, + model=ema.ema, + single_cls=single_cls, + dataloader=val_loader, + save_dir=save_dir, + plots=False, + callbacks=callbacks, + compute_loss=compute_loss, + mask_downsample_ratio=mask_ratio, + 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, - 'best_fitness': best_fitness, - 'model': deepcopy(de_parallel(model)).half(), - 'ema': deepcopy(ema.ema).half(), - 'updates': ema.updates, - 'optimizer': optimizer.state_dict(), - 'opt': vars(opt), - 'git': GIT_INFO, # {remote, branch, commit} if a git repo - 'date': datetime.now().isoformat()} + "epoch": epoch, + "best_fitness": best_fitness, + "model": deepcopy(de_parallel(model)).half(), + "ema": deepcopy(ema.ema).half(), + "updates": ema.updates, + "optimizer": optimizer.state_dict(), + "opt": vars(opt), + "git": GIT_INFO, # {remote, branch, commit} if a git repo + "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') - logger.log_model(w / f'epoch{epoch}.pt') + torch.save(ckpt, 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 - files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))] + 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 = [(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(sorted(save_dir.glob('val*.jpg')), 'Validation', epoch + 1) + logger.log_images(files, "Results", 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('--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('--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('--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('--rect', action='store_true', help='rectangular 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('--noval', action='store_true', help='only validate final epoch') - 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('--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('--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('--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('--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('--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('--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('--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('--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('--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('--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('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify') + 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("--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("--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("--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("--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("--noval", action="store_true", help="only validate final epoch") + 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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--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('--no-overlap', action='store_true', help='Overlap masks train faster at slightly less mAP') + 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") 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 = \ - check_file(opt.data), 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' + 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 + 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 - opt.project = str(ROOT / '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.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' - assert not opt.image_weights, f'--image-weights {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 % 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' + msg = "is not compatible with YOLOv5 Multi-GPU DDP training" + assert not opt.image_weights, f"--image-weights {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 % 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" torch.cuda.set_device(LOCAL_RANK) - device = torch.device('cuda', LOCAL_RANK) - dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo') + device = torch.device("cuda", LOCAL_RANK) + 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) - 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) - 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 - 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay - 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok) - 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum - 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr - 'box': (1, 0.02, 0.2), # box loss gain - 'cls': (1, 0.2, 4.0), # cls loss gain - 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight - 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) - 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight - 'iou_t': (0, 0.1, 0.7), # IoU training threshold - 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold - '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) - '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_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) - 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) - 'translate': (1, 0.0, 0.9), # image translation (+/- fraction) - 'scale': (1, 0.0, 0.9), # image scale (+/- gain) - 'shear': (1, 0.0, 10.0), # image shear (+/- deg) - 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 - 'flipud': (1, 0.0, 1.0), # image flip up-down (probability) - 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) - 'mosaic': (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) - - with open(opt.hyp, errors='ignore') as f: + "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) + "momentum": (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 + "weight_decay": (1, 0.0, 0.001), # optimizer weight decay + "warmup_epochs": (1, 0.0, 5.0), # warmup epochs (fractions ok) + "warmup_momentum": (1, 0.0, 0.95), # warmup initial momentum + "warmup_bias_lr": (1, 0.0, 0.2), # warmup initial bias lr + "box": (1, 0.02, 0.2), # box loss gain + "cls": (1, 0.2, 4.0), # cls loss gain + "cls_pw": (1, 0.5, 2.0), # cls BCELoss positive_weight + "obj": (1, 0.2, 4.0), # obj loss gain (scale with pixels) + "obj_pw": (1, 0.5, 2.0), # obj BCELoss positive_weight + "iou_t": (0, 0.1, 0.7), # IoU training threshold + "anchor_t": (1, 2.0, 8.0), # anchor-multiple threshold + "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) + "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_v": (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) + "degrees": (1, 0.0, 45.0), # image rotation (+/- deg) + "translate": (1, 0.0, 0.9), # image translation (+/- fraction) + "scale": (1, 0.0, 0.9), # image scale (+/- gain) + "shear": (1, 0.0, 10.0), # image shear (+/- deg) + "perspective": (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 + "flipud": (1, 0.0, 1.0), # image flip up-down (probability) + "fliplr": (0, 0.0, 1.0), # image flip left-right (probability) + "mosaic": (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) + + 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 - hyp['anchors'] = 3 + if "anchors" not in hyp: # anchors commented in hyp.yaml + 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([ - 'gsutil', - 'cp', - f'gs://{opt.bucket}/evolve.csv', - str(evolve_csv), ]) + subprocess.run( + [ + "gsutil", + "cp", + 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' - x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1) + parent = "single" # parent selection method: 'single' or 'weighted' + 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) - if parent == 'single' or len(x) == 1: + w = fitness(x) - fitness(x).min() + 1e-6 # weights (sum > 0) + 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' - f"Results saved to {colorstr('bold', save_dir)}\n" - f'Usage example: $ python train.py --hyp {evolve_yaml}') + 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}' + ) def run(**kwargs): @@ -661,6 +741,6 @@ def run(**kwargs): return opt -if __name__ == '__main__': +if __name__ == "__main__": opt = parse_opt() main(opt) diff --git a/segment/val.py b/segment/val.py index 304d0c751314..1e5159c710ed 100644 --- a/segment/val.py +++ b/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 @@ 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, - check_requirements, check_yaml, coco80_to_coco91_class, colorstr, increment_path, - non_max_suppression, print_args, scale_boxes, xywh2xyxy, xyxy2xywh) +from utils.general import ( + LOGGER, + 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: - f.write(('%g ' * len(line)).rstrip() % line + '\n') + with open(file, "a") as f: + 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['counts'] = rle['counts'].decode('utf-8') + rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0] + 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 single_encode(x): 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({ - 'image_id': image_id, - 'category_id': class_map[int(p[5])], - 'bbox': [round(x, 3) for x in b], - 'score': round(p[4], 5), - 'segmentation': rles[i]}) + jdict.append( + { + "image_id": image_id, + "category_id": class_map[int(p[5])], + "bbox": [round(x, 3) for x in b], + "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 @@ -128,39 +146,39 @@ def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, over @smart_inference_mode() def run( - data, - weights=None, # model.pt path(s) - batch_size=32, # batch size - imgsz=640, # inference size (pixels) - conf_thres=0.001, # confidence threshold - iou_thres=0.6, # NMS IoU threshold - max_det=300, # maximum detections per image - task='val', # train, val, test, speed or study - device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu - workers=8, # max dataloader workers (per RANK in DDP mode) - single_cls=False, # treat as single-class dataset - augment=False, # augmented inference - verbose=False, # verbose output - save_txt=False, # save results to *.txt - 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 - 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(''), - plots=True, - overlap=False, - mask_downsample_ratio=1, - compute_loss=None, - callbacks=Callbacks(), + data, + weights=None, # model.pt path(s) + batch_size=32, # batch size + imgsz=640, # inference size (pixels) + conf_thres=0.001, # confidence threshold + iou_thres=0.6, # NMS IoU threshold + max_det=300, # maximum detections per image + task="val", # train, val, test, speed or study + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + workers=8, # max dataloader workers (per RANK in DDP mode) + single_cls=False, # treat as single-class dataset + augment=False, # augmented inference + verbose=False, # verbose output + save_txt=False, # save results to *.txt + 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 + 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(""), + plots=True, + overlap=False, + mask_downsample_ratio=1, + compute_loss=None, + 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' - 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 + cuda = device.type != "cpu" + 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 iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for mAP@0.5:0.95 niou = iouv.numel() @@ -208,31 +226,46 @@ 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} ' \ - f'classes). Pass correct combination of --weights and --data that are trained together.' + assert ncm == nc, ( + 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 - task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images - dataloader = create_dataloader(data[task], - imgsz, - batch_size, - stride, - single_cls, - pad=pad, - rect=rect, - workers=workers, - prefix=colorstr(f'{task}: '), - overlap_mask=overlap, - mask_downsample_ratio=mask_downsample_ratio)[0] + 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 + dataloader = create_dataloader( + data[task], + imgsz, + batch_size, + stride, + single_cls, + pad=pad, + rect=rect, + workers=workers, + prefix=colorstr(f"{task}: "), + overlap_mask=overlap, + 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', - 'mAP50', 'mAP50-95)') + s = ("%22s" + "%11s" * 10) % ( + "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, - conf_thres, - iou_thres, - labels=lb, - multi_label=True, - agnostic=single_cls, - max_det=max_det, - nm=nm) + preds = non_max_suppression( + preds, conf_thres, 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.permute(1, 2, 0).contiguous().cpu().numpy(), shape, shapes[si][1]) + pred_masks = scale_image( + 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, output_to_target(preds, max_det=15), plot_masks, paths, - save_dir / f'val_batch{batch_i}_pred.jpg', names) # pred + 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, + 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 - LOGGER.info(pf % ('all', seen, nt.sum(), *metrics.mean_results())) + pf = "%22s" + "%11i" * 2 + "%11.3g" * 8 # print format + 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 - anno_json = str(Path('../datasets/coco/annotations/instances_val2017.json')) # annotations - pred_json = str(save_dir / f'{w}_predictions.json') # predictions - LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...') - with open(pred_json, 'w') as f: + 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 + pred_json = str(save_dir / f"{w}_predictions.json") # predictions + LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...") + 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('--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('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') - parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold') - parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold') - parser.add_argument('--max-det', type=int, default=300, help='maximum detections per image') - parser.add_argument('--task', default='val', help='train, val, test, speed or study') - parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)') - 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('--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-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-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('--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('--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("--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("--batch-size", type=int, default=32, help="batch size") + parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)") + parser.add_argument("--conf-thres", type=float, default=0.001, help="confidence threshold") + parser.add_argument("--iou-thres", type=float, default=0.6, help="NMS IoU threshold") + parser.add_argument("--max-det", type=int, default=300, help="maximum detections per image") + parser.add_argument("--task", default="val", help="train, val, test, speed or study") + parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)") + 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("--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-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-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("--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("--half", action="store_true", help="use FP16 half-precision 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 - if opt.task == 'speed': # speed benchmarks + opt.half = torch.cuda.is_available() and opt.device != "cpu" # FP16 for fastest results + 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 - subprocess.run(['zip', '-r', 'study.zip', 'study_*.txt']) + np.savetxt(f, y, fmt="%10.4g") # save + 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) diff --git a/train.py b/train.py index 378ff4bd30ea..73297d204393 100644 --- a/train.py +++ b/train.py @@ -1,7 +1,6 @@ # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license """ -Train a YOLOv5 model on a custom dataset. -Models and datasets download automatically from the latest YOLOv5 release. +Train a YOLOv5 model on a custom dataset. 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.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, - 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.general import ( + LOGGER, + TQDM_BAR_FORMAT, + 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, + 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, - smart_resume, torch_distributed_zero_first) - -LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html -RANK = int(os.getenv('RANK', -1)) -WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) +from utils.torch_utils import ( + 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 +RANK = int(os.getenv("RANK", -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 = \ - 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 - callbacks.run('on_pretrain_routine_start') + 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, + 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 / 'opt.yaml', vars(opt)) + yaml_save(save_dir / "hyp.yaml", hyp) + 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): - include_loggers.append('ndjson_console') - if getattr(opt, 'ndjson_file', False): - include_loggers.append('ndjson_file') + if getattr(opt, "ndjson_console", False): + include_loggers.append("ndjson_console") + if getattr(opt, "ndjson_file", False): + 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'] - 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 - is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset + train_path, val_path = data_dict["train"], data_dict["val"] + 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 + is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset # Model - check_suffix(weights, '.pt') # check weights - pretrained = weights.endswith('.pt') + check_suffix(weights, ".pt") # check weights + 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 - 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 - csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 + 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 + 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 = 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 - optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['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"]) # 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,58 +234,62 @@ 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' - 'See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started.' + "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." ) 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, - imgsz, - batch_size // WORLD_SIZE, - gs, - single_cls, - hyp=hyp, - augment=True, - 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: '), - shuffle=True, - seed=opt.seed) + 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, + rect=opt.rect, + rank=LOCAL_RANK, + workers=workers, + image_weights=opt.image_weights, + quad=opt.quad, + prefix=colorstr("train: "), + shuffle=True, + 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, - imgsz, - batch_size // WORLD_SIZE * 2, - gs, - single_cls, - hyp=hyp, - cache=None if noval else opt.cache, - rect=True, - rank=-1, - workers=workers * 2, - pad=0.5, - prefix=colorstr('val: '))[0] + val_loader = create_dataloader( + val_path, + imgsz, + batch_size // WORLD_SIZE * 2, + gs, + single_cls, + hyp=hyp, + cache=None if noval else opt.cache, + rect=True, + rank=-1, + workers=workers * 2, + pad=0.5, + 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['cls'] *= nc / 80 * 3 / nl # scale to classes and layers - hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers - hyp['label_smoothing'] = opt.label_smoothing + hyp["box"] *= 3 / nl # scale to 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["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') - 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...') + callbacks.run("on_train_start") + 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...' + ) 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)]) - if 'momentum' in x: - x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) + 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: + 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,35 +397,39 @@ 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) - pbar.set_description(('%11s' * 2 + '%11.4g' * 5) % - (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)) + 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) + % (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) - ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights']) + callbacks.run("on_train_epoch_end", epoch=epoch) + 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, - batch_size=batch_size // WORLD_SIZE * 2, - imgsz=imgsz, - half=amp, - model=ema.ema, - single_cls=single_cls, - dataloader=val_loader, - save_dir=save_dir, - plots=False, - callbacks=callbacks, - compute_loss=compute_loss) + results, maps, _ = validate.run( + data_dict, + batch_size=batch_size // WORLD_SIZE * 2, + imgsz=imgsz, + half=amp, + model=ema.ema, + single_cls=single_cls, + dataloader=val_loader, + save_dir=save_dir, + plots=False, + callbacks=callbacks, + 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, - 'best_fitness': best_fitness, - 'model': deepcopy(de_parallel(model)).half(), - 'ema': deepcopy(ema.ema).half(), - 'updates': ema.updates, - 'optimizer': optimizer.state_dict(), - 'opt': vars(opt), - 'git': GIT_INFO, # {remote, branch, commit} if a git repo - 'date': datetime.now().isoformat()} + "epoch": epoch, + "best_fitness": best_fitness, + "model": deepcopy(de_parallel(model)).half(), + "ema": deepcopy(ema.ema).half(), + "updates": ema.updates, + "optimizer": optimizer.state_dict(), + "opt": vars(opt), + "git": GIT_INFO, # {remote, branch, commit} if a git repo + "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('--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('--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('--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('--rect', action='store_true', help='rectangular 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('--noval', action='store_true', help='only validate final epoch') - 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('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations') - parser.add_argument('--evolve_population', - type=str, - 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('--bucket', type=str, default='', help='gsutil bucket') - 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('--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('--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('--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('--project', default=ROOT / 'runs/train', 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('--quad', action='store_true', help='quad dataloader') - 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('--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('--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('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify') + 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("--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("--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("--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("--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("--noval", action="store_true", help="only validate final epoch") + 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("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations") + parser.add_argument( + "--evolve_population", type=str, 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("--bucket", type=str, default="", help="gsutil bucket") + 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("--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("--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("--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("--project", default=ROOT / "runs/train", 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("--quad", action="store_true", help="quad dataloader") + 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("--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("--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("--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('--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('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use') + 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("--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") # NDJSON logging - 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-console", action="store_true", help="Log ndjson to console") + 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 = \ - check_file(opt.data), 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' + 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 + 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 - opt.project = str(ROOT / 'runs/evolve') + if opt.project == str(ROOT / "runs/train"): # if default project name, rename to 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' - assert not opt.image_weights, f'--image-weights {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 % 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' + msg = "is not compatible with YOLOv5 Multi-GPU DDP training" + assert not opt.image_weights, f"--image-weights {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 % 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" torch.cuda.set_device(LOCAL_RANK) - device = torch.device('cuda', LOCAL_RANK) - dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo', - timeout=timedelta(seconds=10800)) + device = torch.device("cuda", LOCAL_RANK) + dist.init_process_group( + 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) - 'lrf': (False, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) - 'momentum': (False, 0.6, 0.98), # SGD momentum/Adam beta1 - 'weight_decay': (False, 0.0, 0.001), # optimizer weight decay - 'warmup_epochs': (False, 0.0, 5.0), # warmup epochs (fractions ok) - 'warmup_momentum': (False, 0.0, 0.95), # warmup initial momentum - 'warmup_bias_lr': (False, 0.0, 0.2), # warmup initial bias lr - 'box': (False, 0.02, 0.2), # box loss gain - 'cls': (False, 0.2, 4.0), # cls loss gain - 'cls_pw': (False, 0.5, 2.0), # cls BCELoss positive_weight - 'obj': (False, 0.2, 4.0), # obj loss gain (scale with pixels) - 'obj_pw': (False, 0.5, 2.0), # obj BCELoss positive_weight - 'iou_t': (False, 0.1, 0.7), # IoU training threshold - 'anchor_t': (False, 2.0, 8.0), # anchor-multiple threshold - '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) - '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_v': (True, 0.0, 0.9), # image HSV-Value augmentation (fraction) - 'degrees': (True, 0.0, 45.0), # image rotation (+/- deg) - 'translate': (True, 0.0, 0.9), # image translation (+/- fraction) - 'scale': (True, 0.0, 0.9), # image scale (+/- gain) - 'shear': (True, 0.0, 10.0), # image shear (+/- deg) - 'perspective': (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 - 'flipud': (True, 0.0, 1.0), # image flip up-down (probability) - 'fliplr': (True, 0.0, 1.0), # image flip left-right (probability) - 'mosaic': (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) + "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) + "momentum": (False, 0.6, 0.98), # SGD momentum/Adam beta1 + "weight_decay": (False, 0.0, 0.001), # optimizer weight decay + "warmup_epochs": (False, 0.0, 5.0), # warmup epochs (fractions ok) + "warmup_momentum": (False, 0.0, 0.95), # warmup initial momentum + "warmup_bias_lr": (False, 0.0, 0.2), # warmup initial bias lr + "box": (False, 0.02, 0.2), # box loss gain + "cls": (False, 0.2, 4.0), # cls loss gain + "cls_pw": (False, 0.5, 2.0), # cls BCELoss positive_weight + "obj": (False, 0.2, 4.0), # obj loss gain (scale with pixels) + "obj_pw": (False, 0.5, 2.0), # obj BCELoss positive_weight + "iou_t": (False, 0.1, 0.7), # IoU training threshold + "anchor_t": (False, 2.0, 8.0), # anchor-multiple threshold + "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) + "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_v": (True, 0.0, 0.9), # image HSV-Value augmentation (fraction) + "degrees": (True, 0.0, 45.0), # image rotation (+/- deg) + "translate": (True, 0.0, 0.9), # image translation (+/- fraction) + "scale": (True, 0.0, 0.9), # image scale (+/- gain) + "shear": (True, 0.0, 10.0), # image shear (+/- deg) + "perspective": (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 + "flipud": (True, 0.0, 1.0), # image flip up-down (probability) + "fliplr": (True, 0.0, 1.0), # image flip left-right (probability) + "mosaic": (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) # 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 - hyp['anchors'] = 3 + if "anchors" not in hyp: # anchors commented in hyp.yaml + 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([ - 'gsutil', - 'cp', - f'gs://{opt.bucket}/evolve.csv', - str(evolve_csv), ]) + subprocess.run( + [ + "gsutil", + "cp", + 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!' - with open(ROOT / opt.resume_evolve, errors='ignore') as f: + assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!" + 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', - 'val/box_loss', 'val/obj_loss', 'val/cls_loss') + keys = ( + "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), - int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10)))) + tournament_size = max( + 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, - min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve))) + crossover_rate = max( + 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, - min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve))) + mutation_rate = max( + 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' - f"Results saved to {colorstr('bold', save_dir)}\n" - f'Usage example: $ python train.py --hyp {evolve_yaml}') + 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}' + ) 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) diff --git a/utils/__init__.py b/utils/__init__.py index 4c7379c87466..eff756e2b90e 100644 --- a/utils/__init__.py +++ b/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): - os.system('pip uninstall -y wandb') # eliminate unexpected account creation prompt with infinite hang + if check_requirements("wandb", install=False): + 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 diff --git a/utils/activations.py b/utils/activations.py index e4d4bbde5ec8..616002f06a73 100644 --- a/utils/activations.py +++ b/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 @@ def forward(x): class MemoryEfficientMish(nn.Module): # Mish activation memory-efficient class F(torch.autograd.Function): - @staticmethod def forward(ctx, x): ctx.save_for_backward(x) @@ -62,7 +59,7 @@ def forward(self, x): class AconC(nn.Module): - r""" ACON activation (activate or not) + r"""ACON activation (activate or not) AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter according to "Activate or Not: Learning Customized Activation" . """ @@ -79,7 +76,7 @@ def forward(self, x): class MetaAconC(nn.Module): - r""" ACON activation (activate or not) + r"""ACON activation (activate or not) MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network according to "Activate or Not: Learning Customized Activation" . """ diff --git a/utils/augmentations.py b/utils/augmentations.py index 1e609303e209..b3b9524320d0 100644 --- a/utils/augmentations.py +++ b/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 @@ def __init__(self, size=640): A.CLAHE(p=0.01), A.RandomBrightnessContrast(p=0.0), A.RandomGamma(p=0.0), - A.ImageCompression(quality_lower=75, p=0.0)] # transforms - self.transform = A.Compose(T, bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'])) + A.ImageCompression(quality_lower=75, p=0.0), + ] # 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 @@ -97,7 +97,7 @@ def replicate(im, labels): boxes = labels[:, 1:].astype(int) x1, y1, x2, y2 = boxes.T s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels) - for i in s.argsort()[:round(s.size * 0.5)]: # smallest indices + for i in s.argsort()[: round(s.size * 0.5)]: # smallest indices x1b, y1b, x2b, y2b = boxes[i] bh, bw = y2b - y1b, x2b - x1b yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y @@ -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, - targets=(), - segments=(), - degrees=10, - translate=.1, - scale=.1, - shear=10, - perspective=0.0, - border=(0, 0)): +def random_perspective( + im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.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] @@ -303,50 +297,52 @@ def box_candidates(box1, box2, wh_thr=2, ar_thr=100, area_thr=0.1, eps=1e-16): def classify_albumentations( - augment=True, - size=224, - scale=(0.08, 1.0), - ratio=(0.75, 1.0 / 0.75), # 0.75, 1.33 - hflip=0.5, - vflip=0.0, - jitter=0.4, - mean=IMAGENET_MEAN, - std=IMAGENET_STD, - auto_aug=False): + augment=True, + size=224, + scale=(0.08, 1.0), + ratio=(0.75, 1.0 / 0.75), # 0.75, 1.33 + hflip=0.5, + vflip=0.0, + jitter=0.4, + mean=IMAGENET_MEAN, + std=IMAGENET_STD, + 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)] if vflip > 0: T += [A.VerticalFlip(p=vflip)] if jitter > 0: - color_jitter = (float(jitter), ) * 3 # repeat value for brightness, contrast, satuaration, 0 hue + color_jitter = (float(jitter),) * 3 # repeat value for brightness, contrast, satuaration, 0 hue T += [A.ColorJitter(*color_jitter, 0)] 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)]) @@ -366,7 +362,7 @@ def __call__(self, im): # im = np.array HWC hs, ws = (math.ceil(x / self.stride) * self.stride for x in (h, w)) if self.auto else self.h, self.w top, left = round((hs - h) / 2 - 0.1), round((ws - w) / 2 - 0.1) im_out = np.full((self.h, self.w, 3), 114, dtype=im.dtype) - im_out[top:top + h, left:left + w] = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR) + im_out[top : top + h, left : left + w] = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR) return im_out @@ -380,7 +376,7 @@ def __call__(self, im): # im = np.array HWC imh, imw = im.shape[:2] m = min(imh, imw) # min dimension top, left = (imh - m) // 2, (imw - m) // 2 - return cv2.resize(im[top:top + m, left:left + m], (self.w, self.h), interpolation=cv2.INTER_LINEAR) + return cv2.resize(im[top : top + m, left : left + m], (self.w, self.h), interpolation=cv2.INTER_LINEAR) class ToTensor: diff --git a/utils/autoanchor.py b/utils/autoanchor.py index 4c11ab3decec..89e4d97fdcd5 100644 --- a/utils/autoanchor.py +++ b/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 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 metric(k): # compute metric 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,28 +56,29 @@ def metric(k): # compute metric 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): - """ Creates kmeans-evolved anchors from training dataset +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. - Arguments: - dataset: path to data.yaml, or a loaded dataset - n: number of anchors - img_size: image size used for training - thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 - gen: generations to evolve anchors using genetic algorithm - verbose: print all results + Arguments: + dataset: path to data.yaml, or a loaded dataset + n: number of anchors + img_size: image size used for training + thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 + gen: generations to evolve anchors using genetic algorithm + verbose: print all results - Return: - k: kmeans evolved anchors + Return: + k: kmeans evolved anchors - Usage: - from utils.autoanchor import *; _ = kmean_anchors() + Usage: + from utils.autoanchor import *; _ = kmean_anchors() """ from scipy.cluster.vq import kmeans @@ -100,20 +99,23 @@ def print_results(k, verbose=True): 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' \ - 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: ' + s = ( + f"{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n" + 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 print_results(k, verbose=True): # 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 print_results(k, verbose=True): 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) diff --git a/utils/autobatch.py b/utils/autobatch.py index aa763b888462..396dbed1dda4 100644 --- a/utils/autobatch.py +++ b/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: ') - LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}') + prefix = colorstr("AutoBatch: ") + LOGGER.info(f"{prefix}Computing optimal batch size for --imgsz {imgsz}") device = next(model.parameters()).device # get model device - if device.type == 'cpu': - LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}') + if device.type == "cpu": + 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,11 +51,11 @@ 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] - p = np.polyfit(batch_sizes[:len(y)], y, deg=1) # first degree polynomial fit + p = np.polyfit(batch_sizes[: len(y)], y, deg=1) # first degree polynomial fit b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size) if None in results: # some sizes failed i = results.index(None) # first fail index @@ -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 diff --git a/utils/aws/resume.py b/utils/aws/resume.py index b21731c979a1..4525ba96749a 100644 --- a/utils/aws/resume.py +++ b/utils/aws/resume.py @@ -14,27 +14,27 @@ sys.path.append(str(ROOT)) # add ROOT to PATH port = 0 # --master_port -path = Path('').resolve() -for last in path.rglob('*/**/last.pt'): +path = Path("").resolve() +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) diff --git a/utils/callbacks.py b/utils/callbacks.py index c90fa824cdb4..ab7befdb0b3b 100644 --- a/utils/callbacks.py +++ b/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_end': [], - 'on_train_start': [], - 'on_train_epoch_start': [], - 'on_train_batch_start': [], - 'optimizer_step': [], - 'on_before_zero_grad': [], - 'on_train_batch_end': [], - 'on_train_epoch_end': [], - 'on_val_start': [], - 'on_val_batch_start': [], - 'on_val_image_end': [], - 'on_val_batch_end': [], - 'on_val_end': [], - 'on_fit_epoch_end': [], # fit = train + val - 'on_model_save': [], - 'on_train_end': [], - 'on_params_update': [], - 'teardown': [], } + "on_pretrain_routine_start": [], + "on_pretrain_routine_end": [], + "on_train_start": [], + "on_train_epoch_start": [], + "on_train_batch_start": [], + "optimizer_step": [], + "on_before_zero_grad": [], + "on_train_batch_end": [], + "on_train_epoch_end": [], + "on_val_start": [], + "on_val_batch_start": [], + "on_val_image_end": [], + "on_val_batch_end": [], + "on_val_end": [], + "on_fit_epoch_end": [], # fit = train + val + "on_model_save": [], + "on_train_end": [], + "on_params_update": [], + "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 @@ def register_action(self, hook, name='', callback=None): """ 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 @@ def get_registered_actions(self, hook=None): 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 @@ def run(self, hook, *args, thread=False, **kwargs): 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) diff --git a/utils/dataloaders.py b/utils/dataloaders.py index d422ef0711cb..c821e917ed38 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -1,7 +1,5 @@ # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license -""" -Dataloaders and dataset utils -""" +"""Dataloaders and dataset utils.""" import contextlib import glob @@ -28,25 +26,49 @@ from torch.utils.data import DataLoader, Dataset, dataloader, distributed from tqdm import tqdm -from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste, - letterbox, mixup, random_perspective) -from utils.general import (DATASETS_DIR, LOGGER, NUM_THREADS, TQDM_BAR_FORMAT, check_dataset, check_requirements, - check_yaml, clean_str, cv2, is_colab, is_kaggle, segments2boxes, unzip_file, xyn2xy, - xywh2xyxy, xywhn2xyxy, xyxy2xywhn) +from utils.augmentations import ( + Albumentations, + augment_hsv, + classify_albumentations, + classify_transforms, + copy_paste, + letterbox, + mixup, + random_perspective, +) +from utils.general import ( + DATASETS_DIR, + LOGGER, + NUM_THREADS, + TQDM_BAR_FORMAT, + check_dataset, + check_requirements, + check_yaml, + clean_str, + cv2, + is_colab, + is_kaggle, + segments2boxes, + unzip_file, + xyn2xy, + xywh2xyxy, + xywhn2xyxy, + xyxy2xywhn, +) from utils.torch_utils import torch_distributed_zero_first # Parameters -HELP_URL = 'See https://docs.ultralytics.com/yolov5/tutorials/train_custom_data' -IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm' # include image suffixes -VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' # include video suffixes -LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html -RANK = int(os.getenv('RANK', -1)) -WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) -PIN_MEMORY = str(os.getenv('PIN_MEMORY', True)).lower() == 'true' # global pin_memory for dataloaders +HELP_URL = "See https://docs.ultralytics.com/yolov5/tutorials/train_custom_data" +IMG_FORMATS = "bmp", "dng", "jpeg", "jpg", "mpo", "png", "tif", "tiff", "webp", "pfm" # include image suffixes +VID_FORMATS = "asf", "avi", "gif", "m4v", "mkv", "mov", "mp4", "mpeg", "mpg", "ts", "wmv" # include video suffixes +LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html +RANK = int(os.getenv("RANK", -1)) +WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1)) +PIN_MEMORY = str(os.getenv("PIN_MEMORY", True)).lower() == "true" # global pin_memory for dataloaders # Get orientation exif tag for orientation in ExifTags.TAGS.keys(): - if ExifTags.TAGS[orientation] == 'Orientation': + if ExifTags.TAGS[orientation] == "Orientation": break @@ -54,7 +76,7 @@ def get_hash(paths): # Returns a single hash value of a list of paths (files or dirs) size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes h = hashlib.sha256(str(size).encode()) # hash sizes - h.update(''.join(paths).encode()) # hash paths + h.update("".join(paths).encode()) # hash paths return h.hexdigest() # return hash @@ -86,17 +108,18 @@ def exif_transpose(image): 5: Image.TRANSPOSE, 6: Image.ROTATE_270, 7: Image.TRANSVERSE, - 8: Image.ROTATE_90}.get(orientation) + 8: Image.ROTATE_90, + }.get(orientation) if method is not None: image = image.transpose(method) del exif[0x0112] - image.info['exif'] = exif.tobytes() + image.info["exif"] = exif.tobytes() return image def seed_worker(worker_id): # Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader - worker_seed = torch.initial_seed() % 2 ** 32 + worker_seed = torch.initial_seed() % 2**32 np.random.seed(worker_seed) random.seed(worker_seed) @@ -104,7 +127,6 @@ def seed_worker(worker_id): # Inherit from DistributedSampler and override iterator # https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py class SmartDistributedSampler(distributed.DistributedSampler): - def __iter__(self): # deterministically shuffle based on epoch and seed g = torch.Generator() @@ -118,7 +140,7 @@ def __iter__(self): idx = idx.tolist() if self.drop_last: - idx = idx[:self.num_samples] + idx = idx[: self.num_samples] else: padding_size = self.num_samples - len(idx) if padding_size <= len(idx): @@ -129,25 +151,27 @@ def __iter__(self): return iter(idx) -def create_dataloader(path, - imgsz, - batch_size, - stride, - single_cls=False, - hyp=None, - augment=False, - cache=False, - pad=0.0, - rect=False, - rank=-1, - workers=8, - image_weights=False, - quad=False, - prefix='', - shuffle=False, - seed=0): +def create_dataloader( + path, + imgsz, + batch_size, + stride, + single_cls=False, + hyp=None, + augment=False, + cache=False, + pad=0.0, + rect=False, + rank=-1, + workers=8, + image_weights=False, + quad=False, + prefix="", + shuffle=False, + 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 = LoadImagesAndLabels( @@ -163,7 +187,8 @@ def create_dataloader(path, pad=pad, image_weights=image_weights, prefix=prefix, - rank=rank) + rank=rank, + ) batch_size = min(batch_size, len(dataset)) nd = torch.cuda.device_count() # number of CUDA devices @@ -172,26 +197,29 @@ def create_dataloader(path, loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates generator = torch.Generator() generator.manual_seed(6148914691236517205 + seed + RANK) - return loader(dataset, - batch_size=batch_size, - shuffle=shuffle and sampler is None, - num_workers=nw, - sampler=sampler, - pin_memory=PIN_MEMORY, - collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn, - worker_init_fn=seed_worker, - generator=generator), dataset + return loader( + dataset, + batch_size=batch_size, + shuffle=shuffle and sampler is None, + num_workers=nw, + sampler=sampler, + pin_memory=PIN_MEMORY, + collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn, + worker_init_fn=seed_worker, + generator=generator, + ), dataset class InfiniteDataLoader(dataloader.DataLoader): - """ Dataloader that reuses workers + """ + Dataloader that reuses workers. Uses same syntax as vanilla DataLoader """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler)) + object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler)) self.iterator = super().__iter__() def __len__(self): @@ -203,7 +231,8 @@ def __iter__(self): class _RepeatSampler: - """ Sampler that repeats forever + """ + Sampler that repeats forever. Args: sampler (Sampler) @@ -221,7 +250,7 @@ class LoadScreenshots: # YOLOv5 screenshot dataloader, i.e. `python detect.py --source "screen 0 100 100 512 256"` def __init__(self, source, img_size=640, stride=32, auto=True, transforms=None): # source = [screen_number left top width height] (pixels) - check_requirements('mss') + check_requirements("mss") import mss source, *params = source.split() @@ -236,17 +265,17 @@ def __init__(self, source, img_size=640, stride=32, auto=True, transforms=None): self.stride = stride self.transforms = transforms self.auto = auto - self.mode = 'stream' + self.mode = "stream" self.frame = 0 self.sct = mss.mss() # Parse monitor shape monitor = self.sct.monitors[self.screen] - self.top = monitor['top'] if top is None else (monitor['top'] + top) - self.left = monitor['left'] if left is None else (monitor['left'] + left) - self.width = width or monitor['width'] - self.height = height or monitor['height'] - self.monitor = {'left': self.left, 'top': self.top, 'width': self.width, 'height': self.height} + self.top = monitor["top"] if top is None else (monitor["top"] + top) + self.left = monitor["left"] if left is None else (monitor["left"] + left) + self.width = width or monitor["width"] + self.height = height or monitor["height"] + self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height} def __iter__(self): return self @@ -254,7 +283,7 @@ def __iter__(self): def __next__(self): # mss screen capture: get raw pixels from the screen as np array im0 = np.array(self.sct.grab(self.monitor))[:, :, :3] # [:, :, :3] BGRA to BGR - s = f'screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: ' + s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: " if self.transforms: im = self.transforms(im0) # transforms @@ -269,22 +298,22 @@ def __next__(self): class LoadImages: # YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4` def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vid_stride=1): - if isinstance(path, str) and Path(path).suffix == '.txt': # *.txt file with img/vid/dir on each line + if isinstance(path, str) and Path(path).suffix == ".txt": # *.txt file with img/vid/dir on each line path = Path(path).read_text().rsplit() files = [] for p in sorted(path) if isinstance(path, (list, tuple)) else [path]: p = str(Path(p).resolve()) - if '*' in p: + if "*" in p: files.extend(sorted(glob.glob(p, recursive=True))) # glob elif os.path.isdir(p): - files.extend(sorted(glob.glob(os.path.join(p, '*.*')))) # dir + files.extend(sorted(glob.glob(os.path.join(p, "*.*")))) # dir elif os.path.isfile(p): files.append(p) # files else: - raise FileNotFoundError(f'{p} does not exist') + raise FileNotFoundError(f"{p} does not exist") - images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS] - videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS] + images = [x for x in files if x.split(".")[-1].lower() in IMG_FORMATS] + videos = [x for x in files if x.split(".")[-1].lower() in VID_FORMATS] ni, nv = len(images), len(videos) self.img_size = img_size @@ -292,7 +321,7 @@ def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vi self.files = images + videos self.nf = ni + nv # number of files self.video_flag = [False] * ni + [True] * nv - self.mode = 'image' + self.mode = "image" self.auto = auto self.transforms = transforms # optional self.vid_stride = vid_stride # video frame-rate stride @@ -300,8 +329,10 @@ def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vi self._new_video(videos[0]) # new video else: self.cap = None - assert self.nf > 0, f'No images or videos found in {p}. ' \ - f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}' + assert self.nf > 0, ( + f"No images or videos found in {p}. " + f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}" + ) def __iter__(self): self.count = 0 @@ -314,7 +345,7 @@ def __next__(self): if self.video_flag[self.count]: # Read video - self.mode = 'video' + self.mode = "video" for _ in range(self.vid_stride): self.cap.grab() ret_val, im0 = self.cap.retrieve() @@ -329,14 +360,14 @@ def __next__(self): self.frame += 1 # im0 = self._cv2_rotate(im0) # for use if cv2 autorotation is False - s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ' + s = f"video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: " else: # Read image self.count += 1 im0 = cv2.imread(path) # BGR - assert im0 is not None, f'Image Not Found {path}' - s = f'image {self.count}/{self.nf} {path}: ' + assert im0 is not None, f"Image Not Found {path}" + s = f"image {self.count}/{self.nf} {path}: " if self.transforms: im = self.transforms(im0) # transforms @@ -371,9 +402,9 @@ def __len__(self): class LoadStreams: # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams` - def __init__(self, sources='file.streams', img_size=640, stride=32, auto=True, transforms=None, vid_stride=1): + def __init__(self, sources="file.streams", img_size=640, stride=32, auto=True, transforms=None, vid_stride=1): torch.backends.cudnn.benchmark = True # faster for fixed-size inference - self.mode = 'stream' + self.mode = "stream" self.img_size = img_size self.stride = stride self.vid_stride = vid_stride # video frame-rate stride @@ -383,29 +414,30 @@ def __init__(self, sources='file.streams', img_size=640, stride=32, auto=True, t self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n for i, s in enumerate(sources): # index, source # Start thread to read frames from video stream - st = f'{i + 1}/{n}: {s}... ' - if urlparse(s).hostname in ('www.youtube.com', 'youtube.com', 'youtu.be'): # if source is YouTube video + st = f"{i + 1}/{n}: {s}... " + if urlparse(s).hostname in ("www.youtube.com", "youtube.com", "youtu.be"): # if source is YouTube video # YouTube format i.e. 'https://www.youtube.com/watch?v=Zgi9g1ksQHc' or 'https://youtu.be/LNwODJXcvt4' - check_requirements(('pafy', 'youtube_dl==2020.12.2')) + check_requirements(("pafy", "youtube_dl==2020.12.2")) import pafy - s = pafy.new(s).getbest(preftype='mp4').url # YouTube URL + + s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL s = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam if s == 0: - assert not is_colab(), '--source 0 webcam unsupported on Colab. Rerun command in a local environment.' - assert not is_kaggle(), '--source 0 webcam unsupported on Kaggle. Rerun command in a local environment.' + assert not is_colab(), "--source 0 webcam unsupported on Colab. Rerun command in a local environment." + assert not is_kaggle(), "--source 0 webcam unsupported on Kaggle. Rerun command in a local environment." cap = cv2.VideoCapture(s) - assert cap.isOpened(), f'{st}Failed to open {s}' + assert cap.isOpened(), f"{st}Failed to open {s}" w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) # warning: may return 0 or nan - self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback + self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float("inf") # infinite stream fallback self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30 # 30 FPS fallback _, self.imgs[i] = cap.read() # guarantee first frame self.threads[i] = Thread(target=self.update, args=([i, cap, s]), daemon=True) - LOGGER.info(f'{st} Success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)') + LOGGER.info(f"{st} Success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)") self.threads[i].start() - LOGGER.info('') # newline + LOGGER.info("") # newline # check for common shapes s = np.stack([letterbox(x, img_size, stride=stride, auto=auto)[0].shape for x in self.imgs]) @@ -413,7 +445,7 @@ def __init__(self, sources='file.streams', img_size=640, stride=32, auto=True, t self.auto = auto and self.rect self.transforms = transforms # optional if not self.rect: - LOGGER.warning('WARNING ⚠️ Stream shapes differ. For optimal performance supply similarly-shaped streams.') + LOGGER.warning("WARNING ⚠️ Stream shapes differ. For optimal performance supply similarly-shaped streams.") def update(self, i, cap, stream): # Read stream `i` frames in daemon thread @@ -426,7 +458,7 @@ def update(self, i, cap, stream): if success: self.imgs[i] = im else: - LOGGER.warning('WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.') + LOGGER.warning("WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.") self.imgs[i] = np.zeros_like(self.imgs[i]) cap.open(stream) # re-open stream if signal was lost time.sleep(0.0) # wait time @@ -437,7 +469,7 @@ def __iter__(self): def __next__(self): self.count += 1 - if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit + if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord("q"): # q to quit cv2.destroyAllWindows() raise StopIteration @@ -449,7 +481,7 @@ def __next__(self): im = im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW im = np.ascontiguousarray(im) # contiguous - return self.sources, im, im0, None, '' + return self.sources, im, im0, None, "" def __len__(self): return len(self.sources) # 1E12 frames = 32 streams at 30 FPS for 30 years @@ -457,8 +489,8 @@ def __len__(self): def img2label_paths(img_paths): # Define label paths as a function of image paths - sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}' # /images/, /labels/ substrings - return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths] + sa, sb = f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}" # /images/, /labels/ substrings + return [sb.join(x.rsplit(sa, 1)).rsplit(".", 1)[0] + ".txt" for x in img_paths] class LoadImagesAndLabels(Dataset): @@ -466,22 +498,24 @@ class LoadImagesAndLabels(Dataset): cache_version = 0.6 # dataset labels *.cache version rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4] - def __init__(self, - path, - img_size=640, - batch_size=16, - augment=False, - hyp=None, - rect=False, - image_weights=False, - cache_images=False, - single_cls=False, - stride=32, - pad=0.0, - min_items=0, - prefix='', - rank=-1, - seed=0): + def __init__( + self, + path, + img_size=640, + batch_size=16, + augment=False, + hyp=None, + rect=False, + image_weights=False, + cache_images=False, + single_cls=False, + stride=32, + pad=0.0, + min_items=0, + prefix="", + rank=-1, + seed=0, + ): self.img_size = img_size self.augment = augment self.hyp = hyp @@ -498,46 +532,46 @@ def __init__(self, for p in path if isinstance(path, list) else [path]: p = Path(p) # os-agnostic if p.is_dir(): # dir - f += glob.glob(str(p / '**' / '*.*'), recursive=True) + f += glob.glob(str(p / "**" / "*.*"), recursive=True) # f = list(p.rglob('*.*')) # pathlib elif p.is_file(): # file with open(p) as t: t = t.read().strip().splitlines() parent = str(p.parent) + os.sep - f += [x.replace('./', parent, 1) if x.startswith('./') else x for x in t] # to global path + f += [x.replace("./", parent, 1) if x.startswith("./") else x for x in t] # to global path # f += [p.parent / x.lstrip(os.sep) for x in t] # to global path (pathlib) else: - raise FileNotFoundError(f'{prefix}{p} does not exist') - self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS) + raise FileNotFoundError(f"{prefix}{p} does not exist") + self.im_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in IMG_FORMATS) # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib - assert self.im_files, f'{prefix}No images found' + assert self.im_files, f"{prefix}No images found" except Exception as e: - raise Exception(f'{prefix}Error loading data from {path}: {e}\n{HELP_URL}') from e + raise Exception(f"{prefix}Error loading data from {path}: {e}\n{HELP_URL}") from e # Check cache self.label_files = img2label_paths(self.im_files) # labels - cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') + cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix(".cache") try: cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict - assert cache['version'] == self.cache_version # matches current version - assert cache['hash'] == get_hash(self.label_files + self.im_files) # identical hash + assert cache["version"] == self.cache_version # matches current version + assert cache["hash"] == get_hash(self.label_files + self.im_files) # identical hash except Exception: cache, exists = self.cache_labels(cache_path, prefix), False # run cache ops # Display cache - nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupt, total + nf, nm, ne, nc, n = cache.pop("results") # found, missing, empty, corrupt, total if exists and LOCAL_RANK in {-1, 0}: - d = f'Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt' + d = f"Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt" tqdm(None, desc=prefix + d, total=n, initial=n, bar_format=TQDM_BAR_FORMAT) # display cache results - if cache['msgs']: - LOGGER.info('\n'.join(cache['msgs'])) # display warnings - assert nf > 0 or not augment, f'{prefix}No labels found in {cache_path}, can not start training. {HELP_URL}' + if cache["msgs"]: + LOGGER.info("\n".join(cache["msgs"])) # display warnings + assert nf > 0 or not augment, f"{prefix}No labels found in {cache_path}, can not start training. {HELP_URL}" # Read cache - [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items + [cache.pop(k) for k in ("hash", "version", "msgs")] # remove items labels, shapes, self.segments = zip(*cache.values()) nl = len(np.concatenate(labels, 0)) # number of labels - assert nl > 0 or not augment, f'{prefix}All labels empty in {cache_path}, can not start training. {HELP_URL}' + assert nl > 0 or not augment, f"{prefix}All labels empty in {cache_path}, can not start training. {HELP_URL}" self.labels = list(labels) self.shapes = np.array(shapes) self.im_files = list(cache.keys()) # update @@ -546,7 +580,7 @@ def __init__(self, # Filter images if min_items: include = np.array([len(x) >= min_items for x in self.labels]).nonzero()[0].astype(int) - LOGGER.info(f'{prefix}{n - len(include)}/{n} images filtered from dataset') + LOGGER.info(f"{prefix}{n - len(include)}/{n} images filtered from dataset") self.im_files = [self.im_files[i] for i in include] self.label_files = [self.label_files[i] for i in include] self.labels = [self.labels[i] for i in include] @@ -603,52 +637,56 @@ def __init__(self, self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(int) * stride # Cache images into RAM/disk for faster training - if cache_images == 'ram' and not self.check_cache_ram(prefix=prefix): + if cache_images == "ram" and not self.check_cache_ram(prefix=prefix): cache_images = False self.ims = [None] * n - self.npy_files = [Path(f).with_suffix('.npy') for f in self.im_files] + self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files] if cache_images: b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes self.im_hw0, self.im_hw = [None] * n, [None] * n - fcn = self.cache_images_to_disk if cache_images == 'disk' else self.load_image + fcn = self.cache_images_to_disk if cache_images == "disk" else self.load_image results = ThreadPool(NUM_THREADS).imap(lambda i: (i, fcn(i)), self.indices) pbar = tqdm(results, total=len(self.indices), bar_format=TQDM_BAR_FORMAT, disable=LOCAL_RANK > 0) for i, x in pbar: - if cache_images == 'disk': + if cache_images == "disk": b += self.npy_files[i].stat().st_size else: # 'ram' self.ims[i], self.im_hw0[i], self.im_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i) b += self.ims[i].nbytes * WORLD_SIZE - pbar.desc = f'{prefix}Caching images ({b / gb:.1f}GB {cache_images})' + pbar.desc = f"{prefix}Caching images ({b / gb:.1f}GB {cache_images})" pbar.close() - def check_cache_ram(self, safety_margin=0.1, prefix=''): + def check_cache_ram(self, safety_margin=0.1, prefix=""): # Check image caching requirements vs available memory b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes n = min(self.n, 30) # extrapolate from 30 random images for _ in range(n): im = cv2.imread(random.choice(self.im_files)) # sample image ratio = self.img_size / max(im.shape[0], im.shape[1]) # max(h, w) # ratio - b += im.nbytes * ratio ** 2 + b += im.nbytes * ratio**2 mem_required = b * self.n / n # GB required to cache dataset into RAM mem = psutil.virtual_memory() cache = mem_required * (1 + safety_margin) < mem.available # to cache or not to cache, that is the question if not cache: - LOGGER.info(f'{prefix}{mem_required / gb:.1f}GB RAM required, ' - f'{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, ' - f"{'caching images ✅' if cache else 'not caching images ⚠️'}") + LOGGER.info( + f'{prefix}{mem_required / gb:.1f}GB RAM required, ' + f'{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, ' + f"{'caching images ✅' if cache else 'not caching images ⚠️'}" + ) return cache - def cache_labels(self, path=Path('./labels.cache'), prefix=''): + def cache_labels(self, path=Path("./labels.cache"), prefix=""): # Cache dataset labels, check images and read shapes x = {} # dict nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages - desc = f'{prefix}Scanning {path.parent / path.stem}...' + desc = f"{prefix}Scanning {path.parent / path.stem}..." with Pool(NUM_THREADS) as pool: - pbar = tqdm(pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))), - desc=desc, - total=len(self.im_files), - bar_format=TQDM_BAR_FORMAT) + pbar = tqdm( + pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))), + desc=desc, + total=len(self.im_files), + bar_format=TQDM_BAR_FORMAT, + ) for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar: nm += nm_f nf += nf_f @@ -658,23 +696,23 @@ def cache_labels(self, path=Path('./labels.cache'), prefix=''): x[im_file] = [lb, shape, segments] if msg: msgs.append(msg) - pbar.desc = f'{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt' + pbar.desc = f"{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt" pbar.close() if msgs: - LOGGER.info('\n'.join(msgs)) + LOGGER.info("\n".join(msgs)) if nf == 0: - LOGGER.warning(f'{prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}') - x['hash'] = get_hash(self.label_files + self.im_files) - x['results'] = nf, nm, ne, nc, len(self.im_files) - x['msgs'] = msgs # warnings - x['version'] = self.cache_version # cache version + LOGGER.warning(f"{prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}") + x["hash"] = get_hash(self.label_files + self.im_files) + x["results"] = nf, nm, ne, nc, len(self.im_files) + x["msgs"] = msgs # warnings + x["version"] = self.cache_version # cache version try: np.save(path, x) # save cache for next time - path.with_suffix('.cache.npy').rename(path) # remove .npy suffix - LOGGER.info(f'{prefix}New cache created: {path}') + path.with_suffix(".cache.npy").rename(path) # remove .npy suffix + LOGGER.info(f"{prefix}New cache created: {path}") except Exception as e: - LOGGER.warning(f'{prefix}WARNING ⚠️ Cache directory {path.parent} is not writeable: {e}') # not writeable + LOGGER.warning(f"{prefix}WARNING ⚠️ Cache directory {path.parent} is not writeable: {e}") # not writeable return x def __len__(self): @@ -690,14 +728,14 @@ def __getitem__(self, index): 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"] if mosaic: # Load mosaic img, labels = self.load_mosaic(index) shapes = None # MixUp augmentation - if random.random() < hyp['mixup']: + if random.random() < hyp["mixup"]: img, labels = mixup(img, labels, *self.load_mosaic(random.choice(self.indices))) else: @@ -714,17 +752,19 @@ def __getitem__(self, index): labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1]) if self.augment: - img, labels = random_perspective(img, - labels, - degrees=hyp['degrees'], - translate=hyp['translate'], - scale=hyp['scale'], - shear=hyp['shear'], - perspective=hyp['perspective']) + img, labels = random_perspective( + img, + labels, + degrees=hyp["degrees"], + translate=hyp["translate"], + scale=hyp["scale"], + shear=hyp["shear"], + 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) + labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3) if self.augment: # Albumentations @@ -732,16 +772,16 @@ def __getitem__(self, index): 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] # Flip left-right - if random.random() < hyp['fliplr']: + if random.random() < hyp["fliplr"]: img = np.fliplr(img) if nl: labels[:, 1] = 1 - labels[:, 1] @@ -762,13 +802,17 @@ def __getitem__(self, index): def load_image(self, i): # Loads 1 image from dataset index 'i', returns (im, original hw, resized hw) - im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i], + im, f, fn = ( + self.ims[i], + self.im_files[i], + self.npy_files[i], + ) if im is None: # not cached in RAM if fn.exists(): # load npy im = np.load(fn) else: # read image im = cv2.imread(f) # BGR - assert im is not None, f'Image Not Found {f}' + assert im is not None, f"Image Not Found {f}" h0, w0 = im.shape[:2] # orig hw r = self.img_size / max(h0, w0) # ratio if r != 1: # if sizes are not equal @@ -828,16 +872,18 @@ def load_mosaic(self, index): # img4, labels4 = replicate(img4, labels4) # replicate # Augment - img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste']) - img4, labels4 = random_perspective(img4, - labels4, - segments4, - degrees=self.hyp['degrees'], - translate=self.hyp['translate'], - scale=self.hyp['scale'], - shear=self.hyp['shear'], - perspective=self.hyp['perspective'], - border=self.mosaic_border) # border to remove + img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"]) + img4, labels4 = random_perspective( + img4, + labels4, + segments4, + degrees=self.hyp["degrees"], + translate=self.hyp["translate"], + scale=self.hyp["scale"], + shear=self.hyp["shear"], + perspective=self.hyp["perspective"], + border=self.mosaic_border, + ) # border to remove return img4, labels4 @@ -886,12 +932,12 @@ def load_mosaic9(self, index): segments9.extend(segments) # Image - img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax] + img9[y1:y2, x1:x2] = img[y1 - pady :, x1 - padx :] # img9[ymin:ymax, xmin:xmax] hp, wp = h, w # height, width previous # Offset yc, xc = (int(random.uniform(0, s)) for _ in self.mosaic_border) # mosaic center x, y - img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s] + img9 = img9[yc : yc + 2 * s, xc : xc + 2 * s] # Concat/clip labels labels9 = np.concatenate(labels9, 0) @@ -905,16 +951,18 @@ def load_mosaic9(self, index): # img9, labels9 = replicate(img9, labels9) # replicate # Augment - img9, labels9, segments9 = copy_paste(img9, labels9, segments9, p=self.hyp['copy_paste']) - img9, labels9 = random_perspective(img9, - labels9, - segments9, - degrees=self.hyp['degrees'], - translate=self.hyp['translate'], - scale=self.hyp['scale'], - shear=self.hyp['shear'], - perspective=self.hyp['perspective'], - border=self.mosaic_border) # border to remove + img9, labels9, segments9 = copy_paste(img9, labels9, segments9, p=self.hyp["copy_paste"]) + img9, labels9 = random_perspective( + img9, + labels9, + segments9, + degrees=self.hyp["degrees"], + translate=self.hyp["translate"], + scale=self.hyp["scale"], + shear=self.hyp["shear"], + perspective=self.hyp["perspective"], + border=self.mosaic_border, + ) # border to remove return img9, labels9 @@ -937,8 +985,9 @@ def collate_fn4(batch): for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW i *= 4 if random.random() < 0.5: - im1 = F.interpolate(im[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear', - align_corners=False)[0].type(im[i].type()) + im1 = F.interpolate(im[i].unsqueeze(0).float(), scale_factor=2.0, mode="bilinear", align_corners=False)[ + 0 + ].type(im[i].type()) lb = label[i] else: im1 = torch.cat((torch.cat((im[i], im[i + 1]), 1), torch.cat((im[i + 2], im[i + 3]), 1)), 2) @@ -953,21 +1002,21 @@ def collate_fn4(batch): # Ancillary functions -------------------------------------------------------------------------------------------------- -def flatten_recursive(path=DATASETS_DIR / 'coco128'): +def flatten_recursive(path=DATASETS_DIR / "coco128"): # Flatten a recursive directory by bringing all files to top level - new_path = Path(f'{str(path)}_flat') + new_path = Path(f"{str(path)}_flat") if os.path.exists(new_path): shutil.rmtree(new_path) # delete output folder os.makedirs(new_path) # make new output folder - for file in tqdm(glob.glob(f'{str(Path(path))}/**/*.*', recursive=True)): + for file in tqdm(glob.glob(f"{str(Path(path))}/**/*.*", recursive=True)): shutil.copyfile(file, new_path / Path(file).name) -def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders import *; extract_boxes() +def extract_boxes(path=DATASETS_DIR / "coco128"): # from utils.dataloaders import *; extract_boxes() # Convert detection dataset into classification dataset, with one directory per class path = Path(path) # images dir - shutil.rmtree(path / 'classification') if (path / 'classification').is_dir() else None # remove existing - files = list(path.rglob('*.*')) + shutil.rmtree(path / "classification") if (path / "classification").is_dir() else None # remove existing + files = list(path.rglob("*.*")) n = len(files) # number of files for im_file in tqdm(files, total=n): if im_file.suffix[1:] in IMG_FORMATS: @@ -983,7 +1032,7 @@ def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders impo for j, x in enumerate(lb): c = int(x[0]) # class - f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename + f = (path / "classifier") / f"{c}" / f"{path.stem}_{im_file.stem}_{j}.jpg" # new filename if not f.parent.is_dir(): f.parent.mkdir(parents=True) @@ -994,11 +1043,11 @@ def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders impo b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image b[[1, 3]] = np.clip(b[[1, 3]], 0, h) - assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}' + assert cv2.imwrite(str(f), im[b[1] : b[3], b[0] : b[2]]), f"box failure in {f}" -def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False): - """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files +def autosplit(path=DATASETS_DIR / "coco128/images", weights=(0.9, 0.1, 0.0), annotated_only=False): + """Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files Usage: from utils.dataloaders import *; autosplit() Arguments path: Path to images directory @@ -1006,40 +1055,40 @@ def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), ann annotated_only: Only use images with an annotated txt file """ path = Path(path) # images dir - files = sorted(x for x in path.rglob('*.*') if x.suffix[1:].lower() in IMG_FORMATS) # image files only + files = sorted(x for x in path.rglob("*.*") if x.suffix[1:].lower() in IMG_FORMATS) # image files only n = len(files) # number of files random.seed(0) # for reproducibility indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split - txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files + txt = ["autosplit_train.txt", "autosplit_val.txt", "autosplit_test.txt"] # 3 txt files for x in txt: if (path.parent / x).exists(): (path.parent / x).unlink() # remove existing - print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only) + print(f"Autosplitting images from {path}" + ", using *.txt labeled images only" * annotated_only) for i, img in tqdm(zip(indices, files), total=n): if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label - with open(path.parent / txt[i], 'a') as f: - f.write(f'./{img.relative_to(path.parent).as_posix()}' + '\n') # add image to txt file + with open(path.parent / txt[i], "a") as f: + f.write(f"./{img.relative_to(path.parent).as_posix()}" + "\n") # add image to txt file def verify_image_label(args): # Verify one image-label pair im_file, lb_file, prefix = args - nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments + nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, "", [] # number (missing, found, empty, corrupt), message, segments try: # verify images im = Image.open(im_file) im.verify() # PIL verify shape = exif_size(im) # image size - assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels' - assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}' - if im.format.lower() in ('jpg', 'jpeg'): - with open(im_file, 'rb') as f: + assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels" + assert im.format.lower() in IMG_FORMATS, f"invalid image format {im.format}" + if im.format.lower() in ("jpg", "jpeg"): + with open(im_file, "rb") as f: f.seek(-2, 2) - if f.read() != b'\xff\xd9': # corrupt JPEG - ImageOps.exif_transpose(Image.open(im_file)).save(im_file, 'JPEG', subsampling=0, quality=100) - msg = f'{prefix}WARNING ⚠️ {im_file}: corrupt JPEG restored and saved' + if f.read() != b"\xff\xd9": # corrupt JPEG + ImageOps.exif_transpose(Image.open(im_file)).save(im_file, "JPEG", subsampling=0, quality=100) + msg = f"{prefix}WARNING ⚠️ {im_file}: corrupt JPEG restored and saved" # verify labels if os.path.isfile(lb_file): @@ -1053,15 +1102,15 @@ def verify_image_label(args): lb = np.array(lb, dtype=np.float32) nl = len(lb) if nl: - assert lb.shape[1] == 5, f'labels require 5 columns, {lb.shape[1]} columns detected' - assert (lb >= 0).all(), f'negative label values {lb[lb < 0]}' - assert (lb[:, 1:] <= 1).all(), f'non-normalized or out of bounds coordinates {lb[:, 1:][lb[:, 1:] > 1]}' + assert lb.shape[1] == 5, f"labels require 5 columns, {lb.shape[1]} columns detected" + assert (lb >= 0).all(), f"negative label values {lb[lb < 0]}" + assert (lb[:, 1:] <= 1).all(), f"non-normalized or out of bounds coordinates {lb[:, 1:][lb[:, 1:] > 1]}" _, i = np.unique(lb, axis=0, return_index=True) if len(i) < nl: # duplicate row check lb = lb[i] # remove duplicates if segments: segments = [segments[x] for x in i] - msg = f'{prefix}WARNING ⚠️ {im_file}: {nl - len(i)} duplicate labels removed' + msg = f"{prefix}WARNING ⚠️ {im_file}: {nl - len(i)} duplicate labels removed" else: ne = 1 # label empty lb = np.zeros((0, 5), dtype=np.float32) @@ -1071,12 +1120,13 @@ def verify_image_label(args): return im_file, lb, shape, segments, nm, nf, ne, nc, msg except Exception as e: nc = 1 - msg = f'{prefix}WARNING ⚠️ {im_file}: ignoring corrupt image/label: {e}' + msg = f"{prefix}WARNING ⚠️ {im_file}: ignoring corrupt image/label: {e}" return [None, None, None, None, nm, nf, ne, nc, msg] -class HUBDatasetStats(): - """ Class for generating HUB dataset JSON and `-hub` dataset directory +class HUBDatasetStats: + """ + Class for generating HUB dataset JSON and `-hub` dataset directory. Arguments path: Path to data.yaml or data.zip (with data.yaml inside data.zip) @@ -1090,43 +1140,43 @@ class HUBDatasetStats(): stats.process_images() """ - def __init__(self, path='coco128.yaml', autodownload=False): + def __init__(self, path="coco128.yaml", autodownload=False): # Initialize class zipped, data_dir, yaml_path = self._unzip(Path(path)) try: - with open(check_yaml(yaml_path), errors='ignore') as f: + with open(check_yaml(yaml_path), errors="ignore") as f: data = yaml.safe_load(f) # data dict if zipped: - data['path'] = data_dir + data["path"] = data_dir except Exception as e: - raise Exception('error/HUB/dataset_stats/yaml_load') from e + raise Exception("error/HUB/dataset_stats/yaml_load") from e check_dataset(data, autodownload) # download dataset if missing - self.hub_dir = Path(data['path'] + '-hub') - self.im_dir = self.hub_dir / 'images' + self.hub_dir = Path(data["path"] + "-hub") + self.im_dir = self.hub_dir / "images" self.im_dir.mkdir(parents=True, exist_ok=True) # makes /images - self.stats = {'nc': data['nc'], 'names': list(data['names'].values())} # statistics dictionary + self.stats = {"nc": data["nc"], "names": list(data["names"].values())} # statistics dictionary self.data = data @staticmethod def _find_yaml(dir): # Return data.yaml file - files = list(dir.glob('*.yaml')) or list(dir.rglob('*.yaml')) # try root level first and then recursive - assert files, f'No *.yaml file found in {dir}' + files = list(dir.glob("*.yaml")) or list(dir.rglob("*.yaml")) # try root level first and then recursive + assert files, f"No *.yaml file found in {dir}" if len(files) > 1: files = [f for f in files if f.stem == dir.stem] # prefer *.yaml files that match dir name - assert files, f'Multiple *.yaml files found in {dir}, only 1 *.yaml file allowed' - assert len(files) == 1, f'Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}' + assert files, f"Multiple *.yaml files found in {dir}, only 1 *.yaml file allowed" + assert len(files) == 1, f"Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}" return files[0] def _unzip(self, path): # Unzip data.zip - if not str(path).endswith('.zip'): # path is data.yaml + if not str(path).endswith(".zip"): # path is data.yaml return False, None, path - assert Path(path).is_file(), f'Error unzipping {path}, file not found' + assert Path(path).is_file(), f"Error unzipping {path}, file not found" unzip_file(path, path=path.parent) - dir = path.with_suffix('') # dataset directory == zip name - assert dir.is_dir(), f'Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/' + dir = path.with_suffix("") # dataset directory == zip name + assert dir.is_dir(), f"Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/" return True, str(dir), self._find_yaml(dir) # zipped, data_dir, yaml_path def _hub_ops(self, f, max_dim=1920): @@ -1137,9 +1187,9 @@ def _hub_ops(self, f, max_dim=1920): r = max_dim / max(im.height, im.width) # ratio if r < 1.0: # image too large im = im.resize((int(im.width * r), int(im.height * r))) - im.save(f_new, 'JPEG', quality=50, optimize=True) # save + im.save(f_new, "JPEG", quality=50, optimize=True) # save except Exception as e: # use OpenCV - LOGGER.info(f'WARNING ⚠️ HUB ops PIL failure {f}: {e}') + LOGGER.info(f"WARNING ⚠️ HUB ops PIL failure {f}: {e}") im = cv2.imread(f) im_height, im_width = im.shape[:2] r = max_dim / max(im_height, im_width) # ratio @@ -1153,30 +1203,32 @@ def _round(labels): # Update labels to integer class and 6 decimal place floats return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels] - for split in 'train', 'val', 'test': + for split in "train", "val", "test": if self.data.get(split) is None: self.stats[split] = None # i.e. no test set continue dataset = LoadImagesAndLabels(self.data[split]) # load dataset - x = np.array([ - np.bincount(label[:, 0].astype(int), minlength=self.data['nc']) - for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics')]) # shape(128x80) + x = np.array( + [ + np.bincount(label[:, 0].astype(int), minlength=self.data["nc"]) + for label in tqdm(dataset.labels, total=dataset.n, desc="Statistics") + ] + ) # shape(128x80) self.stats[split] = { - 'instance_stats': { - 'total': int(x.sum()), - 'per_class': x.sum(0).tolist()}, - 'image_stats': { - 'total': dataset.n, - 'unlabelled': int(np.all(x == 0, 1).sum()), - 'per_class': (x > 0).sum(0).tolist()}, - 'labels': [{ - str(Path(k).name): _round(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)]} + "instance_stats": {"total": int(x.sum()), "per_class": x.sum(0).tolist()}, + "image_stats": { + "total": dataset.n, + "unlabelled": int(np.all(x == 0, 1).sum()), + "per_class": (x > 0).sum(0).tolist(), + }, + "labels": [{str(Path(k).name): _round(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)], + } # Save, print and return if save: - stats_path = self.hub_dir / 'stats.json' - print(f'Saving {stats_path.resolve()}...') - with open(stats_path, 'w') as f: + stats_path = self.hub_dir / "stats.json" + print(f"Saving {stats_path.resolve()}...") + with open(stats_path, "w") as f: json.dump(self.stats, f) # save stats.json if verbose: print(json.dumps(self.stats, indent=2, sort_keys=False)) @@ -1184,14 +1236,14 @@ def _round(labels): def process_images(self): # Compress images for Ultralytics HUB - for split in 'train', 'val', 'test': + for split in "train", "val", "test": if self.data.get(split) is None: continue dataset = LoadImagesAndLabels(self.data[split]) # load dataset - desc = f'{split} images' + desc = f"{split} images" for _ in tqdm(ThreadPool(NUM_THREADS).imap(self._hub_ops, dataset.im_files), total=dataset.n, desc=desc): pass - print(f'Done. All images saved to {self.im_dir}') + print(f"Done. All images saved to {self.im_dir}") return self.im_dir @@ -1199,6 +1251,7 @@ def process_images(self): class ClassificationDataset(torchvision.datasets.ImageFolder): """ YOLOv5 Classification Dataset. + Arguments root: Dataset path transform: torchvision transforms, used by default @@ -1209,9 +1262,9 @@ def __init__(self, root, augment, imgsz, cache=False): super().__init__(root=root) self.torch_transforms = classify_transforms(imgsz) self.album_transforms = classify_albumentations(augment, imgsz) if augment else None - self.cache_ram = cache is True or cache == 'ram' - self.cache_disk = cache == 'disk' - self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples] # file, index, npy, im + self.cache_ram = cache is True or cache == "ram" + self.cache_disk = cache == "disk" + self.samples = [list(x) + [Path(x[0]).with_suffix(".npy"), None] for x in self.samples] # file, index, npy, im def __getitem__(self, i): f, j, fn, im = self.samples[i] # filename, index, filename.with_suffix('.npy'), image @@ -1224,20 +1277,15 @@ def __getitem__(self, i): else: # read image im = cv2.imread(f) # BGR if self.album_transforms: - sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))['image'] + sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))["image"] else: sample = self.torch_transforms(im) return sample, j -def create_classification_dataloader(path, - imgsz=224, - batch_size=16, - augment=True, - cache=False, - rank=-1, - workers=8, - shuffle=True): +def create_classification_dataloader( + path, imgsz=224, batch_size=16, augment=True, cache=False, rank=-1, workers=8, shuffle=True +): # Returns Dataloader object to be used with YOLOv5 Classifier with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache) @@ -1247,11 +1295,13 @@ def create_classification_dataloader(path, sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) generator = torch.Generator() generator.manual_seed(6148914691236517205 + RANK) - return InfiniteDataLoader(dataset, - batch_size=batch_size, - shuffle=shuffle and sampler is None, - num_workers=nw, - sampler=sampler, - pin_memory=PIN_MEMORY, - worker_init_fn=seed_worker, - generator=generator) # or DataLoader(persistent_workers=True) + return InfiniteDataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle and sampler is None, + num_workers=nw, + sampler=sampler, + pin_memory=PIN_MEMORY, + worker_init_fn=seed_worker, + generator=generator, + ) # or DataLoader(persistent_workers=True) diff --git a/utils/downloads.py b/utils/downloads.py index 9298259d4ab1..ee700acb618b 100644 --- a/utils/downloads.py +++ b/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. - """ - silent_option = 'sS' if silent else '' # silent - proc = subprocess.run([ - 'curl', - '-#', - f'-{silent_option}L', - url, - '--output', - filename, - '--retry', - '9', - '-C', - '-', ]) + """Download a file from a url to a filename using curl.""" + silent_option = "sS" if silent else "" # silent + proc = subprocess.run( + [ + "curl", + "-#", + f"-{silent_option}L", + url, + "--output", + filename, + "--retry", + "9", + "-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('') + LOGGER.info(f"ERROR: {assert_msg}\n{error_msg}") + 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': - version = f'tags/{version}' # i.e. tags/v7.0 - response = requests.get(f'https://github.com/gitapi/repos/{repository}/releases/{version}').json() # github api - return response['tag_name'], [x['name'] for x in response['assets']] # tag, assets + if version != "latest": + version = f"tags/{version}" # i.e. tags/v7.0 + response = requests.get(f"https://github.com/gitapi/repos/{repository}/releases/{version}").json() # github api + 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 - url = str(file).replace(':/', '://') # Pathlib turns :// -> :/ - file = name.split('?')[0] # parse authentication https://url.com/file.txt?auth... + if str(file).startswith(("http:/", "https:/")): # download + url = str(file).replace(":/", "://") # Pathlib turns :// -> :/ + 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 github_assets(repository, version='latest'): 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, - url=f'https://github.com/{repo}/releases/download/{tag}/{name}', - min_bytes=1E5, - error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/{tag}') + safe_download( + file, + url=f"https://github.com/{repo}/releases/download/{tag}/{name}", + min_bytes=1e5, + error_msg=f"{file} missing, try downloading from https://github.com/{repo}/releases/{tag}", + ) return str(file) diff --git a/utils/flask_rest_api/example_request.py b/utils/flask_rest_api/example_request.py index 256ad1319c82..7b850051cca0 100644 --- a/utils/flask_rest_api/example_request.py +++ b/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' -IMAGE = 'zidane.jpg' +DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s" +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) diff --git a/utils/flask_rest_api/restapi.py b/utils/flask_rest_api/restapi.py index ae4756b276e4..e62c7ebd709f 100644 --- a/utils/flask_rest_api/restapi.py +++ b/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 @@ app = Flask(__name__) models = {} -DETECTION_URL = '/v1/object-detection/' +DETECTION_URL = "/v1/object-detection/" -@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__': - parser = argparse.ArgumentParser(description='Flask API exposing YOLOv5 model') - 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') +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model") + 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") 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 diff --git a/utils/general.py b/utils/general.py index 73925ce5fb95..47ab656e5a3f 100644 --- a/utils/general.py +++ b/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 @@ 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 @@ 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 -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 -TQDM_BAR_FORMAT = '{l_bar}{bar:10}{r_bar}' # tqdm bar format -FONT = 'Arial.ttf' # https://ultralytics.com/assets/Arial.ttf - -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 +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 +VERBOSE = str(os.getenv("YOLOv5_VERBOSE", True)).lower() == "true" # global verbose mode +TQDM_BAR_FORMAT = "{l_bar}{bar:10}{r_bar}" # tqdm bar format +FONT = "Arial.ttf" # https://ultralytics.com/assets/Arial.ttf + +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 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['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["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["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. - Verified on Colab, Jupyterlab, Kaggle, Paperspace. + Check if the current script is running inside a Jupyter Notebook. 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: - return any('docker' in line for line in file) + with open("/proc/self/cgroup") as 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({ - 'version': 1, - 'disable_existing_loggers': False, - 'formatters': { - name: { - 'format': '%(message)s'}}, - 'handlers': { - name: { - 'class': 'logging.StreamHandler', - 'formatter': name, - 'level': level, }}, - 'loggers': { - name: { - 'level': level, - 'handlers': [name], - 'propagate': False, }}}) + logging.config.dictConfig( + { + "version": 1, + "disable_existing_loggers": False, + "formatters": {name: {"format": "%(message)s"}}, + "handlers": { + name: { + "class": "logging.StreamHandler", + "formatter": name, + "level": level, + } + }, + "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 - 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 + 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 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 @@ def time(self): 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 @@ def _timeout_handler(self, signum, frame): 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 @@ def __exit__(self, exc_type, exc_val, exc_tb): 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 '') - LOGGER.info(colorstr(s) + ', '.join(f'{k}={v}' for k, v in args.items())) + 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())) 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['PYTHONHASHSEED'] = str(seed) + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + 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) - return max(last_list, key=os.path.getctime) if last_list else '' + last_list = glob.glob(f"{search_dir}/**/last*.pt", recursive=True) + 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 run_once(): 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() - return check_output(f'git -C {path} describe --tags --long --always', shell=True).decode()[:-1] + assert (Path(path) / ".git").is_dir() + 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}' - msg = f', for updates see {url}' - s = colorstr('github: ') # string - assert Path('.git').exists(), s + 'skipping check (not a git repository)' + msg - assert check_online(), s + 'skipping check (offline)' + msg + url = f"https://github.com/{repo}" + msg = f", for updates see {url}" + s = colorstr("github: ") # string + assert Path(".git").exists(), s + "skipping check (not a git repository)" + 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' - check_output(f'git remote add {remote} {url}', shell=True) - 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 - n = int(check_output(f'git rev-list {local_branch}..{remote}/{branch} --count', shell=True)) # commits behind + remote = "ultralytics" + check_output(f"git remote add {remote} {url}", shell=True) + 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 + 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 - assert 'clearml' in sys.modules, "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'." + 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'." return file else: # search files = [] - for d in 'data', 'models', 'utils': # search directories - files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True)) # find file - assert len(files), f'File not found: {file}' # assert file was found + for d in "data", "models", "utils": # search directories + 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) == 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}' - LOGGER.info(f'Downloading {url} to {file}...') + url = f"https://ultralytics.com/assets/{font.name}" + 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) - data = next((DATASETS_DIR / Path(data).stem).rglob('*.yaml')) + 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")) 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 - 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' - data['nc'] = len(data['names']) + if isinstance(data["names"], (list, tuple)): # old array format + 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" + 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 - for k in 'train', 'val', 'test': + data["path"] = path # download scripts + 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 - LOGGER.info(f'Running {s} ...') + elif s.startswith("bash "): # bash script + LOGGER.info(f"Running {s} ...") r = subprocess.run(s, shell=True) else: # python script - r = exec(s, {'yaml': data}) # return None - 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} ❌' - LOGGER.info(f'Dataset download {s}') - check_font('Arial.ttf' if is_ascii(data['names']) else 'Arial.Unicode.ttf', progress=True) # download fonts + r = exec(s, {"yaml": data}) # return None + 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} ❌" + LOGGER.info(f"Dataset download {s}") + 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 amp_allclose(model, im): 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 - im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if check_online() else np.ones((640, 640, 3)) + 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)) try: - assert amp_allclose(deepcopy(model), im) or amp_allclose(DetectMultiBackend('yolov5n.pt', device), im) - LOGGER.info(f'{prefix}checks passed ✅') + assert amp_allclose(deepcopy(model), im) or amp_allclose(DetectMultiBackend("yolov5n.pt", device), im) + LOGGER.info(f"{prefix}checks passed ✅") return True except Exception: - help_url = 'https://github.com/ultralytics/yolov5/issues/7908' - LOGGER.warning(f'{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}') + help_url = "https://github.com/ultralytics/yolov5/issues/7908" + 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 :// -> :/ - return Path(urllib.parse.unquote(url)).name.split('?')[0] # '%2F' to '/', split https://url.com/file.txt?auth + 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 -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_one(url, dir): 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_one(url, dir): 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)): - LOGGER.info(f'Unzipping {f}...') + if unzip and success and (f.suffix == ".gz" or is_zipfile(f) or is_tarfile(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 - elif f.suffix == '.gz': - subprocess.run(['tar', 'xfz', f, '--directory', f.parent], check=True) # unzip + subprocess.run(["tar", "xf", f, "--directory", f.parent], check=True) # unzip + elif f.suffix == ".gz": + 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 - 'red': '\033[31m', - 'green': '\033[32m', - 'yellow': '\033[33m', - 'blue': '\033[34m', - 'magenta': '\033[35m', - 'cyan': '\033[36m', - 'white': '\033[37m', - 'bright_black': '\033[90m', # bright colors - 'bright_red': '\033[91m', - 'bright_green': '\033[92m', - 'bright_yellow': '\033[93m', - 'bright_blue': '\033[94m', - 'bright_magenta': '\033[95m', - 'bright_cyan': '\033[96m', - 'bright_white': '\033[97m', - 'end': '\033[0m', # misc - 'bold': '\033[1m', - 'underline': '\033[4m'} - return ''.join(colors[x] for x in args) + f'{string}' + colors['end'] + "black": "\033[30m", # basic colors + "red": "\033[31m", + "green": "\033[32m", + "yellow": "\033[33m", + "blue": "\033[34m", + "magenta": "\033[35m", + "cyan": "\033[36m", + "white": "\033[37m", + "bright_black": "\033[90m", # bright colors + "bright_red": "\033[91m", + "bright_green": "\033[92m", + "bright_yellow": "\033[93m", + "bright_blue": "\033[94m", + "bright_magenta": "\033[95m", + "bright_cyan": "\033[96m", + "bright_white": "\033[97m", + "end": "\033[0m", # misc + "bold": "\033[1m", + "underline": "\033[4m", + } + 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, - 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] + 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, + 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 @@ -865,30 +953,31 @@ def clip_segments(segments, shape): def non_max_suppression( - prediction, - conf_thres=0.25, - iou_thres=0.45, - classes=None, - agnostic=False, - multi_label=False, - labels=(), - max_det=300, - nm=0, # number of masks + prediction, + conf_thres=0.25, + iou_thres=0.45, + classes=None, + agnostic=False, + multi_label=False, + labels=(), + 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 <= iou_thres <= 1, f'Invalid IoU {iou_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" 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')) - if x.get('ema'): - x['model'] = x['ema'] # replace model with ema - for k in 'optimizer', 'best_fitness', 'ema', 'updates': # keys + x = torch.load(f, map_location=torch.device("cpu")) + if x.get("ema"): + x["model"] = x["ema"] # replace model with ema + for k in "optimizer", "best_fitness", "ema", "updates": # keys x[k] = None - x['epoch'] = -1 - x['model'].half() # to FP16 - for p in x['model'].parameters(): + x["epoch"] = -1 + x["model"].half() # to FP16 + 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: ')): - evolve_csv = save_dir / 'evolve.csv' - evolve_yaml = save_dir / 'hyp_evolve.yaml' +def print_mutation(keys, results, hyp, save_dir, bucket, prefix=colorstr("evolve: ")): + evolve_csv = save_dir / "evolve.csv" + 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 - with open(evolve_csv, 'a') as f: - f.write(s + ('%20.5g,' * n % vals).rstrip(',') + '\n') + s = "" if evolve_csv.exists() else (("%20s," * n % keys).rstrip(",") + "\n") # add header + with open(evolve_csv, "a") as f: + 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'# 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') + f.write( + "# YOLOv5 Hyperparameter Evolution Results\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 + - ', '.join(f'{x.strip():>20s}' for x in keys) + '\n' + prefix + ', '.join(f'{x:20.5g}' - for x in vals) + '\n\n') + LOGGER.info( + prefix + + 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): @@ -1053,7 +1157,7 @@ def apply_classifier(x, model, img, im0): pred_cls1 = d[:, 5].long() ims = [] for a in d: - cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])] + cutout = im0[i][int(a[1]) : int(a[3]), int(a[0]) : int(a[2])] im = cv2.resize(cutout, (224, 224)) # BGR im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 @@ -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: diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index 381d477d127c..df67e45c8221 100644 --- a/utils/loggers/__init__.py +++ b/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.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 -RANK = int(os.getenv('RANK', -1)) +LOGGERS = ("csv", "tb", "wandb", "clearml", "comet") # *.csv, TensorBoard, Weights & Biases, ClearML +RANK = int(os.getenv("RANK", -1)) try: from torch.utils.tensorboard import SummaryWriter @@ -27,8 +25,8 @@ try: import wandb - 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}: + 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}: try: wandb_login_success = wandb.login(timeout=30) except wandb.errors.UsageError: # known non-TTY terminal issue @@ -41,7 +39,7 @@ 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 @@ 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 @@ 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 @@ def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, self.logger = logger # for printing results to console self.include = include self.keys = [ - 'train/box_loss', - 'train/obj_loss', - 'train/cls_loss', # train loss - 'metrics/precision', - 'metrics/recall', - 'metrics/mAP_0.5', - 'metrics/mAP_0.5:0.95', # metrics - 'val/box_loss', - 'val/obj_loss', - 'val/cls_loss', # val loss - 'x/lr0', - 'x/lr1', - 'x/lr2'] # params - self.best_keys = ['best/epoch', 'best/precision', 'best/recall', 'best/mAP_0.5', 'best/mAP_0.5:0.95'] + "train/box_loss", + "train/obj_loss", + "train/cls_loss", # train loss + "metrics/precision", + "metrics/recall", + "metrics/mAP_0.5", + "metrics/mAP_0.5:0.95", # metrics + "val/box_loss", + "val/obj_loss", + "val/cls_loss", # val loss + "x/lr0", + "x/lr1", + "x/lr2", + ] # 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_file = ('ndjson_file' in self.include) # log ndjson to file + self.ndjson_console = "ndjson_console" in self.include # log ndjson to console + 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: - prefix = colorstr('TensorBoard: ') + if "tb" in self.include and not self.opt.evolve: + 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: ') - LOGGER.warning(f'{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging.' - f' See https://docs.ultralytics.com/yolov5/tutorials/clearml_logging_integration#readme') + prefix = colorstr("ClearML: ") + LOGGER.warning( + 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 isinstance(self.opt.resume, str) and self.opt.resume.startswith('comet://'): - run_id = self.opt.resume.split('/')[-1] + if comet_ml and "comet" in self.include: + if isinstance(self.opt.resume, str) and self.opt.resume.startswith("comet://"): + run_id = self.opt.resume.split("/")[-1] self.comet_logger = CometLogger(self.opt, self.hyp, run_id=run_id) else: @@ -170,9 +175,9 @@ def on_pretrain_routine_end(self, labels, names): # 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 @@ def on_train_batch_end(self, model, ni, imgs, targets, paths, vals): # 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 @@ def on_val_batch_end(self, batch_i, im, targets, paths, shapes, out): 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 @@ def on_fit_epoch_end(self, vals, epoch, best_fitness, fi): # 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 - with open(file, 'a') as f: - f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n') + s = "" if file.exists() else (("%20s," * n % tuple(["epoch"] + self.keys)).rstrip(",") + "\n") # add header + with open(file, "a") as f: + 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' - with open(file, 'a') as f: + file = self.save_dir / "results.ndjson" + with open(file, "a") as f: print(json_data, file=f) if self.tb: @@ -279,9 +284,9 @@ def on_model_save(self, last, epoch, final_epoch, best_fitness, fi): 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), - model_name='Latest Model', - auto_delete_file=False) + self.clearml.task.update_output_model( + model_path=str(last), model_name="Latest Model", 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 @@ def on_model_save(self, last, epoch, final_epoch, best_fitness, fi): 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 - files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))] + 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 = [(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), - type='model', - name=f'run_{self.wandb.wandb_run.id}_model', - aliases=['latest', 'best', 'stripped']) + wandb.log_artifact( + str(best if best.exists() else last), + type="model", + 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), - "Best Model" if best.exists() else "Last Model", epoch) + self.clearml.log_model( + 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 - if 'tb' in self.include: - prefix = colorstr('TensorBoard: ') + self.csv = self.save_dir / "results.csv" # CSV logger + if "tb" in self.include: + 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: - self.wandb = wandb.init(project=web_project_name(str(opt.project)), - name=None if opt.name == 'exp' else opt.name, - config=opt) + if wandb and "wandb" in self.include: + self.wandb = wandb.init( + project=web_project_name(str(opt.project)), name=None if opt.name == "exp" else opt.name, 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: ') - LOGGER.warning(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') + prefix = colorstr("ClearML: ") + LOGGER.warning( + 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 @@ def log_metrics(self, metrics, epoch): 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 - with open(self.csv, 'a') as f: - f.write(s + ('%23.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n') + s = "" if self.csv.exists() else (("%23s," * n % tuple(["epoch"] + keys)).rstrip(",") + "\n") # header + with open(self.csv, "a") as f: + f.write(s + ("%23.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n") if self.tb: for k, v in metrics.items(): @@ -393,20 +404,20 @@ def log_metrics(self, metrics, epoch): 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 @@ def log_graph(self, model, imgsz=(640, 640)): 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 '' - return f'YOLOv5{suffix}' + suffix = "-Classify" if project.endswith("-cls") else "-Segment" if project.endswith("-seg") else "" + return f"YOLOv5{suffix}" diff --git a/utils/loggers/clearml/clearml_utils.py b/utils/loggers/clearml/clearml_utils.py index c7627a261186..8b141d177afd 100644 --- a/utils/loggers/clearml/clearml_utils.py +++ b/utils/loggers/clearml/clearml_utils.py @@ -13,55 +13,63 @@ 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. - """ - dataset_id = clearml_info_string.replace('clearml://', '') + """Load in a clearml dataset and fill the internal data_dict with its contents.""" + 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 ' - 'the dataset definition this way.') + raise ValueError( + "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 ' - 'inside the dataset root path.') + raise ValueError( + "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( - {'train', 'test', 'val', 'nc', 'names'} + assert set( + 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( - (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['val'] = 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'] + data_dict["train"] = ( + 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["val"] = ( + 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 information includes hyperparameters, system configuration and metrics, model metrics, code information and - basic data metrics and analyses. + This logger sends information to ClearML at app.clear.ml or to your own hosted server. By default, this information + includes hyperparameters, system configuration and metrics, model metrics, code information and basic data metrics + and analyses. - By providing additional command line arguments to train.py, datasets, - models and predictions can also be logged. + By providing additional command line arguments to train.py, datasets, models and predictions can also be logged. """ def __init__(self, opt, hyp): @@ -81,36 +89,36 @@ def __init__(self, opt, hyp): 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', - task_name=opt.name if opt.name != 'exp' else 'Training', - tags=['YOLOv5'], + project_name=opt.project if not str(opt.project).startswith("runs/") else "YOLOv5", + task_name=opt.name if opt.name != "exp" else "Training", + tags=["YOLOv5"], output_uri=True, reuse_last_task_id=opt.exist_ok, - auto_connect_frameworks={ - 'pytorch': False, - 'matplotlib': False} + auto_connect_frameworks={"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(opt, name='Args') + self.task.connect(hyp, name="Hyperparameters") + 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', - docker_arguments='--ipc=host -e="CLEARML_AGENT_SKIP_PYTHON_ENV_INSTALL=1"', - docker_setup_bash_script='pip install clearml') + 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", + ) # 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 @@ def __init__(self, opt, hyp): 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), - name=model_name, - iteration=epoch, - auto_delete_file=False) + self.task.update_output_model( + model_path=str(model_path), name=model_name, 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 @@ def log_summary(self, metrics): 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 @@ def log_plot(self, title, plot_path): """ 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 @@ def log_debug_samples(self, files, title='Debug Samples'): """ 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, - series=f.name.replace(f"_batch{iteration}", ''), - local_path=str(f), - iteration=iteration) + self.task.get_logger().report_image( + title=title, series=f.name.replace(f"_batch{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 @@ def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_thres 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', - series=image_path.name, - iteration=self.current_epoch, - image=annotated_image) + self.task.get_logger().report_image( + title="Bounding Boxes", series=image_path.name, iteration=self.current_epoch, image=annotated_image + ) self.current_epoch_logged_images.add(image_path) diff --git a/utils/loggers/clearml/hpo.py b/utils/loggers/clearml/hpo.py index ee518b0fbfc8..4e314ea868df 100644 --- a/utils/loggers/clearml/hpo.py +++ b/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_name='YOLOv5', - task_type=Task.TaskTypes.optimizer, - reuse_last_task_id=False) +task = Task.init( + project_name="Hyper-Parameter Optimization", + task_name="YOLOv5", + task_type=Task.TaskTypes.optimizer, + reuse_last_task_id=False, +) # Example use case: optimizer = HyperParameterOptimizer( # This is the experiment we want to optimize - base_task_id='', + base_task_id="", # here we define the hyper-parameters to optimize # Notice: The parameter name should exactly match what you see in the UI: / # For Example, here we see in the base experiment a section Named: "General" @@ -20,39 +23,40 @@ # 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/lrf', min_value=0.01, max_value=1.0), - 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/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_bias_lr', min_value=0.0, 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_pw', min_value=0.5, max_value=2.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/iou_t', min_value=0.1, max_value=0.7), - 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/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_v', min_value=0.0, max_value=0.9), - UniformParameterRange('Hyperparameters/degrees', min_value=0.0, max_value=45.0), - UniformParameterRange('Hyperparameters/translate', 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/perspective', min_value=0.0, max_value=0.001), - UniformParameterRange('Hyperparameters/flipud', 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/mixup', min_value=0.0, max_value=1.0), - UniformParameterRange('Hyperparameters/copy_paste', min_value=0.0, max_value=1.0)], + UniformParameterRange("Hyperparameters/lr0", min_value=1e-5, max_value=1e-1), + UniformParameterRange("Hyperparameters/lrf", min_value=0.01, max_value=1.0), + 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/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_bias_lr", min_value=0.0, 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_pw", min_value=0.5, max_value=2.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/iou_t", min_value=0.1, max_value=0.7), + 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/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_v", min_value=0.0, max_value=0.9), + UniformParameterRange("Hyperparameters/degrees", min_value=0.0, max_value=45.0), + UniformParameterRange("Hyperparameters/translate", 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/perspective", min_value=0.0, max_value=0.001), + UniformParameterRange("Hyperparameters/flipud", 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/mixup", 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_series='mAP_0.5', + objective_metric_title="metrics", + 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 @@ # make sure background optimization stopped optimizer.stop() -print('We are done, good bye') +print("We are done, good bye") diff --git a/utils/loggers/comet/__init__.py b/utils/loggers/comet/__init__.py index c14a5f885696..bdf81f63982e 100644 --- a/utils/loggers/comet/__init__.py +++ b/utils/loggers/comet/__init__.py @@ -17,7 +17,7 @@ # 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.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_PREDICTIONS = os.getenv('COMET_LOG_PREDICTIONS', 'true').lower() == 'true' -COMET_MAX_IMAGE_UPLOADS = int(os.getenv('COMET_MAX_IMAGE_UPLOADS', 100)) +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_MAX_IMAGE_UPLOADS = int(os.getenv("COMET_MAX_IMAGE_UPLOADS", 100)) # Confusion Matrix Settings -CONF_THRES = float(os.getenv('CONF_THRES', 0.001)) -IOU_THRES = float(os.getenv('IOU_THRES', 0.6)) +CONF_THRES = float(os.getenv("CONF_THRES", 0.001)) +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_BATCH_LOGGING_INTERVAL = os.getenv('COMET_BATCH_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_BATCH_METRICS = os.getenv("COMET_LOG_BATCH_METRICS", "false").lower() == "true" +COMET_BATCH_LOGGING_INTERVAL = os.getenv("COMET_BATCH_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" -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 - with Comet - """ + """Log metrics, parameters, source code, models and much more 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 @@ def __init__(self, opt, hyp, run_id=None, job_type='Training', **experiment_kwar # Default parameters to pass to Experiment objects self.default_experiment_kwargs = { - 'log_code': False, - 'log_env_gpu': True, - 'log_env_cpu': True, - 'project_name': COMET_PROJECT_NAME, } + "log_code": False, + "log_env_gpu": True, + "log_env_cpu": True, + "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.num_classes = self.data_dict['nc'] + self.class_names = self.data_dict["names"] + 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', - f'{workspace}/{project_name}/{experiment_id}', + "Run Path", + 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', - metadata={'type': 'hyp-config-file'}, + name="hyperparameters.json", + metadata={"type": "hyp-config-file"}, ) self.log_asset( - f'{self.opt.save_dir}/opt.yaml', - metadata={'type': 'opt-config-file'}, + f"{self.opt.save_dir}/opt.yaml", + 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 @@ def __init__(self, opt, hyp, run_id=None, job_type='Training', **experiment_kwar self.comet_log_per_class_metrics = COMET_LOG_PER_CLASS_METRICS - self.experiment.log_others({ - 'comet_mode': COMET_MODE, - 'comet_max_image_uploads': COMET_MAX_IMAGE_UPLOADS, - 'comet_log_per_class_metrics': COMET_LOG_PER_CLASS_METRICS, - 'comet_log_batch_metrics': COMET_LOG_BATCH_METRICS, - 'comet_log_confusion_matrix': COMET_LOG_CONFUSION_MATRIX, - 'comet_model_name': COMET_MODEL_NAME, }) + self.experiment.log_others( + { + "comet_mode": COMET_MODE, + "comet_max_image_uploads": COMET_MAX_IMAGE_UPLOADS, + "comet_log_per_class_metrics": COMET_LOG_PER_CLASS_METRICS, + "comet_log_batch_metrics": COMET_LOG_BATCH_METRICS, + "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'): - 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_metric', self.opt.comet_optimizer_metric) - self.experiment.log_other('optimizer_parameters', json.dumps(self.hyp)) + if hasattr(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_metric", self.opt.comet_optimizer_metric) + 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 @@ def _get_experiment(self, mode, experiment_id=None): return comet_ml.Experiment(**self.default_experiment_kwargs) except ValueError: - logger.warning('COMET WARNING: ' - 'Comet credentials have not been set. ' - 'Comet will default to offline logging. ' - 'Please set your credentials to enable online logging.') - return self._get_experiment('offline', experiment_id) + logger.warning( + "COMET WARNING: " + "Comet credentials have not been set. " + "Comet will default to offline logging. " + "Please set your credentials to enable online logging." + ) + return self._get_experiment("offline", experiment_id) return @@ -211,12 +217,13 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False): return model_metadata = { - 'fitness_score': fitness_score[-1], - 'epochs_trained': epoch + 1, - 'save_period': opt.save_period, - 'total_epochs': opt.epochs, } + "fitness_score": fitness_score[-1], + "epochs_trained": epoch + 1, + "save_period": opt.save_period, + "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 @@ def check_dataset(self, data_file): 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 @@ def log_predictions(self, image, labelsn, path, shape, predn): filtered_detections = detections[mask] filtered_labels = labelsn[mask] - image_id = path.split('/')[-1].split('.')[0] - image_name = f'{image_id}_curr_epoch_{self.experiment.curr_epoch}' + image_id = path.split("/")[-1].split(".")[0] + 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 @@ def log_predictions(self, image, labelsn, path, shape, predn): metadata = [] for cls, *xyxy in filtered_labels.tolist(): - metadata.append({ - 'label': f'{self.class_names[int(cls)]}-gt', - 'score': 100, - 'box': { - 'x': xyxy[0], - 'y': xyxy[1], - 'x2': xyxy[2], - 'y2': xyxy[3]}, }) + metadata.append( + { + "label": f"{self.class_names[int(cls)]}-gt", + "score": 100, + "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]}, + } + ) for *xyxy, conf, cls in filtered_detections.tolist(): - metadata.append({ - 'label': f'{self.class_names[int(cls)]}', - 'score': conf * 100, - 'box': { - 'x': xyxy[0], - 'y': xyxy[1], - 'x2': xyxy[2], - 'y2': xyxy[3]}, }) + metadata.append( + { + "label": f"{self.class_names[int(cls)]}", + "score": conf * 100, + "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]}, + } + ) self.metadata_dict[image_name] = metadata self.logged_images_count += 1 @@ -307,7 +312,7 @@ def preprocess_prediction(self, image, labels, shape, pred): 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 @@ def add_assets_to_artifact(self, artifact, path, asset_path, split): 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(f'COMET ERROR: {e}') + logger.error("COMET ERROR: Error adding file to Artifact. Skipping file.") + logger.error(f"COMET ERROR: {e}") continue return artifact def upload_dataset_artifact(self): - dataset_name = self.data_dict.get('dataset_name', 'yolov5-dataset') - path = str((ROOT / Path(self.data_dict['path'])).resolve()) + dataset_name = self.data_dict.get("dataset_name", "yolov5-dataset") + 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 @@ def download_dataset_artifact(self, artifact_path): 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 [ - f'{path}/{x}' for x in split_path]) + data_dict[split] = ( + f"{path}/{split_path}" if isinstance(split, str) else [f"{path}/{x}" for x in split_path] + ) return data_dict @@ -422,11 +428,11 @@ def on_train_batch_end(self, log_dict, step): 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'{save_dir}/results.csv', metadata={'epoch': epoch}) + self.log_asset(f, 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 @@ def on_train_end(self, files, save_dir, last, best, epoch, results): ) # 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 @@ def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix) class_name = self.class_names[c] self.experiment.log_metrics( { - 'mAP@.5': ap50[i], - 'mAP@.5:.95': ap[i], - 'precision': p[i], - 'recall': r[i], - 'f1': f1[i], - 'true_positives': tp[i], - 'false_positives': fp[i], - 'support': nt[c], }, + "mAP@.5": ap50[i], + "mAP@.5:.95": ap[i], + "precision": p[i], + "recall": r[i], + "f1": f1[i], + "true_positives": tp[i], + "false_positives": fp[i], + "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 @@ def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix) max_categories=num_classes, labels=class_names, epoch=epoch, - column_label='Actual Category', - row_label='Predicted Category', - file_name=f'confusion-matrix-epoch-{epoch}.json', + column_label="Actual Category", + row_label="Predicted Category", + file_name=f"confusion-matrix-epoch-{epoch}.json", ) def on_fit_epoch_end(self, result, epoch): diff --git a/utils/loggers/comet/comet_utils.py b/utils/loggers/comet/comet_utils.py index 27600761ad28..f7b56dddd5f7 100644 --- a/utils/loggers/comet/comet_utils.py +++ b/utils/loggers/comet/comet_utils.py @@ -11,28 +11,28 @@ logger = logging.getLogger(__name__) -COMET_PREFIX = 'comet://' -COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5') -COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv('COMET_DEFAULT_CHECKPOINT_FILENAME', 'last.pt') +COMET_PREFIX = "comet://" +COMET_MODEL_NAME = os.getenv("COMET_MODEL_NAME", "yolov5") +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_download_path = f'{model_dir}/{asset_filename}' - with open(model_download_path, 'wb') as f: + model_binary = experiment.get_asset(asset_id, return_type="binary", stream=False) + model_download_path = f"{model_dir}/{asset_filename}" + 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': - asset_id = asset['assetId'] - asset_binary = experiment.get_asset(asset_id, return_type='binary', stream=False) + if asset["fileName"] == "opt.yaml": + asset_id = asset["assetId"] + 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' - with open(hyp_yaml_path, 'w') as f: + hyp_yaml_path = f"{save_dir}/hyp.yaml" + 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) diff --git a/utils/loggers/comet/hpo.py b/utils/loggers/comet/hpo.py index fc49115c1358..a9e6fabec1cd 100644 --- a/utils/loggers/comet/hpo.py +++ b/utils/loggers/comet/hpo.py @@ -21,77 +21,79 @@ # 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('--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('--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('--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('--rect', action='store_true', help='rectangular 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('--noval', action='store_true', help='only validate final epoch') - 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('--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('--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('--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('--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('--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('--project', default=ROOT / 'runs/train', 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('--quad', action='store_true', help='quad dataloader') - 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('--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('--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('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify') + 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("--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("--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("--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("--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("--noval", action="store_true", help="only validate final epoch") + 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("--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("--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("--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("--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("--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("--project", default=ROOT / "runs/train", 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("--quad", action="store_true", help="quad dataloader") + 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("--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("--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("--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('--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('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use') + 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("--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") # Comet Arguments - 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_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_workers', - type=int, - default=1, - help='Comet: Number of Parallel Workers to use with the Comet Optimizer.') + 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_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_workers", + type=int, + default=1, + 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.epochs = parameters.get('epochs') + opt.batch_size = parameters.get("batch_size") + 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 @@ def run(parameters, opt): 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 @@ def run(parameters, opt): opt.comet_optimizer_id = optimizer.id status = optimizer.status() - opt.comet_optimizer_objective = status['spec']['objective'] - opt.comet_optimizer_metric = status['spec']['metric'] + opt.comet_optimizer_objective = status["spec"]["objective"] + 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) diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py index 4ea32b1d4c6e..f8d49a33d00f 100644 --- a/utils/loggers/wandb/wandb_utils.py +++ b/utils/loggers/wandb/wandb_utils.py @@ -15,34 +15,35 @@ 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)) -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.' +RANK = int(os.getenv("RANK", -1)) +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(): - """Log training runs, datasets, models, and predictions to Weights & Biases. +class WandbLogger: + """ + Log training runs, datasets, models, and predictions to Weights & Biases. - This logger sends information to W&B at wandb.ai. By default, this information - includes hyperparameters, system configuration and metrics, model metrics, - and basic data metrics and analyses. + This logger sends information to W&B at wandb.ai. By default, this information includes hyperparameters, system + configuration and metrics, model metrics, and basic data metrics and analyses. - By providing additional command line arguments to train.py, datasets, - models and predictions can also be logged. + By providing additional command line arguments to train.py, datasets, 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 @@ -53,7 +54,7 @@ def __init__(self, opt, run_id=None, job_type='Training'): run_id (str) -- Run ID of W&B run to be resumed job_type (str) -- To set the job_type for this run - """ + """ # Pre-training routine -- self.job_type = job_type self.wandb, self.wandb_run = wandb, wandb.run if wandb else None @@ -64,17 +65,23 @@ def __init__(self, opt, run_id=None, job_type='Training'): self.max_imgs_to_log = 16 self.data_dict = None if self.wandb: - self.wandb_run = wandb.init(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, - job_type=job_type, - id=run_id, - allow_val_change=True) if not wandb.run else wandb.run + self.wandb_run = ( + wandb.init( + 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, + job_type=job_type, + id=run_id, + 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 @@ def setup_training(self, opt): 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( - self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \ - config.hyp, config.imgsz + opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp, opt.imgsz = ( + str(self.weights), + 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 @@ def setup_training(self, opt): 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 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False): 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', - type='model', - metadata={ - 'original_url': str(path), - 'epochs_trained': epoch + 1, - 'save period': opt.save_period, - 'project': opt.project, - 'total_epochs': opt.epochs, - '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 '']) - LOGGER.info(f'Saving model artifact on epoch {epoch + 1}') + model_artifact = wandb.Artifact( + "run_" + wandb.run.id + "_model", + type="model", + metadata={ + "original_url": str(path), + "epochs_trained": epoch + 1, + "save period": opt.save_period, + "project": opt.project, + "total_epochs": opt.epochs, + "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 ""] + ) + 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 @@ def log(self, log_dict): 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 @@ def end_epoch(self): 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(): @@ -180,7 +195,7 @@ def finish_run(self): @contextmanager def all_logging_disabled(highest_level=logging.CRITICAL): - """ source - https://gist.github.com/simon-weber/7853144 + """source - https://gist.github.com/simon-weber/7853144 A context manager that will prevent any logging messages triggered during the body from being processed. :param highest_level: the maximum logging level in use. This would only need to be changed if a custom level greater than CRITICAL is defined. diff --git a/utils/loss.py b/utils/loss.py index 26cca8797315..26b8c06bf333 100644 --- a/utils/loss.py +++ b/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 @@ def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 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 @@ def forward(self, pred, true): 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 @@ def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 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 @@ def forward(self, pred, true): 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 @@ def __init__(self, model, autobalance=False): h = model.hyp # hyperparameters # Define criteria - BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) - BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_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)) # 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 @@ def __call__(self, p, targets): # predictions, targets if self.autobalance: self.balance = [x / self.balance[self.ssi] for x in self.balance] - lbox *= self.hyp['box'] - lobj *= self.hyp['obj'] - lcls *= self.hyp['cls'] + lbox *= self.hyp["box"] + lobj *= self.hyp["obj"] + lcls *= self.hyp["cls"] bs = tobj.shape[0] # batch size return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach() @@ -183,16 +181,20 @@ def build_targets(self, p, targets): targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2) # append anchor indices g = 0.5 # bias - off = torch.tensor( - [ - [0, 0], - [1, 0], - [0, 1], - [-1, 0], - [0, -1], # j,k,l,m - # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm - ], - device=self.device).float() * g # offsets + off = ( + torch.tensor( + [ + [0, 0], + [1, 0], + [0, 1], + [-1, 0], + [0, -1], # j,k,l,m + # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm + ], + device=self.device, + ).float() + * g + ) # offsets for i in range(self.nl): anchors, shape = self.anchors[i], p[i].shape @@ -203,7 +205,7 @@ def build_targets(self, p, targets): 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 diff --git a/utils/metrics.py b/utils/metrics.py index 5646f40e9860..5f45621dc372 100644 --- a/utils/metrics.py +++ b/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=''): - """ Compute the average precision, given the recall and precision curves. +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. + 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_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, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall') + 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, 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") i = smooth(f1.mean(0), 0.1).argmax() # max F1 index p, r, f1 = p[:, i], r[:, i], f1[:, i] @@ -96,7 +96,7 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names def compute_ap(recall, precision): - """ Compute the average precision, given the recall and precision curves + """Compute the average precision, given the recall and precision curves # Arguments recall: The recall curve (list) precision: The precision curve (list) @@ -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' - if method == 'interp': + method = "interp" # methods: 'continuous', '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 @@ def __init__(self, nc, conf=0.25, iou_thres=0.45): 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 @@ def tp_fp(self): # fn = self.matrix.sum(0) - tp # false negatives (missed detections) return tp[:-1], fp[:-1] # remove background class - @TryExcept('WARNING ⚠️ ConfusionMatrix plot failure') - def plot(self, normalize=True, save_dir='', names=()): + @TryExcept("WARNING ⚠️ ConfusionMatrix plot failure") + 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 - sn.heatmap(array, - ax=ax, - annot=nc < 30, - annot_kws={ - 'size': 8}, - cmap='Blues', - fmt='.2f', - square=True, - vmin=0.0, - xticklabels=ticklabels, - yticklabels=ticklabels).set_facecolor((1, 1, 1)) - ax.set_xlabel('True') - ax.set_ylabel('Predicted') - ax.set_title('Confusion Matrix') - fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) + warnings.simplefilter("ignore") # suppress empty matrix RuntimeWarning: All-NaN slice encountered + sn.heatmap( + array, + ax=ax, + annot=nc < 30, + annot_kws={"size": 8}, + cmap="Blues", + fmt=".2f", + square=True, + vmin=0.0, + xticklabels=ticklabels, + yticklabels=ticklabels, + ).set_facecolor((1, 1, 1)) + ax.set_xlabel("True") + ax.set_ylabel("Predicted") + 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) * \ - (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).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) # Union Area union = w1 * h1 + w2 * h2 - inter + eps @@ -247,10 +250,10 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7 cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1) # convex (smallest enclosing box) width ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1) # convex height if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 - c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared + c2 = cw**2 + ch**2 + eps # convex diagonal squared rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center dist ** 2 if CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 - v = (4 / math.pi ** 2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2) + v = (4 / math.pi**2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2) with torch.no_grad(): alpha = v / (v - iou + (1 + eps)) return iou - (rho2 / c2 + v * alpha) # CIoU @@ -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) * \ - (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).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) # 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.set_xlabel('Recall') - ax.set_ylabel('Precision') + 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_ylabel("Precision") ax.set_xlim(0, 1) ax.set_ylim(0, 1) - ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left') - ax.set_title('Precision-Recall Curve') + ax.legend(bbox_to_anchor=(1.04, 1), loc="upper left") + 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.set_title(f'{ylabel}-Confidence Curve') + ax.legend(bbox_to_anchor=(1.04, 1), loc="upper left") + ax.set_title(f"{ylabel}-Confidence Curve") fig.savefig(save_dir, dpi=250) plt.close(fig) diff --git a/utils/plots.py b/utils/plots.py index 5901ca2dbfaa..11c96a6372c3 100644 --- a/utils/plots.py +++ b/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.metrics import fitness # Settings -RANK = int(os.getenv('RANK', -1)) -matplotlib.rc('font', **{'size': 11}) -matplotlib.use('Agg') # for writing to files only +RANK = int(os.getenv("RANK", -1)) +matplotlib.rc("font", **{"size": 11}) +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', - '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7') - self.palette = [self.hex2rgb(f'#{c}') for c in hexs] + hexs = ( + "FF3838", + "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): @@ -45,13 +63,13 @@ def __call__(self, i, bgr=False): @staticmethod def hex2rgb(h): # rgb order (PIL) - return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) + return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4)) 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' - not in module_type) and ('Segment' - not in module_type): # 'Detect' for Object Detect task,'Segment' for Segment task + if ("Detect" not in module_type) and ( + "Segment" not in module_type + ): # '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})') - plt.savefig(f, dpi=300, bbox_inches='tight') + LOGGER.info(f"Saving {f}... ({n}/{channels})") + 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() @@ -125,7 +143,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None): max_subplots = 16 # max image subplots, i.e. 4x4 bs, _, h, w = images.shape # batch size, _, height, width bs = min(bs, max_subplots) # limit plot images - ns = np.ceil(bs ** 0.5) # number of subplots (square) + ns = np.ceil(bs**0.5) # number of subplots (square) if np.max(images[0]) <= 1: images *= 255 # de-normalise (optional) @@ -136,7 +154,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None): break x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin im = im.transpose(1, 2, 0) - mosaic[y:y + h, x:x + w, :] = im + mosaic[y : y + h, x : x + w, :] = im # Resize (optional) scale = max_size / ns / max(h, w) @@ -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']) - plt.plot(y, '.-', label='LR') - plt.xlabel('epoch') - plt.ylabel('LR') + y.append(optimizer.param_groups[0]["lr"]) + plt.plot(y, ".-", label="LR") + plt.xlabel("epoch") + 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') - plt.savefig('hist2d.png', dpi=300) + ax.set_aspect("equal") + 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 - s = ['x targets', 'y targets', 'width targets', 'height targets'] + x = np.loadtxt("targets.txt", dtype=np.float32).T + 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], - y[3, 1:j] * 1E2, - '.-', - linewidth=2, - markersize=8, - 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], - 'k.-', - linewidth=2, - markersize=8, - alpha=.25, - label='EfficientDet') + ax2.plot( + y[5, 1:j], + y[3, 1:j] * 1e2, + ".-", + linewidth=2, + markersize=8, + 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], + "k.-", + linewidth=2, + markersize=8, + alpha=0.25, + 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_ylabel('COCO AP val') - ax2.legend(loc='lower right') - f = save_dir / 'study.png' - print(f'Saving {f}...') + ax2.set_xlabel("GPU Speed (ms/img)") + ax2.set_ylabel("COCO AP val") + ax2.legend(loc="lower right") + f = save_dir / "study.png" + 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)) - plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200) + 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.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') - 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) + ax[0].set_xlabel("classes") + 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) # rectangles labels[:, 1:3] = 0.5 # center @@ -305,47 +328,48 @@ 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) - matplotlib.use('Agg') + plt.savefig(save_dir / "labels.jpg", dpi=200) + 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)] - blocks = torch.chunk(denormalize(im.clone()).cpu().float(), len(im), - dim=0) # select batch index 0, block by channels + names = names or [f"class{i}" for i in range(1000)] + blocks = torch.chunk( + 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 + 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 ax = ax.ravel() if m > 1 else [ax] # 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 '') - ax[i].set_title(s, fontsize=8, verticalalignment='top') - plt.savefig(f, dpi=300, bbox_inches='tight') + s = names[labels[i]] + (f"—{names[pred[i]]}" if pred is not None else "") + ax[i].set_title(s, fontsize=8, verticalalignment="top") + 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}) - print(f'Best results from row {j} of {evolve_csv}:') + matplotlib.rc("font", **{"size": 8}) + 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.plot(mu, f.max(), 'k+', markersize=15) - plt.title(f'{k} = {mu:.3g}', fontdict={'size': 9}) # limit to 40 characters + 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.title(f"{k} = {mu:.3g}", fontdict={"size": 9}) # limit to 40 characters if i % 5 != 0: plt.yticks([]) - print(f'{k:>15}: {mu:.3g}') - f = evolve_csv.with_suffix('.png') # filename + print(f"{k:>15}: {mu:.3g}") + 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')) - assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.' + files = list(save_dir.glob("results*.csv")) + assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot." for f in files: try: data = pd.read_csv(f) 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, gaussian_filter1d(y, sigma=3), ':', label='smooth', linewidth=2) # smoothing line + 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].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'] - files = list(Path(save_dir).glob('frames*.txt')) + 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")) 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_', '') - a.plot(t, results[i], marker='.', label=label, linewidth=1, markersize=5) + label = labels[fi] if len(labels) else f.stem.replace("frames_", "") + 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 @@ -439,10 +463,10 @@ def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad xyxy = xywh2xyxy(b).long() clip_boxes(xyxy, im.shape) - crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)] + 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 diff --git a/utils/segment/augmentations.py b/utils/segment/augmentations.py index f8154b834869..56636b65d93a 100644 --- a/utils/segment/augmentations.py +++ b/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, - targets=(), - segments=(), - degrees=10, - translate=.1, - scale=.1, - shear=10, - perspective=0.0, - border=(0, 0)): +def random_perspective( + im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.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[1, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * height) # y 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) # 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) diff --git a/utils/segment/dataloaders.py b/utils/segment/dataloaders.py index 5398617eef68..b0b3a7424216 100644 --- a/utils/segment/dataloaders.py +++ b/utils/segment/dataloaders.py @@ -1,7 +1,5 @@ # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license -""" -Dataloaders -""" +"""Dataloaders.""" import os import random @@ -17,30 +15,32 @@ from ..torch_utils import torch_distributed_zero_first from .augmentations import mixup, random_perspective -RANK = int(os.getenv('RANK', -1)) - - -def create_dataloader(path, - imgsz, - batch_size, - stride, - single_cls=False, - hyp=None, - augment=False, - cache=False, - pad=0.0, - rect=False, - rank=-1, - workers=8, - image_weights=False, - quad=False, - prefix='', - shuffle=False, - mask_downsample_ratio=1, - overlap_mask=False, - seed=0): +RANK = int(os.getenv("RANK", -1)) + + +def create_dataloader( + path, + imgsz, + batch_size, + stride, + single_cls=False, + hyp=None, + augment=False, + cache=False, + pad=0.0, + rect=False, + rank=-1, + workers=8, + image_weights=False, + quad=False, + prefix="", + shuffle=False, + mask_downsample_ratio=1, + overlap_mask=False, + 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 @@ def __init__( 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, - stride, pad, min_items, prefix, rank, seed) + super().__init__( + 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 @@ def __getitem__(self, index): 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 @@ def __getitem__(self, index): 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 @@ def __getitem__(self, index): 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, - labels, - segments=segments, - degrees=hyp['degrees'], - translate=hyp['translate'], - scale=hyp['scale'], - shear=hyp['shear'], - perspective=hyp['perspective']) + img, labels, segments = random_perspective( + img, + labels, + segments=segments, + degrees=hyp["degrees"], + translate=hyp["translate"], + scale=hyp["scale"], + shear=hyp["shear"], + 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], - segments, - downsample_ratio=self.downsample_ratio) + masks, sorted_idx = polygons2masks_overlap( + img.shape[:2], segments, 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] // - self.downsample_ratio, img.shape[1] // - self.downsample_ratio)) + masks = ( + torch.from_numpy(masks) + 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 @@ def __getitem__(self, index): 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 @@ def load_mosaic(self, index): # img4, labels4 = replicate(img4, labels4) # replicate # Augment - img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste']) - img4, labels4, segments4 = random_perspective(img4, - labels4, - segments4, - degrees=self.hyp['degrees'], - translate=self.hyp['translate'], - scale=self.hyp['scale'], - shear=self.hyp['shear'], - perspective=self.hyp['perspective'], - border=self.mosaic_border) # border to remove + img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"]) + img4, labels4, segments4 = random_perspective( + img4, + labels4, + segments4, + degrees=self.hyp["degrees"], + translate=self.hyp["translate"], + scale=self.hyp["scale"], + shear=self.hyp["shear"], + perspective=self.hyp["perspective"], + 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), - dtype=np.int32 if len(segments) > 255 else np.uint8) + masks = np.zeros( + (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)): diff --git a/utils/segment/general.py b/utils/segment/general.py index f1b2f1dd120f..8cbc745b4a90 100644 --- a/utils/segment/general.py +++ b/utils/segment/general.py @@ -6,8 +6,7 @@ def crop_mask(masks, boxes): """ - "Crop" predicted masks by zeroing out everything not in the predicted bbox. - Vectorized by Chong (thanks Chong). + "Crop" predicted masks by zeroing out everything not in the predicted bbox. 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 diff --git a/utils/segment/loss.py b/utils/segment/loss.py index caeff3cad586..1e007271fa9c 100644 --- a/utils/segment/loss.py +++ b/utils/segment/loss.py @@ -18,14 +18,14 @@ def __init__(self, model, autobalance=False, overlap=False): h = model.hyp # hyperparameters # Define criteria - BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) - BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_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)) # 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 @@ def __call__(self, preds, targets, masks): # predictions, targets, model # 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 @@ def __call__(self, preds, targets, masks): # predictions, targets, model if self.autobalance: self.balance = [x / self.balance[self.ssi] for x in self.balance] - lbox *= self.hyp['box'] - lobj *= self.hyp['obj'] - lcls *= self.hyp['cls'] - lseg *= self.hyp['box'] / bs + lbox *= self.hyp["box"] + lobj *= self.hyp["obj"] + lcls *= self.hyp["cls"] + lseg *= self.hyp["box"] / bs loss = lbox + lobj + lcls + lseg return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach() @@ -111,7 +111,7 @@ def __call__(self, preds, targets, masks): # predictions, targets, model 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,16 +132,20 @@ def build_targets(self, p, targets): targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None], ti[..., None]), 2) # append anchor indices g = 0.5 # bias - off = torch.tensor( - [ - [0, 0], - [1, 0], - [0, 1], - [-1, 0], - [0, -1], # j,k,l,m - # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm - ], - device=self.device).float() * g # offsets + off = ( + torch.tensor( + [ + [0, 0], + [1, 0], + [0, 1], + [-1, 0], + [0, -1], # j,k,l,m + # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm + ], + device=self.device, + ).float() + * g + ) # offsets for i in range(self.nl): anchors, shape = self.anchors[i], p[i].shape @@ -152,7 +156,7 @@ def build_targets(self, p, targets): 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 diff --git a/utils/segment/metrics.py b/utils/segment/metrics.py index 787961bee1bf..222a749b5986 100644 --- a/utils/segment/metrics.py +++ b/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 @@ -15,14 +13,14 @@ def fitness(x): def ap_per_class_box_and_mask( - tp_m, - tp_b, - conf, - pred_cls, - target_cls, - plot=False, - save_dir='.', - names=(), + tp_m, + tp_b, + conf, + pred_cls, + target_cls, + plot=False, + save_dir=".", + names=(), ): """ Args: @@ -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, - conf, - pred_cls, - target_cls, - plot=plot, - save_dir=save_dir, - names=names, - prefix='Box')[2:] - results_masks = ap_per_class(tp_m, - conf, - pred_cls, - target_cls, - plot=plot, - save_dir=save_dir, - names=names, - prefix='Mask')[2:] + results_boxes = ap_per_class( + tp_b, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Box" + )[2:] + results_masks = ap_per_class( + tp_m, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Mask" + )[2:] results = { - 'boxes': { - 'p': results_boxes[0], - 'r': results_boxes[1], - 'ap': results_boxes[3], - 'f1': results_boxes[2], - 'ap_class': results_boxes[4]}, - 'masks': { - 'p': results_masks[0], - 'r': results_masks[1], - 'ap': results_masks[3], - 'f1': results_masks[2], - 'ap_class': results_masks[4]}} + "boxes": { + "p": results_boxes[0], + "r": results_boxes[1], + "ap": results_boxes[3], + "f1": results_boxes[2], + "ap_class": results_boxes[4], + }, + "masks": { + "p": results_masks[0], + "r": results_masks[1], + "ap": results_masks[3], + "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 @@ def __init__(self) -> None: @property def ap50(self): - """AP@0.5 of all classes. + """ + AP@0.5 of all classes. + Return: (nc, ) or []. """ @@ -90,7 +82,9 @@ def ap(self): @property def mp(self): - """mean precision of all classes. + """ + Mean precision of all classes. + Return: float. """ @@ -98,7 +92,9 @@ def mp(self): @property def mr(self): - """mean recall of all classes. + """ + Mean recall of all classes. + Return: float. """ @@ -106,7 +102,9 @@ def mr(self): @property def map50(self): - """Mean AP@0.5 of all classes. + """ + Mean AP@0.5 of all classes. + Return: float. """ @@ -114,18 +112,20 @@ def map50(self): @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 @@ def update(self, results): Args: results: Dict{'boxes': Dict{}, 'masks': Dict{}} """ - self.metric_box.update(list(results['boxes'].values())) - self.metric_mask.update(list(results['masks'].values())) + self.metric_box.update(list(results["boxes"].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 @@ def ap_class_index(self): KEYS = [ - 'train/box_loss', - 'train/seg_loss', # train loss - 'train/obj_loss', - 'train/cls_loss', - 'metrics/precision(B)', - 'metrics/recall(B)', - 'metrics/mAP_0.5(B)', - 'metrics/mAP_0.5:0.95(B)', # metrics - 'metrics/precision(M)', - 'metrics/recall(M)', - 'metrics/mAP_0.5(M)', - 'metrics/mAP_0.5:0.95(M)', # metrics - 'val/box_loss', - 'val/seg_loss', # val loss - 'val/obj_loss', - 'val/cls_loss', - 'x/lr0', - 'x/lr1', - 'x/lr2', ] + "train/box_loss", + "train/seg_loss", # train loss + "train/obj_loss", + "train/cls_loss", + "metrics/precision(B)", + "metrics/recall(B)", + "metrics/mAP_0.5(B)", + "metrics/mAP_0.5:0.95(B)", # metrics + "metrics/precision(M)", + "metrics/recall(M)", + "metrics/mAP_0.5(M)", + "metrics/mAP_0.5:0.95(M)", # metrics + "val/box_loss", + "val/seg_loss", # val loss + "val/obj_loss", + "val/cls_loss", + "x/lr0", + "x/lr1", + "x/lr2", +] BEST_KEYS = [ - 'best/epoch', - 'best/precision(B)', - 'best/recall(B)', - 'best/mAP_0.5(B)', - 'best/mAP_0.5:0.95(B)', - 'best/precision(M)', - 'best/recall(M)', - 'best/mAP_0.5(M)', - 'best/mAP_0.5:0.95(M)', ] + "best/epoch", + "best/precision(B)", + "best/recall(B)", + "best/mAP_0.5(B)", + "best/mAP_0.5:0.95(B)", + "best/precision(M)", + "best/recall(M)", + "best/mAP_0.5(M)", + "best/mAP_0.5:0.95(M)", +] diff --git a/utils/segment/plots.py b/utils/segment/plots.py index f9938cd1b06a..0e30c61be66f 100644 --- a/utils/segment/plots.py +++ b/utils/segment/plots.py @@ -14,7 +14,7 @@ @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() @@ -27,7 +27,7 @@ def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg' max_subplots = 16 # max image subplots, i.e. 4x4 bs, _, h, w = images.shape # batch size, _, height, width bs = min(bs, max_subplots) # limit plot images - ns = np.ceil(bs ** 0.5) # number of subplots (square) + ns = np.ceil(bs**0.5) # number of subplots (square) if np.max(images[0]) <= 1: images *= 255 # de-normalise (optional) @@ -38,7 +38,7 @@ def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg' break x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin im = im.transpose(1, 2, 0) - mosaic[y:y + h, x:x + w, :] = im + mosaic[y : y + h, x : x + w, :] = im # Resize (optional) scale = max_size / ns / max(h, w) @@ -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')) - assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.' + files = list(save_dir.glob("results*.csv")) + assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot." for f in files: try: data = pd.read_csv(f) - index = np.argmax(0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] + - 0.1 * data.values[:, 11]) + index = np.argmax( + 0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] + 0.1 * data.values[:, 11] + ) s = [x.strip() for x in data.columns] x = data.values[:, 0] for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]): y = data.values[:, j] # y[y == 0] = np.nan # don't show zero values - ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=2) + ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=2) if best: # best - ax[i].scatter(index, y[index], color='r', label=f'best:{index}', marker='*', linewidth=3) - ax[i].set_title(s[j] + f'\n{round(y[index], 5)}') + ax[i].scatter(index, y[index], color="r", label=f"best:{index}", marker="*", linewidth=3) + ax[i].set_title(s[j] + f"\n{round(y[index], 5)}") else: # last - ax[i].scatter(x[-1], y[-1], color='r', label='last', marker='*', linewidth=3) - ax[i].set_title(s[j] + f'\n{round(y[-1], 5)}') + ax[i].scatter(x[-1], y[-1], color="r", label="last", marker="*", linewidth=3) + ax[i].set_title(s[j] + f"\n{round(y[-1], 5)}") # if j in [8, 9, 10]: # share train and val loss y axes # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) except Exception as e: - print(f'Warning: Plotting error for {f}: {e}') + 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() diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 13a356f3238c..6bc4b4c7fd04 100644 --- a/utils/torch_utils.py +++ b/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 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 -RANK = int(os.getenv('RANK', -1)) -WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) +LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html +RANK = int(os.getenv("RANK", -1)) +WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1)) try: import thop # for FLOPs computation @@ -31,11 +29,11 @@ thop = None # Suppress PyTorch warnings -warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling') -warnings.filterwarnings('ignore', category=UserWarning) +warnings.filterwarnings("ignore", message="User provided device_type of 'cuda', but CUDA is not available. Disabling") +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 decorate(fn): 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), \ - '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' - if check_version(torch.__version__, '1.11.0'): + 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. " + "Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395" + ) + 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__} ' - device = str(device).strip().lower().replace('cuda:', '').replace('none', '') # to string, 'cuda:0' to '0' - cpu = device == 'cpu' - mps = device == 'mps' # Apple Metal Performance Shaders (MPS) + 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' + cpu = device == "cpu" + 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() - assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \ - f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)" + 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(",", "") + ), 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}' - space = ' ' * (len(s) + 1) + assert batch_size % n == 0, f"batch-size {batch_size} not multiple of GPU count {n}" + 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' - elif mps and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available(): # prefer MPS if available - s += 'MPS\n' - arg = 'mps' + arg = "cuda:0" + elif mps and getattr(torch, "has_mps", False) and torch.backends.mps.is_available(): # prefer MPS if available + s += "MPS\n" + arg = "mps" else: # revert to CPU - s += 'CPU\n' - arg = 'cpu' + s += "CPU\n" + arg = "cpu" if not newline: s = s.rstrip() @@ -149,7 +150,7 @@ def time_sync(): def profile(input, ops, n=10, device=None): - """ YOLOv5 speed/memory/FLOPs profiler + """YOLOv5 speed/memory/FLOPs profiler Usage: input = torch.randn(16, 3, 640, 640) m1 = lambda x: x * torch.sigmoid(x) @@ -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}" - f"{'input':>24s}{'output':>24s}") + print( + 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.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m + 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 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) - s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' for x in (x, y)) # shapes + 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 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.remove(m, 'weight') # make permanent - LOGGER.info(f'Model pruned to {sparsity(model):.3g} global sparsity') + prune.l1_unstructured(m, name="weight", amount=amount) # prune + prune.remove(m, "weight") # make permanent + 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, - 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) + 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) + ) # 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.', '') - print('%5g %40s %9s %12g %20s %10.3g %10.3g' % - (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) + name = name.replace("module_list.", "") + print( + "%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' - LOGGER.info(f'{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}') + 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}") 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.') - - 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) - LOGGER.info(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') + 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[1], "weight_decay": 0.0}) # add g1 (BatchNorm2d weights) + LOGGER.info( + 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'): - kwargs['skip_validation'] = True # validation causes GitHub API rate limit errors - if check_version(torch.__version__, '1.12.0'): - kwargs['trust_repo'] = True # argument required starting in torch 0.12 + if check_version(torch.__version__, "1.9.1"): + kwargs["skip_validation"] = True # validation causes GitHub API rate limit errors + if check_version(torch.__version__, "1.12.0"): + 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 - if ckpt['optimizer'] is not None: - optimizer.load_state_dict(ckpt['optimizer']) # optimizer - best_fitness = ckpt['best_fitness'] - if ema and ckpt.get('ema'): - ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) # EMA - ema.updates = ckpt['updates'] + start_epoch = ckpt["epoch"] + 1 + if ckpt["optimizer"] is not None: + optimizer.load_state_dict(ckpt["optimizer"]) # optimizer + best_fitness = ckpt["best_fitness"] + if ema and ckpt.get("ema"): + ema.ema.load_state_dict(ckpt["ema"].float().state_dict()) # EMA + ema.updates = ckpt["updates"] if resume: - 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') + 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") 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,15 +410,17 @@ def __call__(self, epoch, fitness): 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. ' - f'Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n' - 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.') + LOGGER.info( + f"Stopping training early as no improvement observed in last {self.patience} epochs. " + f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n" + 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 class ModelEMA: - """ Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models + """Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models Keeps a moving average of everything in the model state_dict (parameters and buffers) For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage """ @@ -427,6 +445,6 @@ def update(self, model): 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) diff --git a/utils/triton.py b/utils/triton.py index b5153dad940d..9584d07fbcf0 100644 --- a/utils/triton.py +++ b/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 @@ 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 - as input and returns them as outputs. + """ + 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 as input and returns them as + outputs. """ def __init__(self, url: str): @@ -21,7 +22,7 @@ def __init__(self, url: str): """ 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 @@ def __init__(self, url: str): 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""" - return self.metadata.get('backend', self.metadata.get('platform')) + """Returns the model runtime.""" + 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. - kwargs are matched with the model input names. + """ + 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. 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']: - tensor = torch.as_tensor(response.as_numpy(output['name'])) + for output in self.metadata["outputs"]: + 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: diff --git a/val.py b/val.py index 1a4219c38962..6cc1d37a0a26 100644 --- a/val.py +++ b/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 @@ 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, - check_yaml, coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, - print_args, scale_boxes, xywh2xyxy, xyxy2xywh) +from utils.general import ( + LOGGER, + 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: - f.write(('%g ' * len(line)).rstrip() % line + '\n') + with open(file, "a") as f: + 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({ - 'image_id': image_id, - 'category_id': class_map[int(p[5])], - 'bbox': [round(x, 3) for x in b], - 'score': round(p[4], 5)}) + jdict.append( + { + "image_id": image_id, + "category_id": class_map[int(p[5])], + "bbox": [round(x, 3) for x in b], + "score": round(p[4], 5), + } + ) def process_batch(detections, labels, iouv): @@ -98,47 +115,47 @@ def process_batch(detections, labels, iouv): @smart_inference_mode() def run( - data, - weights=None, # model.pt path(s) - batch_size=32, # batch size - imgsz=640, # inference size (pixels) - conf_thres=0.001, # confidence threshold - iou_thres=0.6, # NMS IoU threshold - max_det=300, # maximum detections per image - task='val', # train, val, test, speed or study - device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu - workers=8, # max dataloader workers (per RANK in DDP mode) - single_cls=False, # treat as single-class dataset - augment=False, # augmented inference - verbose=False, # verbose output - save_txt=False, # save results to *.txt - 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 - 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(''), - plots=True, - callbacks=Callbacks(), - compute_loss=None, + data, + weights=None, # model.pt path(s) + batch_size=32, # batch size + imgsz=640, # inference size (pixels) + conf_thres=0.001, # confidence threshold + iou_thres=0.6, # NMS IoU threshold + max_det=300, # maximum detections per image + task="val", # train, val, test, speed or study + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + workers=8, # max dataloader workers (per RANK in DDP mode) + single_cls=False, # treat as single-class dataset + augment=False, # augmented inference + verbose=False, # verbose output + save_txt=False, # save results to *.txt + 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 + 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(""), + plots=True, + callbacks=Callbacks(), + compute_loss=None, ): # Initialize/load model and set device 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' - 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 + cuda = device.type != "cpu" + 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 iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for mAP@0.5:0.95 niou = iouv.numel() @@ -168,36 +185,40 @@ 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} ' \ - f'classes). Pass correct combination of --weights and --data that are trained together.' + assert ncm == nc, ( + 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 - task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images - dataloader = create_dataloader(data[task], - imgsz, - batch_size, - stride, - single_cls, - pad=pad, - rect=rect, - workers=workers, - prefix=colorstr(f'{task}: '))[0] + 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 + dataloader = create_dataloader( + data[task], + imgsz, + batch_size, + stride, + single_cls, + pad=pad, + rect=rect, + workers=workers, + 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, - conf_thres, - iou_thres, - labels=lb, - multi_label=True, - agnostic=single_cls, - max_det=max_det) + preds = non_max_suppression( + preds, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls, max_det=max_det + ) # Metrics for si, pred in enumerate(preds): @@ -259,18 +276,18 @@ def run( # Save/log if save_txt: - (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_dir / "labels").mkdir(parents=True, exist_ok=True) + 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, output_to_target(preds), paths, save_dir / f'val_batch{batch_i}_pred.jpg', names) # pred + 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 - 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 - LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) + pf = "%22s" + "%11i" * 2 + "%11.3g" * 4 # print format + 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 - anno_json = str(Path('../datasets/coco/annotations/instances_val2017.json')) # annotations + 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 if not os.path.exists(anno_json): - anno_json = os.path.join(data['path'], 'annotations', 'instances_val2017.json') - pred_json = str(save_dir / f'{w}_predictions.json') # predictions - LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...') - with open(pred_json, 'w') as f: + anno_json = os.path.join(data["path"], "annotations", "instances_val2017.json") + pred_json = str(save_dir / f"{w}_predictions.json") # predictions + LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...") + 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('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)') - parser.add_argument('--batch-size', type=int, default=32, help='batch size') - parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') - parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold') - parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold') - parser.add_argument('--max-det', type=int, default=300, help='maximum detections per image') - parser.add_argument('--task', default='val', help='train, val, test, speed or study') - parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)') - 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('--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-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-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('--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('--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("--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("--batch-size", type=int, default=32, help="batch size") + parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)") + parser.add_argument("--conf-thres", type=float, default=0.001, help="confidence threshold") + parser.add_argument("--iou-thres", type=float, default=0.6, help="NMS IoU threshold") + parser.add_argument("--max-det", type=int, default=300, help="maximum detections per image") + parser.add_argument("--task", default="val", help="train, val, test, speed or study") + parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)") + 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("--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-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-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("--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("--half", action="store_true", help="use FP16 half-precision 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 - if opt.task == 'speed': # speed benchmarks + opt.half = torch.cuda.is_available() and opt.device != "cpu" # FP16 for fastest results + 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 - subprocess.run(['zip', '-r', 'study.zip', 'study_*.txt']) + np.savetxt(f, y, fmt="%10.4g") # save + 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)