feat: 特殊车牌生成管线 + C++/多进程/色温校正

CCPD_SPE_GEN/gan_pipeline/dataset.py

@@ -0,0 +1,131 @@
+# -*- coding: utf-8 -*-
+"""
+L通道车牌数据集加载器
+=====================
+Domain A: 合成车牌 (140×440) → 提取 L 通道
+Domain B: 真实车牌 (140×440) → 提取 L 通道
+"""
+
+import os
+import random
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+class LChannelPlateDataset(Dataset):
+    """
+    车牌 L 通道数据集。
+
+    Args:
+        root:     图片根目录
+        img_size: 统一输出尺寸 (W, H)
+        is_real:  True=真实车牌域, False=合成车牌域
+    """
+
+    def __init__(self, root: str, img_size=(440, 140), is_real=False):
+        self.root = Path(root)
+        self.img_size = img_size
+        self.is_real = is_real
+
+        self.paths = []
+        if self.root.is_dir():
+            for ext in ('*.jpg', '*.jpeg', '*.png'):
+                self.paths.extend(self.root.glob(ext))
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx) -> torch.Tensor:
+        path = self.paths[idx]
+        img = cv2.imread(str(path))
+        if img is None:
+            return self.__getitem__((idx + 1) % len(self))
+
+        # 统一尺寸
+        img = cv2.resize(img, self.img_size, interpolation=cv2.INTER_LINEAR)
+
+        # 提取 L 通道: BGR → LAB → L
+        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+        l_channel = lab[:, :, 0].astype(np.float32)  # (H, W)  [0, 255]
+
+        # 归一化到 [-1, 1]
+        l_channel = (l_channel / 127.5) - 1.0
+
+        return torch.from_numpy(l_channel).unsqueeze(0)  # (1, H, W)
+
+
+def make_loader(root: str, batch_size=4, img_size=(440, 140),
+                is_real=False, shuffle=True, num_workers=2):
+    """创建 DataLoader"""
+    from torch.utils.data import DataLoader
+
+    ds = LChannelPlateDataset(root, img_size, is_real)
+    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle,
+                      num_workers=num_workers, pin_memory=True,
+                      drop_last=True)
+
+
+class LChannelInference:
+    """
+    L通道 CUT 推理封装：合成车牌 → 真实风格车牌。
+
+    用法:
+        infer = LChannelInference('checkpoint.pth')
+        result_bgr = infer(synthetic_plate_bgr)
+    """
+
+    def __init__(self, checkpoint_path: str, device=None, n_res=6):
+        from models import LChannelGenerator, init_weights
+
+        self.device = device or torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        self.G = LChannelGenerator(n_res=n_res).to(self.device)
+        ckpt = torch.load(checkpoint_path, map_location=self.device)
+        self.G.load_state_dict(ckpt.get('G_state', ckpt))
+        self.G.eval()
+
+        # 期望输入尺寸
+        self.img_size = (440, 140)
+
+    def __call__(self, plate_bgr: np.ndarray) -> np.ndarray:
+        """
+        Args:
+            plate_bgr: (H, W, 3) 合成车牌 BGR
+
+        Returns:
+            真实风格车牌 BGR (H, W, 3)
+        """
+        h, w = plate_bgr.shape[:2]
+        img = cv2.resize(plate_bgr, self.img_size, interpolation=cv2.INTER_LINEAR)
+
+        # BGR → LAB
+        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
+
+        # L → tensor → G → tensor → L
+        l_tensor = torch.from_numpy(lab[:, :, 0] / 127.5 - 1.0)
+        l_tensor = l_tensor.unsqueeze(0).unsqueeze(0).to(self.device)  # (1,1,H,W)
+
+        with torch.no_grad():
+            l_out = self.G(l_tensor).squeeze().cpu().numpy()  # (H, W)
+
+        # 还原 L 到 [0, 255]
+        l_out = (l_out + 1.0) * 127.5
+        l_out = np.clip(l_out, 0, 255)
+
+        lab[:, :, 0] = l_out
+        # a, b 通道不变 → 保护颜色
+
+        result = cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
+
+        # 还原原始尺寸
+        if (h, w) != (self.img_size[1], self.img_size[0]):
+            result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LINEAR)
+
+        return result