使用Python脚本转换YOLOv5配置文件到https://github.com/ultralytics/ultralytics：一个详细的指南

在深度学习领域，YOLO（You Only Look Once）系列模型因其高效和准确性而广受欢迎。然而，随着项目需求的变化，有时我们需要对预训练模型的配置文件进行调整。本文将详细介绍如何使用Python脚本自动转换YOLOv5的配置文件到https://github.com/ultralytics/ultralytics，并解释每个步骤的目的和实现方法。

准备工作

首先，确保安装了必要的库：

pip install pyyaml

代码详解

定义格式化和过滤函数

def format_and_filter_list(lst):
    formatted_items = []
    for item in lst:
        if isinstance(item, list):
            processed_item = format_and_filter_list(item)
            if processed_item:  # 只添加非空列表
                formatted_items.append(processed_item)
        elif isinstance(item, str):
            if item == 'nearest':
                formatted_items.append(f"'{item}'")
            elif item != 'anchors':  # 过滤掉 'anchors'
                formatted_items.append(item)
        else:
            formatted_items.append(str(item))
    return f"[{', '.join(formatted_items)}]"

定义转换配置函数

def transform_config(original_config):
    transformed = {
        'nc': original_config['nc'],  # number of classes
        'scales': {
            'n': [original_config['depth_multiple'], original_config['width_multiple'], 1024]
        },
        'backbone': [],
        'head': []
    }

    # Transform backbone
    for layer in original_config['backbone']:
        if not isinstance(layer, list):
            layer = [layer]
        transformed['backbone'].append(layer)

    # Adjust the head to match the desired format, also updating nc and anchors usage
    for i, layer in enumerate(original_config['head']):
        if not isinstance(layer, list):
            layer = [layer]

        if isinstance(layer[0], list) and len(layer[0]) > 3 and layer[0][3] == 'Detect':
            layer[0][3] = ['nc', 'anchors']
            transformed['head'].append([layer[0]])
        else:
            transformed['head'].append(layer)

    # Special case: update the final Detect layer configuration
    if transformed['head'] and isinstance(transformed['head'][-1], list) and transformed['head'][-1]:
        last_layer = transformed['head'][-1]
        if isinstance(last_layer[0], list) and len(last_layer[0]) > 3:
            last_layer[0][3] = ['nc']

    return transformed

定义主函数以更改YAML文件

def change_yaml(yolov5_config):
    # Check if input file exists
    if not os.path.exists(yolov5_config):
        print(f"Error: Input file '{yolov5_config}' does not exist.")
    else:
        old_path = os.path.dirname(yolov5_config)
        save_path = os.path.dirname(old_path)
        new_yaml_name_path = os.path.join(save_path, "yolov5_" + os.path.basename(yolov5_config))
        # Load your YAML file
        with open(yolov5_config, 'r') as file:
            try:
                original_config = yaml.safe_load(file)
            except yaml.YAMLError as exc:
                print(f"Error: Failed to parse YAML file. {exc}")
                exit(1)

        transformed_config = transform_config(original_config)

        # Ensure that each element in backbone and head is a list
        def ensure_list_format(config_part):
            new_config_part = []
            for item in config_part:
                if not isinstance(item, list):
                    item = [item]
                new_config_part.append(item)
            return new_config_part

        transformed_config['backbone'] = ensure_list_format(transformed_config['backbone'])
        transformed_config['head'] = ensure_list_format(transformed_config['head'])

        # Generate the formatted YAML content
        formatted_yaml_content = f"""# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# Ultralytics YOLOv5 object detection model with P3/8 - P5/32 outputs
# Model docs: https://docs.ultralytics.com/models/yolov5
# Task docs: https://docs.ultralytics.com/tasks/detect

# Parameters
nc: {transformed_config['nc']} # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov5n.yaml' will call yolov5.yaml with scale 'n'
  # [depth, width, max_channels]
  n: {transformed_config['scales']['n']}
# YOLOv5 v6.0 backbone
backbone:
"""

    for layer in transformed_config['backbone']:
        layer = format_and_filter_list(layer)
        print(layer)
        formatted_yaml_content += f"  - {layer}\n"

    formatted_yaml_content += """
# YOLOv5 v6.0 head with (P2, P3, P4, P5) outputs
head:
"""

    for layer in transformed_config['head']:

        print(layer)
        if isinstance(layer, list) and layer[2] == 'Conv':
            formatted_yaml_content += '\n'

        layer = format_and_filter_list(layer)
        print(layer)
        formatted_yaml_content += f"  - {layer}\n"

    # Save the transformed configuration to a new YAML file
    with open(new_yaml_name_path, 'w') as file:
        file.write(formatted_yaml_content)

    print(f"Configuration has been transformed and saved to '{new_yaml_name_path}'.")

批量处理多个配置文件

input_path = "/media/lindsay/data/ultralytics-main/ultralytics/cfg/models/v5/yolov5"
input_path_list = [os.path.join(input_path, i) for i in os.listdir(input_path)]
for i in input_path_list:
    change_yaml(i)

总结

通过上述步骤，您可以轻松地将YOLOv5的配置文件转换为ultralytics-yolov8所需的格式，并且可以批量处理多个配置文件。这个脚本不仅可以帮助您节省时间，还能减少手动操作带来的错误。希望这篇博客能为您提供有价值的指导，让您更高效地管理深度学习项目的配置文件。

完整代码

import yaml
import os

# 函数用于格式化列表为所需形式
def format_and_filter_list(lst):
    formatted_items = []
    for item in lst:
        if isinstance(item, list):
            # 如果是列表类型，则递归调用此函数进行格式化和过滤
            processed_item = format_and_filter_list(item)
            if processed_item:  # 只添加非空列表
                formatted_items.append(processed_item)
        elif isinstance(item, str):
            # 特殊处理需要保留引号的字符串
            if item == 'nearest':
                formatted_items.append(f"'{item}'")
            elif item != 'anchors':  # 过滤掉 'anchors'
                formatted_items.append(item)
        else:
            # 其他类型直接转换成字符串表示
            formatted_items.append(str(item))

    # 返回格式化后的列表字符串表示
    return f"[{', '.join(formatted_items)}]"


def transform_config(original_config):
    transformed = {
        'nc': original_config['nc'],  # number of classes
        'scales': {
            'n': [original_config['depth_multiple'], original_config['width_multiple'], 1024]
        },
        'backbone': [],
        'head': []
    }

    # Transform backbone
    for layer in original_config['backbone']:
        if not isinstance(layer, list):
            layer = [layer]
        transformed['backbone'].append(layer)

    # Adjust the head to match the desired format, also updating nc and anchors usage
    for i, layer in enumerate(original_config['head']):
        if not isinstance(layer, list):
            layer = [layer]

        if isinstance(layer[0], list) and len(layer[0]) > 3 and layer[0][3] == 'Detect':
            # Update Detect layer to use 'nc' instead of specific numbers for classes
            layer[0][3] = ['nc', 'anchors']
            transformed['head'].append([layer[0]])
        else:
            transformed['head'].append(layer)

    # Special case: update the final Detect layer configuration
    if transformed['head'] and isinstance(transformed['head'][-1], list) and transformed['head'][-1]:
        last_layer = transformed['head'][-1]
        if isinstance(last_layer[0], list) and len(last_layer[0]) > 3:
            last_layer[0][3] = ['nc']

    return transformed


def change_yaml(yolov5_config):
    # Check if input file exists
    if not os.path.exists(yolov5_config):
        print(f"Error: Input file '{yolov5_config}' does not exist.")
    else:
        old_path = os.path.dirname(yolov5_config)
        save_path = os.path.dirname(old_path)
        new_yaml_name_path = os.path.join(save_path, "yolov5_" + os.path.basename(yolov5_config))
        # Load your YAML file
        with open(yolov5_config, 'r') as file:
            try:
                original_config = yaml.safe_load(file)
            except yaml.YAMLError as exc:
                print(f"Error: Failed to parse YAML file. {exc}")
                exit(1)

        transformed_config = transform_config(original_config)

        # Ensure that each element in backbone and head is a list
        def ensure_list_format(config_part):
            new_config_part = []
            for item in config_part:
                if not isinstance(item, list):
                    item = [item]
                new_config_part.append(item)
            return new_config_part

        transformed_config['backbone'] = ensure_list_format(transformed_config['backbone'])
        transformed_config['head'] = ensure_list_format(transformed_config['head'])

        # Generate the formatted YAML content
        formatted_yaml_content = f"""# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# Ultralytics YOLOv5 object detection model with P3/8 - P5/32 outputs
# Model docs: https://docs.ultralytics.com/models/yolov5
# Task docs: https://docs.ultralytics.com/tasks/detect

# Parameters
nc: {transformed_config['nc']} # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov5n.yaml' will call yolov5.yaml with scale 'n'
  # [depth, width, max_channels]
  n: {transformed_config['scales']['n']}
# YOLOv5 v6.0 backbone
backbone:
"""

    for layer in transformed_config['backbone']:
        layer = format_and_filter_list(layer)
        print(layer)
        formatted_yaml_content += f"  - {layer}\n"

    formatted_yaml_content += """
# YOLOv5 v6.0 head with (P2, P3, P4, P5) outputs
head:
"""

    for layer in transformed_config['head']:

        print(layer)
        if isinstance(layer, list) and layer[2] == 'Conv':
            formatted_yaml_content += '\n'

        layer = format_and_filter_list(layer)
        print(layer)
        formatted_yaml_content += f"  - {layer}\n"

    # Save the transformed configuration to a new YAML file
    with open(new_yaml_name_path, 'w') as file:
        file.write(formatted_yaml_content)

    print(f"Configuration has been transformed and saved to '{new_yaml_name_path}'.")


input_path = "/media/lindsay/data/ultralytics-main/ultralytics/cfg/models/v5/yolov5"
input_path_list = [os.path.join(input_path, i) for i in os.listdir(input_path)]
for i in input_path_list:
    change_yaml(i)