First commit to app

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+target/1.jpg filter=lfs diff=lfs merge=lfs -text

main.py

@@ -0,0 +1,174 @@
+from PIL import Image
+from ultralytics import YOLO
+
+import os
+import torch
+import re
+import cv2
+
+import gradio as gr
+import torchvision.transforms as T
+import albumentations as A
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# Folders
+input_folder = "./target"
+output_folder = "./target_output"
+
+os.makedirs(output_folder, exist_ok=True)
+
+# Detector model
+license_plate_detector = YOLO("./models/yolo11x.pt")
+
+# SuperResolution model
+sr = cv2.dnn_superres.DnnSuperResImpl_create()
+sr.readModel("./models/FSRCNN_x3.pb")
+sr.setModel("fsrcnn", 3)
+
+
+class App:
+    models = ['parseq', 'parseq_tiny', 'abinet', 'crnn', 'trba', 'vitstr']
+
+    def __init__(self):
+        self._model_cache = {}
+        self._preprocess = T.Compose([
+            T.Resize((32, 128), T.InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(0.5, 0.5)
+        ])
+
+    def _get_model(self, name):
+        if name in self._model_cache:
+            return self._model_cache[name]
+        model = torch.hub.load('baudm/parseq', name, pretrained=True).eval().to(device)
+        self._model_cache[name] = model
+        return model
+
+    @torch.inference_mode()
+    def __call__(self, model_name, image):
+        if image is None:
+            return '', []
+        model = self._get_model(model_name)
+        image = self._preprocess(image.convert('RGB')).unsqueeze(0).to(device)
+        pred = model(image).softmax(-1)
+        label, _ = model.tokenizer.decode(pred)
+        raw_label, raw_confidence = model.tokenizer.decode(pred, raw=True)
+        max_len = 25 if model_name == 'crnn' else len(label[0]) + 1
+        conf = list(map('{:0.1f}'.format, raw_confidence[0][:max_len].tolist()))
+        return label[0], [raw_label[0][:max_len], conf]
+
+
+p = App()
+
+
+def detect_license_plates(model, image):
+    plate_image_np = pil_to_np(image)
+
+    transform = A.Compose([
+        A.ToGray(p=1.0),
+        A.CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=1.0),
+    ])
+
+    transformed = transform(image=plate_image_np)['image']
+
+    if len(transformed.shape) == 2:
+        transformed = cv2.cvtColor(transformed, cv2.COLOR_GRAY2RGB)
+
+    image = np_to_pil(transformed)
+
+    results = model(image)
+    plates = []
+    for result in results:
+        for box in result.boxes.xyxy.cpu().numpy():
+            x1, y1, x2, y2 = map(int, box)
+            plate = image.crop((x1, y1, x2, y2))
+            plates.append((plate, (x1, y1, x2, y2)))
+
+    return plates
+
+
+def pil_to_np(image):
+    return np.array(image)
+
+
+def np_to_pil(image_np):
+    return Image.fromarray(image_np)
+
+
+def preprocess_license_plate(plate_image: Image):
+    plate_image_np = pil_to_np(plate_image)
+    if not(plate_image_np.ndim == 2 or plate_image_np.shape[-1] == 1):
+        plate_image_np = A.ToGray(p=1.0, num_output_channels=1)(image=plate_image_np)['image']
+    super_resolved = sr.upsample(plate_image_np)
+    augmented = A.Compose([
+        A.CLAHE(clip_limit=2, tile_grid_size=(1, 1), p=1.0),
+        A.Morphological(p=1.0, scale=(4, 4), operation="erosion"),
+    ])(image=super_resolved)['image']
+
+    super_resolved_pil = np_to_pil(augmented)
+    return super_resolved_pil
+
+
+def process_image(image_path: Image):
+    image_np = np.array(image_path)
+
+    fig, ax = plt.subplots(1, figsize=(10, 6))
+    ax.imshow(image_np)
+
+    plates = detect_license_plates(license_plate_detector, image_path)
+    recognized_text = ""
+
+    for i, (plate, bbox) in enumerate(plates):
+        preprocessed_plate = preprocess_license_plate(plate)
+        recognized_text, raw_output = p.__call__("parseq", preprocessed_plate)
+
+        if recognized_text and len(recognized_text) > 5:
+            recognized_text = re.sub(r"[^A-Za-z0-9]", "", recognized_text).upper()
+            recognized_text = recognized_text.replace('V', 'Y').replace('I', '')
+            recognized_text = recognized_text.replace('8', 'В', 1) if recognized_text[0] == "8" else recognized_text
+            recognized_text = recognized_text.replace('7', 'T', 1) if recognized_text[0] == "7" else recognized_text
+            recognized_text = recognized_text.replace('0', 'O', 1) if recognized_text[0] == "0" else recognized_text
+            recognized_text = recognized_text[:9] if len(recognized_text) >= 9 else recognized_text
+
+        x1, y1, x2, y2 = bbox
+        rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor='r', facecolor='none')
+        ax.add_patch(rect)
+        ax.text(x1, y1 - 10, recognized_text, color='red', fontsize=12, bbox=dict(facecolor='white', alpha=0.5))
+
+    plt.axis('off')
+
+    # Saving image to buffer
+    output_buffer = "processed_image.png"
+    plt.savefig(output_buffer, bbox_inches='tight')
+    plt.close()
+
+    return Image.open(output_buffer), recognized_text.strip()
+
+
+# Gradio UI
+
+target_folder = "./target"
+example_images = [
+    os.path.join(target_folder, file) for file in os.listdir(target_folder) if file.lower().endswith(("jpg", "png", "bmp"))
+]
+
+interface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil", label="Загрузите фото машины с номером 📤"),
+    outputs=[
+        gr.Image(type="pil", label="📸 Выход 0 - Обработанное изображение"),
+        gr.Text(label="🔍 Выход 1 - Распознанный номер"),
+    ],
+    title="Распознавание российских номеров",
+    description="🔎 **Загрузите изображение с автомобильным номером** и модель автоматически **определит госномер!** 🔥\n\n📸 **Форматы:** JPG, PNG, BMP",
+    examples=example_images,
+    allow_flagging="never",
+    theme="compact",
+)
+
+if __name__ == "__main__":
+    interface.launch(share=True)