Create app.py

app.py

@@ -0,0 +1,151 @@
+# -----------------------------
+# Full Vehicle Detection & Dashboard
+# -----------------------------
+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 datetime import datetime
+
+# -----------------------------
+# Load Models
+# -----------------------------
+processor = YolosImageProcessor.from_pretrained(
+    "nickmuchi/yolos-small-finetuned-license-plate-detection"
+)
+model = YolosForObjectDetection.from_pretrained(
+    "nickmuchi/yolos-small-finetuned-license-plate-detection"
+)
+model.eval()
+
+reader = easyocr.Reader(['en'])
+
+# -----------------------------
+# Database Setup
+# -----------------------------
+conn = sqlite3.connect("vehicle_data.db", check_same_thread=False)
+cursor = conn.cursor()
+cursor.execute('''
+CREATE TABLE IF NOT EXISTS vehicles (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    timestamp TEXT,
+    license_plate 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:
+        return "EV"
+    elif y > g and y > w:
+        return "Commercial"
+    else:
+        return "Personal"
+
+# -----------------------------
+# Process Image
+# -----------------------------
+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
+    )[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))
+    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
+# -----------------------------
+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.")
+    
+    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])
+
+demo.launch()