mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 21:44:22 +08:00
README and Docs updates with A100 TensorRT times (#270)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
216cf2ddb6
commit
e18ae9d8e1
42
README.md
42
README.md
@ -121,13 +121,13 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on fi
|
|||||||
|
|
||||||
<details open><summary>Detection</summary>
|
<details open><summary>Detection</summary>
|
||||||
|
|
||||||
| Model | size<br><sup>(pixels) | mAP<sup>val<br>50-95 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) |
|
| Model | size<br><sup>(pixels) | mAP<sup>val<br>50-95 | Speed<br><sup>CPU ONNX<br>(ms) | Speed<br><sup>A100 TensorRT<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) |
|
||||||
| ------------------------------------------------------------------------------------ | --------------------- | -------------------- | ------------------------- | ---------------------------- | ------------------ | ----------------- |
|
| ------------------------------------------------------------------------------------ | --------------------- | -------------------- | ------------------------------ | ----------------------------------- | ------------------ | ----------------- |
|
||||||
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt) | 640 | 37.3 | - | - | 3.2 | 8.7 |
|
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt) | 640 | 37.3 | - | 0.99 | 3.2 | 8.7 |
|
||||||
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt) | 640 | 44.9 | - | - | 11.2 | 28.6 |
|
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt) | 640 | 44.9 | - | 1.20 | 11.2 | 28.6 |
|
||||||
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt) | 640 | 50.2 | - | - | 25.9 | 78.9 |
|
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt) | 640 | 50.2 | - | 1.83 | 25.9 | 78.9 |
|
||||||
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l.pt) | 640 | 52.9 | - | - | 43.7 | 165.2 |
|
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l.pt) | 640 | 52.9 | - | 2.39 | 43.7 | 165.2 |
|
||||||
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt) | 640 | 53.9 | - | - | 68.2 | 257.8 |
|
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt) | 640 | 53.9 | - | 3.53 | 68.2 | 257.8 |
|
||||||
|
|
||||||
- **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.
|
- **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.
|
||||||
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
|
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
|
||||||
@ -138,13 +138,13 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on fi
|
|||||||
|
|
||||||
<details><summary>Segmentation</summary>
|
<details><summary>Segmentation</summary>
|
||||||
|
|
||||||
| Model | size<br><sup>(pixels) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) |
|
| Model | size<br><sup>(pixels) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | Speed<br><sup>CPU ONNX<br>(ms) | Speed<br><sup>A100 TensorRT<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) |
|
||||||
| ---------------------------------------------------------------------------------------- | --------------------- | -------------------- | --------------------- | ------------------------- | ---------------------------- | ------------------ | ----------------- |
|
| ---------------------------------------------------------------------------------------- | --------------------- | -------------------- | --------------------- | ------------------------------ | ----------------------------------- | ------------------ | ----------------- |
|
||||||
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt) | 640 | 36.7 | 30.5 | - | - | 3.4 | 12.6 |
|
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt) | 640 | 36.7 | 30.5 | - | - | 3.4 | 12.6 |
|
||||||
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-seg.pt) | 640 | 44.6 | 36.8 | - | - | 11.8 | 42.6 |
|
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-seg.pt) | 640 | 44.6 | 36.8 | - | - | 11.8 | 42.6 |
|
||||||
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-seg.pt) | 640 | 49.9 | 40.8 | - | - | 27.3 | 110.2 |
|
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-seg.pt) | 640 | 49.9 | 40.8 | - | - | 27.3 | 110.2 |
|
||||||
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-seg.pt) | 640 | 52.3 | 42.6 | - | - | 46.0 | 220.5 |
|
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-seg.pt) | 640 | 52.3 | 42.6 | - | - | 46.0 | 220.5 |
|
||||||
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-seg.pt) | 640 | 53.4 | 43.4 | - | - | 71.8 | 344.1 |
|
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-seg.pt) | 640 | 53.4 | 43.4 | - | - | 71.8 | 344.1 |
|
||||||
|
|
||||||
- **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.
|
- **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.
|
||||||
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
|
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
|
||||||
@ -155,13 +155,13 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on fi
|
|||||||
|
|
||||||
<details><summary>Classification</summary>
|
<details><summary>Classification</summary>
|
||||||
|
|
||||||
| Model | size<br><sup>(pixels) | acc<br><sup>top1 | acc<br><sup>top5 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) at 640 |
|
| Model | size<br><sup>(pixels) | acc<br><sup>top1 | acc<br><sup>top5 | Speed<br><sup>CPU ONNX<br>(ms) | Speed<br><sup>A100 TensorRT<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) at 640 |
|
||||||
| ---------------------------------------------------------------------------------------- | --------------------- | ---------------- | ---------------- | ------------------------- | ---------------------------- | ------------------ | ------------------------ |
|
| ---------------------------------------------------------------------------------------- | --------------------- | ---------------- | ---------------- | ------------------------------ | ----------------------------------- | ------------------ | ------------------------ |
|
||||||
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt) | 224 | 66.6 | 87.0 | - | - | 2.7 | 4.3 |
|
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt) | 224 | 66.6 | 87.0 | - | - | 2.7 | 4.3 |
|
||||||
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-cls.pt) | 224 | 72.3 | 91.1 | - | - | 6.4 | 13.5 |
|
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-cls.pt) | 224 | 72.3 | 91.1 | - | - | 6.4 | 13.5 |
|
||||||
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-cls.pt) | 224 | 76.4 | 93.2 | - | - | 17.0 | 42.7 |
|
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-cls.pt) | 224 | 76.4 | 93.2 | - | - | 17.0 | 42.7 |
|
||||||
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-cls.pt) | 224 | 78.0 | 94.1 | - | - | 37.5 | 99.7 |
|
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-cls.pt) | 224 | 78.0 | 94.1 | - | - | 37.5 | 99.7 |
|
||||||
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-cls.pt) | 224 | 78.4 | 94.3 | - | - | 57.4 | 154.8 |
|
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-cls.pt) | 224 | 78.4 | 94.3 | - | - | 57.4 | 154.8 |
|
||||||
|
|
||||||
- **mAP<sup>val</sup>** values are for single-model single-scale on [ImageNet](https://www.image-net.org/) dataset.
|
- **mAP<sup>val</sup>** values are for single-model single-scale on [ImageNet](https://www.image-net.org/) dataset.
|
||||||
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
|
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
|
||||||
|
@ -115,13 +115,13 @@ success = YOLO("yolov8n.pt").export(format="onnx") # 将模型导出为 ONNX
|
|||||||
|
|
||||||
<details open><summary>目标检测</summary>
|
<details open><summary>目标检测</summary>
|
||||||
|
|
||||||
| 模型 | 尺寸<br><sup>(像素) | mAP<sup>val<br>50-95 | 推理速度<br><sup>CPU<br>(ms) | 推理速度<br><sup>T4 GPU<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>(B) |
|
| 模型 | 尺寸<br><sup>(像素) | mAP<sup>val<br>50-95 | 推理速度<br><sup>CPU ONNX<br>(ms) | 推理速度<br><sup>A100 TensorRT<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>(B) |
|
||||||
| ----------------------------------------------------------------------------------------- | --------------- | -------------------- | ------------------------ | --------------------------- | --------------- | ----------------- |
|
| ------------------------------------------------------------------------------------ | --------------- | -------------------- | ----------------------------- | ---------------------------------- | --------------- | ----------------- |
|
||||||
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n.pt) | 640 | 37.3 | - | - | 3.2 | 8.7 |
|
| [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt) | 640 | 37.3 | - | 0.99 | 3.2 | 8.7 |
|
||||||
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s.pt) | 640 | 44.9 | - | - | 11.2 | 28.6 |
|
| [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt) | 640 | 44.9 | - | 1.20 | 11.2 | 28.6 |
|
||||||
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m.pt) | 640 | 50.2 | - | - | 25.9 | 78.9 |
|
| [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt) | 640 | 50.2 | - | 1.83 | 25.9 | 78.9 |
|
||||||
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l.pt) | 640 | 52.9 | - | - | 43.7 | 165.2 |
|
| [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l.pt) | 640 | 52.9 | - | 2.39 | 43.7 | 165.2 |
|
||||||
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x.pt) | 640 | 53.9 | - | - | 68.2 | 257.8 |
|
| [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt) | 640 | 53.9 | - | 3.53 | 68.2 | 257.8 |
|
||||||
|
|
||||||
- **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。
|
- **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。
|
||||||
<br>复现命令 `yolo mode=val task=detect data=coco.yaml device=0`
|
<br>复现命令 `yolo mode=val task=detect data=coco.yaml device=0`
|
||||||
@ -132,13 +132,13 @@ success = YOLO("yolov8n.pt").export(format="onnx") # 将模型导出为 ONNX
|
|||||||
|
|
||||||
<details><summary>实例分割</summary>
|
<details><summary>实例分割</summary>
|
||||||
|
|
||||||
| 模型 | 尺寸<br><sup>(像素) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | 推理速度<br><sup>CPU<br>(ms) | 推理速度<br><sup>T4 GPU<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>(B) |
|
| 模型 | 尺寸<br><sup>(像素) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | 推理速度<br><sup>CPU ONNX<br>(ms) | 推理速度<br><sup>A100 TensorRT<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>(B) |
|
||||||
| --------------------------------------------------------------------------------------------- | --------------- | -------------------- | --------------------- | ------------------------ | --------------------------- | --------------- | ----------------- |
|
| --------------------------------------------------------------------------------------------- | --------------- | -------------------- | --------------------- | ----------------------------- | ---------------------------------- | --------------- | ----------------- |
|
||||||
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n-seg.pt) | 640 | 36.7 | 30.5 | - | - | 3.4 | 12.6 |
|
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n-seg.pt) | 640 | 36.7 | 30.5 | - | - | 3.4 | 12.6 |
|
||||||
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s-seg.pt) | 640 | 44.6 | 36.8 | - | - | 11.8 | 42.6 |
|
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s-seg.pt) | 640 | 44.6 | 36.8 | - | - | 11.8 | 42.6 |
|
||||||
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m-seg.pt) | 640 | 49.9 | 40.8 | - | - | 27.3 | 110.2 |
|
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m-seg.pt) | 640 | 49.9 | 40.8 | - | - | 27.3 | 110.2 |
|
||||||
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l-seg.pt) | 640 | 52.3 | 42.6 | - | - | 46.0 | 220.5 |
|
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l-seg.pt) | 640 | 52.3 | 42.6 | - | - | 46.0 | 220.5 |
|
||||||
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x-seg.pt) | 640 | 53.4 | 43.4 | - | - | 71.8 | 344.1 |
|
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x-seg.pt) | 640 | 53.4 | 43.4 | - | - | 71.8 | 344.1 |
|
||||||
|
|
||||||
- **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。
|
- **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。
|
||||||
<br>复现命令 `yolo mode=val task=detect data=coco.yaml device=0`
|
<br>复现命令 `yolo mode=val task=detect data=coco.yaml device=0`
|
||||||
@ -149,13 +149,13 @@ success = YOLO("yolov8n.pt").export(format="onnx") # 将模型导出为 ONNX
|
|||||||
|
|
||||||
<details><summary>分类</summary>
|
<details><summary>分类</summary>
|
||||||
|
|
||||||
| 模型 | 尺寸<br><sup>(像素) | acc<br><sup>top1 | acc<br><sup>top5 | 推理速度<br><sup>CPU<br>(ms) | 推理速度<br><sup>T4 GPU<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>(B) at 640 |
|
| 模型 | 尺寸<br><sup>(像素) | acc<br><sup>top1 | acc<br><sup>top5 | 推理速度<br><sup>CPU ONNX<br>(ms) | 推理速度<br><sup>A100 TensorRT<br>(ms) | 参数量<br><sup>(M) | FLOPs<br><sup>(B) at 640 |
|
||||||
| --------------------------------------------------------------------------------------------- | --------------- | ---------------- | ---------------- | ------------------------ | --------------------------- | --------------- | ------------------------ |
|
| --------------------------------------------------------------------------------------------- | --------------- | ---------------- | ---------------- | ----------------------------- | ---------------------------------- | --------------- | ------------------------ |
|
||||||
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n-cls.pt) | 224 | 66.6 | 87.0 | - | - | 2.7 | 4.3 |
|
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n-cls.pt) | 224 | 66.6 | 87.0 | - | - | 2.7 | 4.3 |
|
||||||
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s-cls.pt) | 224 | 72.3 | 91.1 | - | - | 6.4 | 13.5 |
|
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s-cls.pt) | 224 | 72.3 | 91.1 | - | - | 6.4 | 13.5 |
|
||||||
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m-cls.pt) | 224 | 76.4 | 93.2 | - | - | 17.0 | 42.7 |
|
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m-cls.pt) | 224 | 76.4 | 93.2 | - | - | 17.0 | 42.7 |
|
||||||
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l-cls.pt) | 224 | 78.0 | 94.1 | - | - | 37.5 | 99.7 |
|
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l-cls.pt) | 224 | 78.0 | 94.1 | - | - | 37.5 | 99.7 |
|
||||||
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x-cls.pt) | 224 | 78.4 | 94.3 | - | - | 57.4 | 154.8 |
|
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x-cls.pt) | 224 | 78.4 | 94.3 | - | - | 57.4 | 154.8 |
|
||||||
|
|
||||||
- **mAP<sup>val</sup>** 都在 [ImageNet](https://www.image-net.org/) 数据集上,使用单模型单尺度测试得到。
|
- **mAP<sup>val</sup>** 都在 [ImageNet](https://www.image-net.org/) 数据集上,使用单模型单尺度测试得到。
|
||||||
<br>复现命令 `yolo mode=val task=detect data=coco.yaml device=0`
|
<br>复现命令 `yolo mode=val task=detect data=coco.yaml device=0`
|
||||||
|
BIN
docs/assets/favicon.ico
Normal file
BIN
docs/assets/favicon.ico
Normal file
Binary file not shown.
After Width: | Height: | Size: 5.3 KiB |
@ -5,7 +5,7 @@ BaseTrainer contains the generic boilerplate training routine. It can be customi
|
|||||||
|
|
||||||
* `get_model(cfg, weights)` - The function that builds a the model to be trained
|
* `get_model(cfg, weights)` - The function that builds a the model to be trained
|
||||||
* `get_dataloder()` - The function that builds the dataloder
|
* `get_dataloder()` - The function that builds the dataloder
|
||||||
More details and source code can be found in [`BaseTrainer` Reference](../reference/base_trainer.md)
|
More details and source code can be found in [`BaseTrainer` Reference](reference/base_trainer.md)
|
||||||
|
|
||||||
## DetectionTrainer
|
## DetectionTrainer
|
||||||
Here's how you can use the YOLOv8 `DetectionTrainer` and customize it.
|
Here's how you can use the YOLOv8 `DetectionTrainer` and customize it.
|
||||||
|
@ -2,9 +2,13 @@
|
|||||||
<a href="https://github.com/ultralytics/ultralytics" target="_blank">
|
<a href="https://github.com/ultralytics/ultralytics" target="_blank">
|
||||||
<img width="1024" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png"></a>
|
<img width="1024" src="https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png"></a>
|
||||||
<br>
|
<br>
|
||||||
<a href="https://console.paperspace.com/github/ultralytics/ultralytics"><img src="https://assets.paperspace.io/img/gradient-badge.svg" alt="Run on Gradient"/></a>
|
<a href="https://github.com/ultralytics/ultralytics/actions/workflows/ci.yaml"><img src="https://github.com/ultralytics/ultralytics/actions/workflows/ci.yaml/badge.svg" alt="Ultralytics CI"></a>
|
||||||
<a href="https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
|
<a href="https://zenodo.org/badge/latestdoi/264818686"><img src="https://zenodo.org/badge/264818686.svg" alt="YOLOv8 Citation"></a>
|
||||||
<a href="https://www.kaggle.com/ultralytics/yolov8"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
<a href="https://hub.docker.com/r/ultralytics/ultralytics"><img src="https://img.shields.io/docker/pulls/ultralytics/ultralytics?logo=docker" alt="Docker Pulls"></a>
|
||||||
|
<br>
|
||||||
|
<a href="https://console.paperspace.com/github/ultralytics/ultralytics"><img src="https://assets.paperspace.io/img/gradient-badge.svg" alt="Run on Gradient"/></a>
|
||||||
|
<a href="https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
|
||||||
|
<a href="https://www.kaggle.com/ultralytics/yolov8"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
||||||
<br>
|
<br>
|
||||||
<br>
|
<br>
|
||||||
</div>
|
</div>
|
||||||
|
27
mkdocs.yml
27
mkdocs.yml
@ -1,25 +1,14 @@
|
|||||||
site_name: Ultralytics Docs
|
site_name: Ultralytics Docs
|
||||||
repo_url: https://github.com/ultralytics/ultralytics
|
repo_url: https://github.com/ultralytics/ultralytics
|
||||||
repo_name: Ultralytics
|
edit_uri: https://github.com/ultralytics/ultralytics/tree/main/docs
|
||||||
|
repo_name: ultralytics/ultralytics
|
||||||
|
|
||||||
theme:
|
theme:
|
||||||
name: "material"
|
name: "material"
|
||||||
logo: https://github.com/ultralytics/assets/raw/main/logo/Ultralytics-logomark-white.png
|
logo: https://github.com/ultralytics/assets/raw/main/logo/Ultralytics-logomark-white.png
|
||||||
icon:
|
favicon: assets/favicon.ico
|
||||||
repo: fontawesome/brands/github
|
font:
|
||||||
admonition:
|
text: Roboto
|
||||||
note: octicons/tag-16
|
|
||||||
abstract: octicons/checklist-16
|
|
||||||
info: octicons/info-16
|
|
||||||
tip: octicons/squirrel-16
|
|
||||||
success: octicons/check-16
|
|
||||||
question: octicons/question-16
|
|
||||||
warning: octicons/alert-16
|
|
||||||
failure: octicons/x-circle-16
|
|
||||||
danger: octicons/zap-16
|
|
||||||
bug: octicons/bug-16
|
|
||||||
example: octicons/beaker-16
|
|
||||||
quote: octicons/quote-16
|
|
||||||
|
|
||||||
palette:
|
palette:
|
||||||
# Palette toggle for light mode
|
# Palette toggle for light mode
|
||||||
@ -34,12 +23,16 @@ theme:
|
|||||||
icon: material/brightness-4
|
icon: material/brightness-4
|
||||||
name: Switch to light mode
|
name: Switch to light mode
|
||||||
features:
|
features:
|
||||||
|
- content.action.edit
|
||||||
- content.code.annotate
|
- content.code.annotate
|
||||||
- content.tooltips
|
- content.tooltips
|
||||||
- search.highlight
|
- search.highlight
|
||||||
- search.share
|
- search.share
|
||||||
- search.suggest
|
- search.suggest
|
||||||
- toc.follow
|
- toc.follow
|
||||||
|
- navigation.top
|
||||||
|
- navigation.expand
|
||||||
|
- navigation.footer
|
||||||
|
|
||||||
extra_css:
|
extra_css:
|
||||||
- stylesheets/style.css
|
- stylesheets/style.css
|
||||||
@ -72,8 +65,10 @@ markdown_extensions:
|
|||||||
- pymdownx.keys
|
- pymdownx.keys
|
||||||
- pymdownx.mark
|
- pymdownx.mark
|
||||||
- pymdownx.tilde
|
- pymdownx.tilde
|
||||||
|
|
||||||
plugins:
|
plugins:
|
||||||
- mkdocstrings
|
- mkdocstrings
|
||||||
|
- search
|
||||||
|
|
||||||
# Primary navigation
|
# Primary navigation
|
||||||
nav:
|
nav:
|
||||||
|
@ -22,32 +22,31 @@ class AutoBackend(nn.Module):
|
|||||||
|
|
||||||
def __init__(self, weights='yolov8n.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
|
def __init__(self, weights='yolov8n.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
|
||||||
"""
|
"""
|
||||||
Ultralytics YOLO MultiBackend class for python inference on various backends
|
MultiBackend class for python inference on various platforms using Ultralytics YOLO.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
weights: the path to the weights file. Defaults to yolov8n.pt
|
weights (str): The path to the weights file. Default: 'yolov8n.pt'
|
||||||
device: The device to run the model on.
|
device (torch.device): The device to run the model on.
|
||||||
dnn: If you want to use OpenCV's DNN module to run the inference, set this to True. Defaults to
|
dnn (bool): Use OpenCV's DNN module for inference if True, defaults to False.
|
||||||
False
|
data (dict): Additional data, optional
|
||||||
data: a dictionary containing the following keys:
|
fp16 (bool): If True, use half precision. Default: False
|
||||||
fp16: If true, will use half precision. Defaults to False
|
fuse (bool): Whether to fuse the model or not. Default: True
|
||||||
fuse: whether to fuse the model or not. Defaults to True
|
|
||||||
|
|
||||||
Supported format and their usage:
|
Supported formats and their usage:
|
||||||
| Platform | weights |
|
Platform | Weights Format
|
||||||
|-----------------------|------------------|
|
-----------------------|------------------
|
||||||
| PyTorch | *.pt |
|
PyTorch | *.pt
|
||||||
| TorchScript | *.torchscript |
|
TorchScript | *.torchscript
|
||||||
| ONNX Runtime | *.onnx |
|
ONNX Runtime | *.onnx
|
||||||
| ONNX OpenCV DNN | *.onnx --dnn |
|
ONNX OpenCV DNN | *.onnx --dnn
|
||||||
| OpenVINO | *.xml |
|
OpenVINO | *.xml
|
||||||
| CoreML | *.mlmodel |
|
CoreML | *.mlmodel
|
||||||
| TensorRT | *.engine |
|
TensorRT | *.engine
|
||||||
| TensorFlow SavedModel | *_saved_model |
|
TensorFlow SavedModel | *_saved_model
|
||||||
| TensorFlow GraphDef | *.pb |
|
TensorFlow GraphDef | *.pb
|
||||||
| TensorFlow Lite | *.tflite |
|
TensorFlow Lite | *.tflite
|
||||||
| TensorFlow Edge TPU | *_edgetpu.tflite |
|
TensorFlow Edge TPU | *_edgetpu.tflite
|
||||||
| PaddlePaddle | *_paddle_model |
|
PaddlePaddle | *_paddle_model
|
||||||
"""
|
"""
|
||||||
super().__init__()
|
super().__init__()
|
||||||
w = str(weights[0] if isinstance(weights, list) else weights)
|
w = str(weights[0] if isinstance(weights, list) else weights)
|
||||||
@ -234,15 +233,16 @@ class AutoBackend(nn.Module):
|
|||||||
|
|
||||||
def forward(self, im, augment=False, visualize=False):
|
def forward(self, im, augment=False, visualize=False):
|
||||||
"""
|
"""
|
||||||
Runs inference on the given model
|
Runs inference on the YOLOv8 MultiBackend model.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
im: the image tensor
|
im (torch.tensor): The image tensor to perform inference on.
|
||||||
augment: whether to augment the image. Defaults to False
|
augment (bool): whether to perform data augmentation during inference, defaults to False
|
||||||
visualize: if True, then the network will output the feature maps of the last convolutional layer.
|
visualize (bool): whether to visualize the output predictions, defaults to False
|
||||||
Defaults to False
|
|
||||||
|
Returns:
|
||||||
|
(tuple): Tuple containing the raw output tensor, and the processed output for visualization (if visualize=True)
|
||||||
"""
|
"""
|
||||||
# YOLOv5 MultiBackend inference
|
|
||||||
b, ch, h, w = im.shape # batch, channel, height, width
|
b, ch, h, w = im.shape # batch, channel, height, width
|
||||||
if self.fp16 and im.dtype != torch.float16:
|
if self.fp16 and im.dtype != torch.float16:
|
||||||
im = im.half() # to FP16
|
im = im.half() # to FP16
|
||||||
@ -325,19 +325,25 @@ class AutoBackend(nn.Module):
|
|||||||
|
|
||||||
def from_numpy(self, x):
|
def from_numpy(self, x):
|
||||||
"""
|
"""
|
||||||
`from_numpy` converts a numpy array to a tensor
|
Convert a numpy array to a tensor.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the numpy array to convert
|
x (numpy.ndarray): The array to be converted.
|
||||||
"""
|
|
||||||
|
Returns:
|
||||||
|
(torch.tensor): The converted tensor
|
||||||
|
"""
|
||||||
return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
|
return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
|
||||||
|
|
||||||
def warmup(self, imgsz=(1, 3, 640, 640)):
|
def warmup(self, imgsz=(1, 3, 640, 640)):
|
||||||
"""
|
"""
|
||||||
Warmup model by running inference once
|
Warm up the model by running one forward pass with a dummy input.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
imgsz: the size of the image you want to run inference on.
|
imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(None): This method runs the forward pass and don't return any value
|
||||||
"""
|
"""
|
||||||
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module
|
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module
|
||||||
if any(warmup_types) and (self.device.type != 'cpu' or self.triton):
|
if any(warmup_types) and (self.device.type != 'cpu' or self.triton):
|
||||||
|
@ -17,35 +17,36 @@ from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_wei
|
|||||||
|
|
||||||
|
|
||||||
class BaseModel(nn.Module):
|
class BaseModel(nn.Module):
|
||||||
'''
|
"""
|
||||||
The BaseModel class is a base class for all the models in the Ultralytics YOLO family.
|
The BaseModel class serves as a base class for all the models in the Ultralytics YOLO family.
|
||||||
'''
|
"""
|
||||||
|
|
||||||
def forward(self, x, profile=False, visualize=False):
|
def forward(self, x, profile=False, visualize=False):
|
||||||
"""
|
"""
|
||||||
> `forward` is a wrapper for `_forward_once` that runs the model on a single scale
|
Forward pass of the model on a single scale.
|
||||||
|
Wrapper for `_forward_once` method.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the input image
|
x (torch.tensor): The input image tensor
|
||||||
profile: whether to profile the model. Defaults to False
|
profile (bool): Whether to profile the model, defaults to False
|
||||||
visualize: if True, will return the intermediate feature maps. Defaults to False
|
visualize (bool): Whether to return the intermediate feature maps, defaults to False
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The output of the network.
|
(torch.tensor): The output of the network.
|
||||||
"""
|
"""
|
||||||
return self._forward_once(x, profile, visualize)
|
return self._forward_once(x, profile, visualize)
|
||||||
|
|
||||||
def _forward_once(self, x, profile=False, visualize=False):
|
def _forward_once(self, x, profile=False, visualize=False):
|
||||||
"""
|
"""
|
||||||
> Forward pass of the network
|
Perform a forward pass through the network.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: input to the model
|
x (torch.tensor): The input tensor to the model
|
||||||
profile: if True, the time taken for each layer will be printed. Defaults to False
|
profile (bool): Print the computation time of each layer if True, defaults to False.
|
||||||
visualize: If True, it will save the feature maps of the model. Defaults to False
|
visualize (bool): Save the feature maps of the model if True, defaults to False
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The last layer of the model.
|
(torch.tensor): The last output of the model.
|
||||||
"""
|
"""
|
||||||
y, dt = [], [] # outputs
|
y, dt = [], [] # outputs
|
||||||
for m in self.model:
|
for m in self.model:
|
||||||
@ -62,13 +63,15 @@ class BaseModel(nn.Module):
|
|||||||
|
|
||||||
def _profile_one_layer(self, m, x, dt):
|
def _profile_one_layer(self, m, x, dt):
|
||||||
"""
|
"""
|
||||||
It takes a model, an input, and a list of times, and it profiles the model on the input, appending
|
Profile the computation time and FLOPs of a single layer of the model on a given input. Appends the results to the provided list.
|
||||||
the time to the list
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
m: the model
|
m (nn.Module): The layer to be profiled.
|
||||||
x: the input image
|
x (torch.Tensor): The input data to the layer.
|
||||||
dt: list of time taken for each layer
|
dt (list): A list to store the computation time of the layer.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
"""
|
"""
|
||||||
c = m == self.model[-1] # is final layer, copy input as inplace fix
|
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
|
||||||
@ -84,10 +87,10 @@ class BaseModel(nn.Module):
|
|||||||
|
|
||||||
def fuse(self):
|
def fuse(self):
|
||||||
"""
|
"""
|
||||||
> It takes a model and fuses the Conv2d() and BatchNorm2d() layers into a single layer
|
Fuse the `Conv2d()` and `BatchNorm2d()` layers of the model into a single layer, in order to improve the computation efficiency.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The model is being returned.
|
(nn.Module): The fused model is returned.
|
||||||
"""
|
"""
|
||||||
LOGGER.info('Fusing layers... ')
|
LOGGER.info('Fusing layers... ')
|
||||||
for m in self.model.modules():
|
for m in self.model.modules():
|
||||||
@ -103,8 +106,8 @@ class BaseModel(nn.Module):
|
|||||||
Prints model information
|
Prints model information
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
verbose: if True, prints out the model information. Defaults to False
|
verbose (bool): if True, prints out the model information. Defaults to False
|
||||||
imgsz: the size of the image that the model will be trained on. Defaults to 640
|
imgsz (int): the size of the image that the model will be trained on. Defaults to 640
|
||||||
"""
|
"""
|
||||||
model_info(self, verbose, imgsz)
|
model_info(self, verbose, imgsz)
|
||||||
|
|
||||||
@ -129,10 +132,10 @@ class BaseModel(nn.Module):
|
|||||||
|
|
||||||
def load(self, weights):
|
def load(self, weights):
|
||||||
"""
|
"""
|
||||||
> This function loads the weights of the model from a file
|
This function loads the weights of the model from a file
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
weights: The weights to load into the model.
|
weights (str): The weights to load into the model.
|
||||||
"""
|
"""
|
||||||
# Force all tasks to implement this function
|
# Force all tasks to implement this function
|
||||||
raise NotImplementedError("This function needs to be implemented by derived classes!")
|
raise NotImplementedError("This function needs to be implemented by derived classes!")
|
||||||
|
@ -84,6 +84,7 @@ class BaseTrainer:
|
|||||||
if overrides is None:
|
if overrides is None:
|
||||||
overrides = {}
|
overrides = {}
|
||||||
self.args = get_config(config, overrides)
|
self.args = get_config(config, overrides)
|
||||||
|
self.device = utils.torch_utils.select_device(self.args.device, self.args.batch)
|
||||||
self.check_resume()
|
self.check_resume()
|
||||||
self.console = LOGGER
|
self.console = LOGGER
|
||||||
self.validator = None
|
self.validator = None
|
||||||
@ -113,7 +114,6 @@ class BaseTrainer:
|
|||||||
print_args(dict(self.args))
|
print_args(dict(self.args))
|
||||||
|
|
||||||
# Device
|
# Device
|
||||||
self.device = utils.torch_utils.select_device(self.args.device, self.batch_size)
|
|
||||||
self.amp = self.device.type != 'cpu'
|
self.amp = self.device.type != 'cpu'
|
||||||
self.scaler = amp.GradScaler(enabled=self.amp)
|
self.scaler = amp.GradScaler(enabled=self.amp)
|
||||||
if self.device.type == 'cpu':
|
if self.device.type == 'cpu':
|
||||||
@ -164,7 +164,15 @@ class BaseTrainer:
|
|||||||
callback(self)
|
callback(self)
|
||||||
|
|
||||||
def train(self):
|
def train(self):
|
||||||
world_size = torch.cuda.device_count()
|
# Allow device='', device=None on Multi-GPU systems to default to device=0
|
||||||
|
if isinstance(self.args.device, int) or self.args.device: # i.e. device=0 or device=[0,1,2,3]
|
||||||
|
world_size = torch.cuda.device_count()
|
||||||
|
elif torch.cuda.is_available(): # i.e. device=None or device=''
|
||||||
|
world_size = 1 # default to device 0
|
||||||
|
else: # i.e. device='cpu' or 'mps'
|
||||||
|
world_size = 0
|
||||||
|
|
||||||
|
# Run subprocess if DDP training, else train normally
|
||||||
if world_size > 1 and "LOCAL_RANK" not in os.environ:
|
if world_size > 1 and "LOCAL_RANK" not in os.environ:
|
||||||
command = generate_ddp_command(world_size, self)
|
command = generate_ddp_command(world_size, self)
|
||||||
try:
|
try:
|
||||||
|
@ -1,5 +1,3 @@
|
|||||||
# Ultralytics YOLO 🚀, GPL-3.0 license
|
|
||||||
|
|
||||||
import contextlib
|
import contextlib
|
||||||
import math
|
import math
|
||||||
import re
|
import re
|
||||||
@ -50,15 +48,15 @@ def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
|
|||||||
|
|
||||||
def segment2box(segment, width=640, height=640):
|
def segment2box(segment, width=640, height=640):
|
||||||
"""
|
"""
|
||||||
> Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to
|
Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to
|
||||||
(xyxy)
|
(xyxy)
|
||||||
Args:
|
Args:
|
||||||
segment: the segment label
|
segment (torch.tensor): the segment label
|
||||||
width: the width of the image. Defaults to 640
|
width (int): the width of the image. Defaults to 640
|
||||||
height: The height of the image. Defaults to 640
|
height (int): The height of the image. Defaults to 640
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the minimum and maximum x and y values of the segment.
|
(np.array): the minimum and maximum x and y values of the segment.
|
||||||
"""
|
"""
|
||||||
# Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
|
# Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
|
||||||
x, y = segment.T # segment xy
|
x, y = segment.T # segment xy
|
||||||
@ -69,17 +67,16 @@ def segment2box(segment, width=640, height=640):
|
|||||||
|
|
||||||
def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
|
def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
|
||||||
"""
|
"""
|
||||||
> Rescale boxes (xyxy) from img1_shape to img0_shape
|
Rescales bounding boxes (in the format of xyxy) from the shape of the image they were originally specified in (img1_shape) to the shape of a different image (img0_shape).
|
||||||
Args:
|
Args:
|
||||||
img1_shape: The shape of the image that the bounding boxes are for.
|
img1_shape (tuple): The shape of the image that the bounding boxes are for, in the format of (height, width).
|
||||||
boxes: the bounding boxes of the objects in the image
|
boxes (torch.tensor): the bounding boxes of the objects in the image, in the format of (x1, y1, x2, y2)
|
||||||
img0_shape: the shape of the original image
|
img0_shape (tuple): the shape of the target image, in the format of (height, width).
|
||||||
ratio_pad: a tuple of (ratio, pad)
|
ratio_pad (tuple): a tuple of (ratio, pad) for scaling the boxes. If not provided, the ratio and pad will be calculated based on the size difference between the two images.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The boxes are being returned.
|
boxes (torch.tensor): The scaled bounding boxes, in the format of (x1, y1, x2, y2)
|
||||||
"""
|
"""
|
||||||
#
|
|
||||||
if ratio_pad is None: # calculate from img0_shape
|
if ratio_pad is None: # calculate from img0_shape
|
||||||
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
||||||
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
||||||
@ -113,7 +110,7 @@ def non_max_suppression(
|
|||||||
nm=0, # number of masks
|
nm=0, # number of masks
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
> Perform non-maximum suppression (NMS) on a set of boxes, with support for masks and multiple labels per box.
|
Perform non-maximum suppression (NMS) on a set of boxes, with support for masks and multiple labels per box.
|
||||||
|
|
||||||
Arguments:
|
Arguments:
|
||||||
prediction (torch.Tensor): A tensor of shape (batch_size, num_boxes, num_classes + 4 + num_masks)
|
prediction (torch.Tensor): A tensor of shape (batch_size, num_boxes, num_classes + 4 + num_masks)
|
||||||
@ -134,7 +131,7 @@ def non_max_suppression(
|
|||||||
nm (int): The number of masks output by the model.
|
nm (int): The number of masks output by the model.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List[torch.Tensor]: A list of length batch_size, where each element is a tensor of
|
(List[torch.Tensor]): A list of length batch_size, where each element is a tensor of
|
||||||
shape (num_boxes, 6 + num_masks) containing the kept boxes, with columns
|
shape (num_boxes, 6 + num_masks) containing the kept boxes, with columns
|
||||||
(x1, y1, x2, y2, confidence, class, mask1, mask2, ...).
|
(x1, y1, x2, y2, confidence, class, mask1, mask2, ...).
|
||||||
"""
|
"""
|
||||||
@ -231,12 +228,12 @@ def non_max_suppression(
|
|||||||
|
|
||||||
def clip_boxes(boxes, shape):
|
def clip_boxes(boxes, shape):
|
||||||
"""
|
"""
|
||||||
> It takes a list of bounding boxes and a shape (height, width) and clips the bounding boxes to the
|
It takes a list of bounding boxes and a shape (height, width) and clips the bounding boxes to the
|
||||||
shape
|
shape
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
boxes: the bounding boxes to clip
|
boxes (torch.tensor): the bounding boxes to clip
|
||||||
shape: the shape of the image
|
shape (tuple): the shape of the image
|
||||||
"""
|
"""
|
||||||
if isinstance(boxes, torch.Tensor): # faster individually
|
if isinstance(boxes, torch.Tensor): # faster individually
|
||||||
boxes[..., 0].clamp_(0, shape[1]) # x1
|
boxes[..., 0].clamp_(0, shape[1]) # x1
|
||||||
@ -262,16 +259,16 @@ def clip_coords(boxes, shape):
|
|||||||
|
|
||||||
def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
|
def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
|
||||||
"""
|
"""
|
||||||
> It takes a mask, and resizes it to the original image size
|
Takes a mask, and resizes it to the original image size
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
im1_shape: model input shape, [h, w]
|
im1_shape (tuple): model input shape, [h, w]
|
||||||
masks: [h, w, num]
|
masks (torch.tensor): [h, w, num]
|
||||||
im0_shape: the original image shape
|
im0_shape (tuple): the original image shape
|
||||||
ratio_pad: the ratio of the padding to the original image.
|
ratio_pad (tuple): the ratio of the padding to the original image.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The masks are being returned.
|
masks (torch.tensor): The masks that are being returned.
|
||||||
"""
|
"""
|
||||||
# Rescale coordinates (xyxy) from im1_shape to im0_shape
|
# Rescale coordinates (xyxy) from im1_shape to im0_shape
|
||||||
if ratio_pad is None: # calculate from im0_shape
|
if ratio_pad is None: # calculate from im0_shape
|
||||||
@ -297,14 +294,12 @@ def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
|
|||||||
|
|
||||||
def xyxy2xywh(x):
|
def xyxy2xywh(x):
|
||||||
"""
|
"""
|
||||||
> It takes a list of bounding boxes, and converts them from the format [x1, y1, x2, y2] to [x, y, w,
|
Convert bounding box coordinates from (x1, y1, x2, y2) format to (x, y, width, height) format.
|
||||||
h] where xy1=top-left, xy2=bottom-right
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the input tensor
|
x (np.ndarray) or (torch.Tensor): The input tensor containing the bounding box coordinates in (x1, y1, x2, y2) format.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the center of the box, the width and the height of the box.
|
y (numpy.ndarray) or (torch.Tensor): The bounding box coordinates in (x, y, width, height) format.
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[..., 0] = (x[..., 0] + x[..., 2]) / 2 # x center
|
y[..., 0] = (x[..., 0] + x[..., 2]) / 2 # x center
|
||||||
@ -316,13 +311,12 @@ def xyxy2xywh(x):
|
|||||||
|
|
||||||
def xywh2xyxy(x):
|
def xywh2xyxy(x):
|
||||||
"""
|
"""
|
||||||
> It converts the bounding box from x,y,w,h to x1,y1,x2,y2 where xy1=top-left, xy2=bottom-right
|
Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is the top-left corner and (x2, y2) is the bottom-right corner.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the input tensor
|
x (np.ndarray) or (torch.Tensor): The input tensor containing the bounding box coordinates in (x, y, width, height) format.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the top left and bottom right coordinates of the bounding box.
|
y (numpy.ndarray) or (torch.Tensor): The bounding box coordinates in (x1, y1, x2, y2) format.
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x
|
y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x
|
||||||
@ -334,17 +328,16 @@ def xywh2xyxy(x):
|
|||||||
|
|
||||||
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
||||||
"""
|
"""
|
||||||
> It converts the normalized coordinates to the actual coordinates [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
Convert normalized bounding box coordinates to pixel coordinates.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the bounding box coordinates
|
x (np.ndarray) or (torch.Tensor): The bounding box coordinates.
|
||||||
w: width of the image. Defaults to 640
|
w (int): Width of the image. Defaults to 640
|
||||||
h: height of the image. Defaults to 640
|
h (int): Height of the image. Defaults to 640
|
||||||
padw: padding width. Defaults to 0
|
padw (int): Padding width. Defaults to 0
|
||||||
padh: height of the padding. Defaults to 0
|
padh (int): Padding height. Defaults to 0
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the xyxy coordinates of the bounding box.
|
y (numpy.ndarray) or (torch.Tensor): The coordinates of the bounding box in the format [x1, y1, x2, y2] where x1,y1 is the top-left corner, x2,y2 is the bottom-right corner of the bounding box.
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[..., 0] = w * (x[..., 0] - x[..., 2] / 2) + padw # top left x
|
y[..., 0] = w * (x[..., 0] - x[..., 2] / 2) + padw # top left x
|
||||||
@ -356,18 +349,16 @@ def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
|||||||
|
|
||||||
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
|
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
|
||||||
"""
|
"""
|
||||||
> It takes in a list of bounding boxes, and returns a list of bounding boxes, but with the x and y
|
Convert bounding box coordinates from (x1, y1, x2, y2) format to (x, y, width, height, normalized) format. x, y, width and height are normalized to image dimensions
|
||||||
coordinates normalized to the width and height of the image
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the bounding box coordinates
|
x (np.ndarray) or (torch.Tensor): The input tensor containing the bounding box coordinates in (x1, y1, x2, y2) format.
|
||||||
w: width of the image. Defaults to 640
|
w (int): The width of the image. Defaults to 640
|
||||||
h: height of the image. Defaults to 640
|
h (int): The height of the image. Defaults to 640
|
||||||
clip: If True, the boxes will be clipped to the image boundaries. Defaults to False
|
clip (bool): If True, the boxes will be clipped to the image boundaries. Defaults to False
|
||||||
eps: the minimum value of the box's width and height.
|
eps (float): The minimum value of the box's width and height. Defaults to 0.0
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the xywhn format of the bounding boxes.
|
y (numpy.ndarray) or (torch.Tensor): The bounding box coordinates in (x, y, width, height, normalized) format
|
||||||
"""
|
"""
|
||||||
if clip:
|
if clip:
|
||||||
clip_boxes(x, (h - eps, w - eps)) # warning: inplace clip
|
clip_boxes(x, (h - eps, w - eps)) # warning: inplace clip
|
||||||
@ -381,17 +372,16 @@ def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
|
|||||||
|
|
||||||
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
||||||
"""
|
"""
|
||||||
> It converts normalized segments into pixel segments of shape (n,2)
|
Convert normalized coordinates to pixel coordinates of shape (n,2)
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the normalized coordinates of the bounding box
|
x (numpy.ndarray) or (torch.Tensor): The input tensor of normalized bounding box coordinates
|
||||||
w: width of the image. Defaults to 640
|
w (int): The width of the image. Defaults to 640
|
||||||
h: height of the image. Defaults to 640
|
h (int): The height of the image. Defaults to 640
|
||||||
padw: padding width. Defaults to 0
|
padw (int): The width of the padding. Defaults to 0
|
||||||
padh: padding height. Defaults to 0
|
padh (int): The height of the padding. Defaults to 0
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the x and y coordinates of the top left corner of the bounding box.
|
y (numpy.ndarray) or (torch.Tensor): The x and y coordinates of the top left corner of the bounding box
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[..., 0] = w * x[..., 0] + padw # top left x
|
y[..., 0] = w * x[..., 0] + padw # top left x
|
||||||
@ -401,13 +391,12 @@ def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
|||||||
|
|
||||||
def xywh2ltwh(x):
|
def xywh2ltwh(x):
|
||||||
"""
|
"""
|
||||||
> It converts the bounding box from [x, y, w, h] to [x1, y1, w, h] where xy1=top-left
|
Convert the bounding box format from [x, y, w, h] to [x1, y1, w, h], where x1, y1 are the top-left coordinates.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the x coordinate of the center of the bounding box
|
x (numpy.ndarray) or (torch.Tensor): The input tensor with the bounding box coordinates in the xywh format
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the top left x and y coordinates of the bounding box.
|
y (numpy.ndarray) or (torch.Tensor): The bounding box coordinates in the xyltwh format
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
|
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
|
||||||
@ -417,13 +406,12 @@ def xywh2ltwh(x):
|
|||||||
|
|
||||||
def xyxy2ltwh(x):
|
def xyxy2ltwh(x):
|
||||||
"""
|
"""
|
||||||
> Convert nx4 boxes from [x1, y1, x2, y2] to [x1, y1, w, h] where xy1=top-left, xy2=bottom-right
|
Convert nx4 bounding boxes from [x1, y1, x2, y2] to [x1, y1, w, h], where xy1=top-left, xy2=bottom-right
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the input tensor
|
x (numpy.ndarray) or (torch.Tensor): The input tensor with the bounding boxes coordinates in the xyxy format
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the xyxy2ltwh function.
|
y (numpy.ndarray) or (torch.Tensor): The bounding box coordinates in the xyltwh format.
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[:, 2] = x[:, 2] - x[:, 0] # width
|
y[:, 2] = x[:, 2] - x[:, 0] # width
|
||||||
@ -433,10 +421,10 @@ def xyxy2ltwh(x):
|
|||||||
|
|
||||||
def ltwh2xywh(x):
|
def ltwh2xywh(x):
|
||||||
"""
|
"""
|
||||||
> Convert nx4 boxes from [x1, y1, w, h] to [x, y, w, h] where xy1=top-left, xy=center
|
Convert nx4 boxes from [x1, y1, w, h] to [x, y, w, h] where xy1=top-left, xy=center
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the input tensor
|
x (torch.tensor): the input tensor
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[:, 0] = x[:, 0] + x[:, 2] / 2 # center x
|
y[:, 0] = x[:, 0] + x[:, 2] / 2 # center x
|
||||||
@ -446,14 +434,13 @@ def ltwh2xywh(x):
|
|||||||
|
|
||||||
def ltwh2xyxy(x):
|
def ltwh2xyxy(x):
|
||||||
"""
|
"""
|
||||||
> It converts the bounding box from [x1, y1, w, h] to [x1, y1, x2, y2] where xy1=top-left,
|
It converts the bounding box from [x1, y1, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
||||||
xy2=bottom-right
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
x: the input image
|
x (numpy.ndarray) or (torch.Tensor): the input image
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the xyxy coordinates of the bounding boxes.
|
y (numpy.ndarray) or (torch.Tensor): the xyxy coordinates of the bounding boxes.
|
||||||
"""
|
"""
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
y[:, 2] = x[:, 2] + x[:, 0] # width
|
y[:, 2] = x[:, 2] + x[:, 0] # width
|
||||||
@ -463,14 +450,13 @@ def ltwh2xyxy(x):
|
|||||||
|
|
||||||
def segments2boxes(segments):
|
def segments2boxes(segments):
|
||||||
"""
|
"""
|
||||||
> It converts segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
|
It converts segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
segments: list of segments, each segment is a list of points, each point is a list of x, y
|
segments (list): list of segments, each segment is a list of points, each point is a list of x, y coordinates
|
||||||
coordinates
|
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the xywh coordinates of the bounding boxes.
|
(np.array): the xywh coordinates of the bounding boxes.
|
||||||
"""
|
"""
|
||||||
boxes = []
|
boxes = []
|
||||||
for s in segments:
|
for s in segments:
|
||||||
@ -481,15 +467,14 @@ def segments2boxes(segments):
|
|||||||
|
|
||||||
def resample_segments(segments, n=1000):
|
def resample_segments(segments, n=1000):
|
||||||
"""
|
"""
|
||||||
> It takes a list of segments (n,2) and returns a list of segments (n,2) where each segment has been
|
It takes a list of segments (n,2) and returns a list of segments (n,2) where each segment has been up-sampled to n points
|
||||||
up-sampled to n points
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
segments: a list of (n,2) arrays, where n is the number of points in the segment.
|
segments (list): a list of (n,2) arrays, where n is the number of points in the segment.
|
||||||
n: number of points to resample the segment to. Defaults to 1000
|
n (int): number of points to resample the segment to. Defaults to 1000
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the resampled segments.
|
segments (list): the resampled segments.
|
||||||
"""
|
"""
|
||||||
for i, s in enumerate(segments):
|
for i, s in enumerate(segments):
|
||||||
s = np.concatenate((s, s[0:1, :]), axis=0)
|
s = np.concatenate((s, s[0:1, :]), axis=0)
|
||||||
@ -501,14 +486,14 @@ def resample_segments(segments, n=1000):
|
|||||||
|
|
||||||
def crop_mask(masks, boxes):
|
def crop_mask(masks, boxes):
|
||||||
"""
|
"""
|
||||||
> It takes a mask and a bounding box, and returns a mask that is cropped to the bounding box
|
It takes a mask and a bounding box, and returns a mask that is cropped to the bounding box
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
masks: [h, w, n] tensor of masks
|
masks (torch.tensor): [h, w, n] tensor of masks
|
||||||
boxes: [n, 4] tensor of bbox coords in relative point form
|
boxes (torch.tensor): [n, 4] tensor of bbox coordinates in relative point form
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The masks are being cropped to the bounding box.
|
(torch.tensor): The masks are being cropped to the bounding box.
|
||||||
"""
|
"""
|
||||||
n, h, w = masks.shape
|
n, h, w = masks.shape
|
||||||
x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n)
|
x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n)
|
||||||
@ -520,17 +505,17 @@ def crop_mask(masks, boxes):
|
|||||||
|
|
||||||
def process_mask_upsample(protos, masks_in, bboxes, shape):
|
def process_mask_upsample(protos, masks_in, bboxes, shape):
|
||||||
"""
|
"""
|
||||||
> It takes the output of the mask head, and applies the mask to the bounding boxes. This produces masks of higher
|
It takes the output of the mask head, and applies the mask to the bounding boxes. This produces masks of higher
|
||||||
quality but is slower.
|
quality but is slower.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
protos: [mask_dim, mask_h, mask_w]
|
protos (torch.tensor): [mask_dim, mask_h, mask_w]
|
||||||
masks_in: [n, mask_dim], n is number of masks after nms
|
masks_in (torch.tensor): [n, mask_dim], n is number of masks after nms
|
||||||
bboxes: [n, 4], n is number of masks after nms
|
bboxes (torch.tensor): [n, 4], n is number of masks after nms
|
||||||
shape: the size of the input image
|
shape (tuple): the size of the input image (h,w)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
mask
|
(torch.tensor): The upsampled masks.
|
||||||
"""
|
"""
|
||||||
c, mh, mw = protos.shape # CHW
|
c, mh, mw = protos.shape # CHW
|
||||||
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
|
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
|
||||||
@ -541,17 +526,17 @@ def process_mask_upsample(protos, masks_in, bboxes, shape):
|
|||||||
|
|
||||||
def process_mask(protos, masks_in, bboxes, shape, upsample=False):
|
def process_mask(protos, masks_in, bboxes, shape, upsample=False):
|
||||||
"""
|
"""
|
||||||
> It takes the output of the mask head, and applies the mask to the bounding boxes. This is faster but produces
|
It takes the output of the mask head, and applies the mask to the bounding boxes. This is faster but produces
|
||||||
downsampled quality of mask
|
downsampled quality of mask
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
protos: [mask_dim, mask_h, mask_w]
|
protos (torch.tensor): [mask_dim, mask_h, mask_w]
|
||||||
masks_in: [n, mask_dim], n is number of masks after nms
|
masks_in (torch.tensor): [n, mask_dim], n is number of masks after nms
|
||||||
bboxes: [n, 4], n is number of masks after nms
|
bboxes (torch.tensor): [n, 4], n is number of masks after nms
|
||||||
shape: the size of the input image
|
shape (tuple): the size of the input image (h,w)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
mask
|
(torch.tensor): The processed masks.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
c, mh, mw = protos.shape # CHW
|
c, mh, mw = protos.shape # CHW
|
||||||
@ -572,16 +557,16 @@ def process_mask(protos, masks_in, bboxes, shape, upsample=False):
|
|||||||
|
|
||||||
def process_mask_native(protos, masks_in, bboxes, shape):
|
def process_mask_native(protos, masks_in, bboxes, shape):
|
||||||
"""
|
"""
|
||||||
> It takes the output of the mask head, and crops it after upsampling to the bounding boxes.
|
It takes the output of the mask head, and crops it after upsampling to the bounding boxes.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
protos: [mask_dim, mask_h, mask_w]
|
protos (torch.tensor): [mask_dim, mask_h, mask_w]
|
||||||
masks_in: [n, mask_dim], n is number of masks after nms
|
masks_in (torch.tensor): [n, mask_dim], n is number of masks after nms
|
||||||
bboxes: [n, 4], n is number of masks after nms
|
bboxes (torch.tensor): [n, 4], n is number of masks after nms
|
||||||
shape: input_image_size, (h, w)
|
shape (tuple): the size of the input image (h,w)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
masks: [h, w, n]
|
masks (torch.tensor): The returned masks with dimensions [h, w, n]
|
||||||
"""
|
"""
|
||||||
c, mh, mw = protos.shape # CHW
|
c, mh, mw = protos.shape # CHW
|
||||||
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
|
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
|
||||||
@ -598,17 +583,17 @@ def process_mask_native(protos, masks_in, bboxes, shape):
|
|||||||
|
|
||||||
def scale_segments(img1_shape, segments, img0_shape, ratio_pad=None, normalize=False):
|
def scale_segments(img1_shape, segments, img0_shape, ratio_pad=None, normalize=False):
|
||||||
"""
|
"""
|
||||||
> Rescale segment coords (xyxy) from img1_shape to img0_shape
|
Rescale segment coordinates (xyxy) from img1_shape to img0_shape
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
img1_shape: The shape of the image that the segments are from.
|
img1_shape (tuple): The shape of the image that the segments are from.
|
||||||
segments: the segments to be scaled
|
segments (torch.tensor): the segments to be scaled
|
||||||
img0_shape: the shape of the image that the segmentation is being applied to
|
img0_shape (tuple): the shape of the image that the segmentation is being applied to
|
||||||
ratio_pad: the ratio of the image size to the padded image size.
|
ratio_pad (tuple): the ratio of the image size to the padded image size.
|
||||||
normalize: If True, the coordinates will be normalized to the range [0, 1]. Defaults to False
|
normalize (bool): If True, the coordinates will be normalized to the range [0, 1]. Defaults to False
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
the segmented image.
|
segments (torch.tensor): the segmented image.
|
||||||
"""
|
"""
|
||||||
if ratio_pad is None: # calculate from img0_shape
|
if ratio_pad is None: # calculate from img0_shape
|
||||||
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
||||||
@ -629,11 +614,11 @@ def scale_segments(img1_shape, segments, img0_shape, ratio_pad=None, normalize=F
|
|||||||
|
|
||||||
def masks2segments(masks, strategy='largest'):
|
def masks2segments(masks, strategy='largest'):
|
||||||
"""
|
"""
|
||||||
> It takes a list of masks(n,h,w) and returns a list of segments(n,xy)
|
It takes a list of masks(n,h,w) and returns a list of segments(n,xy)
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
masks: the output of the model, which is a tensor of shape (batch_size, 160, 160)
|
masks (torch.tensor): the output of the model, which is a tensor of shape (batch_size, 160, 160)
|
||||||
strategy: 'concat' or 'largest'. Defaults to largest
|
strategy (str): 'concat' or 'largest'. Defaults to largest
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
segments (List): list of segment masks
|
segments (List): list of segment masks
|
||||||
@ -654,12 +639,12 @@ def masks2segments(masks, strategy='largest'):
|
|||||||
|
|
||||||
def clip_segments(segments, shape):
|
def clip_segments(segments, shape):
|
||||||
"""
|
"""
|
||||||
> It takes a list of line segments (x1,y1,x2,y2) and clips them to the image shape (height, width)
|
It takes a list of line segments (x1,y1,x2,y2) and clips them to the image shape (height, width)
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
segments: a list of segments, each segment is a list of points, each point is a list of x,y
|
segments (list): a list of segments, each segment is a list of points, each point is a list of x,y
|
||||||
coordinates
|
coordinates
|
||||||
shape: the shape of the image
|
shape (tuple): the shape of the image
|
||||||
"""
|
"""
|
||||||
if isinstance(segments, torch.Tensor): # faster individually
|
if isinstance(segments, torch.Tensor): # faster individually
|
||||||
segments[:, 0].clamp_(0, shape[1]) # x
|
segments[:, 0].clamp_(0, shape[1]) # x
|
||||||
@ -670,5 +655,13 @@ def clip_segments(segments, shape):
|
|||||||
|
|
||||||
|
|
||||||
def clean_str(s):
|
def clean_str(s):
|
||||||
# Cleans a string by replacing special characters with underscore _
|
"""
|
||||||
|
Cleans a string by replacing special characters with underscore _
|
||||||
|
|
||||||
|
Args:
|
||||||
|
s (str): a string needing special characters replaced
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(str): a string with special characters replaced by an underscore _
|
||||||
|
"""
|
||||||
return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
|
return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
|
||||||
|
Loading…
x
Reference in New Issue
Block a user