使用YOLOv8进行汽车零部件分割，汽车部位分割数据集 3000多张，yolo格式（txt标签），分辨率640x640， 22类

使用YOLOv8进行汽车零部件分割和汽车部位分割。代码包括数据集准备、加载预训练模型、设置数据集路径、数据集配置、开始训练、模型评估、可视化训练结果、清理临时文件以及推理和显示结果。

import os
import random
from pathlib import Path
import shutil
import torch
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import tempfile
import matplotlib.pyplot as plt
import yaml

# 1. 创建数据集配置文件 (data.yaml)
def create_data_yaml(train_dir, val_dir, nc, names):
    data_config = {
        'train': train_dir,
        'val': val_dir,
        'nc': nc,
        'names': names
    }
    
    with open('data.yaml', 'w') as f:
        yaml.dump(data_config, f)

# 2. 分割数据集
def split_dataset(data_dir, train_ratio=0.8):
    images = list(Path(data_dir).glob('*.jpg'))
    random.shuffle(images)
    
    num_train = int(len(images) * train_ratio)
    train_images = images[:num_train]
    val_images = images[num_train:]
    
    train_dir = Path(data_dir).parent / 'train'
    val_dir = Path(data_dir).parent / 'val'
    
    train_img_dir = train_dir / 'images'
    train_label_dir = train_dir / 'labels'
    val_img_dir = val_dir / 'images'
    val_label_dir = val_dir / 'labels'
    
    train_img_dir.mkdir(parents=True, exist_ok=True)
    train_label_dir.mkdir(parents=True, exist_ok=True)
    val_img_dir.mkdir(parents=True, exist_ok=True)
    val_label_dir.mkdir(parents=True, exist_ok=True)
    
    for img in train_images:
        label_path = img.with_suffix('.txt')
        shutil.copy(img, train_img_dir / img.name)
        shutil.copy(label_path, train_label_dir / label_path.name)
    
    for img in val_images:
        label_path = img.with_suffix('.txt')
        shutil.copy(img, val_img_dir / img.name)
        shutil.copy(label_path, val_label_dir / label_path.name)

# 使用示例
split_dataset('./datasets/car_parts_dataset/images')

# 3. 创建数据集配置文件 (data.yaml)
create_data_yaml(
    train_dir='./datasets/car_parts_dataset/train/images',
    val_dir='./datasets/car_parts_dataset/val/images',
    nc=5,  # 假设有5个类别
    names=['part1', 'part2', 'part3', 'part4', 'part5']  # 替换为实际类别名称
)

# 4. 训练脚本
def train_model(dataset_config, epochs=100, imgsz=640, batch=16, name='car_parts_detection', project='runs/train'):
    # 设置随机种子以保证可重复性
    torch.manual_seed(42)
    
    # 加载预训练的YOLOv8n模型
    model = YOLO('yolov8n.pt')
    
    # 训练模型
    results = model.train(
        data=dataset_config,
        epochs=epochs,
        imgsz=imgsz,
        batch=batch,
        name=name,
        project=project
    )
    
    # 评估模型
    metrics = model.val()
    
    # 保存最佳模型权重
    best_model_weights = f'{project}/{name}/weights/best.pt'
    print(f"Best model weights saved to {best_model_weights}")
    
    return best_model_weights

# 使用示例
best_model_weights = train_model('data.yaml')

# 5. 可视化训练结果
def visualize_training_results(project_dir, run_name):
    # 加载训练好的模型
    model = YOLO(f'{project_dir}/{run_name}/weights/best.pt')
    
    # 可视化训练结果
    model.plot_results(save=True, save_dir=f'{project_dir}/{run_name}')

# 使用示例
visualize_training_results('runs/train', 'car_parts_detection')

# 6. 清理临时文件
def clean_temp_files(project_dir):
    temp_dirs = [
        f'{project_dir}/wandb',
        f'{project_dir}/cache'
    ]
    
    for dir_path in temp_dirs:
        if os.path.exists(dir_path):
            shutil.rmtree(dir_path)
            print(f"Removed directory: {dir_path}")

# 使用示例
clean_temp_files('runs/train/car_parts_detection')

# 7. 推理脚本
def detect_image(model, image_path, conf_threshold=0.5):
    results = model.predict(image_path, conf=conf_threshold)[0]
    annotated_frame = annotate_image(image_path, results, model)
    return annotated_frame

def detect_video(model, video_path, conf_threshold):
    cap = cv2.VideoCapture(video_path)
    frames = []
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model.predict(frame, conf=conf_threshold)[0]
        annotated_frame = annotate_image(frame, results, model)
        frames.append(annotated_frame)
    cap.release()
    return frames

def detect_camera(model, conf_threshold):
    cap = cv2.VideoCapture(0)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        results = model.predict(frame, conf=conf_threshold)[0]
        annotated_frame = annotate_image(frame, results, model)
        frames.append(annotated_frame)
        cv2.imshow('Camera Detection', annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
    return frames

def annotate_image(image_path, results, model):
    frame = cv2.imread(image_path)
    for result in results.boxes.cpu().numpy():
        r = result.xyxy[0].astype(int)
        cls = int(result.cls[0])
        conf = result.conf[0]
        
        label = f"{model.names[cls]} {conf:.2f}"
        color = (0, 255, 0)
        cv2.rectangle(frame, (r[0], r[1]), (r[2], r[3]), color, 2)
        cv2.putText(frame, label, (r[0], r[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
    return frame

# 使用示例
def main():
    # 加载模型
    model = YOLO(best_model_weights)
    
    # 动态调整置信度阈值
    conf_threshold = 0.5
    
    # 输入方式选择
    input_type = "Image"  # 可选: "Image", "Video", "Camera"
    
    if input_type == "Image":
        test_image_path = './datasets/car_parts_dataset/images/test_image.jpg'
        annotated_image = detect_image(model, test_image_path, conf_threshold)
        cv2.imwrite('annotated_test_image.jpg', annotated_image)
        plt.imshow(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
        plt.axis('off')
        plt.show()
        
        # 统计检测到的物体数量
        results = model.predict(test_image_path, conf=conf_threshold)[0]
        class_counts = {}
        for result in results.boxes.cpu().numpy():
            cls = int(result.cls[0])
            class_name = model.names[cls]
            if class_name in class_counts:
                class_counts[class_name] += 1
            else:
                class_counts[class_name] = 1
        
        print("Detection Summary:")
        for class_name, count in class_counts.items():
            print(f"{class_name}: {count}")
    
    elif input_type == "Video":
        test_video_path = './datasets/car_parts_dataset/videos/test_video.mp4'
        frames = detect_video(model, test_video_path, conf_threshold)
        for i, frame in enumerate(frames):
            cv2.imwrite(f'frame_{i}.jpg', frame)
            plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            plt.axis('off')
            plt.show()
    
    elif input_type == "Camera":
        frames = detect_camera(model, conf_threshold)
        for i, frame in enumerate(frames):
            cv2.imwrite(f'camera_frame_{i}.jpg', frame)
            plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            plt.axis('off')
            plt.show()

if __name__ == "__main__":
    main()

![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/25a3678459de49c482db6e473f6a14d3.png)



```![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/dd8d0005ef5a4ad18f95ffe0fe598ed5.png)

![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/f4fee7f527294fa7a1d25157e34b3144.png)
好的，以下是完整的代码，用于使用YOLOv8进行汽车零部件分割和汽车部位分割。代码包括数据集准备、加载预训练模型、设置数据集路径、数据集配置、开始训练、模型评估、可视化训练结果、清理临时文件以及推理和显示结果。

### 完整项目代码



### 文件结构

car_parts_segmentation/
├── main.py
├── datasets/
│ └── car_parts_dataset/
│ ├── yolo/
│ │ ├── labels/
│ │ │ ├── image1.txt
│ │ │ ├── image2.txt
│ │ │ └── …
│ │ ├── train/
│ │ │ ├── images/
│ │ │ │ ├── image1.jpg
│ │ │ │ ├── image2.jpg
│ │ │ │ └── …
│ │ │ └── labels/
│ │ │ ├── image1.txt
│ │ │ ├── image2.txt
│ │ │ └── …
│ │ └── val/
│ │ ├── images/
│ │ │ ├── image1.jpg
│ │ │ ├── image2.jpg
│ │ │ └── …
│ │ └── labels/
│ │ ├── image1.txt
│ │ ├── image2.txt
│ │ └── …
│ └── images/
│ ├── image1.jpg
│ ├── image2.jpg
│ └── …
├── best_car_parts_yolov8.pt
└── requirements.txt

### 安装依赖项

首先，确保你已经安装了所有必要的依赖项。你可以通过以下命令安装：

```bash
pip install -r requirements.txt

requirements.txt 内容如下：

opencv-python
torch==2.0
ultralytics
matplotlib
pyyaml

运行步骤总结

1. 克隆项目仓库（如果有的话）：

   git clone https://github.com/yourusername/car_parts_segmentation.git
   cd car_parts_segmentation
   

2. 安装依赖项：

   pip install -r requirements.txt
   

3. 数据集准备：

   • 如果数据集尚未分割，请运行数据集分割脚本。

     python main.py
     

   • 确保数据集目录结构正确。

4. 训练模型：

   python main.py
   

5. 评估模型：
   在训练脚本中，模型会在训练结束后自动进行评估。

6. 可视化训练结果：

   python main.py
   

7. 清理临时文件：

   python main.py
   

8. 推理测试：

   python main.py
   

总结

以上是完整的基于YOLOv8的汽车零部件分割系统的代码实现。代码涵盖了从数据集准备、模型训练、评估、可视化训练结果、清理临时文件到推理和显示结果的所有步骤。希望这些详细的信息和代码能够帮助你顺利实施和优化你的汽车零部件分割系统。