Update app.py

app.py

@@ -1,4 +1,5 @@
-# app.py 
+# app.py (Clean Final Version for HF Spaces)
+
 import io
 import os
 import cv2
@@ -6,28 +7,26 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import requests
 import torch
-import pathlib
 import numpy as np
 from urllib.parse import urlparse
-from transformers import AutoImageProcessor, YolosForObjectDetection, DetrForObjectDetection
-import sqlite3
-import pandas as pd
-import matplotlib.pyplot as plt
-from PIL import Image, ImageDraw
+from PIL import Image
 from transformers import YolosImageProcessor, YolosForObjectDetection
 import easyocr
-from datetime import datetime
-
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 
+# ---------------- Globals (lazy loaded) ----------------
+processor = None
+model = None
+reader = easyocr.Reader(["en"], gpu=False)
+
 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]
+    [0.301, 0.745, 0.933],
 ]
 
 # ---------------- Utilities ----------------
@@ -44,34 +43,28 @@ def get_original_image(url_input):
     if url_input and is_valid_url(url_input):
         image = Image.open(requests.get(url_input, stream=True).raw).convert("RGB")
         return image
+    return None
 
 
+# ---------------- Model Loader ----------------
 
-# -------------------- 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()
+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"
+        )
+        model = YolosForObjectDetection.from_pretrained(
+            "nickmuchi/yolos-small-finetuned-license-plate-detection",
+            use_safetensors=True,
+            torch_dtype=torch.float32,
+        )
+        model.eval()
+    return processor, model
 
-reader = easyocr.Reader(['en'], gpu=False)
 
+# ---------------- Plate Color Classifier ----------------
 
-# -------------------- Plate Color Classifier --------------------
 def classify_plate_color(plate_img):
     img = np.array(plate_img)
     hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
@@ -87,39 +80,20 @@ def classify_plate_color(plate_img):
     else:
         return "Personal"
 
-# -------------------- OCR --------------------
+
+# ---------------- 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")
-
-    return fig
-
 
 # ---------------- Core Inference ----------------
 
-def make_prediction(img, processor, model):
+def make_prediction(img):
+    processor, model = load_model()
     inputs = processor(images=img, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
@@ -131,6 +105,8 @@ def make_prediction(img, processor, model):
     return processed_outputs[0]
 
 
+# ---------------- Visualization ----------------
+
 def fig2img(fig):
     buf = io.BytesIO()
     fig.savefig(buf)
@@ -138,33 +114,17 @@ def fig2img(fig):
     pil_img = Image.open(buf)
 
     basewidth = 750
-    wpercent = (basewidth / float(pil_img.size[0]))
-    hsize = int((float(pil_img.size[1]) * float(wpercent)))
+    wpercent = basewidth / float(pil_img.size[0])
+    hsize = int(float(pil_img.size[1]) * float(wpercent))
     img = pil_img.resize((basewidth, hsize), Image.Resampling.LANCZOS)
 
     plt.close(fig)
     return img
 
 
-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 green > yellow and green > white:
-        return "EV"
-    elif yellow > green and yellow > white:
-        return "Commercial"
-    else:
-        return "Personal"
-
-
-# ---------------- Visualization ----------------
-
 def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
+    BASE_TOLL = 100  # base amount for all vehicles
+
     keep = output_dict["scores"] > threshold
     boxes = output_dict["boxes"][keep].tolist()
     scores = output_dict["scores"][keep].tolist()
@@ -178,10 +138,18 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     ax = plt.gca()
     colors = COLORS * 100
 
-    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
-        if label == 'license-plates':
+    for score, (xmin, ymin, xmax, ymax), label, color in zip(
+        scores, boxes, labels, colors
+    ):
+        if "plate" in label.lower():
             crop = img.crop((int(xmin), int(ymin), int(xmax), int(ymax)))
             plate_type = classify_plate_color(crop)
+            # Apply 10% discount for EV vehicles
+            if plate_type == "EV":
+                discounted_amount = BASE_TOLL * 0.9
+                price_text = f"EV | ₹{discounted_amount:.0f} (10% off)"
+            else:
+                price_text = f"{plate_type} | ₹{BASE_TOLL}"
 
             ax.add_patch(
                 plt.Rectangle(
@@ -191,9 +159,9 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
             )
             ax.text(
                 xmin, ymin - 10,
-                f"{plate_type} | {score:0.2f}",
+                f"{price_text} | {score:0.2f}",
                 fontsize=12,
-                bbox=dict(facecolor="yellow", alpha=0.8)
+                bbox=dict(facecolor="yellow", alpha=0.8),
             )
 
     plt.axis("off")
@@ -202,9 +170,7 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
 
 # ---------------- Image Detection ----------------
 
-def detect_objects_image(model_name, url_input, image_input, webcam_input, threshold):
-    processor, model = load_model(model_name)
-
+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)
     elif image_input is not None:
@@ -214,24 +180,26 @@ def detect_objects_image(model_name, url_input, image_input, webcam_input, thres
     else:
         return None
 
-    processed_outputs = make_prediction(image, processor, model)
-    viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
+    processed_outputs = make_prediction(image)
+    viz_img = visualize_prediction(
+        image, processed_outputs, threshold, load_model()[1].config.id2label
+    )
 
     return viz_img
 
 
 # ---------------- Video Detection ----------------
 
-def detect_objects_video(model_name, video_input, threshold):
+def detect_objects_video(video_input, threshold):
     if video_input is None:
         return None
 
-    processor, model = load_model(model_name)
+    processor, model = load_model()
 
     cap = cv2.VideoCapture(video_input)
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
 
-    output_path = "/mnt/data/output_detected.mp4"
+    output_path = "/tmp/output_detected.mp4"
     fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -246,7 +214,7 @@ def detect_objects_video(model_name, video_input, threshold):
         rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         pil_img = Image.fromarray(rgb_frame)
 
-        processed_outputs = make_prediction(pil_img, processor, model)
+        processed_outputs = make_prediction(pil_img)
 
         keep = processed_outputs["scores"] > threshold
         boxes = processed_outputs["boxes"][keep].tolist()
@@ -265,7 +233,7 @@ def detect_objects_video(model_name, video_input, threshold):
                     (int(xmin), int(ymin)),
                     (int(xmax), int(ymax)),
                     (0, 255, 0),
-                    2
+                    2,
                 )
                 cv2.putText(
                     frame,
@@ -274,7 +242,7 @@ def detect_objects_video(model_name, video_input, threshold):
                     cv2.FONT_HERSHEY_SIMPLEX,
                     0.6,
                     (0, 255, 0),
-                    2
+                    2,
                 )
 
         out.write(frame)
@@ -287,21 +255,17 @@ def detect_objects_video(model_name, video_input, threshold):
 
 # ---------------- UI ----------------
 
-title = """<h1 id="title">Smart Vehicle Clssification (Image + Video)</h1>"""
+title = """<h1 id="title">Smart Vehicle Classification (Image + Video)</h1>"""
 
 description = """
-Detect license plates using YOLOS or DETR and Vehicle classification.
-Supports:Image URL, Image Upload, Webcam, Video Upload 
+Smart Vehicle Classification system to Promote EV by applying discount on Toll,
+Tax, parking. 
+Supports:Image URL, Image Upload, Webcam, Video Upload,Vehicle type classification by plate color
 """
 
-#models = [
-#    "nickmuchi/yolos-small-finetuned-license-plate-detection"
-#]
-css = '''
-h1#title {
-  text-align: center;
-}
-'''
+css = """
+h1#title { text-align: center; }
+"""
 
 demo = gr.Blocks()
 
@@ -309,70 +273,67 @@ with demo:
     gr.Markdown(title)
     gr.Markdown(description)
 
-    options = gr.Dropdown(choices=models, label='Object Detection Model', value=models[0])
-    slider_input = gr.Slider(minimum=0.2, maximum=1, value=0.5, step=0.1, label='Prediction Threshold')
+    slider_input = gr.Slider(
+        minimum=0.2, maximum=1, value=0.5, step=0.1, label="Prediction Threshold"
+    )
 
     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 valid image URL here..")
                 original_image = gr.Image(height=750, width=750)
                 url_input.change(get_original_image, url_input, original_image)
                 img_output_from_url = gr.Image(height=750, width=750)
-            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=750, width=750)
+                img_input = gr.Image(type="pil", height=750, width=750)
                 img_output_from_upload = gr.Image(height=750, width=750)
-            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"],
-                    type="pil",
-                    height=750,
-                    width=750,
-                    streaming=True
+                    sources=["webcam"], type="pil", height=750, width=750, streaming=True
                 )
                 img_output_from_webcam = gr.Image(height=750, width=750)
-            cam_but = gr.Button('Detect')
+            cam_but = gr.Button("Detect")
 
-        with gr.TabItem('Video Upload'):
+        with gr.TabItem("Video Upload"):
             with gr.Row():
                 video_input = gr.Video(label="Upload Video")
                 video_output = gr.Video(label="Detected Video")
-            vid_but = gr.Button('Detect Video')
+            vid_but = gr.Button("Detect Video")
 
     url_but.click(
         detect_objects_image,
-        inputs=[options, url_input, img_input, web_input, slider_input],
+        inputs=[url_input, img_input, web_input, slider_input],
         outputs=[img_output_from_url],
-        queue=True
+        queue=True,
     )
 
     img_but.click(
         detect_objects_image,
-        inputs=[options, url_input, img_input, web_input, slider_input],
+        inputs=[url_input, img_input, web_input, slider_input],
         outputs=[img_output_from_upload],
-        queue=True
+        queue=True,
     )
 
     cam_but.click(
         detect_objects_image,
-        inputs=[options, url_input, img_input, web_input, slider_input],
+        inputs=[url_input, img_input, web_input, slider_input],
         outputs=[img_output_from_webcam],
-        queue=True
+        queue=True,
     )
 
     vid_but.click(
         detect_objects_video,
-        inputs=[options, video_input, slider_input],
+        inputs=[video_input, slider_input],
         outputs=[video_output],
-        queue=True
+        queue=True,
     )
 
 
 demo.queue()
-demo.launch(debug=True)
+demo.launch(debug=True, ssr_mode=False)