采用YOLOv5模型使用深度学习目标检测方法训练工地航拍工程车装备车辆数据集 识别航拍状态下的自卸车混凝土车挖掘机起重机等
文章目录
- 1. 安装依赖
- 2. 配置数据集
- 3. 加载预训练模型并开始训练
- 4. 评估模型
- 5. 加载训练好的模型进行推理
如何训练自己的数据集——
工地工程车装备车辆数据集，无人机拍摄，从高空视角pai，标注为YOLO风格。数据集包含6700+条目，9大类别，涵盖挖掘机、塔吊、翻斗车等
使用深度学习目标检测方法训练工地工程车装备车辆数据集（标注为YOLO风格），采用YOLOv5模型，因为它具有良好的性能和易用性。
以下是详细的步骤和代码示例，包括配置数据集、加载预训练模型、训练模型、评估模型以及如何加载训练好的模型进行推理并显示结果。
9个label。
0.叉车 1.打桩机 2.异物
3.起重机 4.自卸车 5.鸟骸
6.混凝土泵车 7.塔吊 8.挖掘机
1. 安装依赖
首先确保安装了必要的库：
pipinstalltorch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113gitclone https://github.com/ultralytics/yolov5cdyolov5 pipinstall-rrequirements.txt2. 配置数据集
假设你的数据集目录结构如下：
dataset/ ├── images/ │ ├── train/ │ │ ├── img1.jpg │ │ └── ... │ └── val/ │ ├── img1.jpg │ └── ... └── labels/ ├── train/ │ ├── img1.txt │ └── ... └── val/ ├── img1.txt └── ...同时，你需要创建一个data.yaml文件来描述你的数据集路径和类别信息：
train:./dataset/images/train/val:./dataset/images/val/nc:9# number of classesnames:['excavator','tower_crane','dump_truck',...]# class names3. 加载预训练模型并开始训练
你可以从YOLOv5的官方仓库中选择一个预训练模型开始训练。下面是如何定义训练命令的例子：
python train.py--img640--batch16--epochs50--datadata.yaml--weightsyolov5s.pt--cache这里我们选择了yolov5s.pt作为预训练权重，它代表YOLOv5的小型版本，适合快速实验和资源受限环境。你可以根据需要选择更大的模型如yolov5m,yolov5l, 或yolov5x。
4. 评估模型
训练完成后，可以使用以下命令对验证集进行评估：
python val.py--weightsruns/train/exp/weights/best.pt--datadata.yaml--img640--taskval这将输出包括精确率、召回率和mAP在内的多个评估指标。
5. 加载训练好的模型进行推理
为了演示如何加载训练好的模型并对新图像进行推理，下面提供了一个Python脚本示例：
importtorchfromPILimportImageimportcv2importmatplotlib.pyplotasplt# Load the modelmodel=torch.hub.load('ultralytics/yolov5','custom',path_or_model='runs/train/exp/weights/best.pt')defdetect_image(image_path):# Inferenceresults=model(image_path)# Resultsdf=results.pandas().xyxy[0]# or .show(), .save(), .crop()# Read image with OpenCVimg=cv2.imread(image_path)img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)# Draw bounding boxes on the imageforindex,rowindf.iterrows():xmin,ymin,xmax,ymax=int(row['xmin']),int(row['ymin']),int(row['xmax']),int(row['ymax'])label=row['name']confidence=row['confidence']cv2.rectangle(img,(xmin,ymin),(xmax,ymax),(255,0,0),2)cv2.putText(img,f'{label}{confidence:.2f}',(xmin,ymin-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,(255,0,0),2)# Display the resultplt.figure(figsize=(12,8))plt.imshow(img)plt.axis('off')plt.show()# 使用函数detect_image('path/to/your/image.jpg')这段代码首先加载训练好的YOLOv5模型，然后定义了一个名为detect_image的函数用于读取图像、执行推理并在原图上绘制预测框，最后显示处理后的图像。
