Update app.py

app.py

@@ -1,5 +1,4 @@
 import io
-import os
 import cv2
 import gradio as gr
 import matplotlib
@@ -45,10 +44,7 @@ def load_model():
     global processor, model
     if processor is None or model is None:
         processor = YolosImageProcessor.from_pretrained(MODEL_NAME)
-        model = YolosForObjectDetection.from_pretrained(
-            MODEL_NAME,
-            torch_dtype=torch.float32
-        )
+        model = YolosForObjectDetection.from_pretrained(MODEL_NAME)
         model.eval()
     return processor, model
 
@@ -62,14 +58,15 @@ def is_valid_url(url):
         return False
 
 def get_original_image(url):
-    return Image.open(requests.get(url, stream=True).raw).convert("RGB")
+    response = requests.get(url, stream=True)
+    return Image.open(response.raw).convert("RGB")
 
 # -------------------- DISCOUNT LOGIC --------------------
 
 def compute_discount(vehicle_type):
     if vehicle_type == "EV":
-        return BASE_AMT * 0.9, "10% EV discount applied"
-    return BASE_AMT, "No discount"
+        return BASE_AMT * 0.9
+    return BASE_AMT
 
 # -------------------- PLATE COLOR CLASSIFICATION --------------------
 
@@ -86,7 +83,7 @@ def classify_plate_color(plate_img):
         return "Commercial"
     return "Personal"
 
-# -------------------- OCR (LIGHTWEIGHT) --------------------
+# -------------------- OCR --------------------
 
 def read_plate(plate_img):
     gray = cv2.cvtColor(np.array(plate_img), cv2.COLOR_RGB2GRAY)
@@ -103,38 +100,40 @@ def read_plate(plate_img):
 def make_prediction(img):
     processor, model = load_model()
     inputs = processor(images=img, return_tensors="pt")
+
     with torch.no_grad():
         outputs = model(**inputs)
 
-    img_size = torch.tensor([tuple(reversed(img.size))])
+    img_size = torch.tensor([img.size[::-1]])
     results = processor.post_process_object_detection(
         outputs, threshold=0.3, target_sizes=img_size
     )
-    return results[0]
+
+    return results[0], model.config.id2label
 
 # -------------------- VISUALIZATION --------------------
 
 def fig_to_img(fig):
     buf = io.BytesIO()
-    fig.savefig(buf)
+    fig.savefig(buf, bbox_inches="tight")
     buf.seek(0)
     img = Image.open(buf)
     plt.close(fig)
     return img
 
-def visualize(img, output, threshold):
+def visualize(img, output, id2label, threshold):
     keep = output["scores"] > threshold
     boxes = output["boxes"][keep]
     labels = output["labels"][keep]
 
-    plt.figure(figsize=(10,10))
-    plt.imshow(img)
-    ax = plt.gca()
+    fig, ax = plt.subplots(figsize=(8,8))
+    ax.imshow(img)
 
-    results = []
+    results_text = []
 
     for box, label in zip(boxes, labels):
-        if "plate" not in load_model()[1].config.id2label[label.item()].lower():
+        label_name = id2label[label.item()].lower()
+        if "plate" not in label_name:
             continue
 
         x1,y1,x2,y2 = map(int, box.tolist())
@@ -142,7 +141,7 @@ def visualize(img, output, threshold):
 
         plate = read_plate(plate_img)
         vtype = classify_plate_color(plate_img)
-        toll, msg = compute_discount(vtype)
+        toll = compute_discount(vtype)
 
         cursor.execute(
             "INSERT INTO vehicles VALUES (?, ?, ?, datetime('now'))",
@@ -150,15 +149,21 @@ def visualize(img, output, threshold):
         )
         conn.commit()
 
-        results.append(f"{plate} | {vtype} | ₹{int(toll)}")
+        results_text.append(f"{plate} | {vtype} | ₹{int(toll)}")
 
         ax.add_patch(
-            plt.Rectangle((x1,y1), x2-x1, y2-y1, fill=False, color="red", linewidth=2)
+            plt.Rectangle((x1,y1), x2-x1, y2-y1,
+                          fill=False, color="red", linewidth=2)
         )
-        ax.text(x1, y1-5, f"{plate} ({vtype})", color="yellow")
+        ax.text(x1, y1-5, f"{plate} ({vtype})",
+                color="yellow", fontsize=10)
+
+    ax.axis("off")
 
-    plt.axis("off")
-    return fig_to_img(plt.gcf()), "\n".join(results) if results else "No plate detected"
+    if not results_text:
+        return fig_to_img(fig), "No plate detected"
+
+    return fig_to_img(fig), "\n".join(results_text)
 
 # -------------------- DASHBOARD --------------------
 
@@ -177,48 +182,57 @@ def get_dashboard():
 
 # -------------------- MAIN CALLBACK --------------------
 
-def detect(url, img, cam, threshold):
-    if url and is_valid_url(url):
-        image = get_original_image(url)
-    elif img is not None:
-        image = img
-    elif cam is not None:
-        image = cam
-    else:
-        return None, "No input"
+def detect_from_url(url, threshold):
+    if not url or not is_valid_url(url):
+        return None, "Invalid URL"
+    img = get_original_image(url)
+    output, id2label = make_prediction(img)
+    return visualize(img, output, id2label, threshold)
 
-    output = make_prediction(image)
-    return visualize(image, output, threshold)
+def detect_from_image(img, threshold):
+    if img is None:
+        return None, "No image provided"
+    output, id2label = make_prediction(img)
+    return visualize(img, output, id2label, threshold)
 
 # -------------------- UI --------------------
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🚦 Smart Vehicle Classification System")
+    gr.Markdown("## Smart Vehicle Classification System")
 
+    slider = gr.Slider(0.3, 1.0, 0.5, label="Confidence Threshold")
     result_box = gr.Textbox(label="Result", lines=4)
-    slider = gr.Slider(0.3,1.0,0.5,label="Confidence Threshold")
 
     with gr.Tabs():
+
         with gr.Tab("Image URL"):
-            url = gr.Textbox(label="Image URL")
-            out1 = gr.Image()
-            btn1 = gr.Button("Detect")
+            url_input = gr.Textbox(label="Image URL")
+            url_output = gr.Image()
+            url_btn = gr.Button("Detect")
+            url_btn.click(detect_from_url,
+                          inputs=[url_input, slider],
+                          outputs=[url_output, result_box])
 
         with gr.Tab("Upload"):
-            img = gr.Image(type="pil")
-            out2 = gr.Image()
-            btn2 = gr.Button("Detect")
+            img_input = gr.Image(type="pil")
+            img_output = gr.Image()
+            img_btn = gr.Button("Detect")
+            img_btn.click(detect_from_image,
+                          inputs=[img_input, slider],
+                          outputs=[img_output, result_box])
 
         with gr.Tab("Webcam"):
-            cam = gr.Image(source="webcam", type="pil")
-            out3 = gr.Image()
-            btn3 = gr.Button("Detect")
-
-    btn1.click(detect, [url, img, cam, slider], [out1, result_box])
-    btn2.click(detect, [url, img, cam, slider], [out2, result_box])
-    btn3.click(detect, [url, img, cam, slider], [out3, result_box])
-
-    gr.Markdown("### 📊 Dashboard")
+            cam_input = gr.Image(
+                sources=["webcam"],   # using web camera
+                type="pil"
+            )
+            cam_output = gr.Image()
+            cam_btn = gr.Button("Detect")
+            cam_btn.click(detect_from_image,
+                          inputs=[cam_input, slider],
+                          outputs=[cam_output, result_box])
+
+    gr.Markdown("###Dashboard")
     gr.Plot(get_dashboard)
 
 demo.launch()