Update app.py

app.py

@@ -1,155 +1,331 @@
+# app_with_video.py
+import io
+import os
 import cv2
-import numpy as np
-from ultralytics import YOLO
-import easyocr
 import gradio as gr
-import tempfile
-import os
-
-# Load YOLOv8 plate detection model
-model = YOLO("best.pt")   # <-- your trained plate model
-
-# Initialize OCR
-reader = easyocr.Reader(['en'], gpu=False)
-
-def preprocess_plate(plate_img):
-    gray = cv2.cvtColor(plate_img, cv2.COLOR_BGR2GRAY)
-    blur = cv2.GaussianBlur(gray, (5, 5), 0)
-    thresh = cv2.adaptiveThreshold(
-        blur, 255,
-        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
-        cv2.THRESH_BINARY, 11, 2
+import matplotlib.pyplot as plt
+import requests, validators
+import torch
+import pathlib
+from PIL import Image
+from transformers import AutoImageProcessor, YolosForObjectDetection, DetrForObjectDetection
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+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]
+]
+
+# ---------- 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
     )
-    return thresh
-
-def recognize_plate(plate_img):
-    processed = preprocess_plate(plate_img)
-    ocr_result = reader.readtext(processed)
-
-    plate_text = ""
-    for (bbox, text, prob) in ocr_result:
-        if prob > 0.4:
-            plate_text += text + " "
-
-    return plate_text.strip()
-
-def process_frame(frame):
-    detected_plates = []
+    return processed_outputs[0]
+
+
+def fig2img(fig):
+    buf = io.BytesIO()
+    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
+
+
+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)
+    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:
+        return "Blue Plate (Diplomatic)"
+    if avg_h > 35 and avg_h < 85:
+        return "Green Plate (Electric)"
+
+    return "Unknown Plate"
+
+
+def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
+    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':
+            crop = img.crop((int(xmin), int(ymin), int(xmax), int(ymax)))
+            plate_type = classify_plate_color(crop)
+
+            ax.add_patch(
+                plt.Rectangle(
+                    (xmin, ymin), xmax - xmin, ymax - ymin,
+                    fill=False, color=color, linewidth=4
+                )
+            )
+            ax.text(
+                xmin, ymin - 10,
+                f"{plate_type} | {score:0.2f}",
+                fontsize=12,
+                bbox=dict(facecolor="yellow", alpha=0.8)
+            )
+
+    plt.axis("off")
+    return fig2img(plt.gcf())(img, output_dict, threshold=0.5, id2label=None):
+    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 ----------
+
+def detect_objects_image(model_name, url_input, image_input, webcam_input, threshold):
+    processor, model = load_model(model_name)
+
+    if validators.url(url_input):
+        image = get_original_image(url_input)
+    elif image_input is not None:
+        image = image_input
+    elif webcam_input is not None:
+        image = webcam_input
+    else:
+        return None
+
+    processed_outputs = make_prediction(image, processor, model)
+    viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
+
+    return viz_img
+
+
+# ---------- Video Detection ----------
+
+def detect_objects_video(model_name, video_input, threshold):
+    if video_input is None:
+        return None
 
-    results = model(frame)
+    processor, model = load_model(model_name)
+
+    cap = cv2.VideoCapture(video_input)
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
 
-    for r in results:
-        if r.boxes is None:
-            continue
-
-        boxes = r.boxes.xyxy.cpu().numpy()
-        confs = r.boxes.conf.cpu().numpy()
-
-        for box, conf in zip(boxes, confs):
-            x1, y1, x2, y2 = map(int, box)
-
-            plate_img = frame[y1:y2, x1:x2]
-            if plate_img.size == 0:
-                continue
-
-            plate_text = recognize_plate(plate_img)
-
-            detected_plates.append({
-                "plate_text": plate_text,
-                "confidence": float(conf)
-            })
+    output_path = "/mnt/data/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))
 
-            # Draw bounding box
-            cv2.rectangle(frame, (x1, y1), (x2, y2),
-                          (0, 255, 0), 2)
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
 
-            # Draw plate text
-            label = plate_text if plate_text else "Plate"
-            cv2.putText(frame, label,
-                        (x1, y1 - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX,
-                        0.8, (255, 0, 0), 2)
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
 
-    return frame, detected_plates
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_img = Image.fromarray(rgb_frame)
+
+        processed_outputs = make_prediction(pil_img, processor, model)
+
+        keep = processed_outputs["scores"] > threshold
+        boxes = processed_outputs["boxes"][keep].tolist()
+        scores = processed_outputs["scores"][keep].tolist()
+        labels = processed_outputs["labels"][keep].tolist()
+
+        labels = [model.config.id2label[x] for x in labels]
+
+        for score, (xmin, ymin, xmax, ymax), label in zip(scores, boxes, labels):
+            if label == 'license-plates':
+                cv2.rectangle(
+                    frame,
+                    (int(xmin), int(ymin)),
+                    (int(xmax), int(ymax)),
+                    (0, 255, 0),
+                    2
+                )
+                cv2.putText(
+                    frame,
+                    f"{label}: {score:.2f}",
+                    (int(xmin), int(ymin) - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.6,
+                    (0, 255, 0),
+                    2
+                )
+
+        out.write(frame)
 
-# =========================
-# IMAGE MODE
-# =========================
-def process_image(image):
-    frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    annotated_frame, plates = process_frame(frame)
-    annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
+    cap.release()
+    out.release()
 
-    plate_texts = [p["plate_text"] for p in plates if p["plate_text"]]
-    result_text = "\n".join(plate_texts) if plate_texts else "No plates detected."
+    return output_path
 
-    return annotated_frame, result_text
 
-# =========================
-# VIDEO MODE
-# =========================
-def process_video(video_file):
-    cap = cv2.VideoCapture(video_file)
+# ---------- UI ----------
 
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
+title = """<h1 id="title">License Plate Detection (Image + Video)</h1>"""
 
-    temp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    out_path = temp_out.name
-    temp_out.close()
+description = """
+Detect license plates using YOLOS or DETR.
+Supports:
+- Image URL
+- Image Upload
+- Webcam
+- Video Upload
+"""
 
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+models = [
+    "nickmuchi/yolos-small-finetuned-license-plate-detection",
+    "nickmuchi/detr-resnet50-license-plate-detection"
+]
 
-    all_detected = set()
+css = '''
+h1#title {
+  text-align: center;
+}
+'''
 
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+demo = gr.Blocks(css=css)
 
-        annotated_frame, plates = process_frame(frame)
+with demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
 
-        for p in plates:
-            if p["plate_text"]:
-                all_detected.add(p["plate_text"])
+    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')
 
-        out.write(annotated_frame)
+    with gr.Tabs():
+        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))
+                url_input.change(get_original_image, url_input, original_image)
+                img_output_from_url = gr.Image(shape=(750, 750))
+            url_but = gr.Button('Detect')
+
+        with gr.TabItem('Image Upload'):
+            with gr.Row():
+                img_input = gr.Image(type='pil', shape=(750, 750))
+                img_output_from_upload = gr.Image(shape=(750, 750))
+            img_but = gr.Button('Detect')
+
+        with gr.TabItem('WebCam'):
+            with gr.Row():
+                web_input = gr.Image(source='webcam', type='pil', shape=(750, 750), streaming=True)
+                img_output_from_webcam = gr.Image(shape=(750, 750))
+            cam_but = gr.Button('Detect')
+
+        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')
+
+    url_but.click(
+        detect_objects_image,
+        inputs=[options, url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_url],
+        queue=True
+    )
 
-    cap.release()
-    out.release()
+    img_but.click(
+        detect_objects_image,
+        inputs=[options, url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_upload],
+        queue=True
+    )
 
-    result_text = "\n".join(all_detected) if all_detected else "No plates detected."
+    cam_but.click(
+        detect_objects_image,
+        inputs=[options, url_input, img_input, web_input, slider_input],
+        outputs=[img_output_from_webcam],
+        queue=True
+    )
 
-    return out_path, result_text
+    vid_but.click(
+        detect_objects_video,
+        inputs=[options, video_input, slider_input],
+        outputs=[video_output],
+        queue=True
+    )
 
-# =========================
-# GRADIO UI
-# =========================
-with gr.Blocks() as demo:
-    gr.Markdown("## Smart Traffic & EV Analytics System")
-    gr.Markdown("Upload an image or video to detect multiple vehicle number plates.")
 
-    with gr.Tabs():
-        with gr.Tab("Image"):
-            image_input = gr.Image(type="numpy", label="Upload Image")
-            image_output = gr.Image(label="Detected Plates")
-            image_text = gr.Textbox(label="Recognized Plate Numbers")
-
-            image_button = gr.Button("Detect Plates")
-            image_button.click(process_image,
-                               inputs=image_input,
-                               outputs=[image_output, image_text])
-
-        with gr.Tab("Video"):
-            video_input = gr.Video(label="Upload Video")
-            video_output = gr.Video(label="Processed Video")
-            video_text = gr.Textbox(label="Recognized Plate Numbers")
-
-            video_button = gr.Button("Detect Plates")
-            video_button.click(process_video,
-                               inputs=video_input,
-                               outputs=[video_output, video_text])
-
-demo.launch()
+demo.launch(debug=True, enable_queue=True)