Yolo整合

helmet检测

准备数据集

kaggle中yolo格式的helmet数据集

导入yolo-v8

1
import ultralytics

安装cuda、torch…

….

[error]https://zzggo.com/2024/09/05/Yolo-fix/

准备yaml文件

1
2
3
4
5
train: ...
val: ...
test ...
nc = 10
names: [Hardhat, Mask, NO-Hardhat, NO-Mask, NO-Safety Vest, Person, Safety Cone, Safety Vest, machinery, vehicle]

基于yolo-v8的coco模型训练helmet模型

1
2
3
4
5
6
7
8
9
10
from ultralytics import YOLO

def main():
    model = YOLO('yolov8n.pt')
    model.to('cuda')
    model.train(data='safehat.yaml', epochs=100)
    model.val()

if __name__ == '__main__':
    main()

编写测试文件

1
2
3
from ultralytics import YOLO
model = YOLO('best.pt')
model.predict('testfile',save=True)

结果

原始图片
运行后图片

附：

检查显卡可用性

1
2
3
4
5
6
7
8
torch.cuda.is_available()
cuda是否可用；
torch.cuda.device_count()
返回gpu数量；
torch.cuda.get_device_name(0)
返回gpu名字，设备索引默认从0开始；
torch.cuda.current_device()
返回当前设备索引；

torch.cuda.is_available()
False

nvcc -V 检查cuda版本

Known issue with inference using CPU on Windows for PyTorch 2.4.0

Fixes

  1. Downgrade PyTorch (CPU) to a previous version
1
2
pip uninstall torch torchvision torchaudio -y
pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cpu
  1. Install PyTorch with CUDA support, requires NVIDIA GPU with CUDA support
1
2
pip uninstall torch torchvision torchaudio -y
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

ERROR: pip’s dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
ultralytics 8.2.87 requires torch<2.4.0,>=1.8.0; sys_platform == “win32”, but you have torch 2.4.1+cu121 which is incompatible.

1
2
3
4
pip uninstall torch torchvision torchaudio -y
pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121
pip uninstall torch torchvision torchaudio -y
pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121

yolo.document

**### [Yolo-v8_train]<**https://docs.ultralytics.com/modes/train/