本地训练YOLO模型实现App图标识别全流程指南

本文档完整记录了从零开始搭建环境、标注数据、训练模型到最终实现图标裁剪的完整流程，以及过程中遇到的所有问题及解决方案。

📋 目录

• 一、项目目标

• 二、环境搭建

• 三、数据准备与标注

• 四、训练YOLO模型

• 五、模型测试与图标裁剪

• 六、问题汇总与解决方案

• 七、后续优化建议

一、项目目标

识别手机App截图中的UI元素，具体实现：

1. 检测并框出 icon（App图标）和 app_name（App名称）

2. 将 icon 区域裁剪出来，保存为独立的图片文件

3. （进阶）通过OCR识别 app_name 文字，用作图标文件名

二、环境搭建

2.1 硬件要求

2.2 软件安装

安装Anaconda

从 Anaconda官网 下载并安装。

创建项目环境

bash

# 创建环境（指定Python 3.10）
conda create -n ui_detection python=3.10

# 激活环境
conda activate ui_detection

安装核心依赖

bash

# 安装PyTorch（CPU版本，有GPU请去官网选择对应CUDA版本）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

# 克隆YOLOv5代码
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装OCR和图像处理库
pip install paddlepaddle paddleocr opencv-python pillow

安装标注工具

bash

# 通过conda安装PyQt5（稳定）
conda install -c conda-forge pyqt

# 安装labelImg
pip install labelImg

三、数据准备与标注

3.1 项目目录结构

text

D:\AI Trainer\UI_Detection/
├── datasets/
│   ├── images/          # 原始截图
│   ├── labels/          # 标注文件（.txt）
│   ├── train/           # 训练集
│   │   ├── images/
│   │   └── labels/
│   └── val/             # 验证集
│       ├── images/
│       └── labels/
├── yolov5/              # YOLOv5源码
├── export/              # Label Studio导出的数据
├── cropped_results/     # 裁剪结果输出
├── data.yaml            # 数据集配置文件
└── crop_icons.py        # 图标裁剪脚本

3.2 收集图片

• 至少收集 50-100张 不同App、不同界面的截图

• 图片格式支持：.jpg、.jpeg、.png

• 存放在 datasets/images/ 目录下

3.3 数据标注

使用Label Studio（推荐）

bash

# 启动Label Studio
label-studio start
# 浏览器访问 http://localhost:8080

标注步骤：

1. 创建项目，选择 Object Detection with Bounding Boxes

2. 添加标签：icon、app_name

3. 导入图片，开始标注

4. 导出时选择 YOLO 格式

使用labelImg（备选）

bash

# 启动labelImg
labelImg

操作快捷键：

3.4 整理数据集

Label Studio导出的文件名与原始图片名不一致，需要重命名：

python

# rename_labels.py - 将导出的标签文件重命名为与图片一致
import os
import shutil

export_folder = r"D:\AI Trainer\UI_Detection\export"
labels_folder = os.path.join(export_folder, "labels")
source_images = r"D:\AI Trainer\UI_Detection\datasets\images"
target_folder = r"D:\AI Trainer\UI_Detection\datasets\labels"

os.makedirs(target_folder, exist_ok=True)

image_files = [f for f in os.listdir(source_images) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

for image_file in image_files:
    base_name = os.path.splitext(image_file)[0]
    for label_file in os.listdir(labels_folder):
        if label_file.endswith('.txt') and base_name in label_file:
            shutil.copy(
                os.path.join(labels_folder, label_file),
                os.path.join(target_folder, f"{base_name}.txt")
            )
            break

3.5 划分训练集和验证集

python

# split_dataset.py
import os
import random
import shutil

random.seed(42)

images_dir = r"D:\AI Trainer\UI_Detection\datasets\images"
labels_dir = r"D:\AI Trainer\UI_Detection\datasets\labels"
base_dir = r"D:\AI Trainer\UI_Detection\datasets"

# 创建目录
for sub in ['train', 'val']:
    for kind in ['images', 'labels']:
        os.makedirs(os.path.join(base_dir, sub, kind), exist_ok=True)

# 获取并打乱图片
image_files = [f for f in os.listdir(images_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
random.shuffle(image_files)

# 8:2 划分
split_idx = int(len(image_files) * 0.8)
train_files, val_files = image_files[:split_idx], image_files[split_idx:]

# 复制文件
for phase, files in [('train', train_files), ('val', val_files)]:
    for f in files:
        base = os.path.splitext(f)[0]
        shutil.copy(os.path.join(images_dir, f), os.path.join(base_dir, phase, 'images', f))
        label_file = os.path.join(labels_dir, f"{base}.txt")
        if os.path.exists(label_file):
            shutil.copy(label_file, os.path.join(base_dir, phase, 'labels', f"{base}.txt"))

3.6 创建数据配置文件

创建 data.yaml：

yaml

train: D:\AI Trainer\UI_Detection\datasets\train\images
val: D:\AI Trainer\UI_Detection\datasets\val\images

nc: 2  # 类别数：icon, app_name
names: ['icon', 'app_name']

四、训练YOLO模型

4.1 下载预训练权重

bash

cd D:\AI Trainer\UI_Detection\yolov5
# 手动下载 yolov5s.pt 放到此目录
# 下载地址: https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt

4.2 启动训练

bash

python train.py --img 640 --batch 8 --epochs 50 --data ../data.yaml --weights yolov5s.pt

参数说明：

4.3 监控训练

训练过程中关注：

• box_loss / cls_loss：应逐渐下降

• mAP@.5：应逐渐上升（0.8以上为佳）

• mAP@.5:.95：综合精度指标

4.4 训练输出

训练完成后，模型保存在：

text

runs/train/exp/weights/best.pt  # 最佳模型
runs/train/exp/weights/last.pt  # 最后模型

五、模型测试与图标裁剪

5.1 测试模型

bash

python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.25 --source ../datasets/images/test.jpg

5.2 图标裁剪脚本

python

# crop_icons.py
import torch
import cv2
import os
from pathlib import Path

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'custom', 
                       path='yolov5/runs/train/exp/weights/best.pt')
model.conf = 0.4  # 置信度阈值

# 输入输出
image_path = "datasets/images/test.jpg"
output_dir = "cropped_results"
os.makedirs(output_dir, exist_ok=True)

# 检测并裁剪
img = cv2.imread(image_path)
results = model(img)
df = results.pandas().xyxy[0]

class_names = ['icon', 'app_name']

for idx, row in df.iterrows():
    x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
    cls_name = class_names[int(row['class'])]
    conf = row['confidence']
    
    crop = img[y1:y2, x1:x2]
    base_name = Path(image_path).stem
    cv2.imwrite(os.path.join(output_dir, f"{base_name}_{cls_name}_{conf:.2f}.png"), crop)
    print(f"✅ 已保存: {cls_name} (置信度: {conf:.2f})")

六、问题汇总与解决方案

6.1 Python版本问题

6.2 安装问题

6.3 Label Studio问题

6.4 训练问题

6.5 路径问题

七、后续优化建议

7.1 提升检测精度

7.2 功能增强

• OCR识别app_name：集成PaddleOCR，识别检测到的 app_name 文字区域，自动用作文件名

• 批量处理：脚本支持文件夹批量处理

• GUI界面：使用PyQt或Gradio构建可视化工具

📚 参考命令速查

bash

# 激活环境
conda activate ui_detection

# 启动Label Studio
label-studio start

# 启动labelImg
labelImg

# 训练模型
cd D:\AI Trainer\UI_Detection\yolov5
python train.py --img 640 --batch 8 --epochs 50 --data ../data.yaml --weights yolov5s.pt

# 测试模型
python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.25 --source test.jpg

# 裁剪图标
cd D:\AI Trainer\UI_Detection
python crop_icons.py

📞 技术支持

如遇到未覆盖的问题，建议：

1. 查看YOLOv5官方文档：https://docs.ultralytics.com/

2. 检查报错信息中的路径是否正确

3. 确认 (ui_detection) 环境已激活

4. 使用 pip list 检查依赖版本