Create app.py

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageEnhance
+from ultralytics import YOLO
+import cv2
+
+# Load YOLO model
+model_path = "./best.pt"
+modelY = YOLO(model_path)
+modelY.to('cpu')
+
+# Preprocessing function
+def preprocessing(image):
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    image = ImageEnhance.Sharpness(image).enhance(2.0)
+    image = ImageEnhance.Contrast(image).enhance(1.5)
+    image = ImageEnhance.Brightness(image).enhance(0.8)
+    width = 448
+    aspect_ratio = image.height / image.width
+    height = int(width * aspect_ratio)
+    return image.resize((width, height))
+
+# YOLO document detection and cropping
+def detect_and_crop_document(image):
+    image_np = np.array(image)
+    results = modelY(image_np, conf=0.85, device='cpu')
+    cropped_images = []
+    predictions = []
+    for result in results:
+        for box in result.boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            conf = int(box.conf[0] * 100)  # Convert confidence to percentage
+            cls = int(box.cls[0])
+            class_name = modelY.names[cls].capitalize()  # Ensure class names are capitalized
+            cropped_image_np = image_np[y1:y2, x1:x2]
+            cropped_image = Image.fromarray(cropped_image_np)
+            cropped_images.append(cropped_image)
+            predictions.append(f"STNK {class_name} -- (Confidence: {conf}%)")
+    if not cropped_images:
+        return None, "No document detected"
+    return cropped_images[0], predictions[0]
+
+# Gradio interface
+def process_image(image):
+    preprocessed_image = preprocessing(image)
+    cropped_image, prediction = detect_and_crop_document(preprocessed_image)
+    return cropped_image, prediction
+
+with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-size: 16px; padding: 10px 20px; border-radius: 8px;}") as demo:
+    gr.Markdown(
+        """
+        <h1 style="text-align: center; color: #4caf50;">📜 License Registration Classification</h1>
+        <p style="text-align: center; font-size: 18px;">Upload an image and let the YOLO model detect and crop license documents automatically.</p>
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=1, min_width=300):
+            input_image = gr.Image(type="pil", label="Upload License Image", interactive=True)
+            with gr.Row():
+                clear_btn = gr.Button("Clear")
+                submit_btn = gr.Button("Detect Document")
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="Cropped Document", interactive=False)
+            output_text = gr.Textbox(label="Detection Result", interactive=False)
+
+    submit_btn.click(process_image, inputs=input_image, outputs=[output_image, output_text])
+    clear_btn.click(lambda: (None, ""), outputs=[output_image, output_text])
+
+demo.launch()