Update app.py

app.py

@@ -1,174 +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)
+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,
+    flagging_mode="never",
+    theme="compact",
+)
+
+if __name__ == "__main__":
+    interface.launch(share=True)