Dataset: add READMEs, organize scripts, update .gitignore
- Add README.md to all 16 dataset directories with script references - Centralize generation scripts under scripts/ with full README - Update .gitignore to cover all data directories, keep only metadata - Remove redundant CCPD/.gitignore (parent covers same rules) - Remove stale scripts/test/ (254MB pipeline residue) - Fix label parser reference: scripts are authoritative source of truth
This commit is contained in:
kirl committed
1 parent
7867c3bfaf
commit
a64690477a
43 files changed
+2691
-205
No files matched your search
@@ -0,0 +1,188 @@
|
||||
#!/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()
|
||||
Reference in new issue
Block a user