Create app.py

app.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.models import convnext_tiny
+from ultralytics import YOLO
+from PIL import Image
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+import gradio as gr
+
+# ---------- 1. Class labels ----------
+class_names = [
+    'beige', 'black', 'blue', 'brown', 'gold',
+    'green', 'grey', 'orange', 'pink', 'purple',
+    'red', 'silver', 'tan', 'white', 'yellow'
+]
+
+# ---------- 2. Load ConvNeXt-Tiny 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))
+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(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+# ---------- 4. Load YOLOv8 Model ----------
+yolo_model = YOLO("yolo11x.pt")
+
+# ---------- Gradio Inference Function ----------
+def detect_vehicle_color(input_img):
+    img_original = input_img.convert("RGB")
+    img_cv2 = cv2.cvtColor(np.array(img_original), cv2.COLOR_RGB2BGR)
+
+    results = yolo_model(img_cv2)
+    boxes = results[0].boxes
+
+    # Vehicle class IDs: car, motorcycle, bus, truck
+    vehicle_class_ids = {2, 3, 5, 7}
+    vehicle_boxes = [box for box in boxes if int(box.cls.item()) in vehicle_class_ids]
+
+    if len(vehicle_boxes) == 0:
+        return "No vehicle detected", img_original, img_original
+
+    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))
+
+    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()
+
+    # Draw bounding box on original image
+    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)
+
+    return f"{pred_class} ({confidence*100:.1f}%)", img_with_box_pil, cropped
+
+# ---------- Gradio UI ----------
+demo = gr.Interface(
+    fn=detect_vehicle_color,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Text(label="Predicted Vehicle Color"),
+        gr.Image(label="Detected Vehicle in Original"),
+        gr.Image(label="Cropped Vehicle Region")
+    ],
+    title="🚗 Vehicle Color Detection",
+    description="Upload an image to detect the most prominent vehicle and its predicted color."
+)
+
+# ---------- Launch ----------
+if __name__ == "__main__":
+    demo.launch()