Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import io
 import os
 import cv2
@@ -6,7 +5,6 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import requests
 import torch
-import pathlib
 import numpy as np
 import sqlite3
 import pandas as pd
@@ -16,23 +14,17 @@ from transformers import YolosImageProcessor, YolosForObjectDetection
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 
-base_amt = 100
-#--- calculate discount 
+MODEL_NAME = "nickmuchi/yolos-small-finetuned-license-plate-detection"
+BASE_AMT = 100
+
+# -------------------- Discount --------------------
+
 def compute_discount(vehicle_type):
     if vehicle_type == "EV":
-        return base_amt * 0.9, "10% discount applied (EV)"
-    return toll_parking_amt, "No discount"
-#------------
-COLORS = [
-    [0.000, 0.447, 0.741],
-    [0.850, 0.325, 0.098],
-    [0.929, 0.694, 0.125],
-    [0.494, 0.184, 0.556],
-    [0.466, 0.674, 0.188],
-    [0.301, 0.745, 0.933]
-]
-
-# ---------------- Utilities ----------------
+        return BASE_AMT * 0.9, "10% discount applied (EV)"
+    return BASE_AMT, "No discount"
+
+# -------------------- Utilities --------------------
 
 def is_valid_url(url):
     try:
@@ -49,6 +41,7 @@ def get_original_image(url_input):
 
 
 # -------------------- Database --------------------
+
 conn = sqlite3.connect("vehicles.db", check_same_thread=False)
 cursor = conn.cursor()
 cursor.execute("""
@@ -61,7 +54,9 @@ CREATE TABLE IF NOT EXISTS vehicles (
 """)
 conn.commit()
 
+
 # -------------------- Lazy Model --------------------
+
 processor = None
 model = None
 
@@ -69,11 +64,9 @@ model = None
 def load_model():
     global processor, model
     if processor is None or model is None:
-        processor = YolosImageProcessor.from_pretrained(
-            "nickmuchi/yolos-small-finetuned-license-plate-detection"
-        )
+        processor = YolosImageProcessor.from_pretrained(MODEL_NAME)
         model = YolosForObjectDetection.from_pretrained(
-            "nickmuchi/yolos-small-finetuned-license-plate-detection",
+            MODEL_NAME,
             use_safetensors=True,
             torch_dtype=torch.float32
         )
@@ -99,11 +92,10 @@ def classify_plate_color(plate_img):
         return "Personal"
 
 
-# ---------------- Dashboard ----------------
+# -------------------- Dashboard --------------------
 
 def get_dashboard():
     df = pd.read_sql("SELECT * FROM vehicles", conn)
-
     fig, ax = plt.subplots(figsize=(7, 5))
 
     if len(df) == 0:
@@ -113,21 +105,16 @@ def get_dashboard():
         return fig
 
     counts = df["type"].value_counts()
-
-    # Use bar chart instead of line for categorical data
     counts.plot(kind="bar", ax=ax, color="steelblue")
 
     ax.set_title("Vehicle Classification Dashboard", fontsize=12)
     ax.set_xlabel("Vehicle Type", fontsize=10)
     ax.set_ylabel("Count", fontsize=10)
 
-    # Ensure labels are fully visible
     ax.set_xticks(range(len(counts.index)))
     ax.set_xticklabels(counts.index, rotation=0, ha="center")
-
     ax.grid(axis="y", linestyle="--", alpha=0.6)
 
-    # Add value labels on top of bars
     for i, v in enumerate(counts.values):
         ax.text(i, v + 0.05, str(v), ha="center", va="bottom", fontsize=10)
 
@@ -135,7 +122,7 @@ def get_dashboard():
     return fig
 
 
-# ---------------- Core Inference ----------------
+# -------------------- YOLOS Inference --------------------
 
 def make_prediction(img):
     processor, model = load_model()
@@ -145,7 +132,7 @@ def make_prediction(img):
 
     img_size = torch.tensor([tuple(reversed(img.size))])
     processed_outputs = processor.post_process_object_detection(
-        outputs, threshold=0.0, target_sizes=img_size
+        outputs, threshold=0.3, target_sizes=img_size
     )
     return processed_outputs[0]
 
@@ -165,7 +152,14 @@ def fig2img(fig):
     return img
 
 
-# ---------------- Visualization ----------------
+# -------------------- OCR Stub --------------------
+
+def read_plate(crop):
+    # Placeholder OCR logic
+    return "KA01AB1234"
+
+
+# -------------------- Visualization --------------------
 
 def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     keep = output_dict["scores"] > threshold
@@ -179,11 +173,10 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     plt.figure(figsize=(20, 20))
     plt.imshow(img)
     ax = plt.gca()
-    colors = COLORS * 100
 
     result_lines = []
 
-    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
+    for score, (xmin, ymin, xmax, ymax), label in zip(scores, boxes, labels):
         if "plate" in label.lower():
             crop = img.crop((int(xmin), int(ymin), int(xmax), int(ymax)))
 
@@ -191,6 +184,12 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
             vehicle_type = classify_plate_color(crop)
             toll, discount_msg = compute_discount(vehicle_type)
 
+            cursor.execute(
+                "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
+                (plate_text, vehicle_type, toll)
+            )
+            conn.commit()
+
             result_lines.append(
                 f"License: {plate_text} | Type: {vehicle_type} | Toll: ₹{int(toll)} | {discount_msg}"
             )
@@ -198,7 +197,7 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
             ax.add_patch(
                 plt.Rectangle(
                     (xmin, ymin), xmax - xmin, ymax - ymin,
-                    fill=False, color=color, linewidth=4
+                    fill=False, color="red", linewidth=3
                 )
             )
 
@@ -219,10 +218,10 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
 
     return final_img, result_text
 
-# ---------------- Image Detection ----------------
 
-def detect_objects_image(model_name, url_input, image_input, webcam_input, threshold):
-    processor, model = load_model(model_name)
+# -------------------- Gradio Callback --------------------
+
+def detect_objects_image(url_input, image_input, webcam_input, threshold):
 
     if url_input and is_valid_url(url_input):
         image = get_original_image(url_input)
@@ -233,73 +232,58 @@ def detect_objects_image(model_name, url_input, image_input, webcam_input, thres
     else:
         return None, "No image provided."
 
-    processed_outputs = make_prediction(image, processor, model)
+    processed_outputs = make_prediction(image)
 
     viz_img, result_text = visualize_prediction(
-        image, processed_outputs, threshold, model.config.id2label
+        image,
+        processed_outputs,
+        threshold,
+        load_model()[1].config.id2label
     )
 
     return viz_img, result_text
 
 
-# ---------------- UI ----------------
-
-title = """<h1 id="title">Smart Vehicle classification</h1>"""
+# -------------------- UI --------------------
 
+title = "<h1>🚦 Smart Vehicle Classification</h1>"
 description = """
 Detect license plates using YOLOS.
 Features:
-- Image URL, Image Upload, Webcam,Vehicle type classification by plate color
-- EV vehicles get 10% discount on Tolls, Tax, parking
+- Image URL, Image Upload, Webcam
+- Vehicle type classification by plate color
+- EV vehicles get 10% discount on Toll / Parking
 """
-result_box = gr.Textbox(
-    label="Detection Result",
-    lines=5,
-    interactive=False
-)
-demo = gr.Blocks()
-  
-with demo:
-    debug=False,
-    share=False,
-    ssr_mode=False
-    gr.Markdown(title)
-    gr.Markdown(description)
-    options = gr.Dropdown(
-    choices=model,
-    label="Object Detection Model",
-    value=model[0]
-    )
-
-    url_input = gr.Textbox(label="Image URL")
-    img_input = gr.Image(type="pil", label="Upload Image")
-    web_input = gr.Image(source="webcam", type="pil", label="Webcam Input")
-    slider_input = gr.Slider(0, 1, value=0.5, step=0.05, label="Confidence Threshold")
 
-    img_output_from_url = gr.Image(label="Detection Output")
+with gr.Blocks() as demo:
 
-    detect_btn = gr.Button("Detect")
+    gr.Markdown(title)
+    gr.Markdown(description)
 
-    slider_input = gr.Slider(minimum=0.2, maximum=1, value=0.5, step=0.1, label='Prediction Threshold')
+    result_box = gr.Textbox(label="Detection Result", lines=5)
 
     with gr.Tabs():
-        with gr.TabItem('Image URL'):
+
+        with gr.TabItem("Image URL"):
             with gr.Row():
-                url_input = gr.Textbox(lines=2, label='Enter valid image URL here..')
+                url_input = gr.Textbox(lines=2, label="Enter Image URL")
                 original_image = gr.Image(height=200)
                 url_input.change(get_original_image, url_input, original_image)
+
                 img_output_from_url = gr.Image(height=200)
+
             dashboard_output_url = gr.Plot()
-            url_but = gr.Button('Detect')
+            url_but = gr.Button("Detect")
 
-        with gr.TabItem('Image Upload'):
+        with gr.TabItem("Image Upload"):
             with gr.Row():
-                img_input = gr.Image(type='pil', height=200)
+                img_input = gr.Image(type="pil", height=200)
                 img_output_from_upload = gr.Image(height=200)
+
             dashboard_output_upload = gr.Plot()
-            img_but = gr.Button('Detect')
+            img_but = gr.Button("Detect")
 
-        with gr.TabItem('WebCam'):
+        with gr.TabItem("Webcam"):
             with gr.Row():
                 web_input = gr.Image(
                     sources=["webcam"],
@@ -308,42 +292,36 @@ with demo:
                     streaming=True
                 )
                 img_output_from_webcam = gr.Image(height=200)
+
             dashboard_output_webcam = gr.Plot()
-            cam_but = gr.Button('Detect')
+            cam_but = gr.Button("Detect")
+
+    slider_input = gr.Slider(0.2, 1.0, value=0.5, step=0.05, label="Confidence Threshold")
 
     url_but.click(
         detect_objects_image,
-        inputs=[options, url_input, img_input, web_input, slider_input],
-        outputs=[img_output_from_url],
+        inputs=[url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_url, result_box],
         queue=True
     )
 
     img_but.click(
         detect_objects_image,
-        inputs=[options, url_input, img_input, web_input, slider_input],
-        outputs=[img_output_from_upload],
+        inputs=[url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_upload, result_box],
         queue=True
     )
 
     cam_but.click(
         detect_objects_image,
-        inputs=[options, url_input, img_input, web_input, slider_input],
-        outputs=[img_output_from_webcam],
+        inputs=[url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_webcam, result_box],
         queue=True
     )
 
-#    vid_but.click(
-#       detect_objects_video,
-#        inputs=[video_input, slider_input],
-#        outputs=[video_output],
-#       queue=True
-#   )
+    url_but.click(get_dashboard, outputs=dashboard_output_url)
+    img_but.click(get_dashboard, outputs=dashboard_output_upload)
+    cam_but.click(get_dashboard, outputs=dashboard_output_webcam)
 
 demo.queue()
-import asyncio
-
-try:
-    asyncio.get_running_loop()
-except RuntimeError:
-    asyncio.set_event_loop(asyncio.new_event_loop())
 demo.launch(debug=True, ssr_mode=False)