From cac7189c3e11de563056e2db08012c2c12b30789 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 3 Jan 2024 08:04:47 +0100 Subject: [PATCH] Create format.yml (#2165) * Create format.yml * Auto-format by Ultralytics actions * Update metrics.py * Auto-format by Ultralytics actions * Update pyproject.toml * Update CONTRIBUTING.md * Auto-format by Ultralytics actions --- .github/workflows/format.yml | 23 + CONTRIBUTING.md | 39 +- README.md | 8 +- README.zh-CN.md | 3 +- benchmarks.py | 116 ++--- classify/predict.py | 127 ++--- classify/train.py | 255 +++++----- classify/val.py | 81 ++-- detect.py | 176 +++---- export.py | 592 ++++++++++++----------- hubconf.py | 72 +-- models/common.py | 312 +++++++------ models/experimental.py | 35 +- models/tf.py | 245 ++++++---- models/yolo.py | 115 +++-- segment/predict.py | 143 +++--- segment/train.py | 522 ++++++++++++--------- segment/val.py | 297 ++++++------ train.py | 546 ++++++++++++---------- 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 | 593 +++++++++++++----------- utils/downloads.py | 91 ++-- utils/flask_rest_api/README.md | 7 +- utils/flask_rest_api/example_request.py | 12 +- utils/flask_rest_api/restapi.py | 28 +- utils/general.py | 568 ++++++++++++++--------- utils/loggers/__init__.py | 170 +++---- utils/loggers/clearml/README.md | 11 +- utils/loggers/clearml/clearml_utils.py | 100 ++-- utils/loggers/clearml/hpo.py | 78 ++-- utils/loggers/comet/README.md | 26 +- 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 | 250 +++++----- utils/segment/augmentations.py | 22 +- utils/segment/dataloaders.py | 147 +++--- 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 | 267 ++++++----- 53 files changed, 4136 insertions(+), 3459 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 0000000000..27b2c7d890 --- /dev/null +++ b/.github/workflows/format.yml @@ -0,0 +1,23 @@ +# 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: + branches: [main,master] + +jobs: + format: + runs-on: ubuntu-latest + steps: + - name: Run Ultralytics Formatting + uses: ultralytics/actions@main + with: + python: true + docstrings: true + markdown: true + spelling: true diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7995069f88..0dbea33b38 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -8,8 +8,7 @@ We love your input! We want to make contributing to YOLOv5 as easy and transpare - Proposing a new feature - Becoming a maintainer -YOLOv5 works so well due to our combined community effort, and for every small improvement you contribute you will be -helping push the frontiers of what's possible in AI 😃! +YOLOv5 works so well due to our combined community effort, and for every small improvement you contribute you will be helping push the frontiers of what's possible in AI 😃! ## Submitting a Pull Request (PR) 🛠️ @@ -35,9 +34,7 @@ Change the `matplotlib` version from `3.2.2` to `3.3`. ### 4. Preview Changes and Submit PR -Click on the **Preview changes** tab to verify your updates. At the bottom of the screen select 'Create a **new branch** -for this commit', assign your branch a descriptive name such as `fix/matplotlib_version` and click the green **Propose -changes** button. All done, your PR is now submitted to YOLOv5 for review and approval 😃! +Click on the **Preview changes** tab to verify your updates. At the bottom of the screen select 'Create a **new branch** for this commit', assign your branch a descriptive name such as `fix/matplotlib_version` and click the green **Propose changes** button. All done, your PR is now submitted to YOLOv5 for review and approval 😃!

PR_step4

@@ -45,8 +42,7 @@ changes** button. All done, your PR is now submitted to YOLOv5 for review and ap To allow your work to be integrated as seamlessly as possible, we advise you to: -- ✅ Verify your PR is **up-to-date** with `ultralytics/yolov5` `master` branch. If your PR is behind you can update - your code by clicking the 'Update branch' button or by running `git pull` and `git merge master` locally. +- ✅ Verify your PR is **up-to-date** with `ultralytics/yolov5` `master` branch. If your PR is behind you can update your code by clicking the 'Update branch' button or by running `git pull` and `git merge master` locally.

Screenshot 2022-08-29 at 22 47 15

@@ -54,40 +50,27 @@ To allow your work to be integrated as seamlessly as possible, we advise you to:

Screenshot 2022-08-29 at 22 47 03

-- ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase - but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ — Bruce Lee +- ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ — Bruce Lee ## Submitting a Bug Report 🐛 If you spot a problem with YOLOv5 please submit a Bug Report! -For us to start investigating a possible problem we need to be able to reproduce it ourselves first. We've created a few -short guidelines below to help users provide what we need to get started. +For us to start investigating a possible problem we need to be able to reproduce it ourselves first. We've created a few short guidelines below to help users provide what we need to get started. -When asking a question, people will be better able to provide help if you provide **code** that they can easily -understand and use to **reproduce** the problem. This is referred to by community members as creating -a [minimum reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/). Your code that reproduces -the problem should be: +When asking a question, people will be better able to provide help if you provide **code** that they can easily understand and use to **reproduce** the problem. This is referred to by community members as creating a [minimum reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/). Your code that reproduces the problem should be: - ✅ **Minimal** – Use as little code as possible that still produces the same problem - ✅ **Complete** – Provide **all** parts someone else needs to reproduce your problem in the question itself - ✅ **Reproducible** – Test the code you're about to provide to make sure it reproduces the problem -In addition to the above requirements, for [Ultralytics](https://ultralytics.com/) to provide assistance your code -should be: +In addition to the above requirements, for [Ultralytics](https://ultralytics.com/) to provide assistance your code should be: -- ✅ **Current** – Verify that your code is up-to-date with the current - GitHub [master](https://github.com/ultralytics/yolov5/tree/master), and if necessary `git pull` or `git clone` a new - copy to ensure your problem has not already been resolved by previous commits. -- ✅ **Unmodified** – Your problem must be reproducible without any modifications to the codebase in this - repository. [Ultralytics](https://ultralytics.com/) does not provide support for custom code ⚠️. +- ✅ **Current** – Verify that your code is up-to-date with the current GitHub [master](https://github.com/ultralytics/yolov5/tree/master), and if necessary `git pull` or `git clone` a new copy to ensure your problem has not already been resolved by previous commits. +- ✅ **Unmodified** – Your problem must be reproducible without any modifications to the codebase in this repository. [Ultralytics](https://ultralytics.com/) does not provide support for custom code ⚠️. -If you believe your problem meets all of the above criteria, please close this issue and raise a new one using the 🐛 -**Bug Report** [template](https://github.com/ultralytics/yolov5/issues/new/choose) and provide -a [minimum reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/) to help us better -understand and diagnose your problem. +If you believe your problem meets all of the above criteria, please close this issue and raise a new one using the 🐛 **Bug Report** [template](https://github.com/ultralytics/yolov5/issues/new/choose) and provide a [minimum reproducible example](https://docs.ultralytics.com/help/minimum_reproducible_example/) to help us better understand and diagnose your problem. ## License -By contributing, you agree that your contributions will be licensed under -the [AGPL-3.0 license](https://choosealicense.com/licenses/agpl-3.0/) +By contributing, you agree that your contributions will be licensed under the [AGPL-3.0 license](https://choosealicense.com/licenses/agpl-3.0/) diff --git a/README.md b/README.md index 87577ded1d..00cc1abf22 100644 --- a/README.md +++ b/README.md @@ -65,9 +65,7 @@ See the [YOLOv3 Docs](https://docs.ultralytics.com) for full documentation on tr
Install -Clone repo and install [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) in a -[**Python>=3.7.0**](https://www.python.org/) environment, including -[**PyTorch>=1.7**](https://pytorch.org/get-started/locally/). +Clone repo and install [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) in a [**Python>=3.7.0**](https://www.python.org/) environment, including [**PyTorch>=1.7**](https://pytorch.org/get-started/locally/). ```bash git clone https://github.com/ultralytics/yolov3 # clone @@ -123,9 +121,7 @@ python detect.py --weights yolov5s.pt --source 0 #
Training -The commands below reproduce YOLOv3 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) -results. [Models](https://github.com/ultralytics/yolov5/tree/master/models) -and [datasets](https://github.com/ultralytics/yolov5/tree/master/data) download automatically from the latest YOLOv3 [release](https://github.com/ultralytics/yolov5/releases). Training times for YOLOv5n/s/m/l/x are 1/2/4/6/8 days on a V100 GPU ([Multi-GPU](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training) times faster). Use the largest `--batch-size` possible, or pass `--batch-size -1` for YOLOv3 [AutoBatch](https://github.com/ultralytics/yolov5/pull/5092). Batch sizes shown for V100-16GB. +The commands below reproduce YOLOv3 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) results. [Models](https://github.com/ultralytics/yolov5/tree/master/models) and [datasets](https://github.com/ultralytics/yolov5/tree/master/data) download automatically from the latest YOLOv3 [release](https://github.com/ultralytics/yolov5/releases). Training times for YOLOv5n/s/m/l/x are 1/2/4/6/8 days on a V100 GPU ([Multi-GPU](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training) times faster). Use the largest `--batch-size` possible, or pass `--batch-size -1` for YOLOv3 [AutoBatch](https://github.com/ultralytics/yolov5/pull/5092). Batch sizes shown for V100-16GB. ```bash python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml --batch-size 128 diff --git a/README.zh-CN.md b/README.zh-CN.md index 355173bf98..e9defd022e 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -121,8 +121,7 @@ python detect.py --weights yolov5s.pt --source 0 #
训练 -下面的命令重现 YOLOv3 在 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) 数据集上的结果。 最新的 [模型](https://github.com/ultralytics/yolov5/tree/master/models) 和 [数据集](https://github.com/ultralytics/yolov5/tree/master/data) -将自动的从 YOLOv3 [release](https://github.com/ultralytics/yolov5/releases) 中下载。 YOLOv5n/s/m/l/x 在 V100 GPU 的训练时间为 1/2/4/6/8 天（ [多GPU](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training) 训练速度更快）。 尽可能使用更大的 `--batch-size` ，或通过 `--batch-size -1` 实现 YOLOv3 [自动批处理](https://github.com/ultralytics/yolov5/pull/5092) 。下方显示的 batchsize 适用于 V100-16GB。 +下面的命令重现 YOLOv3 在 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) 数据集上的结果。 最新的 [模型](https://github.com/ultralytics/yolov5/tree/master/models) 和 [数据集](https://github.com/ultralytics/yolov5/tree/master/data) 将自动的从 YOLOv3 [release](https://github.com/ultralytics/yolov5/releases) 中下载。 YOLOv5n/s/m/l/x 在 V100 GPU 的训练时间为 1/2/4/6/8 天（ [多GPU](https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training) 训练速度更快）。 尽可能使用更大的 `--batch-size` ，或通过 `--batch-size -1` 实现 YOLOv3 [自动批处理](https://github.com/ultralytics/yolov5/pull/5092) 。下方显示的 batchsize 适用于 V100-16GB。 ```bash python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml --batch-size 128 diff --git a/benchmarks.py b/benchmarks.py index cbc5b20499..2f7d58bd3a 100644 --- a/benchmarks.py +++ b/benchmarks.py @@ -1,6 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Run YOLOv3 benchmarks on all supported export formats +Run YOLOv3 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 / 'yolov3-tiny.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 / "yolov3-tiny.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 692cd607b8..17dce7cd16 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 1a3e0511f4..9b6aed86a8 100644 --- a/classify/train.py +++ b/classify/train.py @@ -1,6 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Train a YOLOv3 classifier model on a classification dataset +Train a YOLOv3 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 YOLOv3 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 6e2deca78e..3875001e85 100644 --- a/classify/val.py +++ b/classify/val.py @@ -1,6 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Validate a trained YOLOv3 classification model on a classification dataset +Validate a trained YOLOv3 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(), Profile(), Profile()) 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 8ef95761f0..4502e5fc35 100644 --- a/detect.py +++ b/detect.py @@ -46,53 +46,67 @@ 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_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_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) @@ -140,14 +154,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 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)) @@ -165,20 +179,20 @@ def run( if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format - with open(f'{txt_path}.txt', 'a') as f: - 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]) @@ -187,7 +201,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 @@ -200,18 +214,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) @@ -219,37 +233,35 @@ def run( def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--weights', - nargs='+', - type=str, - default=ROOT / 'yolov3-tiny.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-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 / "yolov3-tiny.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-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)) @@ -257,10 +269,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 93f2b888e5..8cacd97201 100644 --- a/export.py +++ b/export.py @@ -64,30 +64,42 @@ ROOT = FILE.parents[0] # YOLOv3 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(): # YOLOv3 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:")): # YOLOv3 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:")): # YOLOv3 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 = file.with_suffix('.onnx') + LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") + f = 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:")): # YOLOv3 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:")): # YOLOv3 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:")): # YOLOv3 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:")): # YOLOv3 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:"), +): # YOLOv3 TensorFlow SavedModel export try: import tensorflow as tf @@ -396,8 +410,8 @@ def export_saved_model(model, from models.tf import TFModel - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - f = str(file).replace('.pt', '_saved_model') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + 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) @@ -409,7 +423,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 @@ -418,21 +432,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:")): # YOLOv3 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)) @@ -443,13 +460,13 @@ def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')): @try_export -def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')): +def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")): # YOLOv3 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] @@ -457,84 +474,95 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c 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 = [] converter.inference_input_type = tf.uint8 # or tf.int8 converter.inference_output_type = tf.uint8 # or tf.int8 converter.experimental_new_quantizer = 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:")): # YOLOv3 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:")): # YOLOv3 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 @@ -547,8 +575,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() @@ -572,22 +600,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:")): # YOLOv3 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() # YOLOv3 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) @@ -597,7 +625,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) @@ -631,8 +659,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): @@ -648,10 +676,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 @@ -659,10 +687,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) @@ -673,72 +705,76 @@ 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 YOLOv3 Detect() inplace=True - keras=False, # use Keras - optimize=False, # TorchScript: optimize for mobile - int8=False, # CoreML/TF INT8 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 YOLOv3 Detect() inplace=True + keras=False, # use Keras + optimize=False, # TorchScript: optimize for mobile + int8=False, # CoreML/TF INT8 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) @@ -758,12 +794,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 @@ -777,19 +813,21 @@ 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: @@ -807,57 +845,65 @@ 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 / 'yolov3-tiny.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 YOLOv3 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('--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 / "yolov3-tiny.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 YOLOv3 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("--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 89fa522200..729eed0600 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 YOLOv3 model + """ + Creates or loads a YOLOv3 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 ⚠️ YOLOv3 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 ⚠️ YOLOv3 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 ⚠️ YOLOv3 SegmentationModel is not yet AutoShape compatible. ' - 'You will not be able to run inference with this model.') + LOGGER.warning( + "WARNING ⚠️ YOLOv3 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): # YOLOv3 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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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): # YOLOv3-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 152b861f60..d45b99cfd9 100644 --- a/models/common.py +++ b/models/common.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license -""" -Common modules -""" +"""Common modules.""" import ast import contextlib @@ -27,9 +25,23 @@ 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 @@ -211,7 +223,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)) @@ -227,7 +239,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)) @@ -266,9 +278,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) @@ -297,9 +311,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): @@ -314,7 +328,7 @@ def forward(self, x): class DetectMultiBackend(nn.Module): # YOLOv3 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 @@ -336,65 +350,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() @@ -416,22 +433,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)) @@ -440,46 +459,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: YOLOv3 TF.js inference is not supported') + raise NotImplementedError("ERROR: YOLOv3 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) @@ -487,19 +510,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 @@ -526,26 +550,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) @@ -564,17 +588,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] @@ -591,32 +615,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 @@ -633,8 +658,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() @@ -670,7 +695,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 @@ -679,13 +704,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 @@ -706,13 +731,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]) @@ -735,40 +762,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]) @@ -780,22 +811,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 @@ -806,9 +837,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 @@ -832,7 +863,7 @@ def __str__(self): # override print(results) return self._run(pprint=True) # print results def __repr__(self): - return f'YOLOv3 {self.__class__} instance\n' + self.__str__() + return f"YOLOv3 {self.__class__} instance\n" + self.__str__() class Proto(nn.Module): @@ -840,7 +871,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) @@ -850,14 +881,9 @@ def forward(self, x): class Classify(nn.Module): # YOLOv3 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 dd41155a8d..37b5940d65 100644 --- a/models/experimental.py +++ b/models/experimental.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 compatibility 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 # torch 1.7.0 compatibility 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 a933918e14..e59b00055a 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,12 +401,12 @@ 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 = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] + anchors, nc, gd, gw = d["anchors"], d["nc"], d["depth_multiple"], d["width_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) 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): @@ -396,8 +417,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, 8) if c2 != no else c2 @@ -419,15 +452,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) @@ -436,30 +472,33 @@ def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) class TFModel: # TF YOLOv3 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: @@ -479,14 +518,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 @@ -505,36 +540,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 @@ -548,7 +589,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): @@ -563,14 +604,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() @@ -584,15 +625,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 / 'yolov3-tiny.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 / "yolov3-tiny.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)) @@ -603,6 +644,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 0cc5dfc1ba..9c9962e682 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -1,6 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -YOLO-specific modules +YOLO-specific modules. Usage: $ python models/yolo.py --cfg yolov5s.yaml @@ -18,7 +18,7 @@ ROOT = FILE.parents[1] # YOLOv3 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 * # noqa @@ -26,8 +26,15 @@ from utils.autoanchor import check_anchor_order from utils.general import LOGGER, check_version, check_yaml, 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 @@ -49,7 +56,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) @@ -76,14 +83,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 @@ -126,23 +133,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 @@ -164,27 +171,28 @@ def _apply(self, fn): class DetectionModel(BaseModel): # YOLOv3 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() @@ -201,7 +209,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: @@ -242,9 +250,9 @@ def _descale_pred(self, p, flips, scale, img_size): def _clip_augmented(self, y): # Clip YOLOv3 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 @@ -257,7 +265,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) @@ -266,7 +276,7 @@ def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class SegmentationModel(DetectionModel): # YOLOv3 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) @@ -282,9 +292,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 @@ -299,7 +309,7 @@ def _from_yaml(self, cfg): def parse_model(d, ch): # model_dict, input_channels(3) # Parse a YOLOv3 model.yaml dictionary LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") - anchors, nc, gd, gw, act = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation') + anchors, nc, gd, gw, act = d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"], d.get("activation") if act: Conv.default_act = eval(act) # redefine default activation, i.e. Conv.default_act = nn.SiLU() LOGGER.info(f"{colorstr('activation:')} {act}") # print @@ -307,7 +317,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): @@ -315,8 +325,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, 8) @@ -344,10 +371,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: @@ -356,14 +383,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)) @@ -381,11 +408,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 ae3f5d960f..cbd79be552 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 ea7e7bf99b..3535d5bcb7 100644 --- a/segment/train.py +++ b/segment/train.py @@ -1,7 +1,7 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Train a YOLOv3 segment model on a segment dataset -Models and datasets download automatically from the latest YOLOv3 release. +Train a YOLOv3 segment model on a segment dataset Models and datasets download automatically from the latest YOLOv3 +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 YOLOv3 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 YOLOv3 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 87d6ca791c..28f7e5125e 100644 --- a/segment/val.py +++ b/segment/val.py @@ -1,6 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Validate a trained YOLOv3 segment model on a segment dataset +Validate a trained YOLOv3 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(), Profile(), Profile() 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 fa9bae5ad8..146c5cc8a7 100644 --- a/train.py +++ b/train.py @@ -1,7 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Train a YOLOv3 model on a custom dataset. -Models and datasets download automatically from the latest YOLOv3 release. +Train a YOLOv3 model on a custom dataset. Models and datasets download automatically from the latest YOLOv3 release. Usage - Single-GPU training: $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended) @@ -53,47 +52,88 @@ 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 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 @@ -111,39 +151,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 @@ -153,19 +193,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 @@ -181,58 +221,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: @@ -240,10 +284,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 @@ -252,7 +296,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 @@ -261,13 +305,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) @@ -284,12 +330,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 @@ -300,9 +346,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: @@ -310,7 +356,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): @@ -319,7 +365,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() @@ -338,35 +384,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] @@ -374,29 +424,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 @@ -410,12 +461,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, @@ -429,11 +480,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 @@ -441,46 +493,46 @@ 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 / 'yolov3-tiny.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('--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 / "yolov3-tiny.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("--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") return parser.parse_known_args()[0] if known else parser.parse_args() @@ -490,46 +542,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 (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 YOLOv3 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 YOLOv3 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: @@ -539,65 +596,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 @@ -622,15 +683,24 @@ 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) # 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): @@ -642,6 +712,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 1c05ec0d54..8966abbc30 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -1,21 +1,19 @@ # YOLOv3 🚀 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): # YOLOv3 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 f4435ad6db..cb193862ff 100644 --- a/utils/activations.py +++ b/utils/activations.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 64bd7a50a8..d91c84272a 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license -""" -Image augmentation functions -""" +"""Image augmentation functions.""" import math import random @@ -23,10 +21,11 @@ class Albumentations: # YOLOv3 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, +): # YOLOv3 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 b7f902695d..381259ecfb 100644 --- a/utils/autoanchor.py +++ b/utils/autoanchor.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 71faf8783f..4eaf153cb9 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 e34bc2d72c..91643e2c40 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 fff6b0356e..6b5321f1f3 100644 --- a/utils/callbacks.py +++ b/utils/callbacks.py @@ -1,43 +1,40 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license -""" -Callback utils -""" +"""Callback utils.""" import threading class Callbacks: - """" - Handles all registered callbacks for YOLOv3 Hooks - """ + """" Handles all registered callbacks for YOLOv3 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 5d119c2e77..a290111de8 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license -""" -Dataloaders and dataset utils -""" +"""Dataloaders and dataset utils.""" import contextlib import glob @@ -28,24 +26,48 @@ 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)) -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)) +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 @@ -53,7 +75,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 @@ -85,40 +107,43 @@ 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) -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( @@ -133,7 +158,8 @@ def create_dataloader(path, stride=int(stride), pad=pad, image_weights=image_weights, - prefix=prefix) + prefix=prefix, + ) batch_size = min(batch_size, len(dataset)) nd = torch.cuda.device_count() # number of CUDA devices @@ -142,26 +168,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): @@ -173,7 +202,8 @@ def __iter__(self): class _RepeatSampler: - """ Sampler that repeats forever + """ + Sampler that repeats forever. Args: sampler (Sampler) @@ -191,7 +221,7 @@ class LoadScreenshots: # YOLOv3 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() @@ -206,17 +236,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 @@ -224,7 +254,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 @@ -239,22 +269,22 @@ def __next__(self): class LoadImages: # YOLOv3 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 @@ -262,7 +292,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 @@ -270,8 +300,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 @@ -284,7 +316,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() @@ -299,14 +331,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 @@ -341,9 +373,9 @@ def __len__(self): class LoadStreams: # YOLOv3 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 @@ -353,29 +385,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]) @@ -383,7 +416,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 @@ -396,7 +429,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 @@ -407,7 +440,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 @@ -419,7 +452,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 @@ -427,8 +460,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): @@ -436,20 +469,22 @@ 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=''): + 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="", + ): self.img_size = img_size self.augment = augment self.hyp = hyp @@ -466,46 +501,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 @@ -514,7 +549,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] @@ -568,52 +603,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(fcn, range(n)) pbar = tqdm(enumerate(results), total=n, 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 - 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 @@ -623,23 +662,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): @@ -655,14 +694,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.randint(0, self.n - 1))) else: @@ -679,17 +718,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 @@ -697,16 +738,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] @@ -727,13 +768,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 @@ -793,16 +838,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 @@ -851,12 +898,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) @@ -870,16 +917,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 @@ -902,8 +951,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) @@ -918,21 +968,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: @@ -948,7 +998,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) @@ -959,11 +1009,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 @@ -971,40 +1021,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): @@ -1018,15 +1068,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) @@ -1036,12 +1086,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) @@ -1055,43 +1106,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): @@ -1102,9 +1153,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 @@ -1118,30 +1169,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)) @@ -1149,14 +1202,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 @@ -1164,6 +1217,7 @@ def process_images(self): class ClassificationDataset(torchvision.datasets.ImageFolder): """ YOLOv3 Classification Dataset. + Arguments root: Dataset path transform: torchvision transforms, used by default @@ -1174,9 +1228,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 @@ -1189,20 +1243,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 YOLOv3 Classifier with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache) @@ -1212,11 +1261,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 27a879e2c6..59e9325a7e 100644 --- a/utils/downloads.py +++ b/utils/downloads.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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://api.github.com/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://api.github.com/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/README.md b/utils/flask_rest_api/README.md index a726acbd92..b18a3011cf 100644 --- a/utils/flask_rest_api/README.md +++ b/utils/flask_rest_api/README.md @@ -1,8 +1,6 @@ # Flask REST API -[REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are -commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API -created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/). +[REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/). ## Requirements @@ -69,5 +67,4 @@ The model inference results are returned as a JSON response: ] ``` -An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given -in `example_request.py` +An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given in `example_request.py` diff --git a/utils/flask_rest_api/example_request.py b/utils/flask_rest_api/example_request.py index 7dac472e88..55fca5e964 100644 --- a/utils/flask_rest_api/example_request.py +++ b/utils/flask_rest_api/example_request.py @@ -1,19 +1,17 @@ # YOLOv3 🚀 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 b04c295874..62698d848e 100644 --- a/utils/flask_rest_api/restapi.py +++ b/utils/flask_rest_api/restapi.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 YOLOv3 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 YOLOv3 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 6b51d5b48c..ce5e3db069 100644 --- a/utils/general.py +++ b/utils/general.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license -""" -General utils -""" +"""General utils.""" import contextlib import glob @@ -43,67 +41,67 @@ FILE = Path(__file__).resolve() ROOT = FILE.parents[1] # YOLOv3 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 YOLOv3 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 @@ -112,9 +110,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 @@ -122,47 +120,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 @@ -192,7 +195,7 @@ def time(self): class Timeout(contextlib.ContextDecorator): # YOLOv3 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) @@ -201,12 +204,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 @@ -227,7 +230,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): @@ -238,11 +241,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): @@ -253,11 +256,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=()): @@ -271,22 +274,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): @@ -296,7 +299,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 @@ -308,7 +311,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 @@ -319,68 +322,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"): # YOLOv3 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"⚠️ YOLOv3 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="."): # YOLOv3 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.7.0'): +def check_python(minimum="3.7.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 YOLOv3, but {name}{current} is currently installed' # string + s = f"WARNING ⚠️ {name}{minimum} is required by YOLOv3, but {name}{current} is currently installed" # string if hard: assert result, emojis(s) # assert min requirements met if verbose and not result: @@ -396,7 +400,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 @@ -405,18 +409,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): @@ -424,38 +428,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 @@ -465,8 +471,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) @@ -474,10 +480,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) @@ -485,54 +491,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 @@ -548,35 +554,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 @@ -588,11 +594,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 @@ -601,7 +607,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)) @@ -611,18 +617,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 @@ -647,7 +653,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): @@ -657,28 +663,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): @@ -714,9 +721,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): @@ -773,7 +858,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 @@ -854,30 +942,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)): # YOLOv3 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 @@ -948,7 +1037,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 @@ -960,31 +1049,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()) @@ -992,33 +1081,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('# YOLOv3 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( + "# YOLOv3 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): @@ -1042,7 +1146,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 @@ -1056,15 +1160,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) @@ -1099,7 +1203,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 f78d6b5c01..4d4c4fccfd 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license -""" -Logging utils -""" +"""Logging utils.""" 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: @@ -58,7 +56,7 @@ comet_ml = None -class Loggers(): +class Loggers: # YOLOv3 Loggers class def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS): self.save_dir = save_dir @@ -69,60 +67,63 @@ 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 # Messages if not comet_ml: - prefix = colorstr('Comet: ') + prefix = colorstr("Comet: ") s = f"{prefix}run 'pip install comet_ml' to automatically track and visualize YOLOv3 🚀 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: @@ -156,9 +157,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.clearml: # pass # ClearML saves these images automatically using hooks if self.comet_logger: @@ -170,16 +171,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) @@ -210,11 +211,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) @@ -223,18 +224,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.tb: for k, v in x.items(): self.tb.add_scalar(k, v, epoch) elif self.clearml: # log to ClearML if TensorBoard not used for k, v in x.items(): - title, series = k.split('/') + title, series = k.split("/") self.clearml.task.get_logger().report_scalar(title, series, v, epoch) if self.wandb: @@ -258,9 +259,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) @@ -268,30 +269,32 @@ 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.task.update_output_model(model_path=str(best if best.exists() else last), - name='Best Model', - auto_delete_file=False) + self.clearml.task.update_output_model( + model_path=str(best if best.exists() else last), name="Best Model", auto_delete_file=False + ) if self.comet_logger: final_results = dict(zip(self.keys[3:10], results)) @@ -315,22 +318,23 @@ class GenericLogger: include: loggers to include """ - def __init__(self, opt, console_logger, include=('tb', 'wandb')): + def __init__(self, opt, console_logger, include=("tb", "wandb")): # 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 @@ -339,9 +343,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(): @@ -350,14 +354,14 @@ def log_metrics(self, metrics, epoch): if self.wandb: self.wandb.log(metrics, step=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) @@ -370,7 +374,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) @@ -387,15 +391,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'YOLOv3{suffix}' + suffix = "-Classify" if project.endswith("-cls") else "-Segment" if project.endswith("-seg") else "" + return f"YOLOv3{suffix}" diff --git a/utils/loggers/clearml/README.md b/utils/loggers/clearml/README.md index 7dbf6e4263..d228ebfde3 100644 --- a/utils/loggers/clearml/README.md +++ b/utils/loggers/clearml/README.md @@ -38,7 +38,7 @@ Either sign up for free to the [ClearML Hosted Service](https://cutt.ly/yolov5-t pip install clearml ``` -1. Connect the ClearML SDK to the server by [creating credentials](https://app.clear.ml/settings/workspace-configuration) (go right top to Settings -> Workspace -> Create new credentials), then execute the command below and follow the instructions: +2. Connect the ClearML SDK to the server by [creating credentials](https://app.clear.ml/settings/workspace-configuration) (go right top to Settings -> Workspace -> Create new credentials), then execute the command below and follow the instructions: ```bash clearml-init @@ -58,8 +58,7 @@ pip install clearml>=1.2.0 This will enable integration with the YOLOv5 training script. Every training run from now on, will be captured and stored by the ClearML experiment manager. -If you want to change the `project_name` or `task_name`, use the `--project` and `--name` arguments of the `train.py` script, by default the project will be called `YOLOv5` and the task `Training`. -PLEASE NOTE: ClearML uses `/` as a delimiter for subprojects, so be careful when using `/` in your project name! +If you want to change the `project_name` or `task_name`, use the `--project` and `--name` arguments of the `train.py` script, by default the project will be called `YOLOv5` and the task `Training`. PLEASE NOTE: ClearML uses `/` as a delimiter for subprojects, so be careful when using `/` in your project name! ```bash python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt --cache @@ -86,8 +85,7 @@ This will capture: - Validation images per epoch - ... -That's a lot right? 🤯 -Now, we can visualize all of this information in the ClearML UI to get an overview of our training progress. Add custom columns to the table view (such as e.g. mAP_0.5) so you can easily sort on the best performing model. Or select multiple experiments and directly compare them! +That's a lot right? 🤯 Now, we can visualize all of this information in the ClearML UI to get an overview of our training progress. Add custom columns to the table view (such as e.g. mAP_0.5) so you can easily sort on the best performing model. Or select multiple experiments and directly compare them! There even more we can do with all of this information, like hyperparameter optimization and remote execution, so keep reading if you want to see how that works! @@ -181,8 +179,7 @@ python utils/loggers/clearml/hpo.py ## 🤯 Remote Execution (advanced) -Running HPO locally is really handy, but what if we want to run our experiments on a remote machine instead? Maybe you have access to a very powerful GPU machine on-site, or you have some budget to use cloud GPUs. -This is where the ClearML Agent comes into play. Check out what the agent can do here: +Running HPO locally is really handy, but what if we want to run our experiments on a remote machine instead? Maybe you have access to a very powerful GPU machine on-site, or you have some budget to use cloud GPUs. This is where the ClearML Agent comes into play. Check out what the agent can do here: - [YouTube video](https://youtu.be/MX3BrXnaULs) - [Documentation](https://clear.ml/docs/latest/docs/clearml_agent) diff --git a/utils/loggers/clearml/clearml_utils.py b/utils/loggers/clearml/clearml_utils.py index 0c7ddeed31..df4c945202 100644 --- a/utils/loggers/clearml/clearml_utils.py +++ b/utils/loggers/clearml/clearml_utils.py @@ -11,55 +11,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): @@ -84,27 +92,29 @@ def __init__(self, opt, hyp): self.data_dict = None if self.clearml: self.task = Task.init( - project_name=opt.project if opt.project != 'runs/train' else 'YOLOv3', - task_name=opt.name if opt.name != 'exp' else 'Training', - tags=['YOLOv3'], + project_name=opt.project if opt.project != "runs/train" else "YOLOv3", + task_name=opt.name if opt.name != "exp" else "Training", + tags=["YOLOv3"], output_uri=True, reuse_last_task_id=opt.exist_ok, - auto_connect_frameworks={'pytorch': False} + auto_connect_frameworks={"pytorch": 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) @@ -112,7 +122,7 @@ def __init__(self, opt, hyp): # to give it to them opt.data = self.data_dict - 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. @@ -122,12 +132,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(it.group(), ''), - local_path=str(f), - iteration=iteration) + self.task.get_logger().report_image( + title=title, series=f.name.replace(it.group(), ""), local_path=str(f), iteration=iteration + ) def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_threshold=0.25): """ @@ -149,15 +158,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 8cccaa23b6..fcbd6770f6 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='YOLOv3', - task_type=Task.TaskTypes.optimizer, - reuse_last_task_id=False) +task = Task.init( + project_name="Hyper-Parameter Optimization", + task_name="YOLOv3", + 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/README.md b/utils/loggers/comet/README.md index 3ad52b01b4..52f344dba6 100644 --- a/utils/loggers/comet/README.md +++ b/utils/loggers/comet/README.md @@ -8,8 +8,7 @@ This guide will cover how to use YOLOv5 with [Comet](https://bit.ly/yolov5-readm Comet builds tools that help data scientists, engineers, and team leaders accelerate and optimize machine learning and deep learning models. -Track and visualize model metrics in real time, save your hyperparameters, datasets, and model checkpoints, and visualize your model predictions with [Comet Custom Panels](https://www.comet.com/docs/v2/guides/comet-dashboard/code-panels/about-panels/?utm_source=yolov5&utm_medium=partner&utm_campaign=partner_yolov5_2022&utm_content=github)! -Comet makes sure you never lose track of your work and makes it easy to share results and collaborate across teams of all sizes! +Track and visualize model metrics in real time, save your hyperparameters, datasets, and model checkpoints, and visualize your model predictions with [Comet Custom Panels](https://www.comet.com/docs/v2/guides/comet-dashboard/code-panels/about-panels/?utm_source=yolov5&utm_medium=partner&utm_campaign=partner_yolov5_2022&utm_content=github)! Comet makes sure you never lose track of your work and makes it easy to share results and collaborate across teams of all sizes! # Getting Started @@ -84,8 +83,7 @@ By default, Comet will log the following items # Configure Comet Logging -Comet can be configured to log additional data either through command line flags passed to the training script -or through environment variables. +Comet can be configured to log additional data either through command line flags passed to the training script or through environment variables. ```shell export COMET_MODE=online # Set whether to run Comet in 'online' or 'offline' mode. Defaults to online @@ -100,8 +98,7 @@ export COMET_LOG_PREDICTIONS=true # Set this to false to disable logging model p ## Logging Checkpoints with Comet -Logging Models to Comet is disabled by default. To enable it, pass the `save-period` argument to the training script. This will save the -logged checkpoints to Comet based on the interval value provided by `save-period` +Logging Models to Comet is disabled by default. To enable it, pass the `save-period` argument to the training script. This will save the logged checkpoints to Comet based on the interval value provided by `save-period` ```shell python train.py \ @@ -176,14 +173,11 @@ python train.py \ --upload_dataset ``` -You can find the uploaded dataset in the Artifacts tab in your Comet Workspace -artifact-1 +You can find the uploaded dataset in the Artifacts tab in your Comet Workspace artifact-1 -You can preview the data directly in the Comet UI. -artifact-2 +You can preview the data directly in the Comet UI. artifact-2 -Artifacts are versioned and also support adding metadata about the dataset. Comet will automatically log the metadata from your dataset `yaml` file -artifact-3 +Artifacts are versioned and also support adding metadata about the dataset. Comet will automatically log the metadata from your dataset `yaml` file artifact-3 ### Using a saved Artifact @@ -205,8 +199,7 @@ python train.py \ --weights yolov5s.pt ``` -Artifacts also allow you to track the lineage of data as it flows through your Experimentation workflow. Here you can see a graph that shows you all the experiments that have used your uploaded dataset. -artifact-4 +Artifacts also allow you to track the lineage of data as it flows through your Experimentation workflow. Here you can see a graph that shows you all the experiments that have used your uploaded dataset. artifact-4 ## Resuming a Training Run @@ -214,7 +207,7 @@ If your training run is interrupted for any reason, e.g. disrupted internet conn The Run Path has the following format `comet:////`. -This will restore the run to its state before the interruption, which includes restoring the model from a checkpoint, restoring all hyperparameters and training arguments and downloading Comet dataset Artifacts if they were used in the original run. The resumed run will continue logging to the existing Experiment in the Comet UI +This will restore the run to its state before the interruption, which includes restoring the model from a checkpoint, restoring all hyperparameters and training arguments and downloading Comet dataset Artifacts if they were used in the original run. The resumed run will continue logging to the existing Experiment in the Comet UI ```shell python train.py \ @@ -234,8 +227,7 @@ python utils/loggers/comet/hpo.py \ --comet_optimizer_config "utils/loggers/comet/optimizer_config.json" ``` -The `hpo.py` script accepts the same arguments as `train.py`. If you wish to pass additional arguments to your sweep simply add them after -the script. +The `hpo.py` script accepts the same arguments as `train.py`. If you wish to pass additional arguments to your sweep simply add them after the script. ```shell python utils/loggers/comet/hpo.py \ diff --git a/utils/loggers/comet/__init__.py b/utils/loggers/comet/__init__.py index c3e999d81e..36eed15a2c 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', 'YOLOv3') + self.experiment.log_other("Created from", "YOLOv3") 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 dbe7539675..99477f85d9 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 daf21bc5e2..943bd47e00 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 / 'yolov3-tiny.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 / "yolov3-tiny.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 c69e8f3ae7..edbb0b5e2b 100644 --- a/utils/loggers/wandb/wandb_utils.py +++ b/utils/loggers/wandb/wandb_utils.py @@ -15,34 +15,35 @@ ROOT = FILE.parents[3] # YOLOv3 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='YOLOv3' 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="YOLOv3" 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 7fe6ae445b..d533246e60 100644 --- a/utils/loss.py +++ b/utils/loss.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 56a15f45f7..626c290186 100644 --- a/utils/metrics.py +++ b/utils/metrics.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 039c0cdafd..a089a8d6d6 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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,7 +77,7 @@ 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: + if "Detect" not in module_type: 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 @@ -71,12 +89,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): @@ -95,7 +113,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 @@ -112,7 +130,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() @@ -123,7 +141,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) @@ -134,7 +152,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) @@ -154,7 +172,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) @@ -171,59 +189,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 @@ -232,69 +250,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 @@ -303,47 +326,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) @@ -352,83 +376,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 @@ -437,10 +461,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 085ef3a314..f3074d1cdd 100644 --- a/utils/segment/augmentations.py +++ b/utils/segment/augmentations.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 b6daef2c7b..2fa37be554 100644 --- a/utils/segment/dataloaders.py +++ b/utils/segment/dataloaders.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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( @@ -57,7 +57,8 @@ def create_dataloader(path, image_weights=image_weights, prefix=prefix, downsample_ratio=mask_downsample_ratio, - overlap=overlap_mask) + overlap=overlap_mask, + ) batch_size = min(batch_size, len(dataset)) nd = torch.cuda.device_count() # number of CUDA devices @@ -80,7 +81,6 @@ def create_dataloader(path, class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels): # for training/testing - def __init__( self, path, @@ -95,12 +95,25 @@ def __init__( stride=32, pad=0, min_items=0, - prefix='', + prefix="", downsample_ratio=1, overlap=False, ): - super().__init__(path, img_size, batch_size, augment, hyp, rect, image_weights, cache_images, single_cls, - stride, pad, min_items, prefix) + super().__init__( + path, + img_size, + batch_size, + augment, + hyp, + rect, + image_weights, + cache_images, + single_cls, + stride, + pad, + min_items, + prefix, + ) self.downsample_ratio = downsample_ratio self.overlap = overlap @@ -108,7 +121,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 @@ -116,7 +129,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: @@ -144,30 +157,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 @@ -177,17 +196,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] @@ -251,16 +270,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 @@ -309,8 +330,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 f1b2f1dd12..8cbc745b4a 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 caeff3cad5..1e007271fa 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 f9abf3dff1..978017bd31 100644 --- a/utils/segment/metrics.py +++ b/utils/segment/metrics.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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 f9938cd1b0..0e30c61be6 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 940c5ba9e6..000b77b37c 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -1,7 +1,5 @@ # YOLOv3 🚀 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): # YOLOv3 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'YOLOv3 🚀 {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"YOLOv3 🚀 {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): - """ YOLOv3 speed/memory/FLOPs profiler + """YOLOv3 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', 'YOLOv3') 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", "YOLOv3") 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): # YOLOv3 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): # YOLOv3 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 c57ee676e6..f34ec0f323 100644 --- a/utils/triton.py +++ b/utils/triton.py @@ -1,6 +1,5 @@ # YOLOv3 🚀 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 8a82108567..a0a740700b 100644 --- a/val.py +++ b/val.py @@ -1,6 +1,6 @@ # YOLOv3 🚀 by Ultralytics, AGPL-3.0 license """ -Validate a trained YOLOv3 detection model on a detection dataset +Validate a trained YOLOv3 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): @@ -97,47 +114,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) @@ -150,16 +167,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() @@ -167,36 +184,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(), Profile(), Profile() # 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) @@ -217,13 +238,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): @@ -258,17 +275,17 @@ 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: 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 @@ -279,10 +296,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): @@ -290,35 +307,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() @@ -326,12 +343,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): @@ -341,71 +358,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 / 'yolov3-tiny.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 / "yolov3-tiny.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)