#!/usr/bin/env python3 """ CCPD_Random_Blue5K_Green5K → YOLOv5t 推理 → bbox 拓展 → 裁剪 → 角点重映射 ========================================================================== 输出: CCPD_Random_Blue5K_Green5K_EXP_CUT/ 用法: python build_exp_cut.py --full # 全量处理 python build_exp_cut.py --n 100 # 预览 N 张 python build_exp_cut.py --dry-run # 预览 (每类2张, 图上绘制角点) """ import argparse import os import re import sys import time from pathlib import Path import cv2 import numpy as np import torch # Save original imdecode before ultralytics patches it _orig_imdecode = cv2.imdecode from ultralytics import YOLO SCRIPT_DIR = Path(__file__).resolve().parent PROJ_ROOT = SCRIPT_DIR.parents[1] MODEL_PATH = PROJ_ROOT / "ICE_NNv2" / "yolov5t" / "output" / "1_CCPD_19K" / "weights" / "best.pt" SRC_DIR = SCRIPT_DIR DST_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_Random_Blue5K_Green5K_EXP_CUT" IMGSZ = 352 CONF = 0.25 CCPD_W, CCPD_H = 720, 1160 EXPAND_NORMAL = 1.12 EXPAND_SMALL = 1.18 # bw < SMALL_BOX_W or bh < SMALL_BOX_H → small plate SMALL_BOX_W = int(180 * CCPD_W / 1920) # ≈ 68 SMALL_BOX_H = int(48 * CCPD_H / 1080) # ≈ 52 MIN_PAD = 2 CORNER_RE = re.compile(r'-(\d+)&(\d+)_(\d+)&(\d+)_(\d+)&(\d+)_(\d+)&(\d+)-') # ---- Corner remap ---- def parse_ccpd_corners(filename: str): m = CORNER_RE.search(filename) if not m: return None nums = [int(m.group(i)) for i in range(1, 9)] return [(nums[0], nums[1]), (nums[2], nums[3]), (nums[4], nums[5]), (nums[6], nums[7])] def remap_corners_to_crop(corners, x1, y1, x2, y2): crop_w, crop_h = x2 - x1, y2 - y1 if crop_w <= 0 or crop_h <= 0: return None xs, ys = [c[0] for c in corners], [c[1] for c in corners] poly_x1, poly_x2 = min(xs), max(xs) poly_y1, poly_y2 = min(ys), max(ys) inter_w = max(0, min(poly_x2, x2) - max(poly_x1, x1)) inter_h = max(0, min(poly_y2, y2) - max(poly_y1, y1)) poly_area = (poly_x2 - poly_x1) * (poly_y2 - poly_y1) if poly_area <= 0 or (inter_w * inter_h) / poly_area <= 0.25: return None new = [] for px, py in corners: nx = max(0, min(crop_w - 1, int(round(px - x1)))) ny = max(0, min(crop_h - 1, int(round(py - y1)))) new.append((nx, ny)) return new def replace_corners_in_filename(filename: str, new_corners): seg = "-{0}&{1}_{2}&{3}_{4}&{5}_{6}&{7}-".format( new_corners[0][0], new_corners[0][1], new_corners[1][0], new_corners[1][1], new_corners[2][0], new_corners[2][1], new_corners[3][0], new_corners[3][1], ) return CORNER_RE.sub(seg, filename, count=1) # ---- Unicode-safe I/O ---- def _imread(path: str): with open(path, 'rb') as f: data = f.read() if not data: raise ValueError(f"Empty file: {path}") return _orig_imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR) def _imwrite(path: str, img: np.ndarray, params=None): ext = os.path.splitext(path)[1] or '.jpg' ok, buf = cv2.imencode(ext, img, params or []) if ok: with open(path, 'wb') as f: f.write(buf.tobytes()) return ok # ---- Single image ---- def process_one(model, img_path: Path, out_img_path: Path, draw: bool = False): img = _imread(str(img_path)) if img is None: return "miss" orig_h, orig_w = img.shape[:2] results = model.predict(str(img_path), imgsz=IMGSZ, verbose=False, conf=CONF, max_det=5) boxes = results[0].boxes if boxes is None or len(boxes) == 0: return "miss" xywh = boxes.xywh.cpu().numpy() if len(xywh) > 1: return "multi" cx, cy, bw, bh = xywh[0] expand = EXPAND_SMALL if (bw < SMALL_BOX_W or bh < SMALL_BOX_H) else EXPAND_NORMAL pad_x = max(MIN_PAD, int(round(bw * (expand - 1.0) / 2.0))) pad_y = max(MIN_PAD, int(round(bh * (expand - 1.0) / 2.0))) x1 = max(int(cx - bw/2) - pad_x, 0) y1 = max(int(cy - bh/2) - pad_y, 0) x2 = min(int(cx + bw/2) + pad_x + 1, orig_w) y2 = min(int(cy + bh/2) + pad_y + 1, orig_h) if x2 - x1 <= 0 or y2 - y1 <= 0: return "miss" crop = img[y1:y2, x1:x2] stem = img_path.stem corners = parse_ccpd_corners(stem) remapped = None if corners is not None: remapped = remap_corners_to_crop(corners, x1, y1, x2, y2) new_stem = replace_corners_in_filename(stem, remapped) if remapped else stem else: new_stem = stem if draw and remapped: pts_arr = np.array(remapped, dtype=np.int32) cv2.polylines(crop, [pts_arr], isClosed=True, color=(0,255,0), thickness=2) for k, (px, py) in enumerate(remapped): cv2.circle(crop, (px, py), 3, (0,0,255), -1) out_img_path.parent.mkdir(parents=True, exist_ok=True) _imwrite(str(out_img_path.with_stem(new_stem)), crop, [cv2.IMWRITE_JPEG_QUALITY, 92]) return "cropped" if remapped else "cropped_nocorners" # ---- Main ---- def main(): parser = argparse.ArgumentParser() g = parser.add_mutually_exclusive_group(required=True) g.add_argument('--full', action='store_true') g.add_argument('--n', type=int) g.add_argument('--dry-run', action='store_true') args = parser.parse_args() print(f"Model: {MODEL_PATH}") device = 'cuda' if torch.cuda.is_available() else 'cpu' model = YOLO(str(MODEL_PATH)) model.to(device) print(f"Device: {device}\n") classes = ['blue', 'green'] draw = bool(args.dry_run) for cls in classes: src_cls = SRC_DIR / cls dst_cls = DST_DIR / cls dst_cls.mkdir(parents=True, exist_ok=True) images = sorted(src_cls.glob('*.jpg')) if not images: print(f"[WARN] {cls}: no images") continue n = 2 if args.dry_run else (args.n or len(images)) images = images[:n] print(f"[{cls}] {len(images)} images...", flush=True) t0 = time.perf_counter() stats = {'cropped': 0, 'cropped_nocorners': 0, 'miss': 0, 'multi': 0} for img_path in images: out = dst_cls / img_path.name try: r = process_one(model, img_path, out, draw=draw) stats[r] = stats.get(r, 0) + 1 except Exception as e: print(f"\n [ERROR] {img_path.name}: {e}", flush=True) stats['miss'] = stats.get('miss', 0) + 1 elapsed = time.perf_counter() - t0 total = sum(stats.values()) ok = stats['cropped'] + stats['cropped_nocorners'] print(f" → ok:{ok}/{total} miss:{stats['miss']} multi:{stats['multi']} " f"({elapsed:.1f}s, {total/elapsed:.1f} img/s)") print(f"\nDone. Output: {DST_DIR}") if __name__ == '__main__': main()