269 lines
8.2 KiB
Python
269 lines
8.2 KiB
Python
import cv2 as cv
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
import argparse
|
||
import typing
|
||
import logging
|
||
from pathlib import Path
|
||
from dataclasses import dataclass
|
||
|
||
# Reference:
|
||
# - Claude Code
|
||
# - https://www.cnblogs.com/linuxAndMcu/p/19144795
|
||
# - https://www.51halcon.com/forum.php?mod=viewthread&tid=6562
|
||
|
||
UNI_HW: int = 1000
|
||
|
||
|
||
def _uniform_car_plate(img: cv.typing.MatLike) -> cv.typing.MatLike:
|
||
"""
|
||
Uniform the image size to 512x512 while maintaining aspect ratio, padding with black.
|
||
|
||
:param img: The image in BGR format to be uniformed.
|
||
:return: The uniformed image in BGR format.
|
||
"""
|
||
# Calculate the new width and height for given image
|
||
h, w = img.shape[:2]
|
||
scale = min(UNI_HW / w, UNI_HW / h)
|
||
new_w = int(w * scale)
|
||
new_h = int(h * scale)
|
||
|
||
# Resize the image
|
||
resized_img = cv.resize(img, (new_w, new_h))
|
||
|
||
# Create a black canvas of size UNI_HW x UNI_HW
|
||
padded_img = np.zeros((UNI_HW, UNI_HW, 3), dtype=np.uint8)
|
||
|
||
# Calculate position to paste the resized image (centered)
|
||
y_offset = (UNI_HW - new_h) // 2
|
||
x_offset = (UNI_HW - new_w) // 2
|
||
|
||
# Paste the resized image onto the canvas
|
||
padded_img[y_offset : y_offset + new_h, x_offset : x_offset + new_w] = resized_img
|
||
|
||
# Return the padded image
|
||
return padded_img
|
||
|
||
|
||
@dataclass
|
||
class CarPlateHsvBoundary:
|
||
lower_bound: cv.typing.MatLike
|
||
upper_bound: cv.typing.MatLike
|
||
need_revert: bool
|
||
"""是否取反黑白颜色，因为蓝牌和黄牌的操作正好是反的"""
|
||
|
||
|
||
CAR_PLATE_HSV_BOUNDARIES: tuple[CarPlateHsvBoundary, ...] = (
|
||
# 中国蓝牌 HSV 范围
|
||
CarPlateHsvBoundary(np.array([100, 80, 60]), np.array([130, 255, 255]), True),
|
||
# 中国绿牌 HSV 范围
|
||
CarPlateHsvBoundary(np.array([35, 43, 46]), np.array([99, 255, 255]), False),
|
||
# 中国黄牌 HSV 范围
|
||
CarPlateHsvBoundary(np.array([16, 43, 46]), np.array([34, 255, 255]), False),
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class CarPlateMask:
|
||
mask: cv.typing.MatLike
|
||
need_revert: bool
|
||
"""是否对颜色取反，与CarPlateHsvBoundary中的同名字段含义一致"""
|
||
|
||
|
||
def _batchly_mask_car_plate(
|
||
hsv: cv.typing.MatLike,
|
||
) -> typing.Iterator[CarPlateMask]:
|
||
""" """
|
||
for boundary in CAR_PLATE_HSV_BOUNDARIES:
|
||
# 以给定HSV范围检测符合该颜色的位置
|
||
mask = cv.inRange(hsv, boundary.lower_bound, boundary.upper_bound)
|
||
|
||
# 形态学：闭运算填孔 + 开运算去噪
|
||
kernel_close = cv.getStructuringElement(cv.MORPH_RECT, (25, 10))
|
||
kernel_open = cv.getStructuringElement(cv.MORPH_RECT, (5, 5))
|
||
mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel_close)
|
||
mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel_open)
|
||
|
||
# Return value
|
||
yield CarPlateMask(mask, boundary.need_revert)
|
||
|
||
|
||
@dataclass
|
||
class CarPlateRegion:
|
||
x: int
|
||
y: int
|
||
w: int
|
||
h: int
|
||
need_revert: bool
|
||
"""是否对颜色取反，与CarPlateHsvBoundary中的同名字段含义一致"""
|
||
|
||
|
||
MIN_AREA: float = 3000
|
||
MIN_ASPECT_RATIO: float = 1.5
|
||
MAX_ASPECT_RATIO: float = 6.0
|
||
BEST_ASPECT_RATIO: float = 3.5
|
||
|
||
|
||
def _analyse_car_plate_connection(
|
||
mask: CarPlateMask,
|
||
) -> typing.Optional[CarPlateRegion]:
|
||
# 连通域分析，筛选最符合车牌长宽比的区域
|
||
num_labels, labels, stats, _ = cv.connectedComponentsWithStats(
|
||
mask.mask, connectivity=8
|
||
)
|
||
|
||
best: typing.Optional[CarPlateRegion] = None
|
||
best_score = 0
|
||
|
||
for i in range(1, num_labels):
|
||
x, y, w, h, area = stats[i]
|
||
# 检查面积
|
||
if area < MIN_AREA:
|
||
continue
|
||
# 标准车牌宽高比约 3:1 ~ 5:1
|
||
ratio = w / (h + 1e-5)
|
||
if ratio >= MIN_ASPECT_RATIO and ratio <= MAX_ASPECT_RATIO:
|
||
score = area * (1 - abs(ratio - BEST_ASPECT_RATIO) / BEST_ASPECT_RATIO)
|
||
if score > best_score:
|
||
best_score = score
|
||
best = CarPlateRegion(x, y, w, h, mask.need_revert)
|
||
|
||
return best
|
||
|
||
|
||
def extract_car_plate(img: cv.typing.MatLike) -> typing.Optional[cv.typing.MatLike]:
|
||
"""
|
||
Extract the car plate part from given image.
|
||
|
||
:param img: The image containing car plate in BGR format.
|
||
:return: The image of binary car plate in U8 format if succeed, otherwise None.
|
||
"""
|
||
img = _uniform_car_plate(img)
|
||
|
||
# 转换到HSV空间
|
||
hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
|
||
|
||
# 利用车牌颜色在 HSV 空间定位车牌
|
||
candidate: typing.Optional[CarPlateRegion] = None
|
||
for mask in _batchly_mask_car_plate(hsv):
|
||
# 连通域分析，筛选最符合车牌长宽比的区域作为车牌
|
||
candidate = _analyse_car_plate_connection(mask)
|
||
# 找到任意一个就退出
|
||
if candidate is not None:
|
||
break
|
||
|
||
if candidate is None:
|
||
logging.error("Can not find any car plate.")
|
||
return None
|
||
|
||
# 稍微扩边获取最终车牌区域
|
||
pad = 6
|
||
h_img, w_img = img.shape[:2]
|
||
x1 = max(candidate.x - pad, 0)
|
||
y1 = max(candidate.y - pad, 0)
|
||
x2 = min(candidate.x + candidate.w + pad, w_img)
|
||
y2 = min(candidate.y + candidate.h + pad, h_img)
|
||
logging.info(f"车牌区域: x={x1}, y={y1}, w={x2 - x1}, h={y2 - y1}")
|
||
|
||
# # 在原图上标记（仅供调试）
|
||
# debug = img.copy()
|
||
# cv.rectangle(debug, (x1, y1), (x2, y2), (0, 255, 0), 3)
|
||
# cv.imwrite('./debug_detected.jpg', debug)
|
||
|
||
# 二值化：文字/边缘 → 黑色，背景 → 白色
|
||
gray = cv.cvtColor(img[y1:y2, x1:x2], cv.COLOR_BGR2GRAY)
|
||
|
||
# 高斯模糊降噪
|
||
blurred = cv.GaussianBlur(gray, (3, 3), 0)
|
||
|
||
# Otsu 自动阈值，得到白字黑底，再取反 → 黑字白底
|
||
_, binary_otsu = cv.threshold(blurred, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
|
||
# 反转：字符变黑，背景变白
|
||
if candidate.need_revert:
|
||
binary = cv.bitwise_not(binary_otsu)
|
||
else:
|
||
binary = binary_otsu
|
||
|
||
# 去除小噪点（开运算）
|
||
kernel_denoise = cv.getStructuringElement(cv.MORPH_RECT, (2, 2))
|
||
binary = cv.morphologyEx(binary, cv.MORPH_OPEN, kernel_denoise)
|
||
|
||
return binary
|
||
# cv.imwrite('./plate_binary.png', binary)
|
||
# print("二值化结果已保存: plate_binary.png")
|
||
|
||
# 叠加边框轮廓（细化文字边缘，参考效果图）
|
||
# Canny 边缘叠加让效果更接近参考图
|
||
edges = cv.Canny(blurred, 40, 120)
|
||
edges_inv = cv.bitwise_not(edges) # 边缘→黑色
|
||
combined = cv.bitwise_and(binary, edges_inv) # 合并
|
||
# 再做一次轻微腐蚀让字体略粗
|
||
kernel_dilate = cv.getStructuringElement(cv.MORPH_RECT, (2, 2))
|
||
combined = cv.erode(combined, kernel_dilate, iterations=1)
|
||
|
||
return combined
|
||
# cv.imwrite('./plate_final.png', combined)
|
||
# print("最终结果已保存: plate_final.png")
|
||
|
||
|
||
@dataclass
|
||
class Cli:
|
||
input_file: Path
|
||
"""The path to input file"""
|
||
output_file: Path
|
||
"""The path to output file"""
|
||
|
||
@staticmethod
|
||
def from_cmdline() -> "Cli":
|
||
# Build parser
|
||
parser = argparse.ArgumentParser(
|
||
prog="Car Plate Extractor",
|
||
description="Extract the car plate part from given image.",
|
||
)
|
||
parser.add_argument(
|
||
"-i",
|
||
"--in",
|
||
required=True,
|
||
type=str,
|
||
action="store",
|
||
dest="input_file",
|
||
metavar="in.jpg",
|
||
help="""The path to input image containing car plate.""",
|
||
)
|
||
parser.add_argument(
|
||
"-o",
|
||
"--out",
|
||
required=True,
|
||
type=str,
|
||
action="store",
|
||
dest="output_file",
|
||
metavar="out.png",
|
||
help="""The path to output image for extracted car plate.""",
|
||
)
|
||
|
||
# Parse argument from cmdline and return
|
||
args = parser.parse_args()
|
||
return Cli(Path(args.input_file), Path(args.output_file))
|
||
|
||
|
||
def main():
|
||
# Setup logging format
|
||
logging.basicConfig(format="[%(levelname)s] %(message)s", level=logging.DEBUG)
|
||
|
||
# Get user request
|
||
cli = Cli.from_cmdline()
|
||
|
||
# Load file
|
||
in_img = cv.imread(str(cli.input_file), cv.IMREAD_COLOR)
|
||
if in_img is None:
|
||
logging.error(f"Fail to load image {cli.input_file}")
|
||
return
|
||
# Save extracted file if possible
|
||
out_img = extract_car_plate(in_img)
|
||
if out_img is not None:
|
||
cv.imwrite(str(cli.output_file), out_img)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|