Create app.py

app.py

@@ -0,0 +1,80 @@
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image, ImageDraw
+from transformers import YolosImageProcessor, YolosForObjectDetection
+
+# Load model
+processor = YolosImageProcessor.from_pretrained(
+    "nickmuchi/yolos-small-finetuned-license-plate-detection"
+)
+model = YolosForObjectDetection.from_pretrained(
+    "nickmuchi/yolos-small-finetuned-license-plate-detection"
+)
+model.eval()
+
+# -------- Plate Color Classifier -------- #
+def classify_plate_color(plate_img):
+    img = np.array(plate_img)
+    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+
+    green = cv2.inRange(hsv, (35, 40, 40), (85, 255, 255))
+    yellow = cv2.inRange(hsv, (15, 50, 50), (35, 255, 255))
+    white = cv2.inRange(hsv, (0, 0, 200), (180, 30, 255))
+
+    g = np.sum(green)
+    y = np.sum(yellow)
+    w = np.sum(white)
+
+    if g > y and g > w:
+        return "Electric Vehicle (Green Plate)"
+    elif y > g and y > w:
+        return "Commercial Vehicle (Yellow Plate)"
+    else:
+        return "Private Vehicle (White Plate)"
+
+# -------- Main Pipeline -------- #
+def process_image(img):
+    image = Image.fromarray(img)
+
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    target_sizes = torch.tensor([[image.size[1], image.size[0]]])
+    results = processor.post_process_object_detection(
+        outputs, threshold=0.3, target_sizes=target_sizes
+    )[0]
+
+    draw = ImageDraw.Draw(image)
+
+    if len(results["boxes"]) == 0:
+        return image, "No license plate detected"
+
+    box = results["boxes"][0].tolist()
+    x1, y1, x2, y2 = map(int, box)
+
+    plate = image.crop((x1, y1, x2, y2))
+    vehicle_type = classify_plate_color(plate)
+
+    draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
+    draw.text((x1, y1 - 10), vehicle_type, fill="red")
+
+    return image, vehicle_type
+
+# -------- Gradio UI -------- #
+with gr.Blocks() as demo:
+    gr.Markdown("# 🚗 Vehicle Classification using License Plate")
+    gr.Markdown("Upload or take a photo of a car. The AI detects the license plate and classifies the vehicle.")
+
+    with gr.Row():
+        input_img = gr.Image(type="numpy", sources=["upload", "webcam"])
+        output_img = gr.Image()
+    
+    result = gr.Textbox(label="Vehicle Type")
+
+    btn = gr.Button("Detect Vehicle")
+    btn.click(process_image, input_img, [output_img, result])
+
+demo.launch()