Create app.py

app.py

@@ -0,0 +1,224 @@
+import io
+import os
+import cv2
+import gradio as gr
+import matplotlib
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt
+import requests
+import torch
+import numpy as np
+import sqlite3
+import pandas as pd
+import pytesseract
+
+from urllib.parse import urlparse
+from PIL import Image
+from transformers import YolosImageProcessor, YolosForObjectDetection
+
+# -------------------- CONFIG --------------------
+
+MODEL_NAME = "nickmuchi/yolos-small-finetuned-license-plate-detection"
+BASE_AMT = 100
+
+# -------------------- 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()
+
+# -------------------- MODEL (Lazy Load) --------------------
+
+processor = None
+model = None
+
+def load_model():
+    global processor, model
+    if processor is None or model is None:
+        processor = YolosImageProcessor.from_pretrained(MODEL_NAME)
+        model = YolosForObjectDetection.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.float32
+        )
+        model.eval()
+    return processor, model
+
+# -------------------- UTILITIES --------------------
+
+def is_valid_url(url):
+    try:
+        r = urlparse(url)
+        return all([r.scheme, r.netloc])
+    except:
+        return False
+
+def get_original_image(url):
+    return Image.open(requests.get(url, stream=True).raw).convert("RGB")
+
+# -------------------- DISCOUNT LOGIC --------------------
+
+def compute_discount(vehicle_type):
+    if vehicle_type == "EV":
+        return BASE_AMT * 0.9, "10% EV discount applied"
+    return BASE_AMT, "No discount"
+
+# -------------------- PLATE COLOR CLASSIFICATION --------------------
+
+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)))
+
+    if green > yellow:
+        return "EV"
+    elif yellow > green:
+        return "Commercial"
+    return "Personal"
+
+# -------------------- OCR (LIGHTWEIGHT) --------------------
+
+def read_plate(plate_img):
+    gray = cv2.cvtColor(np.array(plate_img), cv2.COLOR_RGB2GRAY)
+    gray = cv2.threshold(gray, 120, 255, cv2.THRESH_BINARY)[1]
+
+    text = pytesseract.image_to_string(
+        gray,
+        config="--psm 7 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+    )
+    return text.strip() if text.strip() else "UNKNOWN"
+
+# -------------------- YOLOS 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))])
+    results = processor.post_process_object_detection(
+        outputs, threshold=0.3, target_sizes=img_size
+    )
+    return results[0]
+
+# -------------------- VISUALIZATION --------------------
+
+def fig_to_img(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    img = Image.open(buf)
+    plt.close(fig)
+    return img
+
+def visualize(img, output, threshold):
+    keep = output["scores"] > threshold
+    boxes = output["boxes"][keep]
+    labels = output["labels"][keep]
+
+    plt.figure(figsize=(10,10))
+    plt.imshow(img)
+    ax = plt.gca()
+
+    results = []
+
+    for box, label in zip(boxes, labels):
+        if "plate" not in load_model()[1].config.id2label[label.item()].lower():
+            continue
+
+        x1,y1,x2,y2 = map(int, box.tolist())
+        plate_img = img.crop((x1,y1,x2,y2))
+
+        plate = read_plate(plate_img)
+        vtype = classify_plate_color(plate_img)
+        toll, msg = compute_discount(vtype)
+
+        cursor.execute(
+            "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
+            (plate, vtype, toll)
+        )
+        conn.commit()
+
+        results.append(f"{plate} | {vtype} | ₹{int(toll)}")
+
+        ax.add_patch(
+            plt.Rectangle((x1,y1), x2-x1, y2-y1, fill=False, color="red", linewidth=2)
+        )
+        ax.text(x1, y1-5, f"{plate} ({vtype})", color="yellow")
+
+    plt.axis("off")
+    return fig_to_img(plt.gcf()), "\n".join(results) if results else "No plate detected"
+
+# -------------------- DASHBOARD --------------------
+
+def get_dashboard():
+    df = pd.read_sql("SELECT * FROM vehicles", conn)
+    fig, ax = plt.subplots()
+
+    if df.empty:
+        ax.text(0.5,0.5,"No data yet",ha="center")
+        ax.axis("off")
+        return fig
+
+    df["type"].value_counts().plot(kind="bar", ax=ax)
+    ax.set_title("Vehicle Types")
+    return fig
+
+# -------------------- MAIN CALLBACK --------------------
+
+def detect(url, img, cam, threshold):
+    if url and is_valid_url(url):
+        image = get_original_image(url)
+    elif img is not None:
+        image = img
+    elif cam is not None:
+        image = cam
+    else:
+        return None, "No input"
+
+    output = make_prediction(image)
+    return visualize(image, output, threshold)
+
+# -------------------- UI --------------------
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🚦 Smart Vehicle Classification System")
+
+    result_box = gr.Textbox(label="Result", lines=4)
+    slider = gr.Slider(0.3,1.0,0.5,label="Confidence Threshold")
+
+    with gr.Tabs():
+        with gr.Tab("Image URL"):
+            url = gr.Textbox(label="Image URL")
+            out1 = gr.Image()
+            btn1 = gr.Button("Detect")
+
+        with gr.Tab("Upload"):
+            img = gr.Image(type="pil")
+            out2 = gr.Image()
+            btn2 = gr.Button("Detect")
+
+        with gr.Tab("Webcam"):
+            cam = gr.Image(source="webcam", type="pil")
+            out3 = gr.Image()
+            btn3 = gr.Button("Detect")
+
+    btn1.click(detect, [url, img, cam, slider], [out1, result_box])
+    btn2.click(detect, [url, img, cam, slider], [out2, result_box])
+    btn3.click(detect, [url, img, cam, slider], [out3, result_box])
+
+    gr.Markdown("### 📊 Dashboard")
+    gr.Plot(get_dashboard)
+
+demo.launch()