Update app.py

app.py

@@ -11,6 +11,9 @@ import numpy as np
 from urllib.parse import urlparse
 from PIL import Image
 from transformers import AutoImageProcessor, YolosForObjectDetection, DetrForObjectDetection
+import easyocr
+from datetime import datetime
+
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 
@@ -39,20 +42,74 @@ def get_original_image(url_input):
         return image
 
 
-# ---------------- Model Loading ----------------
 
-def load_model(model_name):
-    processor = AutoImageProcessor.from_pretrained(model_name)
+# -------------------- 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,
+    time TEXT
+)
+""")
+conn.commit()
+
+# -------------------- 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'], gpu=False)
+
+# -------------------- Plate Color Classifier --------------------
+def classify_plate_color(plate_img):
+    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)))
+    white = np.sum(cv2.inRange(hsv, (0, 0, 200), (180, 30, 255)))
 
-    if "yolos" in model_name:
-        model = YolosForObjectDetection.from_pretrained(model_name)
-    elif "detr" in model_name:
-        model = DetrForObjectDetection.from_pretrained(model_name)
+    if green > yellow and green > white:
+        return "EV"
+    elif yellow > green and yellow > white:
+        return "Commercial"
     else:
-        raise ValueError("Unsupported model")
+        return "Personal"
+
+# -------------------- OCR --------------------
+def read_plate(plate_img):
+    results = reader.readtext(np.array(plate_img))
+    if results:
+        return results[0][1]
+    return "UNKNOWN"
+
+# -------------------- Dashboard --------------------
+def get_dashboard():
+    df = pd.read_sql("SELECT * FROM vehicles", conn)
+
+    fig, ax = plt.subplots(figsize=(8, 5))
+
+    if len(df) == 0:
+        ax.text(0.5, 0.5, "No vehicles scanned yet",
+                ha="center", va="center", fontsize=10)
+        ax.axis("off")
+        return fig
+
+    counts = df["type"].value_counts()
+    counts.plot(kind="bar", ax=ax)
+
+    ax.set_title("Vehicle Classification Dashboard")
+    ax.set_xlabel("Vehicle Type")
+    ax.set_ylabel("Count")
+    ax.grid(axis="y")
 
-    model.eval()
-    return processor, model
+    return fig
 
 
 # ---------------- Core Inference ----------------