145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
"""
|
|
将标注的 55 组样本转换为标准 YOLO 目标检测数据集
|
|
生成:
|
|
tcaptcha/data/yolo/images/train/
|
|
tcaptcha/data/yolo/images/val/
|
|
tcaptcha/data/yolo/labels/train/
|
|
tcaptcha/data/yolo/labels/val/
|
|
tcaptcha/data/yolo/dataset.yaml
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import random
|
|
import shutil
|
|
from pathlib import Path
|
|
import cv2
|
|
import numpy as np
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
logger = logging.getLogger("prepare_yolo")
|
|
|
|
RAW_DIR = Path("tcaptcha/data/raw")
|
|
LABEL_FILE = Path("tcaptcha/data/labels.jsonl")
|
|
YOLO_DIR = Path("tcaptcha/data/yolo")
|
|
|
|
IMG_W = 672
|
|
IMG_H = 480
|
|
BOX_SIZE = 54 # 图标大致大小 (约 54x54 像素)
|
|
|
|
|
|
def find_all_icons_in_bg(bg_path: Path, labeled_points: list) -> list:
|
|
"""
|
|
结合人工标注点与轮廓补充检测，找到大图中所有图标（包含 3 个目标 + 2~3 个干扰图标）
|
|
返回: [(cx, cy, w, h), ...] 绝对像素坐标
|
|
"""
|
|
img = cv2.imread(str(bg_path))
|
|
if img is None:
|
|
return []
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
|
# 提取暗色轮廓找干扰图标
|
|
edges = cv2.Canny(gray, 40, 120)
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
|
|
closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
|
|
contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
all_boxes = []
|
|
# 1. 先加入人工标注的 3 个目标（绝对精准）
|
|
for p in labeled_points:
|
|
cx, cy = p["x"], p["y"]
|
|
all_boxes.append((cx, cy, BOX_SIZE, BOX_SIZE))
|
|
|
|
# 2. 补充检测未标注的干扰图标
|
|
for cnt in contours:
|
|
x, y, w, h = cv2.boundingRect(cnt)
|
|
if 26 <= w <= 85 and 26 <= h <= 85 and (w * h) >= 650:
|
|
cx, cy = x + w // 2, y + h // 2
|
|
# 必须远离已有图标
|
|
if all(np.hypot(cx - lx, cy - ly) > 42 for lx, ly, _, _ in all_boxes):
|
|
# 限制尺寸
|
|
bw = min(68, max(46, w + 8))
|
|
bh = min(68, max(46, h + 8))
|
|
all_boxes.append((cx, cy, bw, bh))
|
|
|
|
return all_boxes
|
|
|
|
|
|
def prepare_dataset(val_ratio: float = 0.2, seed: int = 42):
|
|
random.seed(seed)
|
|
if not LABEL_FILE.exists():
|
|
logger.error("标签文件不存在: %s", LABEL_FILE)
|
|
return
|
|
|
|
# 清理并创建目录
|
|
if YOLO_DIR.exists():
|
|
shutil.rmtree(YOLO_DIR)
|
|
|
|
for split in ["train", "val"]:
|
|
(YOLO_DIR / "images" / split).mkdir(parents=True, exist_ok=True)
|
|
(YOLO_DIR / "labels" / split).mkdir(parents=True, exist_ok=True)
|
|
|
|
with open(LABEL_FILE, "r", encoding="utf-8") as f:
|
|
records = [json.loads(line) for line in f if line.strip()]
|
|
|
|
random.shuffle(records)
|
|
val_count = max(5, int(len(records) * val_ratio))
|
|
train_records = records[val_count:]
|
|
val_records = records[:val_count]
|
|
|
|
logger.info("总标注样本: %d (训练集: %d, 验证集: %d)", len(records), len(train_records), len(val_records))
|
|
|
|
def process_split(record_list, split_name):
|
|
total_boxes = 0
|
|
for r in record_list:
|
|
sid = r["sample_id"]
|
|
src_img = RAW_DIR / f"{sid}_bg.png"
|
|
if not src_img.exists():
|
|
continue
|
|
|
|
dst_img = YOLO_DIR / "images" / split_name / f"{sid}.png"
|
|
shutil.copy(src_img, dst_img)
|
|
|
|
# 提取全部图标坐标
|
|
boxes = find_all_icons_in_bg(src_img, r["points"])
|
|
total_boxes += len(boxes)
|
|
|
|
# 写入 YOLO 格式标注: class_id x_center y_center width height (归一化到 0~1)
|
|
dst_label = YOLO_DIR / "labels" / split_name / f"{sid}.txt"
|
|
lines = []
|
|
for cx, cy, bw, bh in boxes:
|
|
# 边界截断保护
|
|
x1 = max(0, cx - bw / 2)
|
|
y1 = max(0, cy - bh / 2)
|
|
x2 = min(IMG_W, cx + bw / 2)
|
|
y2 = min(IMG_H, cy + bh / 2)
|
|
|
|
norm_cx = (x1 + x2) / 2.0 / IMG_W
|
|
norm_cy = (y1 + y2) / 2.0 / IMG_H
|
|
norm_w = (x2 - x1) / IMG_W
|
|
norm_h = (y2 - y1) / IMG_H
|
|
|
|
lines.append(f"0 {norm_cx:.6f} {norm_cy:.6f} {norm_w:.6f} {norm_h:.6f}")
|
|
|
|
dst_label.write_text("\n".join(lines), encoding="utf-8")
|
|
logger.info(" [%s] 处理完成: %d 张图片, 共 %d 个图标标注", split_name, len(record_list), total_boxes)
|
|
|
|
process_split(train_records, "train")
|
|
process_split(val_records, "val")
|
|
|
|
# 生成 dataset.yaml
|
|
yaml_content = f"""path: {YOLO_DIR.resolve()}
|
|
train: images/train
|
|
val: images/val
|
|
|
|
names:
|
|
0: icon
|
|
"""
|
|
yaml_path = YOLO_DIR / "dataset.yaml"
|
|
yaml_path.write_text(yaml_content, encoding="utf-8")
|
|
logger.info("✓ YOLO 数据集生成完毕: %s", yaml_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
prepare_dataset()
|