132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
# -*- 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
|