Доработки по скрипту для вырезания VIN из СТС.

scripts/detect_vin.py

@@ -4,81 +4,42 @@ exec(open("scripts/detect_vin.py").read())
 
 from paddleocr import PaddleOCR
 import cv2
-import numpy as np
 import re
 
-def auto_rotate_image(image):
-    """Автоматически выравнивает повёрнутое изображение"""
-    # Конвертация в grayscale и бинаризация
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
-    
-    # Нахождение контуров текстовых блоков
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5]
-    
-    # Вычисление угла поворота основных контуров
-    angles = []
-    for cnt in contours:
-        rect = cv2.minAreaRect(cnt)
-        angle = rect[-1]
-        if angle < -45:
-            angle = 90 + angle
-        angles.append(angle)
-    
-    # Медианный угол для устойчивости к выбросам
-    median_angle = np.median(angles)
-    
-    # Поворот изображения
-    (h, w) = image.shape[:2]
-    center = (w // 2, h // 2)
-    M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
-    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
-    
-    return rotated
-
-# Основной процесс
 ocr = PaddleOCR(lang='en', use_textline_orientation=True)
 
-# Загрузка изображения
-image = cv2.imread("scripts/image.jpg")
+def cut_vin(ocr, input_image, output_image):
+    """
+    Вырезать VIN из изображения.
+    """
     
-# Автоматическое выравнивание
-# aligned_image = auto_rotate_image(image)
-aligned_image = image
+    image = cv2.imread(input_image)
 
-# Распознавание текста
-result = ocr.predict(image)
-output_image = result[0]["doc_preprocessor_res"]["output_img"]
-cv2.imwrite('scripts/output_image.jpg', output_image)
+    result = ocr.predict(image)
+    processed_image = result[0]["doc_preprocessor_res"]["output_img"]
 
-# Поиск VIN
-vin_pattern = re.compile(r'^[A-HJ-NPR-Z0-9]{17}$')
-found_vin = None
+    vin_pattern = re.compile(r'^[A-HJ-NPR-Z0-9]{17}$')
+    found_vin = None
+
+    for text, bbox in zip(result[0]["rec_texts"], result[0]["rec_boxes"]):
+        if vin_pattern.match(text):
+            found_vin = text
+            break
 
-for ind, text in enumerate(result[0]["rec_texts"]):
-    if vin_pattern.match(text):
-        found_vin = text
-        bbox = result[0]["rec_boxes"][ind]
-        break
     if found_vin:
-        break
+        x_min, y_min = bbox[0], bbox[1]
+        x_max, y_max = bbox[2], bbox[3]
 
-# Вырезание области VIN
-if found_vin:
-    # xs = [int(p[0]) for p in bbox]
-    # ys = [int(p[1]) for p in bbox]
+        vin_region = processed_image[y_min:y_max, x_min:x_max]
 
-    x_min, y_min = bbox[0], bbox[1]
-    x_max, y_max = bbox[2], bbox[3]
+        cv2.imwrite(output_image, vin_region)
 
-    # Добавление запаса вокруг области
-    vin_region = output_image[
-        y_min:y_max,
-        x_min:x_max
-    ]
+    return found_vin
 
-    cv2.imwrite('scripts/vin.jpg', vin_region)
-    print(f"VIN найден: {found_vin}")
+vin = cut_vin(ocr, "input/image.jpg", "output/vin.jpg")
+
+if vin:
+    print(f"VIN найден: {vin}.")
 else:
-    print("VIN не обнаружен после выравнивания")
+    print("VIN не обнаружен.")
+