Create app.py

app.py

@@ -0,0 +1,293 @@
+# app.py
+import io
+import os
+import cv2
+import gradio as gr
+import matplotlib.pyplot as plt
+import requests
+import torch
+import pathlib
+import numpy as np
+import sqlite3
+import pandas as pd
+from urllib.parse import urlparse
+from PIL import Image
+from transformers import YolosImageProcessor, YolosForObjectDetection
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+BASE_TOLL = 100
+
+COLORS = [
+    [0.000, 0.447, 0.741],
+    [0.850, 0.325, 0.098],
+    [0.929, 0.694, 0.125],
+    [0.494, 0.184, 0.556],
+    [0.466, 0.674, 0.188],
+    [0.301, 0.745, 0.933]
+]
+
+# ---------------- Utilities ----------------
+
+def is_valid_url(url):
+    try:
+        result = urlparse(url)
+        return all([result.scheme, result.netloc])
+    except Exception:
+        return False
+
+
+def get_original_image(url_input):
+    if url_input and is_valid_url(url_input):
+        image = Image.open(requests.get(url_input, stream=True).raw).convert("RGB")
+        return image
+
+
+# -------------------- Database --------------------
+conn = sqlite3.connect("vehicles.db", check_same_thread=False)
+cursor = conn.cursor()
+cursor.execute("""
+CREATE TABLE IF NOT EXISTS vehicles (
+    plate TEXT,
+    type TEXT,
+    amount REAL,
+    time TEXT
+)
+""")
+conn.commit()
+
+# -------------------- Lazy Model --------------------
+processor = None
+model = None
+
+
+def load_model():
+    global processor, model
+    if processor is None or model is None:
+        processor = YolosImageProcessor.from_pretrained(
+            "nickmuchi/yolos-small-finetuned-license-plate-detection"
+        )
+        model = YolosForObjectDetection.from_pretrained(
+            "nickmuchi/yolos-small-finetuned-license-plate-detection",
+            use_safetensors=True,
+            torch_dtype=torch.float32
+        )
+        model.eval()
+    return processor, model
+
+
+# -------------------- Plate Color Classifier --------------------
+
+def classify_plate_color(plate_img):
+    img = np.array(plate_img)
+    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+
+    green = np.sum(cv2.inRange(hsv, (35, 40, 40), (85, 255, 255)))
+    yellow = np.sum(cv2.inRange(hsv, (15, 50, 50), (35, 255, 255)))
+    white = np.sum(cv2.inRange(hsv, (0, 0, 200), (180, 30, 255)))
+
+    if green > yellow and green > white:
+        return "EV"
+    elif yellow > green and yellow > white:
+        return "Commercial"
+    else:
+        return "Personal"
+
+
+# ---------------- Dashboard ----------------
+
+def get_dashboard():
+    df = pd.read_sql("SELECT * FROM vehicles", conn)
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+
+    if len(df) == 0:
+        ax.text(0.5, 0.5, "No vehicles scanned yet",
+                ha="center", va="center", fontsize=12)
+        ax.axis("off")
+        return fig
+
+    counts = df["type"].value_counts()
+    counts.plot(kind="bar", ax=ax)
+
+    ax.set_title("Vehicle Classification Dashboard")
+    ax.set_xlabel("Vehicle Type")
+    ax.set_ylabel("Count")
+    ax.grid(axis="y")
+
+    return fig
+
+
+# ---------------- Core Inference ----------------
+
+def make_prediction(img):
+    processor, model = load_model()
+    inputs = processor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    img_size = torch.tensor([tuple(reversed(img.size))])
+    processed_outputs = processor.post_process_object_detection(
+        outputs, threshold=0.0, target_sizes=img_size
+    )
+    return processed_outputs[0]
+
+
+def fig2img(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    pil_img = Image.open(buf)
+
+    basewidth = 750
+    wpercent = (basewidth / float(pil_img.size[0]))
+    hsize = int((float(pil_img.size[1]) * float(wpercent)))
+    img = pil_img.resize((basewidth, hsize), Image.Resampling.LANCZOS)
+
+    plt.close(fig)
+    return img
+
+
+# ---------------- Visualization ----------------
+
+def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
+    keep = output_dict["scores"] > threshold
+    boxes = output_dict["boxes"][keep].tolist()
+    scores = output_dict["scores"][keep].tolist()
+    labels = output_dict["labels"][keep].tolist()
+
+    if id2label is not None:
+        labels = [id2label[x] for x in labels]
+
+    plt.figure(figsize=(20, 20))
+    plt.imshow(img)
+    ax = plt.gca()
+    colors = COLORS * 100
+
+    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
+        if "plate" in label.lower():
+            crop = img.crop((int(xmin), int(ymin), int(xmax), int(ymax)))
+            plate_type = classify_plate_color(crop)
+
+            if plate_type == "EV":
+                amount = BASE_TOLL * 0.9
+                price_text = f"EV | ₹{amount:.0f} (10% off)"
+            else:
+                amount = BASE_TOLL
+                price_text = f"{plate_type} | ₹{amount:.0f}"
+
+            cursor.execute(
+                "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
+                ("UNKNOWN", plate_type, amount)
+            )
+            conn.commit()
+
+            ax.add_patch(
+                plt.Rectangle(
+                    (xmin, ymin), xmax - xmin, ymax - ymin,
+                    fill=False, color=color, linewidth=4
+                )
+            )
+            ax.text(
+                xmin, ymin - 10,
+                f"{price_text} | {score:0.2f}",
+                fontsize=12,
+                bbox=dict(facecolor="yellow", alpha=0.8)
+            )
+
+    plt.axis("off")
+    return fig2img(plt.gcf())
+
+
+# ---------------- Image Detection ----------------
+
+def detect_objects_image(url_input, image_input, webcam_input, threshold):
+    if url_input and is_valid_url(url_input):
+        image = get_original_image(url_input)
+    elif image_input is not None:
+        image = image_input
+    elif webcam_input is not None:
+        image = webcam_input
+    else:
+        return None, None
+
+    processed_outputs = make_prediction(image)
+    viz_img = visualize_prediction(image, processed_outputs, threshold, load_model()[1].config.id2label)
+    dashboard_fig = get_dashboard()
+
+    return viz_img, dashboard_fig
+
+
+# ---------------- UI ----------------
+
+title = """<h1 id="title">License Plate Detection + Toll Billing</h1>"""
+
+description = """
+Detect license plates using YOLOS.
+Features:
+- Image URL
+- Image Upload
+- Webcam
+- Vehicle type classification by plate color
+- EV vehicles get 10% discount
+- Billing dashboard
+"""
+
+demo = gr.Blocks()
+
+with demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    slider_input = gr.Slider(minimum=0.2, maximum=1, value=0.5, step=0.1, label='Prediction Threshold')
+
+    with gr.Tabs():
+        with gr.TabItem('Image URL'):
+            with gr.Row():
+                url_input = gr.Textbox(lines=2, label='Enter valid image URL here..')
+                original_image = gr.Image(height=400)
+                url_input.change(get_original_image, url_input, original_image)
+                img_output_from_url = gr.Image(height=400)
+            dashboard_output_url = gr.Plot()
+            url_but = gr.Button('Detect')
+
+        with gr.TabItem('Image Upload'):
+            with gr.Row():
+                img_input = gr.Image(type='pil', height=400)
+                img_output_from_upload = gr.Image(height=400)
+            dashboard_output_upload = gr.Plot()
+            img_but = gr.Button('Detect')
+
+        with gr.TabItem('WebCam'):
+            with gr.Row():
+                web_input = gr.Image(
+                    sources=["webcam"],
+                    type="pil",
+                    height=400,
+                    streaming=True
+                )
+                img_output_from_webcam = gr.Image(height=400)
+            dashboard_output_webcam = gr.Plot()
+            cam_but = gr.Button('Detect')
+
+    url_but.click(
+        detect_objects_image,
+        inputs=[url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_url, dashboard_output_url]
+    )
+
+    img_but.click(
+        detect_objects_image,
+        inputs=[url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_upload, dashboard_output_upload]
+    )
+
+    cam_but.click(
+        detect_objects_image,
+        inputs=[url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_webcam, dashboard_output_webcam]
+    )
+
+
+demo.queue()
+demo.launch(debug=True, ssr_mode=False)