Update app.py

app.py

@@ -1,21 +1,15 @@
-# -----------------------------
-# Full Vehicle Detection & Dashboard
-# -----------------------------
+import gradio as gr
 import torch
-from transformers import YolosImageProcessor, YolosForObjectDetection
-from PIL import Image, ImageDraw
-import easyocr
 import numpy as np
 import cv2
 import sqlite3
-import pandas as pd
-import matplotlib.pyplot as plt
-import gradio as gr
+from PIL import Image, ImageDraw
+from transformers import YolosImageProcessor, YolosForObjectDetection
 from datetime import datetime
 
-# -----------------------------
-# Load Models
-# -----------------------------
+# -------------------------------
+# Load YOLOS (lightweight)
+# -------------------------------
 processor = YolosImageProcessor.from_pretrained(
     "nickmuchi/yolos-small-finetuned-license-plate-detection"
 )
@@ -24,128 +18,107 @@ model = YolosForObjectDetection.from_pretrained(
 )
 model.eval()
 
-reader = easyocr.Reader(['en'])
-
-# -----------------------------
-# Database Setup
-# -----------------------------
-conn = sqlite3.connect("vehicle_data.db", check_same_thread=False)
+# -------------------------------
+# SQLite (SAFE)
+# -------------------------------
+conn = sqlite3.connect("vehicles.db", check_same_thread=False)
 cursor = conn.cursor()
-cursor.execute('''
+cursor.execute("""
 CREATE TABLE IF NOT EXISTS vehicles (
     id INTEGER PRIMARY KEY AUTOINCREMENT,
-    timestamp TEXT,
-    license_plate TEXT,
+    time TEXT,
     vehicle_type TEXT
 )
-''')
+""")
 conn.commit()
 
-# -----------------------------
+# -------------------------------
 # Plate Color Classifier
-# -----------------------------
+# -------------------------------
 def classify_plate_color(plate_img):
     img = np.array(plate_img)
     hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
-    green = cv2.inRange(hsv, (35, 40, 40), (85, 255, 255))
-    yellow = cv2.inRange(hsv, (15, 50, 50), (35, 255, 255))
-    white = cv2.inRange(hsv, (0, 0, 200), (180, 30, 255))
-    g = np.sum(green)
-    y = np.sum(yellow)
-    w = np.sum(white)
-    if g > y and g > w:
+
+    green = cv2.inRange(hsv, (35,40,40), (85,255,255))
+    yellow = cv2.inRange(hsv, (15,50,50), (35,255,255))
+
+    if np.sum(green) > np.sum(yellow):
         return "EV"
-    elif y > g and y > w:
+    elif np.sum(yellow) > 0:
         return "Commercial"
-    else:
-        return "Personal"
-
-# -----------------------------
-# Process Image
-# -----------------------------
+    return "Personal"
+
+# -------------------------------
+# Dashboard (NO matplotlib)
+# -------------------------------
+def generate_dashboard():
+    cursor.execute("SELECT vehicle_type, COUNT(*) FROM vehicles GROUP BY vehicle_type")
+    data = cursor.fetchall()
+
+    canvas = np.ones((250, 400, 3), dtype=np.uint8) * 255
+    y = 40
+    for vtype, count in data:
+        cv2.putText(
+            canvas,
+            f"{vtype}: {count}",
+            (40, y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.9,
+            (0,0,0),
+            2
+        )
+        y += 50
+    return canvas
+
+# -------------------------------
+# Main Pipeline
+# -------------------------------
 def process_image(img):
     image = Image.fromarray(img)
     draw = ImageDraw.Draw(image)
-    
+
     inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
 
     target_sizes = torch.tensor([[image.size[1], image.size[0]]])
     results = processor.post_process_object_detection(
-        outputs, threshold=0.3, target_sizes=target_sizes
+        outputs, threshold=0.4, target_sizes=target_sizes
     )[0]
 
     if len(results["boxes"]) == 0:
-        return image, "No license plate detected", update_dashboard()
-
-    # Assume first detected plate
-    box = results["boxes"][0].tolist()
-    x1, y1, x2, y2 = map(int, box)
-    plate_crop = image.crop((x1, y1, x2, y2))
-
-    # OCR
-    ocr_result = reader.readtext(np.array(plate_crop))
-    license_plate = ocr_result[0][1].replace(" ", "") if ocr_result else "UNKNOWN"
-
-    # Vehicle Type
-    vehicle_type = classify_plate_color(plate_crop)
-
-    # Draw rectangle + label
-    draw.rectangle([x1, y1, x2, y2], outline="yellow", width=3)
-    draw.text((x1, y1 - 20), f"{vehicle_type} | {license_plate}", fill="black")
-
-    # -----------------------------
-    # Insert into Database
-    # -----------------------------
-    cursor.execute('''
-    INSERT INTO vehicles (timestamp, license_plate, vehicle_type)
-    VALUES (?, ?, ?)
-    ''', (datetime.now().strftime("%Y-%m-%d %H:%M:%S"), license_plate, vehicle_type))
+        return image, "No Plate Detected", generate_dashboard()
+
+    x1, y1, x2, y2 = map(int, results["boxes"][0].tolist())
+    plate = image.crop((x1, y1, x2, y2))
+
+    vehicle_type = classify_plate_color(plate)
+
+    draw.rectangle([x1,y1,x2,y2], outline="green", width=3)
+    draw.text((x1,y1-10), vehicle_type, fill="black")
+
+    cursor.execute(
+        "INSERT INTO vehicles (time, vehicle_type) VALUES (?,?)",
+        (datetime.now().isoformat(), vehicle_type)
+    )
     conn.commit()
 
-    # -----------------------------
-    # Update Dashboard
-    # -----------------------------
-    return image, f"{vehicle_type} | {license_plate}", update_dashboard()
-
-# -----------------------------
-# Update Dashboard
-# -----------------------------
-def update_dashboard():
-    df = pd.read_sql_query("SELECT * FROM vehicles", conn)
-    if df.empty:
-        return "No vehicles detected yet"
-
-    counts = df['vehicle_type'].value_counts()
-    plt.figure(figsize=(5,3))
-    counts.plot(kind='bar', color=['green','yellow','blue'])
-    plt.title("Vehicle Type Counts")
-    plt.ylabel("Number of Vehicles")
-    plt.xlabel("Type")
-    plt.tight_layout()
-    
-    buf = np.zeros((1,1))  # placeholder
-    plt.savefig("dashboard.png")
-    plt.close()
-    
-    return "dashboard.png"
-
-# -----------------------------
-# Gradio Interface
-# -----------------------------
+    return image, vehicle_type, generate_dashboard()
+
+# -------------------------------
+# Gradio UI
+# -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🚗 Smart Vehicle Classification & EV Dashboard")
-    gr.Markdown("Upload or use webcam to detect vehicles. System logs each vehicle, reads license plate, classifies type, and updates dashboard.")
-    
+    gr.Markdown("## 🚗 Vehicle Classification & Live Dashboard")
+
     with gr.Row():
-        input_img = gr.Image(type="numpy", source="upload")
-        output_img = gr.Image()
-    
-    result_text = gr.Textbox(label="Detected Vehicle")
-    dashboard_img = gr.Image(label="Dashboard")
-    
-    detect_btn = gr.Button("Detect Vehicle")
-    detect_btn.click(process_image, inputs=input_img, outputs=[output_img, result_text, dashboard_img])
+        inp = gr.Image(type="numpy", sources=["upload", "webcam"])
+        out = gr.Image()
+
+    result = gr.Textbox(label="Vehicle Type")
+    dashboard = gr.Image(label="Live Dashboard")
+
+    btn = gr.Button("Detect Vehicle")
+    btn.click(process_image, inp, [out, result, dashboard])
 
 demo.launch()