diff --git a/scripts/detect_vin.py b/scripts/detect_vin.py index f1ee7d0..d188ff6 100644 --- a/scripts/detect_vin.py +++ b/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") - -# Автоматическое выравнивание -# aligned_image = auto_rotate_image(image) -aligned_image = image - -# Распознавание текста -result = ocr.predict(image) -output_image = result[0]["doc_preprocessor_res"]["output_img"] -cv2.imwrite('scripts/output_image.jpg', output_image) - -# Поиск VIN -vin_pattern = re.compile(r'^[A-HJ-NPR-Z0-9]{17}$') -found_vin = None - -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 - -# Вырезание области VIN -if found_vin: - # xs = [int(p[0]) for p in bbox] - # ys = [int(p[1]) for p in bbox] +def cut_vin(ocr, input_image, output_image): + """ + Вырезать VIN из изображения. + """ - x_min, y_min = bbox[0], bbox[1] - x_max, y_max = bbox[2], bbox[3] + image = cv2.imread(input_image) - # Добавление запаса вокруг области - vin_region = output_image[ - y_min:y_max, - x_min:x_max - ] + result = ocr.predict(image) + processed_image = result[0]["doc_preprocessor_res"]["output_img"] - cv2.imwrite('scripts/vin.jpg', vin_region) - print(f"VIN найден: {found_vin}") + 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 + + if found_vin: + x_min, y_min = bbox[0], bbox[1] + x_max, y_max = bbox[2], bbox[3] + + vin_region = processed_image[y_min:y_max, x_min:x_max] + + cv2.imwrite(output_image, vin_region) + + return found_vin + +vin = cut_vin(ocr, "input/image.jpg", "output/vin.jpg") + +if vin: + print(f"VIN найден: {vin}.") else: - print("VIN не обнаружен после выравнивания") + print("VIN не обнаружен.") +