BBIT-Kai
99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
|
|
import cv2
|
|
import torch
|
|
from ultralytics import YOLO
|
|
|
|
from ai.plate.detect_rec_plate import det_rec_plate, draw_result
|
|
from ai.plate.plate_recognition.plate_rec import init_model
|
|
from config.yolo import logger
|
|
|
|
|
|
@dataclass
|
|
class PlateInfo:
|
|
plate_no: str
|
|
plate_color: str
|
|
result_img_path: str
|
|
|
|
|
|
class PlateRecognizer:
|
|
_instance = None
|
|
_ready = False
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
if cls._instance is None:
|
|
cls._instance = super(PlateRecognizer, cls).__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(
|
|
self,
|
|
detect_model_path,
|
|
rec_model_path,
|
|
output_dir="result",
|
|
device="cuda" if torch.cuda.is_available() else "cpu",
|
|
):
|
|
if hasattr(self, "_initialized") and self._initialized:
|
|
return
|
|
self.device = torch.device(device)
|
|
self.output_dir = output_dir
|
|
os.makedirs(self.output_dir, exist_ok=True)
|
|
|
|
try:
|
|
# 模型加载
|
|
self.detect_model = YOLO(detect_model_path)
|
|
self.detect_model.to(self.device)
|
|
self.detect_model.eval()
|
|
self.plate_rec_model = init_model(
|
|
self.device, rec_model_path, is_color=True
|
|
)
|
|
self._initialized = True
|
|
self._ready = True
|
|
except Exception as e:
|
|
self.detect_model = None
|
|
self._ready = False
|
|
logger.error(f"❌车牌 YOLO 模型加载失败: {e}")
|
|
|
|
def analyze_image(self, image_path):
|
|
img = cv2.imread(image_path)
|
|
if not self._ready or img is None:
|
|
return []
|
|
|
|
# 检测与识别
|
|
result_list = det_rec_plate(
|
|
img, self.detect_model, self.plate_rec_model, self.device
|
|
)
|
|
|
|
# 可视化 & 保存
|
|
vis_img, _ = draw_result(img, result_list)
|
|
result_img_path = os.path.join(self.output_dir, os.path.basename(image_path))
|
|
cv2.imwrite(result_img_path, vis_img)
|
|
|
|
# 构建返回结果
|
|
results = [
|
|
PlateInfo(
|
|
plate_no=r["plate_no"],
|
|
plate_color=r["plate_color"],
|
|
result_img_path=result_img_path,
|
|
)
|
|
for r in result_list
|
|
]
|
|
return results
|
|
|
|
|
|
# 初始化一次
|
|
recognizer = PlateRecognizer(
|
|
detect_model_path="/app/models/sentinel/yolo26s-plate-detect.pt",
|
|
rec_model_path="/app/models/sentinel/plate_rec_color.pth",
|
|
output_dir="result",
|
|
)
|
|
|
|
|
|
# base_dir = Path(r"C:\Users\BBIT\Desktop\yolo26-plate-main")
|
|
#
|
|
# recognizer = PlateRecognizer(
|
|
# detect_model_path=str(base_dir / "weights" / "yolo26s-plate-detect.pt"),
|
|
# rec_model_path=str(base_dir / "weights" / "plate_rec_color.pth"),
|
|
# output_dir="result",
|
|
# )
|