app.py

import gradio as gr
from PIL import Image
import torch
from torchvision import transforms

# --------------------------
# 1. Define your model class
# --------------------------
# Make sure this matches the architecture you used to train your model


# Model definition is now handled directly in the loading section

# --------------------------
# 2. Load model weights
# --------------------------
# Load the checkpoint directly as it was saved (a plain ResNet50 with custom fc head)
# state_dict = torch.load("best_stanford_cars_transfer_model.pth", map_location="cpu")
state_dict = torch.load("test_with_YOLO.pth", map_location="cpu")

# Create a ResNet50 and modify its fc to match the checkpoint
from torchvision.models import resnet50
model = resnet50(weights=None)

# Replace fc with the multi-layer head that matches checkpoint
num_classes = 196
in_ch = model.fc.in_features
model.fc = torch.nn.Linear(in_ch, num_classes)
# model.fc = torch.nn.Sequential()
# model.fc.add_module('1', torch.nn.Linear(in_ch, 1024))      # fc.1
# model.fc.add_module('2', torch.nn.ReLU())
# model.fc.add_module('3', torch.nn.BatchNorm1d(1024))        # fc.3  
# model.fc.add_module('4', torch.nn.ReLU())
# model.fc.add_module('5', torch.nn.Linear(1024, 512))        # fc.5
# model.fc.add_module('6', torch.nn.ReLU())
# model.fc.add_module('7', torch.nn.BatchNorm1d(512))         # fc.7
# model.fc.add_module('8', torch.nn.ReLU())
# model.fc.add_module('9', torch.nn.Linear(512, 196))         # fc.9

# Load the state dict
model.load_state_dict(state_dict, strict=True)
model.eval()  # important for inference

# --------------------------
# 3. Define labels
# --------------------------
# Full integer-to-label mapping for Stanford Cars dataset
labels = [
    "AM General Hummer SUV 2000",
    "Acura RL Sedan 2012",
    "Acura TL Sedan 2012",
    "Acura TL Type-S 2008",
    "Acura TSX Sedan 2012",
    "Acura Integra Type R 2001",
    "Acura ZDX Hatchback 2012",
    "Aston Martin V8 Vantage Convertible 2012",
    "Aston Martin V8 Vantage Coupe 2012",
    "Aston Martin Virage Convertible 2012",
    "Aston Martin Virage Coupe 2012",
    "Audi RS 4 Convertible 2008",
    "Audi A5 Coupe 2012",
    "Audi TTS Coupe 2012",
    "Audi R8 Coupe 2012",
    "Audi V8 Sedan 1994",
    "Audi 100 Sedan 1994",
    "Audi 100 Wagon 1994",
    "Audi TT Hatchback 2011",
    "Audi S6 Sedan 2011",
    "Audi S5 Convertible 2012",
    "Audi S5 Coupe 2012",
    "Audi S4 Sedan 2012",
    "Audi S4 Sedan 2007",
    "Audi TT RS Coupe 2012",
    "BMW ActiveHybrid 5 Sedan 2012",
    "BMW 1 Series Convertible 2012",
    "BMW 1 Series Coupe 2012",
    "BMW 3 Series Sedan 2012",
    "BMW 3 Series Wagon 2012",
    "BMW 6 Series Convertible 2007",
    "BMW X5 SUV 2007",
    "BMW X6 SUV 2012",
    None,
    None,
    "BMW M6 Convertible 2010",
    "BMW X3 SUV 2012",
    "BMW Z4 Convertible 2012",
    "Bentley Continental Supersports Conv. Convertible 2012",
    "Bentley Arnage Sedan 2009",
    "Bentley Mulsanne Sedan 2011",
    "Bentley Continental GT Coupe 2012",
    "Bentley Continental GT Coupe 2007",
    "Bentley Continental Flying Spur Sedan 2007",
    "Bugatti Veyron 16.4 Convertible 2009",
    "Bugatti Veyron 16.4 Coupe 2009",
    "Buick Regal GS 2012",
    "Buick Rainier SUV 2007",
    "Buick Verano Sedan 2012",
    "Buick Enclave SUV 2012",
    "Cadillac CTS-V Sedan 2012",
    None,
    "Cadillac Escalade EXT Crew Cab 2007",
    "Chevrolet Silverado 1500 Hybrid Crew Cab 2012",
    "Chevrolet Corvette Convertible 2012",
    "Chevrolet Corvette ZR1 2012",
    "Chevrolet Corvette Ron Fellows Edition Z06 2007",
    "Chevrolet Traverse SUV 2012",
    "Chevrolet Camaro Convertible 2012",
    "Chevrolet HHR SS 2010",
    "Chevrolet Impala Sedan 2007",
    "Chevrolet Tahoe Hybrid SUV 2012",
    "Chevrolet Sonic Sedan 2012",
    "Chevrolet Express Cargo Van 2007",
    "Chevrolet Avalanche Crew Cab 2012",
    "Chevrolet Cobalt SS 2010",
    "Chevrolet Malibu Hybrid Sedan 2010",
    "Chevrolet TrailBlazer SS 2009",
    "Chevrolet Silverado 2500HD Regular Cab 2012",
    "Chevrolet Silverado 1500 Classic Extended Cab 2007",
    "Chevrolet Express Van 2007",
    "Chevrolet Monte Carlo Coupe 2007",
    "Chevrolet Malibu Sedan 2007",
    "Chevrolet Silverado 1500 Extended Cab 2012",
    "Chevrolet Silverado 1500 Regular Cab 2012",
    "Chrysler Aspen SUV 2009",
    "Chrysler Sebring Convertible 2010",
    None,
    "Chrysler 300 SRT-8 2010",
    None,
    None,
    None,
    "Dodge Caliber Wagon 2012",
    "Dodge Caliber Wagon 2007",
    "Dodge Caravan Minivan 1997",
    "Dodge Ram Pickup 3500 Crew Cab 2010",
    "Dodge Ram Pickup 3500 Quad Cab 2009",
    "Dodge Sprinter Cargo Van 2009",
    "Dodge Journey SUV 2012",
    "Dodge Dakota Crew Cab 2010",
    "Dodge Dakota Club Cab 2007",
    None,
    "Dodge Challenger SRT8 2011",
    "Dodge Durango SUV 2012",
    None,
    "Dodge Charger Sedan 2012",
    "Dodge Charger SRT-8 2009",
    None,
    "FIAT 500 Abarth 2012",
    "FIAT 500 Convertible 2012",
    "Ferrari FF Coupe 2012",
    "Ferrari California Convertible 2012",
    "Ferrari 458 Italia Convertible 2012",
    "Ferrari 458 Italia Coupe 2012",
    "Fisker Karma Sedan 2012",
    "Ford F-450 Super Duty Crew Cab 2012",
    "Ford Mustang Convertible 2007",
    "Ford Freestar Minivan 2007",
    "Ford Expedition EL SUV 2009",
    "Ford Edge SUV 2012",
    "Ford Ranger SuperCab 2011",
    "Ford GT Coupe 2006",
    "Ford F-150 Regular Cab 2012",
    "Ford F-150 Regular Cab 2007",
    "Ford Focus Sedan 2007",
    "Ford E-Series Wagon Van 2012",
    "Ford Fiesta Sedan 2012",
    None,
    "GMC Savana Van 2012",
    "GMC Yukon Hybrid SUV 2012",
    "GMC Acadia SUV 2012",
    "GMC Canyon Extended Cab 2012",
    "Geo Metro Convertible 1993",
    "HUMMER H3T Crew Cab 2010",
    "HUMMER H2 SUT Crew Cab 2009",
    "Honda Odyssey Minivan 2012",
    "Honda Odyssey Minivan 2007",
    "Honda Accord Coupe 2012",
    None,
    "Hyundai Veloster Hatchback 2012",
    "Hyundai Santa Fe SUV 2012",
    "Hyundai Tucson SUV 2012",
    "Hyundai Veracruz SUV 2012",
    "Hyundai Sonata Hybrid Sedan 2012",
    "Hyundai Elantra Sedan 2007",
    "Hyundai Accent Sedan 2012",
    "Hyundai Genesis Sedan 2012",
    "Hyundai Sonata Sedan 2012",
    "Hyundai Elantra Touring Hatchback 2012",
    "Hyundai Azera Sedan 2012",
    "Infiniti G Coupe IPL 2012",
    "Infiniti QX56 SUV 2011",
    "Isuzu Ascender SUV 2008",
    "Jaguar XK XKR 2012",
    "Jeep Patriot SUV 2012",
    "Jeep Wrangler SUV 2012",
    "Jeep Liberty SUV 2012",
    "Jeep Grand Cherokee SUV 2012",
    "Jeep Compass SUV 2012",
    None,
    "Lamborghini Aventador Coupe 2012",
    None,
    "Lamborghini Diablo Coupe 2001",
    "Land Rover Range Rover SUV 2012",
    "Land Rover LR2 SUV 2012",
    "Lincoln Town Car Sedan 2011",
    None,
    "Maybach Landaulet Convertible 2012",
    None,
    "McLaren MP4-12C Coupe 2012",
    "Mercedes-Benz 300-Class Convertible 1993",
    "Mercedes-Benz C-Class Sedan 2012",
    "Mercedes-Benz SL-Class Coupe 2009",
    None,
    "Mercedes-Benz S-Class Sedan 2012",
    "Mercedes-Benz Sprinter Van 2012",
    "Mitsubishi Lancer Sedan 2012",
    "Nissan Leaf Hatchback 2012",
    "Nissan NV Passenger Van 2012",
    "Nissan Juke Hatchback 2012",
    "Nissan 240SX Coupe 1998",
    "Plymouth Neon Coupe 1999",
    "Porsche Panamera Sedan 2012",
    "Ram C/V Cargo Van Minivan 2012",
    "Rolls-Royce Phantom Drophead Coupe Convertible 2012",
    "Rolls-Royce Ghost Sedan 2012",
    "Rolls-Royce Phantom Sedan 2012",
    "Scion xD Hatchback 2012",
    "Spyker C8 Convertible 2009",
    "Spyker C8 Coupe 2009",
    "Suzuki Aerio Sedan 2007",
    "Suzuki Kizashi Sedan 2012",
    "Suzuki SX4 Hatchback 2012",
    "Suzuki SX4 Sedan 2012",
    "Tesla Model S Sedan 2012",
    "Toyota Sequoia SUV 2012",
    "Toyota Camry Sedan 2012",
    "Toyota Corolla Sedan 2012",
    "Toyota 4Runner SUV 2012",
    "Volkswagen Golf Hatchback 2012",
    "Volkswagen Golf Hatchback 1991",
    "Volkswagen Beetle Hatchback 2012",
    "Volvo C30 Hatchback 2012",
    "Volvo 240 Sedan 1993",
    "Volvo XC90 SUV 2007",
    "smart fortwo Convertible 2012"
]


from ultralytics import YOLO
import numpy as np

# --------------------------
# Load YOLO model for cropping
# --------------------------
device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
yolo_model = YOLO('yolov8n.pt')  # Using the small 'nano' model
print("YOLOv8 model loaded.")

# --------------------------
# Define YOLO cropping function
# --------------------------
def detect_and_crop_pil(pil_image, model=yolo_model, device=device_str, conf_thresh=0.25, pad_ratio=0.05):
    """
    Run YOLO on a PIL image and return a cropped PIL image around the best car detection.
    If no car is found, it returns the original image.
    """
    results = model(pil_image, imgsz=640, conf=conf_thresh, device=device, verbose=False)
    if len(results) == 0 or results[0].boxes is None or len(results[0].boxes) == 0:
        return pil_image

    r = results[0]
    boxes = r.boxes.xyxy.cpu().numpy()
    try:
        classes = r.boxes.cls.cpu().numpy().astype(int)
    except Exception:
        classes = np.zeros(len(boxes), dtype=int)

    # Prefer COCO car class (index 2)
    car_indices = np.where(classes == 2)[0]
    if len(car_indices) == 0:
        return pil_image # Return original if no car detected

    # Choose the car detection with the largest box area
    areas = (boxes[car_indices, 2] - boxes[car_indices, 0]) * (boxes[car_indices, 3] - boxes[car_indices, 1])
    best_idx = car_indices[np.argmax(areas)]
    x1, y1, x2, y2 = boxes[best_idx].astype(int)

    # Add padding
    w, h = x2 - x1, y2 - y1
    pad = int(max(w, h) * pad_ratio)
    x1, y1 = max(0, x1 - pad), max(0, y1 - pad)
    x2, y2 = min(pil_image.width, x2 + pad), min(pil_image.height, y2 + pad)

    return pil_image.crop((x1, y1, x2, y2))


# --------------------------
# 4. Preprocessing function
# --------------------------
def preprocess_image(img: Image.Image):
    cropped_img = detect_and_crop_pil(img)

    transform = transforms.Compose([
        transforms.Resize((224, 224)),           # match your training input size
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],         # if you used ImageNet pretrained normalization
            std=[0.229, 0.224, 0.225]
        )
    ])
    x = transform(cropped_img).unsqueeze(0)  # add batch dimension
    return x

# --------------------------
# 5. Prediction function
# --------------------------
def predict(img: Image.Image):
    x = preprocess_image(img)
    with torch.no_grad():
        logits = model(x)
        probs = torch.nn.functional.softmax(logits, dim=1)[0]
        top_idx = torch.argmax(probs).item()
        top_label = labels[top_idx] if labels[top_idx] is not None else f"Class {top_idx}"
        confidence = float(probs[top_idx])
        return {top_label: confidence}

# --------------------------
# 6. Gradio interface
# --------------------------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=1),
    title="Stanford Cars Classifier",
    description="Upload a car image and see the most likely predicted make/model and its confidence."
)

if __name__ == "__main__":
    demo.launch()