Update app.py

app.py

@@ -3,20 +3,22 @@ import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import convnext_tiny
 from ultralytics import YOLO
-from PIL import Image, ImageDraw
 import numpy as np
 import cv2
 import gradio as gr
+from PIL import Image, ImageDraw
 from fast_alpr import ALPR
 
-# ---------- 1. Class labels ----------
+# ------------------ Constants and Models ------------------
 class_names = [
     'beige', 'black', 'blue', 'brown', 'gold',
     'green', 'grey', 'orange', 'pink', 'purple',
     'red', 'silver', 'tan', 'white', 'yellow'
 ]
 
-# ---------- 2. Load ConvNeXt-Tiny Model ----------
+DETECTOR_MODEL = "yolo-v9-s-608-license-plate-end2end"
+OCR_MODEL = "global-plates-mobile-vit-v2-model"
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = convnext_tiny(pretrained=False)
 model.classifier[2] = nn.Linear(768, len(class_names))
@@ -24,7 +26,6 @@ model.load_state_dict(torch.load("convnext_best_model.pth", map_location=device)
 model = model.to(device)
 model.eval()
 
-# ---------- 3. Image Transform ----------
 transform = transforms.Compose([
     transforms.Resize((512, 512)),
     transforms.ToTensor(),
@@ -32,91 +33,93 @@ transform = transforms.Compose([
                          [0.229, 0.224, 0.225])
 ])
 
-# ---------- 4. Load YOLOv8 Model ----------
 yolo_model = YOLO("yolo11x.pt")
 
-# ---------- 5. ALPR Setup ----------
-DETECTOR_MODEL = "yolo-v9-s-608-license-plate-end2end"
-OCR_MODEL = "global-plates-mobile-vit-v2-model"
-
-def detect_vehicle_and_plate(input_img):
-    if input_img is None:
-        return None, None, None, None, "Please upload an image."
-
-    # Convert to RGB
-    img_original = input_img.convert("RGB")
-    img_cv2 = cv2.cvtColor(np.array(img_original), cv2.COLOR_RGB2BGR)
-
-    # ---------- Vehicle Detection ----------
-    results = yolo_model(img_cv2)
-    boxes = results[0].boxes
-    vehicle_class_ids = {2, 3, 5, 7}
-    vehicle_boxes = [box for box in boxes if int(box.cls.item()) in vehicle_class_ids]
-
-    if not vehicle_boxes:
-        return "No vehicle detected", img_original, img_original, img_original, "No plate detected."
-
-    def box_area(box):
-        x1, y1, x2, y2 = box.xyxy[0].tolist()
-        return (x2 - x1) * (y2 - y1)
-
-    largest_vehicle = max(vehicle_boxes, key=box_area)
-    x1, y1, x2, y2 = map(int, largest_vehicle.xyxy[0].tolist())
-    cropped = img_original.crop((x1, y1, x2, y2))
+# ------------------ Unified Inference Function ------------------
+def alpr_color_inference(image):
+    if image is None:
+        return None, None, None, "Please upload an image to continue."
 
-    input_tensor = transform(cropped).unsqueeze(0).to(device)
-    with torch.no_grad():
-        output = model(input_tensor)
-        probs = torch.softmax(output, dim=1)[0]
-        pred_idx = torch.argmax(probs).item()
-        pred_class = class_names[pred_idx]
-        confidence = probs[pred_idx].item()
-
-    img_with_box = np.array(img_original).copy()
-    cv2.rectangle(img_with_box, (x1, y1), (x2, y2), (255, 0, 0), 3)
-    img_with_box_pil = Image.fromarray(img_with_box)
-
-    # ---------- License Plate Detection ----------
+    img = image.convert("RGB")
+    img_array = np.array(img)
     alpr = ALPR(detector_model=DETECTOR_MODEL, ocr_model=OCR_MODEL)
-    results = alpr.predict(np.array(img_original))
-    annotated_img = img_original.copy()
+    results = alpr.predict(img_array)
+
+    annotated_img = Image.fromarray(img_array.copy())
     draw = ImageDraw.Draw(annotated_img)
-    final_text = ""
 
+    plate_texts = []
     for result in results:
         detection = getattr(result, 'detection', None)
         ocr = getattr(result, 'ocr', None)
-        if detection and getattr(detection, 'bounding_box', None):
-            bbox = detection.bounding_box
-            box_coords = [int(bbox.x1), int(bbox.y1), int(bbox.x2), int(bbox.y2)]
-            draw.rectangle(box_coords, outline="red", width=3)
-            if ocr:
-                text = ocr.text
-                final_text += text + " "
-                draw.text((box_coords[0], max(box_coords[1] - 10, 0)), text, fill="red")
-
-    plate_result = f"Detected plate(s): {final_text.strip()}" if final_text else "No license plate detected."
-
-    return f"{pred_class} ({confidence*100:.1f}%)", img_with_box_pil, cropped, annotated_img, plate_result
+        if detection is not None:
+            bbox_obj = getattr(detection, 'bounding_box', None)
+            if bbox_obj is not None:
+                bbox = [int(bbox_obj.x1), int(bbox_obj.y1), int(bbox_obj.x2), int(bbox_obj.y2)]
+                draw.rectangle(bbox, outline="red", width=3)
+                if ocr is not None:
+                    text = getattr(ocr, 'text', '')
+                    plate_texts.append(text)
+                    draw.text((bbox[0], max(bbox[1] - 10, 0)), text, fill="red")
+
+    # Color Detection
+    img_cv2 = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    yolo_results = yolo_model(img_cv2)
+    boxes = yolo_results[0].boxes
+    vehicle_class_ids = {2, 3, 5, 7}  # car, motorcycle, bus, truck
+    vehicle_boxes = [box for box in boxes if int(box.cls.item()) in vehicle_class_ids]
 
-# ---------- Gradio UI ----------
+    if not vehicle_boxes:
+        color_text = "No vehicle detected"
+        cropped_img = img
+        vehicle_img = img
+    else:
+        largest_vehicle = max(vehicle_boxes, key=lambda box: (box.xyxy[0][2] - box.xyxy[0][0]) * (box.xyxy[0][3] - box.xyxy[0][1]))
+        x1, y1, x2, y2 = map(int, largest_vehicle.xyxy[0].tolist())
+        cropped_img = img.crop((x1, y1, x2, y2))
+        input_tensor = transform(cropped_img).unsqueeze(0).to(device)
+        with torch.no_grad():
+            output = model(input_tensor)
+            probs = torch.softmax(output, dim=1)[0]
+            pred_idx = torch.argmax(probs).item()
+            pred_class = class_names[pred_idx]
+            confidence = probs[pred_idx].item()
+        vehicle_img = Image.fromarray(cv2.rectangle(np.array(img), (x1, y1), (x2, y2), (255, 0, 0), 3))
+        color_text = f"{pred_class} ({confidence*100:.1f}%)"
+
+    detection_results = (f"Detected {len(results)} license plate(s): {', '.join(plate_texts)}"
+                         if results else "No license plate detected 😔.")
+
+    return annotated_img, vehicle_img, cropped_img, f"{detection_results}\nVehicle Color: {color_text}"
+
+# ------------------ Gradio UI ------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🚗 Vehicle Color + License Plate Detection")
+    gr.Markdown("# License Plate + Vehicle Color Detection")
     with gr.Row():
         with gr.Column():
-            img_input = gr.Image(type="pil", label="Upload Vehicle Image")
-            btn = gr.Button("Run Detection")
+            image_input = gr.Image(type="pil", label="Upload an image")
+            submit_btn = gr.Button("Run Detection")
         with gr.Column():
-            color_result = gr.Text(label="Predicted Vehicle Color")
-            image_with_box = gr.Image(label="Detected Vehicle in Original")
-            cropped_img = gr.Image(label="Cropped Vehicle Region")
-            plate_annotated = gr.Image(label="License Plate Detection")
-            plate_result = gr.Text(label="License Plate OCR")
-
-    btn.click(
-        detect_vehicle_and_plate,
-        inputs=[img_input],
-        outputs=[color_result, image_with_box, cropped_img, plate_annotated, plate_result]
+            plate_output = gr.Image(label="License Plate Detection")
+            vehicle_output = gr.Image(label="Detected Vehicle in Original")
+            cropped_output = gr.Image(label="Cropped Vehicle Region")
+            result_text = gr.Markdown(label="Results")
+
+    submit_btn.click(
+        alpr_color_inference,
+        inputs=[image_input],
+        outputs=[plate_output, vehicle_output, cropped_output, result_text]
+    )
+
+    gr.Examples(
+        examples=[
+            "examples/car1.jpg",
+            "examples/car2.jpg",
+            "examples/car3.jpg",
+            "examples/car4.jpg",
+        ],
+        inputs=[image_input],
+        label="Example Images"
     )
 
 if __name__ == "__main__":