BBIT-Kai
434 lines
14 KiB
Python
434 lines
14 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from ultralytics import YOLO
|
|
|
|
from ai.plate.plate_recognition.double_plate_split_merge import get_split_merge
|
|
from ai.plate.plate_recognition.plate_rec import get_plate_result
|
|
|
|
|
|
@dataclass
|
|
class PlateResult:
|
|
plate: str
|
|
color: str
|
|
|
|
|
|
def collect_files(root_path):
|
|
file_list = []
|
|
for root, _, files in os.walk(root_path):
|
|
for name in files:
|
|
file_list.append(os.path.join(root, name))
|
|
return sorted(file_list)
|
|
|
|
|
|
def four_point_transform(image, pts):
|
|
rect = pts.astype(np.float32)
|
|
(tl, tr, br, bl) = rect
|
|
|
|
width_a = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
|
|
width_b = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
|
|
max_width = max(int(width_a), int(width_b))
|
|
|
|
height_a = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
|
|
height_b = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
|
|
max_height = max(int(height_a), int(height_b))
|
|
|
|
dst = np.array(
|
|
[
|
|
[0, 0],
|
|
[max_width - 1, 0],
|
|
[max_width - 1, max_height - 1],
|
|
[0, max_height - 1],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
matrix = cv2.getPerspectiveTransform(rect, dst)
|
|
return cv2.warpPerspective(image, matrix, (max_width, max_height))
|
|
|
|
|
|
def load_model(weights, device):
|
|
model = YOLO(weights)
|
|
model.to(device)
|
|
return model
|
|
|
|
|
|
def det_rec_plate(img_ori, detect_model, plate_rec_model, device, conf=0.3, iou=0.5):
|
|
result_list = []
|
|
results = detect_model(img_ori, conf=conf, iou=iou, verbose=False)
|
|
|
|
for result in results:
|
|
boxes = result.boxes
|
|
keypoints = result.keypoints
|
|
if len(boxes) == 0 or keypoints is None:
|
|
continue
|
|
|
|
kpts_xy = keypoints.xy
|
|
num_det = min(len(boxes), len(kpts_xy))
|
|
for idx in range(num_det):
|
|
box = boxes.xyxy[idx].cpu().numpy()
|
|
det_conf = float(boxes.conf[idx])
|
|
plate_type = int(boxes.cls[idx])
|
|
|
|
landmarks = kpts_xy[idx].cpu().numpy().astype(np.int64)
|
|
roi_img = four_point_transform(img_ori, landmarks)
|
|
if plate_type == 1:
|
|
roi_img = get_split_merge(roi_img)
|
|
|
|
plate_number, _, plate_color, color_conf = get_plate_result(
|
|
roi_img, device, plate_rec_model, is_color=True
|
|
)
|
|
|
|
result_list.append(
|
|
{
|
|
"plate_no": plate_number,
|
|
"plate_color": plate_color,
|
|
"rect": [int(v) for v in box],
|
|
"detect_conf": det_conf,
|
|
"landmarks": landmarks.tolist(),
|
|
"roi_height": roi_img.shape[0],
|
|
"color_conf": color_conf,
|
|
"plate_type": plate_type, # 0: 单层, 1: 双层
|
|
}
|
|
)
|
|
return result_list
|
|
|
|
|
|
def _clamp(value, low, high):
|
|
return max(low, min(high, value))
|
|
|
|
|
|
def _normalize_rect(rect):
|
|
if not rect or len(rect) < 4:
|
|
return None
|
|
x1 = int(round(float(rect[0])))
|
|
y1 = int(round(float(rect[1])))
|
|
x2 = int(round(float(rect[2])))
|
|
y2 = int(round(float(rect[3])))
|
|
left = min(x1, x2)
|
|
top = min(y1, y2)
|
|
right = max(x1, x2)
|
|
bottom = max(y1, y2)
|
|
if right <= left or bottom <= top:
|
|
return None
|
|
return [left, top, right, bottom]
|
|
|
|
|
|
def _get_plate_theme(plate_color):
|
|
theme_map = {
|
|
"蓝色": {
|
|
"bg": (72, 33, 6),
|
|
"border": (250, 165, 96),
|
|
"text": (254, 242, 224),
|
|
"glow": (246, 130, 59),
|
|
},
|
|
"黄色": {
|
|
"bg": (0, 49, 74),
|
|
"border": (21, 204, 250),
|
|
"text": (195, 249, 254),
|
|
"glow": (8, 179, 234),
|
|
},
|
|
"绿色": {
|
|
"bg": (27, 53, 4),
|
|
"border": (128, 222, 74),
|
|
"text": (231, 252, 220),
|
|
"glow": (94, 197, 34),
|
|
},
|
|
"白色": {
|
|
"bg": (68, 50, 38),
|
|
"border": (240, 232, 226),
|
|
"text": (250, 250, 248),
|
|
"glow": (184, 163, 148),
|
|
},
|
|
"黑色": {
|
|
"bg": (20, 12, 8),
|
|
"border": (184, 163, 148),
|
|
"text": (240, 232, 226),
|
|
"glow": (139, 116, 100),
|
|
},
|
|
}
|
|
return theme_map.get(
|
|
plate_color,
|
|
{
|
|
"bg": (43, 24, 8),
|
|
"border": (248, 189, 56),
|
|
"text": (255, 248, 232),
|
|
"glow": (200, 140, 32),
|
|
},
|
|
)
|
|
|
|
|
|
def _draw_alpha_rect(img, x1, y1, x2, y2, color, alpha=0.75):
|
|
h, w = img.shape[:2]
|
|
x1 = _clamp(x1, 0, w - 1)
|
|
y1 = _clamp(y1, 0, h - 1)
|
|
x2 = _clamp(x2, 0, w)
|
|
y2 = _clamp(y2, 0, h)
|
|
if x2 <= x1 or y2 <= y1:
|
|
return
|
|
roi = img[y1:y2, x1:x2]
|
|
overlay = np.full_like(roi, color, dtype=np.uint8)
|
|
cv2.addWeighted(overlay, alpha, roi, 1 - alpha, 0, roi)
|
|
|
|
|
|
def _measure_text(text, text_size=16):
|
|
try:
|
|
font = ImageFont.truetype(get_font_file_path(), text_size, encoding="utf-8")
|
|
left, top, right, bottom = font.getbbox(text)
|
|
return right - left, bottom - top
|
|
except Exception:
|
|
print("字体文件加载失败")
|
|
size, baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
|
|
return size[0], size[1] + baseline
|
|
|
|
|
|
def _draw_glow_border(img, x1, y1, x2, y2, border_color, glow_color):
|
|
h, w = img.shape[:2]
|
|
x1 = _clamp(x1, 0, w - 1)
|
|
y1 = _clamp(y1, 0, h - 1)
|
|
x2 = _clamp(x2, 0, w - 1)
|
|
y2 = _clamp(y2, 0, h - 1)
|
|
if x2 <= x1 or y2 <= y1:
|
|
return
|
|
glow_layer = np.zeros_like(img)
|
|
cv2.rectangle(glow_layer, (x1, y1), (x2, y2), glow_color, 2)
|
|
glow_layer = cv2.GaussianBlur(glow_layer, (0, 0), sigmaX=1.6, sigmaY=1.6)
|
|
cv2.addWeighted(glow_layer, 0.45, img, 1.0, 0, img)
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), border_color, 1)
|
|
|
|
|
|
def _draw_tech_box(img, x1, y1, x2, y2, border_color, glow_color, track_id=None):
|
|
h, w = img.shape[:2]
|
|
x1 = _clamp(x1, 0, w - 1)
|
|
y1 = _clamp(y1, 0, h - 1)
|
|
x2 = _clamp(x2, 0, w - 1)
|
|
y2 = _clamp(y2, 0, h - 1)
|
|
if x2 <= x1 or y2 <= y1:
|
|
return
|
|
|
|
bw = x2 - x1
|
|
bh = y2 - y1
|
|
diag = float(np.hypot(bw, bh))
|
|
base_thick = _clamp(int(round(diag / 70.0)), 2, 5)
|
|
glow_sigma = _clamp(diag / 55.0, 1.2, 3.6)
|
|
|
|
glow_layer = np.zeros_like(img)
|
|
cv2.rectangle(glow_layer, (x1, y1), (x2, y2), glow_color, max(1, base_thick - 1))
|
|
glow_layer = cv2.GaussianBlur(
|
|
glow_layer, (0, 0), sigmaX=glow_sigma, sigmaY=glow_sigma
|
|
)
|
|
cv2.addWeighted(glow_layer, 0.48, img, 1.0, 0, img)
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), border_color, max(1, base_thick - 1))
|
|
|
|
corner_len = _clamp(int(round(min(bw, bh) * 0.28)), 8, 20)
|
|
t = base_thick
|
|
cv2.line(img, (x1, y1), (x1 + corner_len, y1), border_color, t)
|
|
cv2.line(img, (x1, y1), (x1, y1 + corner_len), border_color, t)
|
|
cv2.line(img, (x2, y1), (x2 - corner_len, y1), border_color, t)
|
|
cv2.line(img, (x2, y1), (x2, y1 + corner_len), border_color, t)
|
|
cv2.line(img, (x1, y2), (x1 + corner_len, y2), border_color, t)
|
|
cv2.line(img, (x1, y2), (x1, y2 - corner_len), border_color, t)
|
|
cv2.line(img, (x2, y2), (x2 - corner_len, y2), border_color, t)
|
|
cv2.line(img, (x2, y2), (x2, y2 - corner_len), border_color, t)
|
|
|
|
if track_id is not None:
|
|
badge = "T%02d" % track_id
|
|
badge_w_txt, badge_h_txt = _measure_text(badge, text_size=12)
|
|
pad_x = 5
|
|
pad_y = 3
|
|
badge_w = badge_w_txt + pad_x * 2
|
|
badge_h = badge_h_txt + pad_y * 2
|
|
bx2 = _clamp(x2, badge_w + 2, w - 2)
|
|
by1 = _clamp(y1 - badge_h - 2, 2, h - badge_h - 2)
|
|
bx1 = bx2 - badge_w
|
|
by2 = by1 + badge_h
|
|
_draw_alpha_rect(img, bx1, by1, bx2, by2, (18, 18, 18), alpha=0.65)
|
|
cv2.rectangle(img, (bx1, by1), (bx2, by2), border_color, 1)
|
|
text_x = bx1 + pad_x
|
|
text_y = by1 + max(1, int((badge_h - badge_h_txt) / 2))
|
|
img[:] = cv2ImgAddText(img, badge, text_x, text_y, border_color, 12)
|
|
|
|
|
|
def cv2ImgAddText(img, text, left, top, textColor=(0, 255, 0), textSize=20):
|
|
if isinstance(img, np.ndarray): # 判断是否OpenCV图片类型
|
|
img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
|
draw = ImageDraw.Draw(img)
|
|
fontText = ImageFont.truetype(get_font_file_path(), textSize, encoding="utf-8")
|
|
draw.text((left, top), text, textColor, font=fontText)
|
|
return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
|
|
|
|
|
|
def get_font_file_path():
|
|
# win
|
|
# font_path = r"C:\Users\BBIT\Desktop\yolo26-plate-main\fonts\platech.ttf"
|
|
# linux
|
|
font_path = r"/app/models/sentinel/platech.ttf"
|
|
|
|
return font_path
|
|
|
|
|
|
def _draw_tech_landmark(img, x, y, border_color, glow_color):
|
|
h, w = img.shape[:2]
|
|
x = _clamp(int(round(x)), 0, w - 1)
|
|
y = _clamp(int(round(y)), 0, h - 1)
|
|
glow_layer = np.zeros_like(img)
|
|
cv2.circle(glow_layer, (x, y), 5, glow_color, -1)
|
|
glow_layer = cv2.GaussianBlur(glow_layer, (0, 0), sigmaX=1.2, sigmaY=1.2)
|
|
cv2.addWeighted(glow_layer, 0.5, img, 1.0, 0, img)
|
|
cv2.circle(img, (x, y), 2, border_color, -1)
|
|
cv2.circle(img, (x, y), 4, border_color, 1)
|
|
|
|
|
|
def _plate_width_from_landmarks(landmarks, fallback_width):
|
|
if not landmarks or len(landmarks) < 4:
|
|
return float(fallback_width)
|
|
try:
|
|
p0 = np.array(landmarks[0], dtype=np.float32)
|
|
p1 = np.array(landmarks[1], dtype=np.float32)
|
|
p2 = np.array(landmarks[2], dtype=np.float32)
|
|
p3 = np.array(landmarks[3], dtype=np.float32)
|
|
top_w = float(np.linalg.norm(p1 - p0))
|
|
bottom_w = float(np.linalg.norm(p2 - p3))
|
|
width = (top_w + bottom_w) / 2.0
|
|
if np.isfinite(width) and width > 1:
|
|
return width
|
|
except Exception:
|
|
pass
|
|
return float(fallback_width)
|
|
|
|
|
|
def _plate_height_from_landmarks(landmarks, fallback_height):
|
|
if not landmarks or len(landmarks) < 4:
|
|
return float(fallback_height)
|
|
try:
|
|
p0 = np.array(landmarks[0], dtype=np.float32)
|
|
p1 = np.array(landmarks[1], dtype=np.float32)
|
|
p2 = np.array(landmarks[2], dtype=np.float32)
|
|
p3 = np.array(landmarks[3], dtype=np.float32)
|
|
left_h = float(np.linalg.norm(p3 - p0))
|
|
right_h = float(np.linalg.norm(p2 - p1))
|
|
height = (left_h + right_h) / 2.0
|
|
if np.isfinite(height) and height > 1:
|
|
return height
|
|
except Exception:
|
|
pass
|
|
return float(fallback_height)
|
|
|
|
|
|
def _fit_font_size(text, max_w, max_h, min_size=10, max_size=24):
|
|
if max_w <= 0 or max_h <= 0:
|
|
return min_size
|
|
for size in range(max_size, min_size - 1, -1):
|
|
tw, th = _measure_text(text, text_size=size)
|
|
if tw <= max_w and th <= max_h:
|
|
return size
|
|
return min_size
|
|
|
|
|
|
def draw_result(orgimg, result_list):
|
|
landmark_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
|
|
result_str = []
|
|
img_h, img_w = orgimg.shape[:2]
|
|
for idx, result in enumerate(result_list, start=1):
|
|
raw_rect = _normalize_rect(result.get("rect"))
|
|
if raw_rect is None:
|
|
continue
|
|
|
|
x1, y1, x2, y2 = raw_rect
|
|
w = x2 - x1
|
|
h = y2 - y1
|
|
padding_w = int(round(0.05 * w))
|
|
padding_h = int(round(0.11 * h))
|
|
rx1 = _clamp(x1 - padding_w, 0, img_w - 1)
|
|
ry1 = _clamp(y1 - padding_h, 0, img_h - 1)
|
|
rx2 = _clamp(x2 + padding_w, 0, img_w - 1)
|
|
ry2 = _clamp(y2 + padding_h, 0, img_h - 1)
|
|
|
|
landmarks = result.get("landmarks", [])
|
|
plate_no = result.get("plate_no", "")
|
|
plate_color = result.get("plate_color", "")
|
|
if result.get("plate_type", 0) == 1:
|
|
result_p = "%s %s双层" % (plate_no, plate_color)
|
|
else:
|
|
result_p = "%s %s" % (plate_no, plate_color)
|
|
result_str.append(result_p)
|
|
|
|
theme = _get_plate_theme(plate_color)
|
|
for i in range(min(4, len(landmarks))):
|
|
point = landmarks[i]
|
|
if len(point) < 2:
|
|
continue
|
|
point_color = landmark_colors[i]
|
|
_draw_tech_landmark(orgimg, point[0], point[1], point_color, point_color)
|
|
|
|
_draw_tech_box(
|
|
orgimg, rx1, ry1, rx2, ry2, theme["border"], theme["glow"], track_id=idx
|
|
)
|
|
|
|
label = "%s | %s" % (plate_no, plate_color)
|
|
plate_w = _plate_width_from_landmarks(landmarks, rx2 - rx1)
|
|
plate_h = _plate_height_from_landmarks(landmarks, ry2 - ry1)
|
|
pre_card_h = _clamp(int(round(plate_h)), 24, min(110, img_h - 4))
|
|
pre_pad_y = _clamp(int(round(pre_card_h * 0.16)), 3, 10)
|
|
pre_inner_h = max(8, pre_card_h - pre_pad_y * 2)
|
|
pre_max_font = _clamp(int(round(pre_card_h * 0.72)), 14, 44)
|
|
pre_min_font = _clamp(int(round(pre_card_h * 0.42)), 10, pre_max_font)
|
|
pre_font_size = _fit_font_size(
|
|
label, 4096, pre_inner_h, min_size=pre_min_font, max_size=pre_max_font
|
|
)
|
|
pre_text_w, _ = _measure_text(label, text_size=pre_font_size)
|
|
|
|
min_w_by_text = pre_text_w + 20
|
|
base_w_by_plate = int(round(plate_w * 1.05))
|
|
card_w = max(90, base_w_by_plate, min_w_by_text)
|
|
card_w = min(card_w, img_w - 8)
|
|
|
|
card_h = pre_card_h
|
|
card_pad_x = _clamp(int(round(card_w * 0.08)), 8, 18)
|
|
card_pad_y = _clamp(int(round(card_h * 0.16)), 3, 10)
|
|
|
|
card_x = int(rx1 + (rx2 - rx1 - card_w) / 2)
|
|
card_x = _clamp(card_x, 4, max(4, img_w - card_w - 4))
|
|
card_y = ry1 - card_h - 2
|
|
if card_y < 2:
|
|
card_y = _clamp(ry1 + 2, 2, max(2, img_h - card_h - 2))
|
|
|
|
_draw_alpha_rect(
|
|
orgimg,
|
|
card_x,
|
|
card_y,
|
|
card_x + card_w,
|
|
card_y + card_h,
|
|
theme["bg"],
|
|
alpha=0.78,
|
|
)
|
|
_draw_glow_border(
|
|
orgimg,
|
|
card_x,
|
|
card_y,
|
|
card_x + card_w,
|
|
card_y + card_h,
|
|
theme["border"],
|
|
theme["glow"],
|
|
)
|
|
|
|
inner_w = max(8, card_w - card_pad_x * 2)
|
|
inner_h = max(8, card_h - card_pad_y * 2)
|
|
dynamic_max_font = _clamp(int(round(card_h * 0.72)), 14, 44)
|
|
dynamic_min_font = _clamp(int(round(card_h * 0.42)), 10, dynamic_max_font)
|
|
font_size = _fit_font_size(
|
|
label,
|
|
inner_w,
|
|
inner_h,
|
|
min_size=dynamic_min_font,
|
|
max_size=dynamic_max_font,
|
|
)
|
|
text_w, text_h = _measure_text(label, text_size=font_size)
|
|
text_x = card_x + max(card_pad_x, int((card_w - text_w) / 2))
|
|
text_y = card_y + max(card_pad_y - 1, int((card_h - text_h) / 2))
|
|
orgimg = cv2ImgAddText(orgimg, label, text_x, text_y, theme["text"], font_size)
|
|
|
|
return orgimg, result_str
|