Yolo迁移训练-带训练源码

目录

下载Git

拉下yolo模型

下载labelimg

准备训练集

迁移训练

继续训练

---

下载Git

Git - Downloading Package

拉下yolo模型

然后用克隆脚本拉下yolo模型

python clone_yolo.py

import os
import subprocess
import sys
import shutil

def check_git_installed():
    """检查Git是否已安装"""
    try:
        result = subprocess.run(['git', '--version'], capture_output=True, text=True)
        if result.returncode == 0:
            print(f"Git已安装: {result.stdout.strip()}")
            return True
        return False
    except:
        return False

def clone_yolo():
    """克隆YOLOv5仓库"""
    # 获取当前工作目录
    work_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 检查是否已存在yolov5目录
    yolo_dir = os.path.join(work_dir, 'yolov5')
    if (os.path.exists(yolo_dir)):
        print(f"YOLO目录已存在({yolo_dir})。自动删除并重新克隆...")
        shutil.rmtree(yolo_dir, ignore_errors=True)
    
    # 克隆YOLOv5仓库 - 使用实时输出而不是捕获输出
    print("正在克隆YOLOv5仓库，这可能需要几分钟时间...")
    print("(下载进度将实时显示，请耐心等待)")
    
    try:
        # 使用实时输出方式运行git clone
        result = subprocess.run(
            ['git', 'clone', 'https://github.com/ultralytics/yolov5.git'],
            cwd=work_dir,
            capture_output=False,  # 不捕获输出，让它直接显示在控制台
            text=True
        )
        
        if result.returncode == 0:
            print("YOLOv5仓库克隆成功!")
            
            # 安装依赖
            print("\n正在安装YOLOv5依赖...")
            requirements_file = os.path.join(yolo_dir, 'requirements.txt')
            if os.path.exists(requirements_file):
                try:
                    subprocess.run(
                        [sys.executable, '-m', 'pip', 'install', '-r', requirements_file],
                        check=True,
                        capture_output=False  # 实时显示安装进度
                    )
                    print("YOLOv5依赖安装完成!")
                except subprocess.CalledProcessError as e:
                    print(f"警告: 安装依赖时出错: {str(e)}")
                    print("您可能需要手动安装依赖。")
            else:
                print(f"警告: 找不到requirements.txt文件({requirements_file})。")
            
            return True
        else:
            print(f"克隆YOLOv5仓库失败，返回代码: {result.returncode}")
            return False
    except Exception as e:
        print(f"执行Git命令时出错: {e}")
        return False

if __name__ == "__main__":
    if not check_git_installed():
        print("Git尚未安装或无法在PATH中找到。")
        print("请先运行 install_git.py 脚本安装Git，然后重新运行此脚本。")
        sys.exit(1)
    
    print("开始克隆YOLOv5仓库，请确保网络连接正常...")
    if clone_yolo():
        print("\nYOLOv5克隆和设置完成!")
        print("您现在可以使用 train_fruits.py 脚本进行水果识别的迁移学习。")
    else:
        print("\n克隆YOLO仓库失败，请检查错误信息并重试。")

下载labelimg

在下载yolo的过程中可以先去https://github.com/tzutalin/labelImg/releases下载给图像打标签的工具 labelimg

准备训练集

在 label 完后会生成 txt 文件（注意同一个水果图片的 .jpg和 .txt前面的名称要一样，比如apple_01.jpg 和 apple_01.txt ）

建立一个 dataset 文件夹，并在该目录下建立 images 和 labels 文件夹，并在这两个文件夹下面创建两个子文件夹，一个叫 train 用于存储训练集，一个叫 val 用于存储验证集

比如标签的训练集.txt文件就放在 label\val

接着我们就用labelimg框选训练集

🧭 一、启动软件

1. 如果你是通过命令行运行的，输入：

   python labelImg.py

2. 或者直接运行 labelImg.exe（Windows 下）。

🖼 二、打开图片文件夹

点击左上角的按钮：

• Open Dir：选择你存放图像的文件夹（比如 images/）

• 如果 Open Dir 找不到的话，也可以直接 Open 打开图片

🧩 三、选择保存标注的格式

点击：

•  YOLO（保存为 .txt 格式，适用于 YOLO 模型）

✏️ 四、开始标注步骤

1. 点击工具栏上的 Create RectBox 按钮（或快捷键 w）

2. 用鼠标在图片上 点击并拖动，画出一个框框（比如把猫框住）

3. 系统会弹出一个窗口，输入标签名（如：cat） → 回车

4. 你可以继续框其他对象，每个对象一个框

💾 五、保存标注

点击左上角的：

• Save 按钮（或快捷键 Ctrl+S），系统会把标注保存为 .xml 或 .txt 文件，放在和图像同目录下。

⏭ 六、切换到下一张图像

• 点箭头按钮 ▶ 或按 d 切换到下一张图像继续标注

• 按 a 切换上一张图像

• 按 Ctrl + S 随时保存标注结果

⌨️ 常用快捷键：

w                框选

Ctrl + s       保存

d                下一张

a                上一张

Delete        删除框

Crtl + e      修改标签名

然后我们写 fruits.yaml # 水果识别数据集配置（注意是文件中写的是相对路径）

# 训练和验证数据集路径（相对于YOLOv5目录）
path: ../dataset  # 数据集根目录
train: images/train  # 训练图像路径（相对于path）
val: images/val  # 验证图像路径（相对于path）

# 类别
nc: 5  # 类别数量
names: ['apple', 'banana', 'orange', 'strawberry', 'kiwi']  # 类别名称

迁移训练

运行训练脚本 train_fruis.py

import os
import argparse
import torch
import shutil

def train_yolo(epochs=100, batch_size=16, img_size=640, weights='yolov5s.pt'):
    """
    使用YOLOv5进行水果识别的迁移学习
    """
    print(f"开始训练水果识别模型，训练轮次：{epochs}，批次大小：{batch_size}，图像尺寸：{img_size}")
    
    # 确保我们在正确的工作目录
    os.chdir(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'yolov5'))
    
    # 训练命令
    train_cmd = f"python train.py --img {img_size} --batch {batch_size} --epochs {epochs} --data ../dataset/fruits.yaml --weights {weights} --cache"
    
    # 执行训练
    print(f"执行命令: {train_cmd}")
    os.system(train_cmd)
    
    # 复制训练结果到models目录
    last_model = os.path.join('runs', 'train', 'exp', 'weights', 'last.pt')
    best_model = os.path.join('runs', 'train', 'exp', 'weights', 'best.pt')
    
    if os.path.exists(last_model):
        print("复制模型到models目录...")
        os.makedirs('../models', exist_ok=True)
        shutil.copy(last_model, '../models/fruits_last.pt')
        
    if os.path.exists(best_model):
        shutil.copy(best_model, '../models/fruits_best.pt')
    
    print("训练完成！模型已保存到models目录。")

def parse_arguments():
    parser = argparse.ArgumentParser(description='YOLO水果识别迁移学习')
    parser.add_argument('--epochs', type=int, default=100, help='训练轮次')
    parser.add_argument('--batch-size', type=int, default=16, help='批次大小')
    parser.add_argument('--img-size', type=int, default=640, help='图像尺寸')
    parser.add_argument('--weights', type=str, default='yolov5s.pt', 
                        help='预训练权重 (yolov5s.pt, yolov5m.pt, yolov5l.pt, yolov5x.pt)')
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()
    train_yolo(epochs=args.epochs, batch_size=args.batch_size, 
               img_size=args.img_size, weights=args.weights)

继续训练

如果新增了训练集的话，直接运行上面的 train_fruits.py 会生成新的模型

如果我们想继续训练我们之前的模型

import os
import argparse
import shutil

def continue_training(epochs=50, batch_size=16, img_size=640, weights='models/fruits_best.pt'):
    """
    使用已训练的水果模型继续训练（增量学习）
    """
    print(f"使用已训练的模型继续训练水果识别模型，训练轮次：{epochs}，批次大小：{batch_size}")
    
    # 确保我们在正确的工作目录
    os.chdir(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'yolov5'))
    
    # 检查指定的权重文件是否存在
    weights_path = os.path.join('..', weights) if not os.path.isabs(weights) else weights
    if not os.path.exists(weights_path):
        print(f"错误：找不到指定的权重文件：{weights_path}")
        print("如果这是您第一次训练，请使用train_fruits.py脚本，或者指定yolov5s.pt作为起点")
        return False
    
    # 训练命令 - 使用之前训练的模型继续训练
    train_cmd = f"python train.py --img {img_size} --batch {batch_size} --epochs {epochs} --data ../dataset/fruits.yaml --weights {weights_path} --cache"
    
    # 执行训练
    print(f"执行命令: {train_cmd}")
    os.system(train_cmd)
    
    # 复制训练结果到models目录
    last_model = os.path.join('runs', 'train', 'exp', 'weights', 'last.pt')
    best_model = os.path.join('runs', 'train', 'exp', 'weights', 'best.pt')
    
    if os.path.exists(last_model):
        print("复制更新后的模型到models目录...")
        os.makedirs('../models', exist_ok=True)
        shutil.copy(last_model, '../models/fruits_continued_last.pt')
        
    if os.path.exists(best_model):
        shutil.copy(best_model, '../models/fruits_continued_best.pt')
    
    print("增量训练完成！更新后的模型已保存到models目录。")
    return True

def parse_arguments():
    parser = argparse.ArgumentParser(description='YOLO水果识别增量训练')
    parser.add_argument('--epochs', type=int, default=50, help='训练轮次')
    parser.add_argument('--batch-size', type=int, default=16, help='批次大小')
    parser.add_argument('--img-size', type=int, default=640, help='图像尺寸')
    parser.add_argument('--weights', type=str, default='models/fruits_best.pt', 
                        help='之前训练好的模型权重文件路径')
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()
    continue_training(epochs=args.epochs, batch_size=args.batch_size, 
                     img_size=args.img_size, weights=args.weights)