Scripts/refine_annotations.py

"""标注精修工具 — 用模型预测辅助修正GT标注

原理:
  1. 用最佳模型预测训练集
  2. 对每个GT框,找最匹配的预测框
  3. 如果IoU在0.5~0.85之间(有偏差但能匹配), 用加权平均修正GT
  4. IoU>0.85的保持不变(标注已经很好)
  5. IoU<0.5的标记为需人工审核

为什么有效:
  模型在大量数据上训练后, 对目标边界的回归比单个标注员更稳定。
  特别是小目标, 标注时2px偏差就会导致IoU从0.9降到0.5,
  而模型预测虽然不完美, 但平均来看比有噪声的标注更一致。

用法:
  python Scripts/refine_annotations.py
  python Scripts/refine_annotations.py --dry-run     # 只统计, 不修改
  python Scripts/refine_annotations.py --ratio 0.7    # GT权重0.7, 预测权重0.3
"""
import sys, os, glob, argparse, shutil
import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def compute_iou(box1, box2):
    x1 = max(box1[0], box2[0]); y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2]); y2 = min(box1[3], box2[3])
    inter = max(0, x2-x1) * max(0, y2-y1)
    a1 = (box1[2]-box1[0]) * (box1[3]-box1[1])
    a2 = (box2[2]-box2[0]) * (box2[3]-box2[1])
    return inter / (a1 + a2 - inter + 1e-7)


def refine_box(gt_xyxy, pred_xyxy, gt_weight=0.6):
    """加权平均修正: 偏向GT但吸收模型的定位优势"""
    pw = 1 - gt_weight
    return [g * gt_weight + p * pw for g, p in zip(gt_xyxy, pred_xyxy)]


def xyxy_to_xywhn(box_xyxy, img_w, img_h):
    """xyxy像素坐标 → 归一化xywh"""
    cx = (box_xyxy[0] + box_xyxy[2]) / 2 / img_w
    cy = (box_xyxy[1] + box_xyxy[3]) / 2 / img_h
    w = (box_xyxy[2] - box_xyxy[0]) / img_w
    h = (box_xyxy[3] - box_xyxy[1]) / img_h
    return [max(0, min(1, v)) for v in [cx, cy, w, h]]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="runs/detect/Detection_experiments/v5_5_s_highres/weights/best.pt")
    parser.add_argument("--data-dir", default="Data/Detection_dataset")
    parser.add_argument("--imgsz", type=int, default=1536)
    parser.add_argument("--ratio", type=float, default=0.6,
                        help="GT权重 (0.5=各一半, 0.7=偏GT, 0.3=偏模型)")
    parser.add_argument("--iou-low", type=float, default=0.4,
                        help="IoU低于此值不修正(可能是漏检/误检)")
    parser.add_argument("--iou-high", type=float, default=0.85,
                        help="IoU高于此值不修正(标注已精确)")
    parser.add_argument("--dry-run", action="store_true",
                        help="只统计不修改")
    parser.add_argument("--split", default="train",
                        help="处理哪个split")
    args = parser.parse_args()

    from ultralytics import YOLO
    model = YOLO(args.model)

    img_dir = os.path.join(args.data_dir, f"images/{args.split}")
    lbl_dir = os.path.join(args.data_dir, f"labels/{args.split}")

    # 备份标注
    if not args.dry_run:
        backup_dir = lbl_dir + "_backup_before_refine"
        if not os.path.exists(backup_dir):
            shutil.copytree(lbl_dir, backup_dir)
            print(f"[Backup] 原始标注已备份到: {backup_dir}")
        else:
            print(f"[Backup] 备份已存在: {backup_dir}")

    imgs = sorted(glob.glob(os.path.join(img_dir, "*")))
    print(f"\n处理 {len(imgs)} 张图片...")

    stats = {"total": 0, "refined": 0, "kept": 0, "unmatched": 0, "need_review": []}

    for idx, img_path in enumerate(imgs):
        basename = os.path.splitext(os.path.basename(img_path))[0]
        lbl_path = os.path.join(lbl_dir, basename + ".txt")
        if not os.path.exists(lbl_path):
            continue

        with open(lbl_path) as f:
            gt_lines = [l.strip() for l in f if l.strip()]
        if not gt_lines:
            continue

        # 模型预测
        results = model.predict(img_path, imgsz=args.imgsz, conf=0.25, iou=0.6, verbose=False)
        pred_boxes = results[0].boxes.xyxy.cpu().numpy()
        pred_confs = results[0].boxes.conf.cpu().numpy() if len(results[0].boxes) > 0 else np.array([])

        img = Image.open(img_path)
        W, H = img.size

        new_lines = []
        review_count = 0

        for line in gt_lines:
            parts = line.split()
            cls_id = parts[0]
            cx, cy, w, h = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
            gt_xyxy = [(cx-w/2)*W, (cy-h/2)*H, (cx+w/2)*W, (cy+h/2)*H]
            area = w * h
            stats["total"] += 1

            if len(pred_boxes) == 0:
                new_lines.append(line)
                stats["unmatched"] += 1
                continue

            # 找最匹配的预测框
            best_iou, best_idx = 0, -1
            for j, pb in enumerate(pred_boxes):
                iou = compute_iou(gt_xyxy, pb)
                if iou > best_iou:
                    best_iou = iou
                    best_idx = j

            if best_iou >= args.iou_high:
                # 标注已经很好, 不修改
                new_lines.append(line)
                stats["kept"] += 1
            elif best_iou >= args.iou_low:
                # 有偏差, 加权修正
                # 小目标更信任模型(标注噪声大), 大目标更信任GT
                if area < 0.003:
                    local_ratio = max(args.ratio - 0.1, 0.3)  # 小目标: 更信模型
                else:
                    local_ratio = min(args.ratio + 0.1, 0.8)  # 大目标: 更信GT

                refined = refine_box(gt_xyxy, pred_boxes[best_idx], local_ratio)
                new_xywhn = xyxy_to_xywhn(refined, W, H)
                new_lines.append(f"{cls_id} {new_xywhn[0]:.6f} {new_xywhn[1]:.6f} {new_xywhn[2]:.6f} {new_xywhn[3]:.6f}")
                stats["refined"] += 1
            else:
                # IoU太低, 可能是漏检或误标, 保持原样但标记
                new_lines.append(line)
                stats["unmatched"] += 1
                review_count += 1

        if review_count > 3:
            stats["need_review"].append(basename)

        # 写入修正后的标注
        if not args.dry_run:
            with open(lbl_path, "w") as f:
                f.write("\n".join(new_lines) + "\n")

        if (idx + 1) % 100 == 0:
            print(f"  [{idx+1}/{len(imgs)}] refined={stats['refined']} kept={stats['kept']} unmatched={stats['unmatched']}")

    print(f"\n{'='*50}")
    print(f"标注精修完成{'(dry-run)' if args.dry_run else ''}!")
    print(f"  总框数:   {stats['total']}")
    print(f"  已修正:   {stats['refined']} ({stats['refined']/max(stats['total'],1)*100:.1f}%)")
    print(f"  已保留:   {stats['kept']} ({stats['kept']/max(stats['total'],1)*100:.1f}%)")
    print(f"  未匹配:   {stats['unmatched']} ({stats['unmatched']/max(stats['total'],1)*100:.1f}%)")
    print(f"  需审核图: {len(stats['need_review'])} 张")

    if stats["need_review"]:
        review_file = "need_review_images.txt"
        with open(review_file, "w") as f:
            f.write("\n".join(stats["need_review"]))
        print(f"  需审核列表: {review_file}")

    if not args.dry_run:
        print(f"\n下一步: 用精修后的标注重新训练")
        print(f"  原始标注备份: {lbl_dir}_backup_before_refine/")
    print(f"{'='*50}")


if __name__ == "__main__":
    main()