Update app.py

app.py

@@ -1,13 +1,14 @@
-# app_with_video.py
+# app.py (FINAL CLEAN VERSION)
 import io
 import os
 import cv2
-import numpy as np
 import gradio as gr
 import matplotlib.pyplot as plt
-import requests, validators
+import requests
 import torch
 import pathlib
+import numpy as np
+from urllib.parse import urlparse
 from PIL import Image
 from transformers import AutoImageProcessor, YolosForObjectDetection, DetrForObjectDetection
 
@@ -22,12 +23,45 @@ COLORS = [
     [0.301, 0.745, 0.933]
 ]
 
-# ---------- Core Inference ----------
+# ---------------- Utilities ----------------
+
+def is_valid_url(url):
+    try:
+        result = urlparse(url)
+        return all([result.scheme, result.netloc])
+    except Exception:
+        return False
+
+
+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
+
+
+# ---------------- Model Loading ----------------
+
+def load_model(model_name):
+    processor = AutoImageProcessor.from_pretrained(model_name)
+
+    if "yolos" in model_name:
+        model = YolosForObjectDetection.from_pretrained(model_name)
+    elif "detr" in model_name:
+        model = DetrForObjectDetection.from_pretrained(model_name)
+    else:
+        raise ValueError("Unsupported model")
+
+    model.eval()
+    return processor, model
+
+
+# ---------------- Core Inference ----------------
 
 def make_prediction(img, processor, model):
     inputs = processor(images=img, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
+
     img_size = torch.tensor([tuple(reversed(img.size))])
     processed_outputs = processor.post_process_object_detection(
         outputs, threshold=0.0, target_sizes=img_size
@@ -40,36 +74,42 @@ def fig2img(fig):
     fig.savefig(buf)
     buf.seek(0)
     pil_img = Image.open(buf)
+
     basewidth = 750
     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
 
 
+# ---------------- Plate Color Classification ----------------
+
 def classify_plate_color(crop_img):
-    # Convert PIL to OpenCV BGR
     img = cv2.cvtColor(np.array(crop_img), cv2.COLOR_RGB2BGR)
     hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
     h, s, v = cv2.split(hsv)
+
     avg_h, avg_s, avg_v = np.mean(h), np.mean(s), np.mean(v)
 
-    # Heuristic thresholds (India-style plates)
+    # Heuristic thresholds (India-style)
     if avg_v < 80:
         return "Black Plate (Commercial)"
     if avg_s < 40 and avg_v > 180:
         return "White Plate (Private)"
     if 15 < avg_h < 35 and avg_s > 80:
         return "Yellow Plate (Commercial)"
-    if avg_h > 80 and avg_h < 130:
+    if 80 < avg_h < 130:
         return "Blue Plate (Diplomatic)"
-    if avg_h > 35 and avg_h < 85:
+    if 35 < avg_h < 85:
         return "Green Plate (Electric)"
 
     return "Unknown Plate"
 
 
+# ---------------- Visualization ----------------
+
 def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     keep = output_dict["scores"] > threshold
     boxes = output_dict["boxes"][keep].tolist()
@@ -104,66 +144,14 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
 
     plt.axis("off")
     return fig2img(plt.gcf())
-    keep = output_dict["scores"] > threshold
-    boxes = output_dict["boxes"][keep].tolist()
-    scores = output_dict["scores"][keep].tolist()
-    labels = output_dict["labels"][keep].tolist()
-
-    if id2label is not None:
-        labels = [id2label[x] for x in labels]
-
-    plt.figure(figsize=(20, 20))
-    plt.imshow(img)
-    ax = plt.gca()
-    colors = COLORS * 100
-
-    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
-        if label == 'license-plates':
-            ax.add_patch(
-                plt.Rectangle(
-                    (xmin, ymin), xmax - xmin, ymax - ymin,
-                    fill=False, color=color, linewidth=4
-                )
-            )
-            ax.text(
-                xmin, ymin,
-                f"{label}: {score:0.2f}",
-                fontsize=12,
-                bbox=dict(facecolor="yellow", alpha=0.8)
-            )
-
-    plt.axis("off")
-    return fig2img(plt.gcf())
-
 
-# ---------- Utilities ----------
-
-def get_original_image(url_input):
-    if validators.url(url_input):
-        image = Image.open(requests.get(url_input, stream=True).raw).convert("RGB")
-        return image
-
-
-def load_model(model_name):
-    processor = AutoImageProcessor.from_pretrained(model_name)
-
-    if "yolos" in model_name:
-        model = YolosForObjectDetection.from_pretrained(model_name)
-    elif "detr" in model_name:
-        model = DetrForObjectDetection.from_pretrained(model_name)
-    else:
-        raise ValueError("Unsupported model")
-
-    model.eval()
-    return processor, model
 
-
-# ---------- Image Detection ----------
+# ---------------- Image Detection ----------------
 
 def detect_objects_image(model_name, url_input, image_input, webcam_input, threshold):
     processor, model = load_model(model_name)
 
-    if validators.url(url_input):
+    if url_input and is_valid_url(url_input):
         image = get_original_image(url_input)
     elif image_input is not None:
         image = image_input
@@ -178,7 +166,7 @@ def detect_objects_image(model_name, url_input, image_input, webcam_input, thres
     return viz_img
 
 
-# ---------- Video Detection ----------
+# ---------------- Video Detection ----------------
 
 def detect_objects_video(model_name, video_input, threshold):
     if video_input is None:
@@ -215,6 +203,9 @@ def detect_objects_video(model_name, video_input, threshold):
 
         for score, (xmin, ymin, xmax, ymax), label in zip(scores, boxes, labels):
             if label == 'license-plates':
+                crop = pil_img.crop((int(xmin), int(ymin), int(xmax), int(ymax)))
+                plate_type = classify_plate_color(crop)
+
                 cv2.rectangle(
                     frame,
                     (int(xmin), int(ymin)),
@@ -224,7 +215,7 @@ def detect_objects_video(model_name, video_input, threshold):
                 )
                 cv2.putText(
                     frame,
-                    f"{label}: {score:.2f}",
+                    f"{plate_type} | {score:.2f}",
                     (int(xmin), int(ymin) - 10),
                     cv2.FONT_HERSHEY_SIMPLEX,
                     0.6,
@@ -240,7 +231,7 @@ def detect_objects_video(model_name, video_input, threshold):
     return output_path
 
 
-# ---------- UI ----------
+# ---------------- UI ----------------
 
 title = """<h1 id="title">License Plate Detection (Image + Video)</h1>"""
 
@@ -251,6 +242,7 @@ Supports:
 - Image Upload
 - Webcam
 - Video Upload
+- Vehicle type classification by plate color
 """
 
 models = [
@@ -264,7 +256,7 @@ h1#title {
 }
 '''
 
-demo = gr.Blocks(css=css)
+demo = gr.Blocks()
 
 with demo:
     gr.Markdown(title)
@@ -277,7 +269,7 @@ with demo:
         with gr.TabItem('Image URL'):
             with gr.Row():
                 url_input = gr.Textbox(lines=2, label='Enter valid image URL here..')
-                original_image = gr.Image(shape=(750, 750))
+                original_image = gr.Image(height=750, width=750)
                 url_input.change(get_original_image, url_input, original_image)
                 img_output_from_url = gr.Image(shape=(750, 750))
             url_but = gr.Button('Detect')