后端更新

2026-03-26 17:48:20 +08:00
parent 4c2bcd7dce
commit 0c2859b0db
22 changed files with 1336 additions and 213 deletions
@@ -0,0 +1,98 @@
+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",
+# )