MLBench/License_Registration_Classification

1# import gradio as gr2# import numpy as np3# from PIL import Image, ImageEnhance4# from ultralytics import YOLO5# import cv26 7# # Load YOLO model8# model_path = "./best.pt"9# modelY = YOLO(model_path)10# modelY.to('cpu')11 12# # Preprocessing function13# def preprocessing(image):14#     if image.mode != 'RGB':15#         image = image.convert('RGB')16#     image = ImageEnhance.Sharpness(image).enhance(2.0)17#     image = ImageEnhance.Contrast(image).enhance(1.5)18#     image = ImageEnhance.Brightness(image).enhance(0.8)19#     width = 44820#     aspect_ratio = image.height / image.width21#     height = int(width * aspect_ratio)22#     return image.resize((width, height))23 24# # YOLO document detection and cropping25# def detect_and_crop_document(image):26#     image_np = np.array(image)27#     results = modelY(image_np, conf=0.80, device='cpu')28#     cropped_images = []29#     predictions = []30    31#     for result in results:32#         for box in result.boxes:33#             x1, y1, x2, y2 = map(int, box.xyxy[0])34#             conf = int(box.conf[0] * 100)  # Convert confidence to percentage35#             cls = int(box.cls[0])36#             class_name = modelY.names[cls].capitalize()  # Capitalize class names37#             cropped_image_np = image_np[y1:y2, x1:x2]38#             cropped_image = Image.fromarray(cropped_image_np)39#             cropped_images.append(cropped_image)40#             predictions.append(f"Detected: STNK {class_name} -- (Confidence: {conf}%)")41    42#     if not cropped_images:43#         return None, "No document detected"44#     return cropped_images, predictions45 46# # Gradio interface47# def process_image(image):48#     preprocessed_image = preprocessing(image)49#     cropped_images, predictions = detect_and_crop_document(preprocessed_image)50    51#     if cropped_images:52#         return cropped_images, '\n'.join(predictions)53#     return None, "No document detected"54 55# with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-size: 16px; padding: 10px 20px; border-radius: 8px;}") as demo:56#     gr.Markdown(57#         """58#         <h1 style="text-align: center; color: #4caf50;">📜 License Registration Classification</h1>59#         <p style="text-align: center; font-size: 18px;">Upload an image and let the YOLO model detect and crop license documents automatically.</p>60#         """61#     )62#     with gr.Row():63#         with gr.Column(scale=1, min_width=300):64#             input_image = gr.Image(type="pil", label="Upload License Image", interactive=True)65#             with gr.Row():66#                 clear_btn = gr.Button("Clear")67#                 submit_btn = gr.Button("Detect Document")68#         with gr.Column(scale=2):69#             output_image = gr.Gallery(label="Cropped Documents", interactive=False)70#             output_text = gr.Textbox(label="Detection Result", interactive=False)71 72#     submit_btn.click(process_image, inputs=input_image, outputs=[output_image, output_text])73#     clear_btn.click(lambda: (None, ""), outputs=[output_image, output_text])74 75# demo.launch()76 77 78 79 80import gradio as gr81import numpy as np82from PIL import Image, ImageEnhance83from ultralytics import YOLO84import cv285import os86 87# --- DOCUMENTATION STRINGS (English Only) ---88 89GUIDELINE_SETUP = """90## 1. Quick Start Guide: Setup and Run Instructions91 92This application uses a YOLO model to automatically detect, classify, and extract specific license registration documents (STNK).93 941.  **Preparation:** Ensure your image clearly shows the target license document.952.  **Upload:** Click the 'Upload License Image' box and select your image (JPG, PNG).963.  **Run:** Click the **"Detect Document"** button.974.  **Review:** The detected documents will appear in the 'Cropped Documents' gallery, and the 'Detection Result' box will show the classification and confidence score.98"""99 100GUIDELINE_INPUT = """101## 2. Expected Inputs and Preprocessing102 103| Input Field | Purpose | Requirement |104| :--- | :--- | :--- |105| **Upload License Image** | The image containing the license document you want to detect and classify. | Must be an image file (e.g., JPG, PNG). |106 107### Automatic Preprocessing Steps:108Before detection, the input image is automatically adjusted to enhance accuracy:1091.  **Sharpness:** Increased sharpness by 2.0.1102.  **Contrast:** Increased contrast by 1.5.1113.  **Brightness:** Slightly reduced brightness by 0.8.1124.  **Resizing:** The image is resized to a width of 448 pixels while maintaining its original aspect ratio.113"""114 115GUIDELINE_OUTPUT = """116## 3. Expected Outputs (Detection and Classification)117 118The application produces two outputs based on a successful detection:119 1201.  **Cropped Documents (Gallery):**121    *   This gallery displays only the regions of the image where a license document was confidently detected (Confidence > 80%).122    *   If multiple documents are found, all cropped images will appear here.123 1242.  **Detection Result (Textbox):**125    *   A text summary listing each detected document, including its specific class name (e.g., 'STNK Class A'), and the model's confidence level (as a percentage).126 127### Failure Modes:128*   If "No document detected" is returned, it means the model did not find a document with a confidence level of 80% or higher, or the image quality was too poor for detection.129"""130 131# --- CORE LOGIC ---132 133# Load YOLO model134# NOTE: Ensure 'best.pt' is available in the execution directory.135model_path = "./best.pt"136try:137    modelY = YOLO(model_path)138    modelY.to('cpu')139except Exception as e:140    print(f"Error loading model: {e}")141    modelY = None142 143# Preprocessing function144def preprocessing(image):145    if image.mode != 'RGB':146        image = image.convert('RGB')147    148    # Enhancement steps149    image = ImageEnhance.Sharpness(image).enhance(2.0)150    image = ImageEnhance.Contrast(image).enhance(1.5)151    image = ImageEnhance.Brightness(image).enhance(0.8)152    153    # Resizing while preserving aspect ratio154    width = 448155    aspect_ratio = image.height / image.width156    height = int(width * aspect_ratio)157    return image.resize((width, height))158 159# YOLO document detection and cropping160def detect_and_crop_document(image):161    if modelY is None:162        return [], ["Model not loaded."]163        164    image_np = np.array(image)165    # Run inference with confidence threshold 0.80166    results = modelY(image_np, conf=0.80, device='cpu', verbose=False)167    168    cropped_images = []169    predictions = []170    171    for result in results:172        for box in result.boxes:173            x1, y1, x2, y2 = map(int, box.xyxy[0])174            conf = int(box.conf[0].item() * 100) # Ensure conversion to scalar for item()175            cls = int(box.cls[0].item())176            class_name = modelY.names.get(cls, "Unknown").capitalize()177            178            cropped_image_np = image_np[y1:y2, x1:x2]179            180            # Check for valid crop size before converting to PIL181            if cropped_image_np.size > 0:182                cropped_image = Image.fromarray(cropped_image_np)183                cropped_images.append(cropped_image)184                predictions.append(f"Detected: STNK {class_name} -- (Confidence: {conf}%)")185    186    return cropped_images, predictions187 188# Gradio interface function189def process_image(image):190    if image is None:191        raise gr.Error("Please upload an image.")192        193    preprocessed_image = preprocessing(image)194    cropped_images, predictions = detect_and_crop_document(preprocessed_image)195    196    if cropped_images:197        return cropped_images, '\n'.join(predictions)198    199    # If no documents are detected with sufficient confidence200    return [], "No document detected (Confidence threshold not met or image is unclear)."201 202# --- GRADIO UI SETUP ---203 204# Define example paths (NOTE: Replace with actual paths if needed)205examples = [206    ["./licence2.jpg"],207    ["./licence.jpg"],208]209 210with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-size: 16px; padding: 10px 20px; border-radius: 8px;}") as demo:211    212    gr.Markdown(213        """214        <h1 style="color: #4caf50;">License Registration Classification</h1>215        <p style="font-size: 18px;">Upload an image and let the YOLO model detect and crop license documents automatically.</p>216        """217    )218    219    # 1. GUIDELINES SECTION220    with gr.Accordion("User Guidelines and Documentation", open=False):221        gr.Markdown(GUIDELINE_SETUP)222        gr.Markdown("---")223        gr.Markdown(GUIDELINE_INPUT)224        gr.Markdown("---")225        gr.Markdown(GUIDELINE_OUTPUT)226        227    gr.Markdown("---")228    229    # 2. APPLICATION INTERFACE230    with gr.Row():231        with gr.Column(scale=1, min_width=300):232            input_image = gr.Image(type="pil", label="Upload License Image", interactive=True)233            with gr.Row():234                clear_btn = gr.Button("Clear")235                submit_btn = gr.Button("Detect Document")236        237        with gr.Column(scale=2):238            output_image = gr.Gallery(label="Cropped Documents", interactive=False, object_fit="contain")239            output_text = gr.Textbox(label="Detection Result", interactive=False, lines=5)240 241    submit_btn.click(process_image, inputs=input_image, outputs=[output_image, output_text])242    clear_btn.click(lambda: (None, ""), outputs=[output_image, output_text, input_image], show_progress=False)243 244    gr.Markdown("---")245 246    # 3. EXAMPLES SECTION247    gr.Markdown("## Sample Data for Testing")248    249    gr.Examples(250        examples=examples,251        inputs=input_image,252        outputs=[output_image, output_text],253        fn=process_image,254        cache_examples=False,255        label="Click to load and run a sample detection.",256    )257 258demo.queue()259demo.launch()