基于改进yolo的桥梁、楼房建筑物墙体裂缝检测识别

程序员uu

540人浏览 · 2025-12-27 13:17:58

程序员uu · 2025-12-27 13:17:58 发布

视频演示：

https://www.bilibili.com/video/BV1S7iwBLEXy/?pop_share=1&spm_id_from=333.40164.0.0&vd_source=66998a92e3158aac4fbba5005183878c

数据：

原始图像没有了请用这个代码

代码：

import os
import shutil
import random
import torch
from ultralytics import YOLO

# ================== 确保使用纯英文路径（终极修复） ==================
# 获取当前工作目录的绝对路径
current_dir = os.path.abspath(os.getcwd())

# 创建英文路径的临时数据集目录（在当前工作目录下创建英文文件夹）
ENGLISH_DATASET_DIR = "crack_dataset_en"  # 纯英文目录名
DATASET_DIR = os.path.join(current_dir, ENGLISH_DATASET_DIR)

# 确保临时目录存在
os.makedirs(DATASET_DIR, exist_ok=True)

# 创建数据集子目录
for sub_dir in ["images/train", "images/val", "labels/train", "labels/val"]:
    os.makedirs(os.path.join(DATASET_DIR, sub_dir), exist_ok=True)
# =================================================================

# ================== 配置区域（请根据实际修改） ==================
# 您的原始数据路径
ORIGINAL_IMAGES_DIR = "JPEGImages"  # 图像文件夹
ORIGINAL_LABELS_DIR = "labels_txt"  # 标签文件夹
NUM_SAMPLES = 500                   # 总样本数
TRAIN_RATIO = 0.9                   # 训练集比例
# ===========================================================

# 获取所有有效图像文件
image_files = [
    f for f in os.listdir(os.path.join(current_dir))
    if f.lower().endswith(('.jpg', '.jpeg', '.png'))
]

# 筛选有对应标签的图像
valid_files = []
for img_file in image_files:
    base_name = os.path.splitext(img_file)[0]
    label_file = f"{base_name}.txt"
    if os.path.exists(os.path.join(current_dir, ORIGINAL_LABELS_DIR, label_file)):
        valid_files.append((img_file, label_file))

# 检查有效样本数
if len(valid_files) < NUM_SAMPLES:
    print(f"警告：仅找到 {len(valid_files)} 个有效样本（需要 {NUM_SAMPLES} 个），将使用全部样本")
    NUM_SAMPLES = len(valid_files)

# 随机选择样本
selected_files = random.sample(valid_files, NUM_SAMPLES)
train_files = selected_files[:int(NUM_SAMPLES * TRAIN_RATIO)]
val_files = selected_files[int(NUM_SAMPLES * TRAIN_RATIO):]

# 复制到英文路径数据集
for img_file, label_file in train_files:
    shutil.copy(os.path.join(current_dir, ORIGINAL_IMAGES_DIR, img_file),
                os.path.join(DATASET_DIR, "images/train", img_file))
    shutil.copy(os.path.join(current_dir, ORIGINAL_LABELS_DIR, label_file),
                os.path.join(DATASET_DIR, "labels/train", label_file))

for img_file, label_file in val_files:
    shutil.copy(os.path.join(current_dir, ORIGINAL_IMAGES_DIR, img_file),
                os.path.join(DATASET_DIR, "images/val", img_file))
    shutil.copy(os.path.join(current_dir, ORIGINAL_LABELS_DIR, label_file),
                os.path.join(DATASET_DIR, "labels/val", label_file))

# 生成数据集配置文件 (dataset.yaml)
yaml_content = f"""
path: {DATASET_DIR}
train: images/train
val: images/val
nc: 1  # 裂缝通常为1类
names: ['crack']  # 类别名称
"""

with open(os.path.join(DATASET_DIR, "dataset.yaml"), "w") as f:
    f.write(yaml_content)

print(f"\n✅ 数据集准备完成！（使用纯英文路径）\n"
      f"  - 总样本: {NUM_SAMPLES}\n"
      f"  - 训练集: {len(train_files)} 个\n"
      f"  - 验证集: {len(val_files)} 个\n"
      f"  - 配置文件: {os.path.join(DATASET_DIR, 'dataset.yaml')}\n")

# ================== 开始训练（自动检测GPU/CPU） ==================
print("\n🚀 正在使用YOLOv8n进行训练（自动检测设备）...")
print("⚠️ 注意：首次运行会自动下载预训练模型（约100MB），请保持网络畅通")

# 自动检测设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"🔍 检测到设备: {device} {'(GPU)' if device == 'cuda' else '(CPU)'}")

# 重要修复：使用绝对路径并确保路径是纯英文
data_path = os.path.join(DATASET_DIR, "dataset.yaml")
print(f"🔍 正在使用数据路径: {data_path}")

# 确保路径是纯英文，没有中文字符
data_path = os.path.normpath(data_path)  # 规范化路径
data_path = data_path.replace("\\", "/")  # 确保使用正斜杠

# 加载预训练模型并训练
model = YOLO('yolov8n.pt')

# 训练参数（小数据集优化）
results = model.train(
    data=data_path,
    epochs=50,        # 小数据集建议多训练
    imgsz=640,         # 图像尺寸
    batch=32,          # 小数据集建议小batch
    name='crack_detection',
    project='runs',
    device=device      # 自动使用CPU/GPU
)

print("\n✅ 训练完成！结果保存在：")
print(f"  - 模型: {os.path.abspath('runs/detect/crack_detection/weights/best.pt')}")
print("  - 训练曲线: runs/detect/crack_detection/results.png")
print("  - 预测示例: runs/detect/crack_detection/val_batch0_pred.jpg")
print("\n💡 提示：训练时间约10-30分钟（CPU环境），请耐心等待")