kadupul-git/BridgeDiseaseBackend-main/scripts/create_demo_dataset.py

# -*- coding: utf-8 -*-
"""
从 seed_assets/medias 生成迷你 YOLO 分割演示数据集（占位标注，仅用于试跑训练流程）。

输出：
  training_templates/bridge_hazard_demo/
  training_templates/bridge_hazard_demo.zip
"""
from __future__ import annotations

import shutil
import zipfile
from pathlib import Path

from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
ASSETS = ROOT / 'seed_assets' / 'medias'
OUT_ROOT = ROOT / 'training_templates' / 'bridge_hazard_demo'
ZIP_PATH = ROOT / 'training_templates' / 'bridge_hazard_demo.zip'

# (源图, 划分, 类别 id, 类别说明)
SAMPLES = [
    ('01_concrete_crack_bridge.jpg', 'train', 0, 'concrete_crack'),
    ('02_bridge_concrete_cracks.jpg', 'train', 0, 'concrete_crack'),
    ('04_concrete_bending_cracks.jpg', 'train', 0, 'concrete_crack'),
    ('06_shrinkage_cracks_concrete.jpg', 'train', 0, 'concrete_crack'),
    ('03_steel_bridge_corrosion.jpg', 'val', 1, 'steel_corrosion'),
    ('08_concrete_rebar_corrosion.jpg', 'val', 1, 'steel_corrosion'),
]


def rect_polygon(class_id: int, cx=0.5, cy=0.5, w=0.35, h=0.35) -> str:
    """生成归一化矩形四顶点分割标注行。"""
    x1, y1 = cx - w / 2, cy - h / 2
    x2, y2 = cx + w / 2, cy - h / 2
    x3, y3 = cx + w / 2, cy + h / 2
    x4, y4 = cx - w / 2, cy + h / 2
    pts = [x1, y1, x2, y2, x3, y3, x4, y4]
    pts = [max(0.0, min(1.0, p)) for p in pts]
    return f"{class_id} " + " ".join(f"{p:.6f}" for p in pts)


def write_data_yaml(root: Path):
    content = """# 检澜演示数据集 — 占位标注，仅用于试跑 YOLOv8 分割训练
path: .
train: train/images
val: val/images

nc: 2
names:
  0: concrete_crack
  1: steel_corrosion
"""
    (root / 'data.yaml').write_text(content, encoding='utf-8')


def main():
    if not ASSETS.is_dir():
        print(f'缺少 {ASSETS}，请先运行 scripts/download_real_medias.py')
        return 1

    if OUT_ROOT.exists():
        shutil.rmtree(OUT_ROOT)
    OUT_ROOT.mkdir(parents=True)

    for split in ('train', 'val'):
        (OUT_ROOT / split / 'images').mkdir(parents=True)
        (OUT_ROOT / split / 'labels').mkdir(parents=True)

    idx = 0
    for src_name, split, class_id, _ in SAMPLES:
        src = ASSETS / src_name
        if not src.is_file():
            print(f'[skip] 缺少 {src_name}')
            continue
        idx += 1
        stem = f'demo_{idx:03d}'
        img_name = f'{stem}.jpg'
        dest_img = OUT_ROOT / split / 'images' / img_name

        with Image.open(src) as im:
            im.convert('RGB').save(dest_img, format='JPEG', quality=90)

        label_line = rect_polygon(class_id)
        (OUT_ROOT / split / 'labels' / f'{stem}.txt').write_text(
            label_line + '\n', encoding='utf-8'
        )
        print(f'[ok] {split} {img_name} class={class_id}')

    if idx < 2:
        print('有效图片不足，无法生成数据集')
        return 1

    write_data_yaml(OUT_ROOT)

    if ZIP_PATH.is_file():
        ZIP_PATH.unlink()
    with zipfile.ZipFile(ZIP_PATH, 'w', zipfile.ZIP_DEFLATED) as zf:
        for file in OUT_ROOT.rglob('*'):
            if file.is_file():
                arc = file.relative_to(OUT_ROOT.parent)
                zf.write(file, arc.as_posix())

    print()
    print('已生成：')
    print(f'  目录  {OUT_ROOT}')
    print(f'  ZIP   {ZIP_PATH}')
    print()
    print('下一步：登录 developer → 模型训练 → 上传 bridge_hazard_demo.zip')
    print('建议试跑：yolov8n-seg，epochs=5，batch=2')
    return 0


if __name__ == '__main__':
    raise SystemExit(main())