Update app.py

app.py

@@ -1,251 +1,249 @@
+import io
+import cv2
 import gradio as gr
+import matplotlib
+matplotlib.use("Agg")
+
 import matplotlib.pyplot as plt
+import torch
+import numpy as np
+import sqlite3
 import pandas as pd
-import os
-import random
-from PIL import ImageDraw
-
-# ===============================
-# GLOBAL VARIABLES
-# ===============================
-total_vehicles = 0
-ev_count = 0
-total_co2_saved = 0
-
-DISTANCE_KM = 10
-CO2_PER_KM_PETROL = 0.150
-CO2_SAVED_PER_EV = DISTANCE_KM * CO2_PER_KM_PETROL
-
-FEEDBACK_FILE = "feedback.csv"
-
-if not os.path.exists(FEEDBACK_FILE):
-    pd.DataFrame(columns=["Predicted_Label", "Feedback"]).to_csv(FEEDBACK_FILE, index=False)
-
-# ===============================
-# SIMULATED NUMBER PLATE COLOUR DETECTION
-# Replace with real color detection later
-# ===============================
+import pytesseract
+
+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
+)
+""")
+
+cursor.execute("""
+CREATE TABLE IF NOT EXISTS feedback (
+    result TEXT,
+    feedback TEXT
+)
+""")
+
+conn.commit()
+
+# ---------------- MODEL (Lazy Load) ----------------
+
+processor = None
+model = None
+
+def load_model():
+    global processor, model
+    if processor is None:
+        processor = YolosImageProcessor.from_pretrained(MODEL_NAME)
+        model = YolosForObjectDetection.from_pretrained(MODEL_NAME)
+        model.eval()
+    return processor, model
+
+# ---------------- LOGIC ----------------
+
+def compute_discount(vehicle_type):
+    if vehicle_type == "EV":
+        return BASE_AMT * 0.9
+    return BASE_AMT
+
 def classify_plate_color(plate_img):
     try:
         img = np.array(plate_img)
-        img = cv2.GaussianBlur(img, (5,5), 0)
         hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
 
-        green_mask = cv2.inRange(hsv, (35, 40, 40), (85, 255, 255))
-        yellow_mask = cv2.inRange(hsv, (15, 50, 50), (35, 255, 255))
-        white_mask = cv2.inRange(hsv, (0, 0, 200), (180, 40, 255))
-
-        green = np.sum(green_mask)
-        yellow = np.sum(yellow_mask)
-        white = np.sum(white_mask)
-
-        if green > yellow and green > white:
-            return "Electric Vehicle (Green Plate)"
-        elif yellow > green and yellow > white:
-            return "Commercial Vehicle (Yellow Plate)"
-        elif white > green and white > yellow:
-            return "Private Vehicle (White Plate)"
-        else:
-            return "Unknown Classification"
+        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"
     except:
-        return "Unknown Classification"
+        return "Unknown"
 
+def read_plate(plate_img):
+    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"
+    except:
+        return "UNKNOWN"
 
+def make_prediction(img):
+    processor, model = load_model()
+    inputs = processor(images=img, return_tensors="pt")
 
-def get_vehicle_type_from_colour(colour):
-    mapping = {
-        "White": "Private Vehicle",
-        "Yellow": "Commercial Vehicle",
-        "Green": "Electric Vehicle",
-        "Black": "Rental Vehicle"
-    }
-    return mapping.get(colour, "Unknown")
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-# ===============================
-# DASHBOARD
-# ===============================
-def generate_dashboard():
-    global total_vehicles, ev_count
-    non_ev = total_vehicles - ev_count
+    img_size = torch.tensor([img.size[::-1]])
+    results = processor.post_process_object_detection(
+        outputs, threshold=0.3, target_sizes=img_size
+    )
 
-    fig, ax = plt.subplots()
-    ax.bar(["EV", "Non-EV"], [ev_count, non_ev])
-    ax.set_title("Vehicle Distribution")
-    ax.set_ylabel("Count")
-    return fig
+    return results[0], model.config.id2label
 
-# ===============================
-# EVALUATION SUMMARY
-# ===============================
-def get_evaluation_summary():
-    df = pd.read_csv(FEEDBACK_FILE)
-    total = len(df)
-    correct = len(df[df["Feedback"] == "Correct"])
-    incorrect = len(df[df["Feedback"] == "Incorrect"])
+# ---------------- VISUALIZATION ----------------
 
-    if total == 0:
-        return "Evaluation Summary:\nNo feedback yet."
+def visualize(img, output, id2label, threshold):
+    try:
+        keep = output["scores"] > threshold
+        boxes = output["boxes"][keep]
+        labels = output["labels"][keep]
 
-    precision = correct / total
+        fig, ax = plt.subplots(figsize=(6,6))
+        ax.imshow(img)
+        results_text = []
 
-    return f"""Evaluation Summary
-Total: {total}
-Correct: {correct}
-Incorrect: {incorrect}
-Accuracy: {precision:.2f}
-"""
+        for box, label in zip(boxes, labels):
+            label_name = id2label[label.item()].lower()
 
-# ===============================
-# DETECTION FUNCTION
-# ===============================
-def detect_image(image):
+            if "plate" not in label_name:
+                continue
 
-    global total_vehicles, ev_count, total_co2_saved
+            x1,y1,x2,y2 = map(int, box.tolist())
+            plate_img = img.crop((x1,y1,x2,y2))
 
-    if image is None:
-        return None, "Upload image first.", "", "", "", "", "", ""
+            plate = read_plate(plate_img)
+            vtype = classify_plate_color(plate_img)
+            toll = compute_discount(vtype)
 
-    plate_colour = detect_plate_colour()
-    vehicle_type = classify_plate_color(plate_colour)
-    plate_number = "KA01AB1234"
+            cursor.execute(
+                "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
+                (plate, vtype, toll)
+            )
+            conn.commit()
 
-    total_vehicles += 1
-    co2_saved_this = 0
+            results_text.append(f"{plate} | {vtype} | ₹{int(toll)}")
 
-    if vehicle_type == "Electric Vehicle":
-        ev_count += 1
-        co2_saved_this = CO2_SAVED_PER_EV
-        total_co2_saved += co2_saved_this
+            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)
 
-    ev_percent = (ev_count / total_vehicles) * 100
+        ax.axis("off")
 
-    # Draw boxes
-    img = image.copy()
-    draw = ImageDraw.Draw(img)
-    w, h = img.size
+        if not results_text:
+            return fig, "No plate detected"
 
-    vehicle_box = [w*0.1, h*0.2, w*0.9, h*0.8]
-    plate_box = [w*0.4, h*0.6, w*0.7, h*0.75]
+        return fig, "\n".join(results_text)
 
-    draw.rectangle(vehicle_box, outline="red", width=4)
-    draw.text((vehicle_box[0], vehicle_box[1]-20),
-              f"{vehicle_type}", fill="red")
+    except Exception as e:
+        return None, f"Error: {str(e)}"
 
-    draw.rectangle(plate_box, outline="green", width=4)
-    draw.text((plate_box[0], plate_box[1]-20),
-              f"{plate_number} ({plate_colour})", fill="green")
+# ---------------- DASHBOARD ----------------
 
+def get_dashboard():
+    df = pd.read_sql("SELECT * FROM vehicles", conn)
     fig, ax = plt.subplots()
-    ax.imshow(img)
-    ax.axis("off")
-
-    result_text = f"""
-Vehicle Type: {vehicle_type}
-Plate Colour: {plate_colour}
-Plate Number: {plate_number}
-CO₂ Saved: {co2_saved_this:.2f} kg
-"""
-
-    total_card = f"### 🚘 Total: {total_vehicles}"
-    ev_card = f"### ⚡ EV: {ev_count}"
-    percent_card = f"### 📊 EV Rate: {ev_percent:.2f}%"
-    co2_card = f"### 🌱 CO₂ Saved: {total_co2_saved:.2f} kg"
-
-    dashboard_fig = generate_dashboard()
-
-    summary = get_evaluation_summary()
-
-    return fig, result_text, total_card, ev_card, percent_card, co2_card, dashboard_fig, vehicle_type, summary
-
-# ===============================
-# FEEDBACK SAVE
-# ===============================
-def save_feedback(predicted_label, feedback):
-    df = pd.read_csv(FEEDBACK_FILE)
-    df = pd.concat([df, pd.DataFrame(
-        [{"Predicted_Label": predicted_label, "Feedback": feedback}]
-    )], ignore_index=True)
-    df.to_csv(FEEDBACK_FILE, index=False)
-    return "Saved!", get_evaluation_summary()
-
-# ===============================
-# RESET DATABASE
-# ===============================
-def reset_database():
-    pd.DataFrame(columns=["Predicted_Label", "Feedback"]).to_csv(FEEDBACK_FILE, index=False)
-    return "Database Reset!", "Evaluation Summary:\nNo feedback yet."
-
-# ===============================
-# UI
-# ===============================
-with gr.Blocks() as demo:
 
-    gr.Markdown("## 🚦 Smart Vehicle Detection Based on Number Plate Colour")
+    if df.empty:
+        ax.text(0.5,0.5,"No data yet",ha="center")
+        ax.axis("off")
+        return fig
 
-    with gr.Row():
-        img_input = gr.Image(type="pil", label="Upload Vehicle Image")
-        img_output = gr.Plot()
+    df["type"].value_counts().plot(kind="bar", ax=ax)
+    ax.set_title("Vehicle Types")
+    return fig
 
-    detect_btn = gr.Button("Detect", size="sm")
+# ---------------- FEEDBACK ----------------
 
-    result_box = gr.Textbox(label="Detection Result", lines=5)
+def submit_feedback(result_text, feedback_choice):
+    if not result_text:
+        return "No result available."
 
-    with gr.Row():
-        total_card = gr.Markdown("### 🚘 Total: 0")
-        ev_card = gr.Markdown("### ⚡ EV: 0")
-        percent_card = gr.Markdown("### 📊 EV Rate: 0%")
-        co2_card = gr.Markdown("### 🌱 CO₂ Saved: 0 kg")
+    cursor.execute(
+        "INSERT INTO feedback VALUES (?, ?)",
+        (result_text, feedback_choice)
+    )
+    conn.commit()
+    return "Feedback recorded!"
 
-    dashboard_plot = gr.Plot()
+def show_accuracy():
+    df = pd.read_sql("SELECT * FROM feedback", conn)
 
-    predicted_label_state = gr.State()
+    if df.empty:
+        return "No feedback yet."
 
-    detect_btn.click(
-        fn=detect_image,
-        inputs=[img_input],
-        outputs=[
-            img_output,
-            result_box,
-            total_card,
-            ev_card,
-            percent_card,
-            co2_card,
-            dashboard_plot,
-            predicted_label_state,
-            gr.Markdown()  # evaluation summary placeholder
-        ]
-    )
+    correct = len(df[df["feedback"] == "Correct"])
+    total = len(df)
 
-    # -------------------------
-    # FEEDBACK + EVALUATION ROW
-    # -------------------------
-    with gr.Row():
+    accuracy = (correct / total) * 100
+    return f"Accuracy (User Feedback Based): {accuracy:.2f}%"
 
-        with gr.Column(scale=1):
-            correct_btn = gr.Button("✔", size="sm")
-            incorrect_btn = gr.Button("✘", size="sm")
-            feedback_status = gr.Textbox(label="Feedback", lines=1)
+# ---------------- CALLBACK ----------------
 
-        with gr.Column(scale=1):
-            evaluation_summary = gr.Textbox(label="Evaluation Summary", lines=6)
+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)
 
-    correct_btn.click(
-        fn=save_feedback,
-        inputs=[predicted_label_state, gr.State("Correct")],
-        outputs=[feedback_status, evaluation_summary]
-    )
+# ---------------- UI ----------------
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Smart Vehicle Classification System")
 
-    incorrect_btn.click(
-        fn=save_feedback,
-        inputs=[predicted_label_state, gr.State("Incorrect")],
-        outputs=[feedback_status, evaluation_summary]
+    slider = gr.Slider(0.3, 1.0, 0.5, label="Confidence Threshold")
+
+    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_image,
+        inputs=[img_input, slider],
+        outputs=[img_output, result_box]
     )
 
-    reset_btn = gr.Button("Reset Database", size="sm")
-    reset_btn.click(
-        fn=reset_database,
-        outputs=[feedback_status, evaluation_summary]
+    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")
+
+    feedback_btn.click(
+        submit_feedback,
+        inputs=[result_box, feedback_radio],
+        outputs=feedback_msg
     )
 
-demo.launch(theme=gr.themes.Soft())
+    gr.Markdown("### Model Accuracy")
+    accuracy_btn = gr.Button("Show Accuracy")
+    accuracy_box = gr.Textbox(label="Accuracy")
+
+    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()