Update app.py

app.py

@@ -7,9 +7,9 @@ from PIL import Image, ImageDraw
 from transformers import YolosImageProcessor, YolosForObjectDetection
 from datetime import datetime
 
-# -------------------------------
-# Load YOLOS (lightweight)
-# -------------------------------
+# -----------------------------
+# Load Model
+# -----------------------------
 processor = YolosImageProcessor.from_pretrained(
     "nickmuchi/yolos-small-finetuned-license-plate-detection"
 )
@@ -18,23 +18,24 @@ model = YolosForObjectDetection.from_pretrained(
 )
 model.eval()
 
-# -------------------------------
-# SQLite (SAFE)
-# -------------------------------
+# -----------------------------
+# Database
+# -----------------------------
 conn = sqlite3.connect("vehicles.db", check_same_thread=False)
 cursor = conn.cursor()
 cursor.execute("""
 CREATE TABLE IF NOT EXISTS vehicles (
     id INTEGER PRIMARY KEY AUTOINCREMENT,
     time TEXT,
+    plate_id 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)
@@ -48,77 +49,23 @@ def classify_plate_color(plate_img):
         return "Commercial"
     return "Personal"
 
-# -------------------------------
-# Dashboard (NO matplotlib)
-# -------------------------------
-def generate_dashboard():
-    cursor.execute("SELECT vehicle_type, COUNT(*) FROM vehicles GROUP BY vehicle_type")
+# -----------------------------
+# Bar Chart (OpenCV)
+# -----------------------------
+def generate_bar_chart():
+    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.4, target_sizes=target_sizes
-    )[0]
-
-    if len(results["boxes"]) == 0:
-        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)
+    chart = np.ones((300, 400, 3), dtype=np.uint8) * 255
+    cv2.putText(chart, "Vehicle Count Report", (60,30),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,0), 2)
 
-    draw.rectangle([x1,y1,x2,y2], outline="green", width=3)
-    draw.text((x1,y1-10), vehicle_type, fill="black")
+    max_count = max([c for _, c in data], default=1)
+    x = 70
 
-    cursor.execute(
-        "INSERT INTO vehicles (time, vehicle_type) VALUES (?,?)",
-        (datetime.now().isoformat(), vehicle_type)
-    )
-    conn.commit()
-
-    return image, vehicle_type, generate_dashboard()
-
-# -------------------------------
-# Gradio UI
-# -------------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## 🚗 Vehicle Classification & Live Dashboard")
-
-    with gr.Row():
-        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()
+    for vtype, count in data:
+        bar_height = int((count / max_count) * 180)
+        cv2.rectangle(chart, (x,2