Update app.py

app.py

@@ -5,23 +5,21 @@ 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 --------------------
+# ---------------- CONFIG ----------------
 
 MODEL_NAME = "nickmuchi/yolos-small-finetuned-license-plate-detection"
 BASE_AMT = 100
 
-# -------------------- DATABASE --------------------
+# ---------------- DATABASE ----------------
 
 conn = sqlite3.connect("vehicles.db", check_same_thread=False)
 cursor = conn.cursor()
@@ -44,33 +42,20 @@ CREATE TABLE IF NOT EXISTS feedback (
 
 conn.commit()
 
-# -------------------- MODEL --------------------
+# ---------------- MODEL (Lazy Load) ----------------
 
 processor = None
 model = None
 
 def load_model():
     global processor, model
-    if processor is None or model is None:
+    if processor is None:
         processor = YolosImageProcessor.from_pretrained(MODEL_NAME)
         model = YolosForObjectDetection.from_pretrained(MODEL_NAME)
         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):
-    response = requests.get(url, stream=True)
-    return Image.open(response.raw).convert("RGB")
-
-# -------------------- LOGIC --------------------
+# ---------------- LOGIC ----------------
 
 def compute_discount(vehicle_type):
     if vehicle_type == "EV":
@@ -78,27 +63,33 @@ def compute_discount(vehicle_type):
     return BASE_AMT
 
 def classify_plate_color(plate_img):
-    img = np.array(plate_img)
-    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    try:
+        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)))
+        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"
+        if green > yellow:
+            return "EV"
+        elif yellow > green:
+            return "Commercial"
+        return "Personal"
+    except:
+        return "Unknown"
 
 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]
+    try:
+        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"
+        text = pytesseract.image_to_string(
+            gray,
+            config="--psm 7 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+        )
+        return text.strip() if text.strip() else "UNKNOWN"
+    except:
+        return "UNKNOWN"
 
 def make_prediction(img):
     processor, model = load_model()
@@ -114,110 +105,126 @@ def make_prediction(img):
 
     return results[0], model.config.id2label
 
-# -------------------- VISUALIZATION --------------------
+# ---------------- VISUALIZATION ----------------
 
 def visualize(img, output, id2label, threshold):
-    keep = output["scores"] > threshold
-    boxes = output["boxes"][keep]
-    labels = output["labels"][keep]
+    try:
+        keep = output["scores"] > threshold
+        boxes = output["boxes"][keep]
+        labels = output["labels"][keep]
 
-    fig, ax = plt.subplots(figsize=(8,8))
-    ax.imshow(img)
-    results_text = []
+        fig, ax = plt.subplots(figsize=(6,6))
+        ax.imshow(img)
+        results_text = []
 
-    for box, label in zip(boxes, labels):
-        label_name = id2label[label.item()].lower()
-        if "plate" not in label_name:
-            continue
+        for box, label in zip(boxes, labels):
+            label_name = id2label[label.item()].lower()
 
-        x1,y1,x2,y2 = map(int, box.tolist())
-        plate_img = img.crop((x1,y1,x2,y2))
+            if "plate" not in label_name:
+                continue
 
-        plate = read_plate(plate_img)
-        vtype = classify_plate_color(plate_img)
-        toll = compute_discount(vtype)
+            x1,y1,x2,y2 = map(int, box.tolist())
+            plate_img = img.crop((x1,y1,x2,y2))
 
-        cursor.execute(
-            "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
-            (plate, vtype, toll)
-        )
-        conn.commit()
+            plate = read_plate(plate_img)
+            vtype = classify_plate_color(plate_img)
+            toll = compute_discount(vtype)
 
-        results_text.append(f"{plate} | {vtype} | ₹{int(toll)}")
+            cursor.execute(
+                "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
+                (plate, vtype, toll)
+            )
+            conn.commit()
 
-        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", fontsize=10)
+            results_text.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", fontsize=8)
+
+        ax.axis("off")
 
-    ax.axis("off")
+        if not results_text:
+            return fig, "No plate detected"
 
-    if not results_text:
-        return fig, "No plate detected"
+        return fig, "\n".join(results_text)
 
-    return fig, "\n".join(results_text)
+    except Exception as e:
+        return None, f"Error: {str(e)}"
 
-# -------------------- FEEDBACK --------------------
+# ---------------- 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
+
+# ---------------- FEEDBACK ----------------
 
 def submit_feedback(result_text, feedback_choice):
     if not result_text:
-        return "No result to evaluate."
+        return "No result available."
 
     cursor.execute(
         "INSERT INTO feedback VALUES (?, ?)",
         (result_text, feedback_choice)
     )
     conn.commit()
-
     return "Feedback recorded!"
 
 def show_accuracy():
     df = pd.read_sql("SELECT * FROM feedback", conn)
 
     if df.empty:
-        return "No feedback data available."
+        return "No feedback yet."
 
     correct = len(df[df["feedback"] == "Correct"])
     total = len(df)
 
     accuracy = (correct / total) * 100
-    return f"Model Accuracy (based on user feedback): {accuracy:.2f}%"
+    return f"Accuracy (User Feedback Based): {accuracy:.2f}%"
 
-# -------------------- CALLBACKS --------------------
+# ---------------- CALLBACK ----------------
 
-def detect_from_image(img, threshold):
+def detect_image(img, threshold):
     if img is None:
         return None, "No image provided"
     output, id2label = make_prediction(img)
     return visualize(img, output, id2label, threshold)
 
-# -------------------- UI --------------------
+# ---------------- UI ----------------
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🚦 Smart Vehicle Classification System")
 
     slider = gr.Slider(0.3, 1.0, 0.5, label="Confidence Threshold")
-    result_box = gr.Textbox(label="Detection Result", lines=4)
 
-    img_input = gr.Image(type="pil")
-    img_output = gr.Plot()
+    with gr.Row():
+        img_input = gr.Image(type="pil")
+        img_output = gr.Plot()
+
+    result_box = gr.Textbox(label="Detection Result", lines=4)
 
     detect_btn = gr.Button("Detect")
     detect_btn.click(
-        detect_from_image,
+        detect_image,
         inputs=[img_input, slider],
         outputs=[img_output, result_box]
     )
 
-    gr.Markdown("### Provide Feedback")
-
-    feedback_radio = gr.Radio(
-        ["Correct", "Incorrect"],
-        label="Was the prediction correct?"
-    )
-
+    gr.Markdown("### Feedback")
+    feedback_radio = gr.Radio(["Correct", "Incorrect"], label="Prediction correct?")
     feedback_btn = gr.Button("Submit Feedback")
     feedback_msg = gr.Textbox(label="Feedback Status")
 
@@ -227,14 +234,16 @@ with gr.Blocks() as demo:
         outputs=feedback_msg
     )
 
-    gr.Markdown("### Model Evaluation")
-
+    gr.Markdown("### Model Accuracy")
     accuracy_btn = gr.Button("Show Accuracy")
     accuracy_box = gr.Textbox(label="Accuracy")
 
-    accuracy_btn.click(
-        show_accuracy,
-        outputs=accuracy_box
-    )
+    accuracy_btn.click(show_accuracy, outputs=accuracy_box)
+
+    gr.Markdown("### 📊 Dashboard")
+    dashboard_plot = gr.Plot()
+
+    refresh_btn = gr.Button("Refresh Dashboard")
+    refresh_btn.click(get_dashboard, outputs=dashboard_plot)
 
 demo.launch()