红外船舶数据集 6000张 yolo格式![]()
红外船舶数据集 6000张 yolo格式
红外船舶数据集包含6000张图片,并且已经标注为YOLO格式。我们可以使用这些数据来训练一个目标检测模型。我们使用YOLOv8架构来进行目标检测任务。
以下是详细的步骤和代码示例,帮助你开始使用这个数据集进行训练和评估。
代码仅供参考!!!
项目结构
ship_detection/ ├── main.py ├── train.py ├── evaluate.py ├── infer.py ├── visualize.py ├── datasets/ │ ├── ship_images/ │ │ ├── images/ │ │ │ ├── train/ │ │ │ └── val/ │ │ ├── labels/ │ │ │ ├── train/ │ │ │ └── val/ │ │ ├── train.txt │ │ └── val.txt ├── best_ship_model.pt ├── requirements.txt └── data.yaml文件内容
requirements.txt
opencv-python==4.5.3.56 torch==2.0.0+cu117 matplotlib numpy pandas albumentations ultralyticsdata.yaml
假设你的数据集自带data.yaml文件,其内容应类似于:
train:./datasets/ship_images/images/trainval:./datasets/ship_images/images/valnc:1# 假设只有一个类别:shipnames:['ship']数据准备
确认数据集目录结构:
确保你的数据集已经按照以下结构组织好:datasets/ └── ship_images/ ├── images/ │ ├── train/ │ └── val/ ├── labels/ │ ├── train/ │ └── val/ ├── train.txt └── val.txt检查
data.yaml文件:
确认data.yaml文件中的路径和类别信息正确无误。
训练脚本
train.py
fromultralyticsimportYOLO# 设置随机种子以保证可重复性importtorch torch.manual_seed(42)# 定义数据集路径dataset_config='data.yaml'# 加载预训练的YOLOv8模型model=YOLO('yolov8n.pt')# 训练模型results=model.train(imgsz=640,batch=16,epochs=50,data=dataset_config,weights='yolov8n.pt',name='ship_detection',project='runs/train')# 打印训练结果print(results)评估脚本
evaluate.py
fromultralyticsimportYOLO# 初始化YOLOv8模型model_path='runs/train/ship_detection/weights/best.pt'# 加载模型model=YOLO(model_path)# 评估模型metrics=model.val(data='data.yaml',imgsz=640)# 打印评估结果print(metrics)推理脚本
infer.py
fromultralyticsimportYOLOimportcv2importnumpyasnpimportmatplotlib.pyplotasplt# 初始化YOLOv8模型model_path='runs/train/ship_detection/weights/best.pt'model=YOLO(model_path)defdetect_and_visualize(image_path):results=model.predict(source=image_path,conf=0.25,iou=0.45,agnostic_nms=False)forresultinresults:boxes=result.boxes.cpu().numpy()im_array=result.orig_imgforboxinboxes:r=box.xyxy[0].astype(int)cls=int(box.cls[0])conf=box.conf[0]label=f'{model.names[cls]}{conf:.2f}'# 绘制边界框cv2.rectangle(im_array,(r[0],r[1]),(r[2],r[3]),(0,255,0),2)# 绘制标签cv2.putText(im_array,label,(r[0],r[1]-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,(0,255,0),2)# 显示图像plt.imshow(cv2.cvtColor(im_array,cv2.COLOR_BGR2RGB))plt.axis('off')plt.show()if__name__=="__main__":image_path='path/to/your/image.jpg'# 替换为你的图像路径detect_and_visualize(image_path)可视化脚本
visualize.py
fromultralyticsimportYOLOimportcv2importnumpyasnpimportmatplotlib.pyplotaspltimportrandom# 初始化YOLOv8模型model_path='runs/train/ship_detection/weights/best.pt'model=YOLO(model_path)defvisualize_dataset(dataset_path,num_samples=5):withopen(dataset_path,'r')asf:lines=f.readlines()sample_paths=random.sample(lines,num_samples)forpathinsample_paths:path=path.strip()results=model.predict(source=path,conf=0.25,iou=0.45,agnostic_nms=False)forresultinresults:boxes=result.boxes.cpu().numpy()im_array=result.orig_imgforboxinboxes:r=box.xyxy[0].astype(int)cls=int(box.cls[0])conf=box.conf[0]label=f'{model.names[cls]}{conf:.2f}'# 绘制边界框cv2.rectangle(im_array,(r[0],r[1]),(r[2],r[3]),(0,255,0),2)# 绘制标签cv2.putText(im_array,label,(r[0],r[1]-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,(0,255,0),2)# 显示图像plt.imshow(cv2.cvtColor(im_array,cv2.COLOR_BGR2RGB))plt.axis('off')plt.show()if__name__=="__main__":dataset_path='datasets/ship_images/val.txt'# 替换为你的验证集路径visualize_dataset(dataset_path)运行步骤总结
克隆项目仓库(如果有的话):
gitclone https://github.com/yourusername/ship_detection.gitcdship_detection安装依赖项:
conda create--nameship_det_envpython=3.8conda activate ship_det_env pipinstall-rrequirements.txt准备数据集:
- 确保你的数据集已经按照上述结构组织好。
- 确认
data.yaml文件中的路径和类别信息正确无误。
训练模型:
python train.py评估模型:
python evaluate.py运行推理:
python infer.py可视化数据集:
python visualize.py
操作界面
- 选择图片进行检测: 修改
infer.py中的image_path变量指向你要检测的图片路径,然后运行python infer.py。 - 批量检测: 在
visualize.py中设置dataset_path为你想要可视化的数据集路径,然后运行python visualize.py。
详细解释
requirements.txt
列出项目所需的所有Python包及其版本。
data.yaml
配置数据集路径和类别信息,用于YOLOv8模型训练。
train.py
加载预训练的YOLOv8模型并使用自定义数据集进行训练。训练完成后打印训练结果。
evaluate.py
加载训练好的YOLOv8模型并对验证集进行评估,打印评估结果。
infer.py
对单张图像进行预测并可视化检测结果。
visualize.py
对数据集中的一些样本进行预测并可视化检测结果。
希望这些详细的信息和代码能够帮助你顺利实施和优化你的红外船舶检测系统。