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,26 @@
|
||||
# 数据生成脚本
|
||||
|
||||
此目录包含所有数据集生成脚本和少量测试样本。
|
||||
|
||||
## 脚本清单
|
||||
|
||||
| 脚本 | 输入 | 输出 | 用途 |
|
||||
|------|------|------|------|
|
||||
| `pipeline_322k.py` | `CCPD_322K/` | `CCPD_322K_EXP_CUT_TRANS/` | 322K 合成全管线（3 步） |
|
||||
| `sample_ccpd_full.py` | `CCPD/` | `CCPD_REAL_TRANS/` | CCPD 真图全量 GT-warp |
|
||||
| `sample_ccpd_real.py` | `CCPD/` | `CCPD_REAL_TRANS/` | CCPD 真图采样（实验用） |
|
||||
| `build_cblprd_trans.py` | CBLPRD zip | `CBLPRD_TRANS/` | CBLPRD COR 矫正 TRANS |
|
||||
| `build_ccpd_combined.py` | 多源目录 | `ccpd_combined_5k5k_spe/` | 合并数据集 |
|
||||
| `rebuild_spe.py` | — | `CCPD_SPE/` | SPE 特殊车牌生成 |
|
||||
| `gen_random_bg.py` | — | `CCPD_Random_*` | 随机背景蓝绿合成 |
|
||||
|
||||
## 目录结构
|
||||
|
||||
```
|
||||
scripts/
|
||||
├── *.py, *.ps1 ← 生成脚本（权威来源）
|
||||
├── train/ ← 空（数据已移至 CCPD_322K_EXP_CUT_TRANS/）
|
||||
└── test/ ← 34004 张测试样本
|
||||
```
|
||||
|
||||
> **以脚本为准**：所有数据集的生成逻辑、标签格式、预处理参数均以本目录中的脚本为唯一权威来源。文件名后缀标签的解析规则见 `ICE_NNv2/ICE-OCR/script/tools/dataset.py`。
|
||||
@@ -0,0 +1,138 @@
|
||||
"""
|
||||
CBLPRD → CCPD_TRANS 数据集构建 (ICE-COR 矫正)
|
||||
输出: ICE_Dataset/CBLPRD_TRANS/train/
|
||||
"""
|
||||
import sys, os, cv2, numpy as np, torch
|
||||
from pathlib import Path
|
||||
import zipfile
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from threading import Lock
|
||||
|
||||
ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
|
||||
COR_DIR = ICEv2 / "ICE_NNv2/ICE-COR"
|
||||
MODEL_COR = COR_DIR / "output/1_CCPD_19K_EXP_CUT/best.pth"
|
||||
ZIP_PATH = r"C:\Users\Kirl\Downloads\CBLPRD-330k_v1.zip"
|
||||
OUT_DIR = ICEv2 / "ICE_Dataset/CBLPRD_TRANS/train"
|
||||
|
||||
sys.path.insert(0, str(COR_DIR))
|
||||
from corner_rectify import load_corner_model, predict_plate_corners, order_corners
|
||||
|
||||
# 类型→文件名标记
|
||||
TYPE_MARKER = {
|
||||
'普通蓝牌': 'blue',
|
||||
'新能源小型车': 'green',
|
||||
'新能源大型车': 'new_energy',
|
||||
'单层黄牌': 'yellow',
|
||||
'黑色车牌': 'black',
|
||||
}
|
||||
SUPPORTED = set(TYPE_MARKER.keys())
|
||||
|
||||
print(f"COR: {MODEL_COR}")
|
||||
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
cor_model, _ = load_corner_model(str(MODEL_COR), dev)
|
||||
|
||||
print(f"Output: {OUT_DIR}")
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 读取 CBLPRD train.txt
|
||||
zf = zipfile.ZipFile(ZIP_PATH, 'r')
|
||||
labels = []
|
||||
for line in zf.read('train.txt').decode('utf-8').strip().split('\n'):
|
||||
parts = line.strip().split(' ')
|
||||
if len(parts) < 3 or parts[2] not in SUPPORTED:
|
||||
continue
|
||||
fname = parts[0].split('/')[-1]
|
||||
labels.append((fname, parts[1], parts[2]))
|
||||
|
||||
# 同时加入 val.txt (合并训练)
|
||||
for line in zf.read('val.txt').decode('utf-8').strip().split('\n'):
|
||||
parts = line.strip().split(' ')
|
||||
if len(parts) < 3 or parts[2] not in SUPPORTED:
|
||||
continue
|
||||
fname = parts[0].split('/')[-1]
|
||||
labels.append((fname, parts[1], parts[2]))
|
||||
|
||||
print(f"Total labels: {len(labels)}")
|
||||
from collections import Counter
|
||||
for t, c in Counter(l[2] for l in labels).most_common():
|
||||
print(f" {t}: {c}")
|
||||
|
||||
# Build name→zip_path index
|
||||
name_idx = {}
|
||||
for n in zf.namelist():
|
||||
if n.endswith('.jpg'):
|
||||
name_idx[n.split('/')[-1]] = n
|
||||
print(f"Zip index: {len(name_idx)} images")
|
||||
|
||||
# 统计已完成的
|
||||
existing = set(f.stem.split('_trans')[0] if '_trans' in f.stem else '' for f in OUT_DIR.glob('*.jpg'))
|
||||
print(f"Already done: {len(existing)}")
|
||||
|
||||
# 过滤已完成 (按原始文件名前缀)
|
||||
pending = []
|
||||
for fname, plate, ptype in labels:
|
||||
if fname in name_idx and fname.rsplit('.',1)[0] not in existing:
|
||||
pending.append((fname, plate, ptype, name_idx[fname]))
|
||||
|
||||
print(f"Pending: {len(pending)}")
|
||||
|
||||
lock = Lock()
|
||||
count = [0]
|
||||
total = len(pending)
|
||||
|
||||
def process_one(item):
|
||||
fname, plate, ptype, zip_path = item
|
||||
try:
|
||||
data = zf.read(zip_path)
|
||||
img = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
|
||||
if img is None:
|
||||
return f"READ_ERR:{fname}"
|
||||
|
||||
# ICE-COR
|
||||
corners = predict_plate_corners(img, cor_model, dev)
|
||||
ordered = order_corners(corners)
|
||||
OW, OH = 192, 64
|
||||
cx, cy = ordered.mean(axis=0)
|
||||
exp = ordered + (ordered - np.array([[cx, cy]])) * np.array([0.08, 0.05])
|
||||
exp[:, 0] = np.clip(exp[:, 0], 0, img.shape[1] - 1)
|
||||
exp[:, 1] = np.clip(exp[:, 1], 0, img.shape[0] - 1)
|
||||
dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
|
||||
M = cv2.getPerspectiveTransform(exp.astype(np.float32), dst)
|
||||
warped = cv2.warpPerspective(img, M, (OW, OH))
|
||||
except Exception as e:
|
||||
return f"COR_ERR:{fname}:{e}"
|
||||
|
||||
# 文件名: {原始名}_{marker}_{plate}_trans.jpg
|
||||
marker = TYPE_MARKER[ptype]
|
||||
stem = fname.rsplit('.', 1)[0]
|
||||
out_name = f"{stem}_{marker}_{plate}_trans.jpg"
|
||||
out_path = OUT_DIR / out_name
|
||||
|
||||
# Unicode safe write
|
||||
ret, enc = cv2.imencode('.jpg', warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
if ret:
|
||||
with open(str(out_path), 'wb') as f:
|
||||
f.write(enc.tobytes())
|
||||
else:
|
||||
return f"WRITE_ERR:{fname}"
|
||||
|
||||
with lock:
|
||||
count[0] += 1
|
||||
if count[0] % 500 == 0:
|
||||
print(f" [{count[0]}/{total}] {fname} -> {marker}:{plate}")
|
||||
return f"OK:{fname}"
|
||||
|
||||
# 多线程处理
|
||||
WORKERS = 8
|
||||
print(f"\nProcessing {total} images with {WORKERS} workers...")
|
||||
with ThreadPoolExecutor(max_workers=WORKERS) as executor:
|
||||
futures = {executor.submit(process_one, item): item for item in pending}
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result.startswith("READ_ERR") or result.startswith("WRITE_ERR"):
|
||||
print(f" [!] {result}")
|
||||
|
||||
zf.close()
|
||||
print(f"\nDone! Total saved: {count[0]}/{total}")
|
||||
print(f"Output: {OUT_DIR}")
|
||||
print(f"Files: {len(list(OUT_DIR.glob('*.jpg')))}")
|
||||
@@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
全量生成 Random 蓝绿车牌
|
||||
blue 150K + green 100K, 覆盖全部 CCPD 背景
|
||||
"""
|
||||
|
||||
import os, sys, time
|
||||
from pathlib import Path
|
||||
|
||||
ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
|
||||
os.chdir(str(ICE_DATASET))
|
||||
|
||||
sys.path.insert(0, str(ICE_DATASET / 'CCPD_SPE_GEN' / 'cv_pipeline'))
|
||||
from plate_swap import PlateSwapPipeline
|
||||
|
||||
CCPD2019 = ICE_DATASET / 'CCPD' / 'CCPD2019'
|
||||
CCPD2020 = ICE_DATASET / 'CCPD' / 'CCPD2020'
|
||||
OUT = ICE_DATASET / 'CCPD_Random_Blue150K_Green100K'
|
||||
|
||||
# 全部可用 CCPD 子集 (含全量, 不限于2019)
|
||||
CCPD_DIRS = [
|
||||
str(CCPD2019 / 'ccpd_base'),
|
||||
str(CCPD2019 / 'ccpd_blur'),
|
||||
str(CCPD2019 / 'ccpd_challenge'),
|
||||
str(CCPD2019 / 'ccpd_fn'),
|
||||
str(CCPD2019 / 'ccpd_rotate'),
|
||||
str(CCPD2019 / 'ccpd_tilt'),
|
||||
str(CCPD2019 / 'ccpd_weather'),
|
||||
]
|
||||
|
||||
# 只保留实际存在的
|
||||
CCPD_DIRS = [d for d in CCPD_DIRS if Path(d).is_dir()]
|
||||
print(f"CCPD 背景目录: {len(CCPD_DIRS)}")
|
||||
|
||||
SEED = 42
|
||||
WORKERS = 24
|
||||
|
||||
def main():
|
||||
import shutil
|
||||
# 断点续传: 不删除已有数据
|
||||
total_start = time.time()
|
||||
|
||||
for plate_type, target in [('blue', 150000), ('green', 100000)]:
|
||||
out_dir = OUT / plate_type
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
existing = len(list(out_dir.glob('*.jpg')))
|
||||
if existing >= target:
|
||||
print(f"\n[{plate_type}] 已有 {existing} >= {target}, 跳过")
|
||||
continue
|
||||
need = target - existing
|
||||
print(f"\n{'='*60}")
|
||||
print(f"[{plate_type}] 已有 {existing}, 需补 {need} → 目标 {target}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
pipeline = PlateSwapPipeline(plate_type=plate_type)
|
||||
|
||||
print(f" 生成中...", end=' ', flush=True)
|
||||
t0 = time.time()
|
||||
pipeline.process_batch(
|
||||
ccpd_dirs=CCPD_DIRS,
|
||||
output_dir=str(out_dir),
|
||||
num_images=need,
|
||||
preview=False,
|
||||
seed=SEED,
|
||||
workers=WORKERS,
|
||||
)
|
||||
elapsed = time.time() - t0
|
||||
print(f"完成 ({elapsed:.0f}s, {need/elapsed:.0f} 张/s)")
|
||||
|
||||
total_t = time.time() - total_start
|
||||
blue_final = len(list((OUT / 'blue').glob('*.jpg')))
|
||||
green_final = len(list((OUT / 'green').glob('*.jpg')))
|
||||
print(f"\nDONE! blue={blue_final} green={green_final}, {total_t/60:.1f}min")
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,38 @@
|
||||
# 322K 全自动: 等生成完 → 构建 → EXP_CUT → TRANS → 训练
|
||||
# 用法: 生成进程跑完后执行此脚本
|
||||
$ErrorActionPreference = "Continue"
|
||||
|
||||
$PY = "C:/Program Files/Python313/python.exe"
|
||||
$ICE = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset"
|
||||
|
||||
Write-Host "===== 322K Pipeline =====" -ForegroundColor Cyan
|
||||
Write-Host "Waiting for SPE + Random generation to complete..."
|
||||
Write-Host "SPE dir: $ICE\CCPD_SPE\ccpd_base"
|
||||
Write-Host "Random dir: $ICE\CCPD_Random_Blue150K_Green100K"
|
||||
|
||||
# Wait for SPE (9 types x 8000)
|
||||
$spe_types = @('police','army','gangao','shiguan_old','shiguan_new','lingguan_old','lingguan_new','coach','new_energy')
|
||||
do {
|
||||
$done = $true
|
||||
foreach ($t in $spe_types) {
|
||||
$c = @(Get-ChildItem "$ICE\CCPD_SPE\ccpd_base\$t" -Filter *.jpg -ErrorAction SilentlyContinue).Count
|
||||
if ($c -lt 8000) { $done = $false; break }
|
||||
}
|
||||
if (-not $done) { Write-Host " SPE not ready, waiting 60s..."; Start-Sleep 60 }
|
||||
} while (-not $done)
|
||||
Write-Host "SPE: OK" -ForegroundColor Green
|
||||
|
||||
# Wait for Random
|
||||
do {
|
||||
$b = @(Get-ChildItem "$ICE\CCPD_Random_Blue150K_Green100K\blue" -Filter *.jpg -ErrorAction SilentlyContinue).Count
|
||||
$g = @(Get-ChildItem "$ICE\CCPD_Random_Blue150K_Green100K\green" -Filter *.jpg -ErrorAction SilentlyContinue).Count
|
||||
if ($b -lt 150000 -or $g -lt 100000) {
|
||||
Write-Host " Random: blue=$b green=$g, waiting 60s..."
|
||||
Start-Sleep 60
|
||||
} else { break }
|
||||
} while ($true)
|
||||
Write-Host "Random: blue=$b green=$g OK" -ForegroundColor Green
|
||||
|
||||
# Now run the pipeline
|
||||
Write-Host "`nAll data ready. Starting pipeline..." -ForegroundColor Cyan
|
||||
& "C:/Program Files/Python313/python.exe" "$PSScriptRoot\pipeline_322k.py"
|
||||
@@ -0,0 +1,85 @@
|
||||
# 322K 全自动流水线: 构建数据集 → EXP_CUT → TRANS → 训练
|
||||
$ErrorActionPreference = "Continue"
|
||||
$PY = "C:/Program Files/Python313/python.exe"
|
||||
$ICE = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset"
|
||||
$NN = "C:\Users\Kirl\Desktop\ICEv2\ICE_NNv2"
|
||||
|
||||
# ===== Step 1: Build CCPD_322K =====
|
||||
Write-Host "`n===== 1/4 构建 CCPD_322K =====" -ForegroundColor Cyan
|
||||
|
||||
$SRC_SPE = "$ICE\CCPD_SPE\ccpd_base"
|
||||
$SRC_RND = "$ICE\CCPD_Random_Blue150K_Green100K"
|
||||
$DST = "$ICE\CCPD_322K"
|
||||
|
||||
if (Test-Path $DST) { Remove-Item -Recurse -Force $DST }
|
||||
New-Item -ItemType Directory "$DST\train" -Force | Out-Null
|
||||
New-Item -ItemType Directory "$DST\test" -Force | Out-Null
|
||||
|
||||
# SPE 9类 (9:1 split)
|
||||
$types = @('police','army','gangao','shiguan_old','shiguan_new','lingguan_old','lingguan_new','coach','new_energy')
|
||||
foreach ($t in $types) {
|
||||
$src = "$SRC_SPE\$t"
|
||||
if (-not (Test-Path $src)) { Write-Host " SKIP $t"; continue }
|
||||
$files = Get-ChildItem $src -Filter *.jpg | Sort-Object Name
|
||||
$n = $files.Count
|
||||
$nTrain = [math]::Floor($n * 0.9)
|
||||
for ($i = 0; $i -lt $n; $i++) {
|
||||
$split = if ($i -lt $nTrain) { "train" } else { "test" }
|
||||
$d = "$DST\$split\$t"; New-Item -ItemType Directory $d -Force | Out-Null
|
||||
Copy-Item $files[$i].FullName "$d\" -Force
|
||||
}
|
||||
Write-Host " $t : $n (train=$nTrain test=$($n-$nTrain))"
|
||||
}
|
||||
|
||||
# Random蓝绿 (9:1 split)
|
||||
foreach ($t in @('blue','green')) {
|
||||
$src = "$SRC_RND\$t"
|
||||
if (-not (Test-Path $src)) { Write-Host " SKIP $t"; continue }
|
||||
$files = Get-ChildItem $src -Filter *.jpg | Sort-Object Name
|
||||
$n = $files.Count
|
||||
$nTrain = [math]::Floor($n * 0.9)
|
||||
for ($i = 0; $i -lt $n; $i++) {
|
||||
$split = if ($i -lt $nTrain) { "train" } else { "test" }
|
||||
$d = "$DST\$split\$t"; New-Item -ItemType Directory $d -Force | Out-Null
|
||||
Copy-Item $files[$i].FullName "$d\" -Force
|
||||
}
|
||||
Write-Host " $t : $n (train=$nTrain test=$($n-$nTrain))"
|
||||
}
|
||||
|
||||
$total = (Get-ChildItem $DST -Recurse -Filter *.jpg).Count
|
||||
Write-Host " TOTAL: $total" -ForegroundColor Green
|
||||
|
||||
# ===== Step 2: EXP_CUT =====
|
||||
Write-Host "`n===== 2/4 YOLO EXP_CUT =====" -ForegroundColor Cyan
|
||||
$EXP_CUT = "$ICE\CCPD_322K_EXP_CUT"
|
||||
if (Test-Path $EXP_CUT) { Remove-Item -Recurse -Force $EXP_CUT }
|
||||
|
||||
# Patch the build script with our paths
|
||||
$script = Get-Content "$ICE\CCPD_19K_EXP_CUT\build_ccpd_19k_exp_cut.py" -Raw -Encoding UTF8
|
||||
$script = $script -replace 'SRC_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K"', "SRC_DIR = Path(r'$DST')"
|
||||
$script = $script -replace 'DST_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT"', "DST_DIR = Path(r'$EXP_CUT')"
|
||||
$tmp = "$ICE\scripts\_exp_cut_322k.py"
|
||||
[System.IO.File]::WriteAllText($tmp, $script, [System.Text.UTF8Encoding]::new($false))
|
||||
& $PY $tmp --full
|
||||
Remove-Item $tmp
|
||||
|
||||
# ===== Step 3: TRANS =====
|
||||
Write-Host "`n===== 3/4 ICE-COR TRANS =====" -ForegroundColor Cyan
|
||||
$TRANS = "$ICE\CCPD_322K_EXP_CUT_TRANS"
|
||||
if (Test-Path $TRANS) { Remove-Item -Recurse -Force $TRANS }
|
||||
|
||||
$script = Get-Content "$ICE\CCPD_19K_EXP_CUT_TRANS\build_ccpd_19k_exp_cut_trans.py" -Raw -Encoding UTF8
|
||||
$script = $script -replace 'SRC_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT"', "SRC_DIR = Path(r'$EXP_CUT')"
|
||||
$script = $script -replace 'DST_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT_TRANS"', "DST_DIR = Path(r'$TRANS')"
|
||||
$tmp = "$ICE\scripts\_trans_322k.py"
|
||||
[System.IO.File]::WriteAllText($tmp, $script, [System.Text.UTF8Encoding]::new($false))
|
||||
& $PY $tmp --full
|
||||
Remove-Item $tmp
|
||||
|
||||
# ===== Step 4: Training =====
|
||||
Write-Host "`n===== 4/4 ICE-OCR Training =====" -ForegroundColor Cyan
|
||||
Write-Host " Data: $TRANS"
|
||||
Write-Host " Launching..."
|
||||
& $PY "$NN\ICE-OCR\script\train.py" --dataset ccpd_trans --epochs 100 --batch 128 --output ccpd_322k_best.pth --workers 2 --data-dir $TRANS
|
||||
|
||||
Write-Host "`n===== DONE =====" -ForegroundColor Green
|
||||
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env python3
|
||||
"""322K Pipeline Step 2-4: EXP_CUT -> TRANS -> Train"""
|
||||
import os, sys, shutil, subprocess
|
||||
from pathlib import Path
|
||||
|
||||
ICE = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
|
||||
PY = r'C:/Program Files/Python313/python.exe'
|
||||
NN = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_NNv2')
|
||||
|
||||
DATASET = ICE / 'CCPD_322K'
|
||||
EXP_CUT = ICE / 'CCPD_322K_EXP_CUT'
|
||||
TRANS = ICE / 'CCPD_322K_EXP_CUT_TRANS'
|
||||
|
||||
def run(cmd):
|
||||
print(f" {str(cmd)[:150]}...")
|
||||
if isinstance(cmd, list):
|
||||
r = subprocess.run(cmd)
|
||||
else:
|
||||
r = subprocess.run(cmd, shell=True)
|
||||
if r.returncode != 0: print(f" FAILED ({r.returncode})"); sys.exit(1)
|
||||
|
||||
# ===== Step 2: EXP_CUT =====
|
||||
print("===== 2/4 EXP_CUT =====")
|
||||
shutil.rmtree(EXP_CUT, ignore_errors=True)
|
||||
EXP_CUT.mkdir(parents=True)
|
||||
|
||||
exp_py = ICE / 'CCPD_19K_EXP_CUT' / 'build_ccpd_19k_exp_cut.py'
|
||||
code = exp_py.read_text(encoding='utf-8')
|
||||
code = code.replace(
|
||||
'SRC_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K" # 源数据集',
|
||||
f'SRC_DIR = Path(r"{DATASET}")')
|
||||
code = code.replace(
|
||||
'DST_DIR = SCRIPT_DIR # 本目录 (CCPD_19K_EXP_CUT)',
|
||||
f'DST_DIR = Path(r"{EXP_CUT}")')
|
||||
tmp = ICE / 'scripts' / '_e.py'
|
||||
tmp.write_text(code, encoding='utf-8')
|
||||
run([PY, str(tmp), '--full'])
|
||||
tmp.unlink()
|
||||
print(f" Images: {sum(1 for _ in EXP_CUT.rglob('*.jpg'))}")
|
||||
|
||||
# ===== Step 3: TRANS =====
|
||||
print("===== 3/4 TRANS =====")
|
||||
shutil.rmtree(TRANS, ignore_errors=True)
|
||||
TRANS.mkdir(parents=True)
|
||||
|
||||
trans_py = ICE / 'CCPD_19K_EXP_CUT_TRANS' / 'build_ccpd_19k_exp_cut_trans.py'
|
||||
code = trans_py.read_text(encoding='utf-8')
|
||||
code = code.replace(
|
||||
'SRC_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT" # 源数据集',
|
||||
f'SRC_DIR = Path(r"{EXP_CUT}")')
|
||||
code = code.replace(
|
||||
'DST_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT_TRANS"',
|
||||
f'DST_DIR = Path(r"{TRANS}")')
|
||||
tmp = ICE / 'scripts' / '_t.py'
|
||||
tmp.write_text(code, encoding='utf-8')
|
||||
run([PY, str(tmp), '--full'])
|
||||
tmp.unlink()
|
||||
print(f" Images: {sum(1 for _ in TRANS.rglob('*.jpg'))}")
|
||||
|
||||
# ===== Step 4: Training =====
|
||||
print("===== 4/4 Training =====")
|
||||
os.chdir(str(NN / 'ICE-OCR'))
|
||||
run([PY, 'script/train.py', '--dataset', 'ccpd_trans', '--epochs', '100', '--batch', '128', '--output', 'ccpd_322k_best.pth', '--workers', '2', '--data-dir', str(TRANS)])
|
||||
print("DONE")
|
||||
@@ -0,0 +1,54 @@
|
||||
# C++ 管线批量重生 CCPD_SPE (24核)
|
||||
$ErrorActionPreference = "Continue"
|
||||
|
||||
$env:Path = "C:\opencv\build\x64\vc16\bin;$env:Path"
|
||||
|
||||
$EXE = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset\CCPD_SPE_GEN\cpp_pipeline\build\Release\plate_swap.exe"
|
||||
$POOL = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset\cpp_pipeline\output\plate_pool"
|
||||
$CCPD = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset\CCPD\CCPD2019"
|
||||
$SPE = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset\CCPD_SPE"
|
||||
|
||||
# 清旧
|
||||
if (Test-Path $SPE) {
|
||||
Write-Host "删除旧 CCPD_SPE..."
|
||||
Remove-Item -Recurse -Force $SPE
|
||||
}
|
||||
|
||||
$DIST = @{
|
||||
'ccpd_base' = 588
|
||||
'ccpd_blur' = 55
|
||||
'ccpd_challenge' = 159
|
||||
'ccpd_fn' = 58
|
||||
'ccpd_rotate' = 30
|
||||
'ccpd_tilt' = 82
|
||||
'ccpd_weather' = 28
|
||||
}
|
||||
|
||||
$TYPES = @('police','army','gangao','shiguan_old','shiguan_new',
|
||||
'lingguan_old','lingguan_new','coach','new_energy')
|
||||
|
||||
$total = 0
|
||||
$sw = [System.Diagnostics.Stopwatch]::StartNew()
|
||||
|
||||
foreach ($type in $TYPES) {
|
||||
Write-Host "`n=== $type ===" -ForegroundColor Cyan
|
||||
foreach ($subdir in $DIST.Keys) {
|
||||
$count = $DIST[$subdir]
|
||||
$ccpdDir = Join-Path $CCPD $subdir
|
||||
$outDir = Join-Path $SPE "$subdir\$type"
|
||||
New-Item -ItemType Directory -Force -Path $outDir | Out-Null
|
||||
|
||||
Write-Host " $subdir : $count 张... " -NoNewline
|
||||
& $EXE --ccpd-dir $ccpdDir --plate-type $type --plate-pool $POOL `
|
||||
--output $outDir --num $count --workers 24 --seed 42 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -eq 0) {
|
||||
Write-Host "OK" -ForegroundColor Green
|
||||
$total += $count
|
||||
} else {
|
||||
Write-Host "FAIL" -ForegroundColor Red
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
$sw.Stop()
|
||||
Write-Host "`nDONE! $total 张, $([math]::Round($sw.Elapsed.TotalMinutes,1))min" -ForegroundColor Green
|
||||
@@ -0,0 +1,70 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
批量重新生成 CCPD_SPE 数据集 — 全量版
|
||||
SPE 9类 × 8000 (ccpd_base) = 72,000 张
|
||||
蓝绿由独立脚本生成 150K+100K
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
|
||||
os.chdir(str(ICE_DATASET))
|
||||
|
||||
sys.path.insert(0, str(ICE_DATASET / 'CCPD_SPE_GEN' / 'cv_pipeline'))
|
||||
from plate_swap import PlateSwapPipeline
|
||||
|
||||
# 只生成 9 类特殊车牌 (蓝绿由独立脚本处理)
|
||||
SPE_TYPES = ['police','army','gangao','shiguan_old','shiguan_new',
|
||||
'lingguan_old','lingguan_new','coach','new_energy']
|
||||
|
||||
COUNT_PER_TYPE = 8000 # 每类 8000, 全部用 ccpd_base 背景
|
||||
|
||||
CCPD2019_DIR = ICE_DATASET / 'CCPD' / 'CCPD2019'
|
||||
SPE_OUTPUT = ICE_DATASET / 'CCPD_SPE'
|
||||
|
||||
SEED = 42
|
||||
WORKERS = 24
|
||||
|
||||
def main():
|
||||
total_start = time.time()
|
||||
total_images = 0
|
||||
|
||||
import shutil
|
||||
if SPE_OUTPUT.exists():
|
||||
print(f"删除旧 CCPD_SPE: {SPE_OUTPUT}")
|
||||
shutil.rmtree(str(SPE_OUTPUT))
|
||||
SPE_OUTPUT.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for plate_type in SPE_TYPES:
|
||||
print(f"\n{'='*60}")
|
||||
print(f"[{plate_type}] {COUNT_PER_TYPE} 张")
|
||||
print(f"{'='*60}")
|
||||
|
||||
pipeline = PlateSwapPipeline(plate_type=plate_type)
|
||||
out_dir = SPE_OUTPUT / 'ccpd_base' / plate_type
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f" 生成中...", end=' ', flush=True)
|
||||
t0 = time.time()
|
||||
pipeline.process_batch(
|
||||
ccpd_dirs=[str(CCPD2019_DIR / 'ccpd_base')],
|
||||
output_dir=str(out_dir),
|
||||
num_images=COUNT_PER_TYPE,
|
||||
preview=False,
|
||||
seed=SEED,
|
||||
workers=WORKERS,
|
||||
)
|
||||
elapsed = time.time() - t0
|
||||
total_images += COUNT_PER_TYPE
|
||||
print(f"完成 ({elapsed:.0f}s, {COUNT_PER_TYPE/elapsed:.0f} 张/s)")
|
||||
|
||||
total_t = time.time() - total_start
|
||||
print(f"\n{'='*60}")
|
||||
print(f"SPE 完成! {total_images} 张, {total_t/60:.1f}min")
|
||||
print(f"{'='*60}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,99 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
C++ 管线批量重生 CCPD_SPE (9 类 x 7 子集 = 9,000 张)
|
||||
比 Python 版快 ~5x，24 线程饱和
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
|
||||
CPP_EXE = ICE_DATASET / 'CCPD_SPE_GEN' / 'cpp_pipeline' / 'build' / 'Release' / 'plate_swap.exe'
|
||||
PLATE_POOL = ICE_DATASET / 'cpp_pipeline' / 'output' / 'plate_pool'
|
||||
CCPD2019 = ICE_DATASET / 'CCPD' / 'CCPD2019'
|
||||
SPE_OUT = ICE_DATASET / 'CCPD_SPE'
|
||||
|
||||
# 原始分布
|
||||
DIST = {
|
||||
'ccpd_base': 588,
|
||||
'ccpd_blur': 55,
|
||||
'ccpd_challenge': 159,
|
||||
'ccpd_fn': 58,
|
||||
'ccpd_rotate': 30,
|
||||
'ccpd_tilt': 82,
|
||||
'ccpd_weather': 28,
|
||||
}
|
||||
|
||||
TYPES = ['police','army','gangao','shiguan_old','shiguan_new',
|
||||
'lingguan_old','lingguan_new','coach','new_energy']
|
||||
|
||||
WORKERS = 24 # 物理核心数
|
||||
SEED = 42
|
||||
|
||||
def main():
|
||||
if not CPP_EXE.exists():
|
||||
print(f"[FATAL] C++ exe not found: {CPP_EXE}")
|
||||
return 1
|
||||
if not PLATE_POOL.is_dir():
|
||||
print(f"[FATAL] Plate pool not found: {PLATE_POOL}")
|
||||
return 1
|
||||
|
||||
# 清空旧数据
|
||||
if SPE_OUT.exists():
|
||||
print(f"删除旧 CCPD_SPE: {SPE_OUT}")
|
||||
shutil.rmtree(str(SPE_OUT))
|
||||
|
||||
total_start = time.time()
|
||||
total = 0
|
||||
|
||||
for plate_type in TYPES:
|
||||
t0 = time.time()
|
||||
type_count = 0
|
||||
print(f"\n{'='*50}")
|
||||
print(f"[{plate_type}]")
|
||||
print(f"{'='*50}")
|
||||
|
||||
for subdir, count in DIST.items():
|
||||
ccpd_dir = CCPD2019 / subdir
|
||||
if not ccpd_dir.is_dir():
|
||||
continue
|
||||
|
||||
out_dir = SPE_OUT / subdir / plate_type
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
cmd = [
|
||||
str(CPP_EXE),
|
||||
'--ccpd-dir', str(ccpd_dir),
|
||||
'--plate-type', plate_type,
|
||||
'--plate-pool', str(PLATE_POOL),
|
||||
'--output', str(out_dir),
|
||||
'--num', str(count),
|
||||
'--workers', str(WORKERS),
|
||||
'--seed', str(SEED + hash(plate_type) % 10000),
|
||||
]
|
||||
|
||||
print(f" {subdir}: {count} 张...", end=' ', flush=True)
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode == 0:
|
||||
print(f"OK")
|
||||
type_count += count
|
||||
else:
|
||||
print(f"FAIL")
|
||||
print(f" stderr: {result.stderr[:200]}")
|
||||
|
||||
elapsed = time.time() - t0
|
||||
print(f" [{plate_type}] {type_count} 张, {elapsed:.0f}s")
|
||||
total += type_count
|
||||
|
||||
total_t = time.time() - total_start
|
||||
print(f"\n{'='*50}")
|
||||
print(f"DONE! {total} 张, {total_t/60:.1f}min")
|
||||
print(f"{'='*50}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
os.chdir(str(ICE_DATASET))
|
||||
main()
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
CCPD原始真图 全量 TRANS (GT顶点外接框 + 12%/18%外扩)
|
||||
输出: ICE_Dataset/CCPD_REAL_TRANS/train/
|
||||
全量 ~338K, 文件名 _mask_real_ 标记 (训练时屏蔽省位+城市位)
|
||||
"""
|
||||
import os, sys, cv2, shutil
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
|
||||
sys.path.insert(0, str(ICEv2 / "ICE_NNv2/ICE-OCR/script/tools"))
|
||||
from dataset import parse_ccpd_filename, imread_unicode
|
||||
|
||||
CCPD_DIRS = [
|
||||
ICEv2 / "ICE_Dataset/CCPD/CCPD2019",
|
||||
ICEv2 / "ICE_Dataset/CCPD/CCPD2020",
|
||||
]
|
||||
OUT_DIR = ICEv2 / "ICE_Dataset/CCPD_REAL_TRANS/train"
|
||||
EXPAND_NORMAL = 1.12
|
||||
EXPAND_SMALL = 1.18
|
||||
SMALL_BOX_W = SMALL_BOX_H = 16
|
||||
MIN_PAD = 2
|
||||
|
||||
|
||||
def process_one(full_path, stem):
|
||||
"""GT顶点→外接框→12/18%外扩→裁切→顶点重映射→warp"""
|
||||
img = imread_unicode(str(full_path))
|
||||
if img is None:
|
||||
return None
|
||||
h, w = img.shape[:2]
|
||||
|
||||
parts = stem.split("-")
|
||||
if len(parts) < 4:
|
||||
return None
|
||||
# GT顶点 (CCPD: br, bl, tl, tr)
|
||||
verts = []
|
||||
for v in parts[3].split("_"):
|
||||
x, y = v.split("&")
|
||||
verts.append([float(x), float(y)])
|
||||
verts = np.float32(verts)
|
||||
|
||||
# 外接框
|
||||
x1, y1 = verts[:, 0].min(), verts[:, 1].min()
|
||||
x2, y2 = verts[:, 0].max(), verts[:, 1].max()
|
||||
bw, bh = x2 - x1, y2 - y1
|
||||
|
||||
# 外扩
|
||||
ef = EXPAND_SMALL if (bw < SMALL_BOX_W or bh < SMALL_BOX_H) else EXPAND_NORMAL
|
||||
px = max(MIN_PAD, int(round(bw * (ef - 1.0) / 2.0)))
|
||||
py = max(MIN_PAD, int(round(bh * (ef - 1.0) / 2.0)))
|
||||
cx1 = max(0, int(x1) - px); cy1 = max(0, int(y1) - py)
|
||||
cx2 = min(w, int(x2) + px); cy2 = min(h, int(y2) + py)
|
||||
if cx2 <= cx1 or cy2 <= cy1:
|
||||
return None
|
||||
|
||||
cropped = img[cy1:cy2, cx1:cx2]
|
||||
|
||||
# 顶点重映射到裁切坐标 + OpenCV顺序 (CCPD: br,bl,tl,tr → tl,tr,br,bl)
|
||||
verts_crop = verts[[2, 3, 0, 1], :] - np.array([cx1, cy1])
|
||||
verts_crop = verts_crop.astype(np.float32)
|
||||
|
||||
OW, OH = 192, 64
|
||||
dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
|
||||
M = cv2.getPerspectiveTransform(verts_crop, dst)
|
||||
return cv2.warpPerspective(cropped, M, (OW, OH))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Step 0: 清空
|
||||
if OUT_DIR.exists():
|
||||
shutil.rmtree(str(OUT_DIR))
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Step 1: 收集
|
||||
print("[1/2] Scanning all CCPD...")
|
||||
all_files = []
|
||||
total = skipped = 0
|
||||
for d in CCPD_DIRS:
|
||||
for root, dirs, files in os.walk(str(d)):
|
||||
for f in files:
|
||||
if not f.lower().endswith(".jpg"):
|
||||
continue
|
||||
indices, plate = parse_ccpd_filename(f)
|
||||
if indices is None:
|
||||
skipped += 1
|
||||
continue
|
||||
total += 1
|
||||
all_files.append((os.path.join(root, f), Path(f).stem, plate))
|
||||
|
||||
print(f" Valid: {total}, Skipped: {skipped}")
|
||||
for prov in sorted(set(p[2][0] for p in all_files),
|
||||
key=lambda p: -sum(1 for x in all_files if x[2][0] == p)):
|
||||
cnt = sum(1 for x in all_files if x[2][0] == prov)
|
||||
print(f" {prov}: {cnt}")
|
||||
|
||||
# Step 2: GT顶点 warp
|
||||
print(f"\n[2/2] GT-vertex warp {len(all_files)} -> {OUT_DIR}...")
|
||||
done = err = 0
|
||||
for i, (full_path, stem, plate) in enumerate(all_files):
|
||||
try:
|
||||
warped = process_one(full_path, stem)
|
||||
if warped is None:
|
||||
err += 1
|
||||
continue
|
||||
|
||||
# _mask_real_ 标记: 训练时屏蔽 pos0(省)+pos1(城市码)
|
||||
out_name = f"{stem}_mask_real_{plate}_trans.jpg"
|
||||
out_path = OUT_DIR / out_name
|
||||
|
||||
ret, enc = cv2.imencode(".jpg", warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
if ret:
|
||||
with open(str(out_path), "wb") as wf:
|
||||
wf.write(enc.tobytes())
|
||||
done += 1
|
||||
else:
|
||||
err += 1
|
||||
except Exception as e:
|
||||
err += 1
|
||||
if err <= 3:
|
||||
print(f" [!] {Path(full_path).name}: {e}")
|
||||
|
||||
if (i + 1) % 10000 == 0:
|
||||
print(f" [{i+1}/{len(all_files)}] done={done}, err={err}")
|
||||
|
||||
print(f"\nDone! Saved: {done}, Errors: {err}")
|
||||
print(f"Output: {OUT_DIR} ({len(list(OUT_DIR.glob('*.jpg')))} files)")
|
||||
@@ -0,0 +1,156 @@
|
||||
"""
|
||||
CCPD原始真图 去偏采样 + YOLO EXP_CUT + 角点TRANS处理
|
||||
输出: ICE_Dataset/CCPD_REAL_TRANS/train/
|
||||
省份均衡: 每省最多 MAX_PER_PROVINCE 张
|
||||
管线: YOLO检出→12%/18%外扩→裁切→角点重映射→warp (与322K EXP_CUT一致)
|
||||
"""
|
||||
import os, sys, cv2, random, shutil
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from collections import defaultdict
|
||||
|
||||
ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
|
||||
OCR_TOOLS = ICEv2 / "ICE_NNv2/ICE-OCR/script/tools"
|
||||
sys.path.insert(0, str(OCR_TOOLS))
|
||||
from dataset import parse_ccpd_filename, imread_unicode
|
||||
|
||||
CCPD_DIRS = [
|
||||
ICEv2 / "ICE_Dataset/CCPD/CCPD2019",
|
||||
ICEv2 / "ICE_Dataset/CCPD/CCPD2020",
|
||||
]
|
||||
OUT_DIR = ICEv2 / "ICE_Dataset/CCPD_REAL_TRANS/train"
|
||||
MAX_PER_PROVINCE = 300
|
||||
SEED = 42
|
||||
|
||||
YOLO_PATH = str(ICEv2 / "ICE_NNv2/yolov5t/output/1_CCPD_19K/weights/best.pt")
|
||||
EXPAND_NORMAL = 1.12 # 常规外扩 12%
|
||||
EXPAND_SMALL = 1.18 # 小车牌外扩 18%
|
||||
SMALL_BOX_W = SMALL_BOX_H = 16
|
||||
MIN_PAD = 2
|
||||
|
||||
random.seed(SEED)
|
||||
|
||||
|
||||
def process_one(img_path, plate, model, img_size=640):
|
||||
"""YOLO→外扩→裁切→角点重映射→warp (仿 EXP_CUT 管线)"""
|
||||
img = imread_unicode(str(img_path))
|
||||
if img is None:
|
||||
return None
|
||||
|
||||
h, w = img.shape[:2]
|
||||
results = model(img_path, imgsz=img_size, verbose=False)
|
||||
if len(results) == 0 or results[0].boxes is None or len(results[0].boxes) == 0:
|
||||
return None
|
||||
|
||||
boxes = results[0].boxes
|
||||
# 取置信度最高的框 (YOLO 输出 xyxy 格式)
|
||||
best_idx = boxes.conf.argmax()
|
||||
x1, y1, x2, y2 = boxes.xyxy[best_idx].cpu().numpy().astype(int)
|
||||
bw, bh = x2 - x1, y2 - y1
|
||||
|
||||
# 外扩 (与 EXP_CUT 完全一致)
|
||||
expand_factor = EXPAND_SMALL if (bw < SMALL_BOX_W or bh < SMALL_BOX_H) else EXPAND_NORMAL
|
||||
pad_x = max(MIN_PAD, int(round(bw * (expand_factor - 1.0) / 2.0)))
|
||||
pad_y = max(MIN_PAD, int(round(bh * (expand_factor - 1.0) / 2.0)))
|
||||
|
||||
cx1 = max(0, int(x1) - pad_x)
|
||||
cy1 = max(0, int(y1) - pad_y)
|
||||
cx2 = min(w, int(x2) + pad_x)
|
||||
cy2 = min(h, int(y2) + pad_y)
|
||||
if cx2 <= cx1 or cy2 <= cy1:
|
||||
return None
|
||||
|
||||
cropped = img[cy1:cy2, cx1:cx2]
|
||||
|
||||
# 角点: 用 YOLO 检出框的四角 (原始图坐标), 重映射到裁切图坐标
|
||||
corners_orig = np.float32([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]) # TL, TR, BR, BL
|
||||
corners_crop = (corners_orig - np.array([cx1, cy1])).astype(np.float32)
|
||||
|
||||
# 透视变换 → 192x64
|
||||
OW, OH = 192, 64
|
||||
dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
|
||||
M = cv2.getPerspectiveTransform(corners_crop, dst)
|
||||
warped = cv2.warpPerspective(cropped, M, (OW, OH))
|
||||
return warped
|
||||
|
||||
|
||||
# ---- Step 0: 清空旧输出 ----
|
||||
if OUT_DIR.exists():
|
||||
shutil.rmtree(str(OUT_DIR))
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
# ---- Step 1: 收集CCPD真图, 按省份分组 ----
|
||||
print("[1/3] Scanning CCPD directories...")
|
||||
prov_files = defaultdict(list) # province -> [(full_path, stem, plate)]
|
||||
total = skipped = 0
|
||||
|
||||
for ccpd_dir in CCPD_DIRS:
|
||||
for root, dirs, files in os.walk(str(ccpd_dir)):
|
||||
for f in files:
|
||||
if not f.lower().endswith(".jpg"):
|
||||
continue
|
||||
indices, plate = parse_ccpd_filename(f)
|
||||
if indices is None:
|
||||
skipped += 1
|
||||
continue
|
||||
total += 1
|
||||
prov = plate[0]
|
||||
full_path = os.path.join(root, f)
|
||||
stem = Path(f).stem
|
||||
prov_files[prov].append((full_path, stem, plate))
|
||||
|
||||
print(f" Valid: {total}, Skipped: {skipped}, Provinces: {len(prov_files)}")
|
||||
for prov in sorted(prov_files, key=lambda p: -len(prov_files[p])):
|
||||
print(f" {prov}: {len(prov_files[prov])}")
|
||||
|
||||
|
||||
# ---- Step 2: 去偏采样 ----
|
||||
print(f"\n[2/3] Sampling max {MAX_PER_PROVINCE} per province...")
|
||||
sampled = []
|
||||
for prov, items in prov_files.items():
|
||||
if len(items) <= MAX_PER_PROVINCE:
|
||||
sampled.extend(items)
|
||||
else:
|
||||
sampled.extend(random.sample(items, MAX_PER_PROVINCE))
|
||||
random.shuffle(sampled)
|
||||
print(f" Sampled: {len(sampled)}")
|
||||
prov_dist = defaultdict(int)
|
||||
for _, _, plate in sampled:
|
||||
prov_dist[plate[0]] += 1
|
||||
for prov in sorted(prov_dist):
|
||||
print(f" {prov}: {prov_dist[prov]}")
|
||||
|
||||
|
||||
# ---- Step 3: YOLO EXP_CUT + warp + 保存 ----
|
||||
print(f"\n[3/3] YOLO EXP_CUT (12%/18%) + warp {len(sampled)} images -> {OUT_DIR}...")
|
||||
print(f" YOLO: {YOLO_PATH}")
|
||||
from ultralytics import YOLO
|
||||
model = YOLO(YOLO_PATH)
|
||||
done = err = 0
|
||||
for i, (full_path, stem, plate) in enumerate(sampled):
|
||||
try:
|
||||
warped = process_one(full_path, plate, model)
|
||||
if warped is None:
|
||||
err += 1; continue
|
||||
|
||||
out_name = f"{stem}_real_blue_{plate}_trans.jpg"
|
||||
out_path = OUT_DIR / out_name
|
||||
|
||||
ret, enc = cv2.imencode(".jpg", warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
if ret:
|
||||
with open(str(out_path), "wb") as wf:
|
||||
wf.write(enc.tobytes())
|
||||
done += 1
|
||||
else:
|
||||
err += 1
|
||||
except Exception as e:
|
||||
err += 1
|
||||
if err <= 5:
|
||||
print(f" [!] {Path(full_path).name}: {e}")
|
||||
|
||||
if (i + 1) % 200 == 0:
|
||||
print(f" [{i+1}/{len(sampled)}] done={done}, err={err}")
|
||||
|
||||
print(f"\nDone! Saved: {done}, Errors: {err}")
|
||||
print(f"Output: {OUT_DIR} ({len(list(OUT_DIR.glob('*.jpg')))} files)")
|
||||
Reference in new issue
Block a user