Spaces:
Sleeping
Sleeping
File size: 2,946 Bytes
11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b f726feb 11cfb1b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
# --------------------
# Class Mapping
# --------------------
class_to_idx = {
'Acura': 0, 'Alfa Romeo': 1, 'Aston Martin': 2, 'Audi': 3, 'BMW': 4,
'Bentley': 5, 'Bugatti': 6, 'Buick': 7, 'Cadillac': 8, 'Chevrolet': 9,
'Chrysler': 10, 'Citroen': 11, 'Daewoo': 12, 'Dodge': 13, 'Ferrari': 14,
'Fiat': 15, 'Ford': 16, 'GMC': 17, 'Genesis': 18, 'Honda': 19,
'Hudson': 20, 'Hyundai': 21, 'Infiniti': 22, 'Jaguar': 23, 'Jeep': 24,
'Kia': 25, 'Land Rover': 26, 'Lexus': 27, 'Lincoln': 28, 'MG': 29,
'Maserati': 30, 'Mazda': 31, 'Mercedes-Benz': 32, 'Mini': 33,
'Mitsubishi': 34, 'Nissan': 35, 'Oldsmobile': 36, 'Peugeot': 37,
'Pontiac': 38, 'Porsche': 39, 'Ram Trucks': 40, 'Renault': 41,
'Saab': 42, 'Studebaker': 43, 'Subaru': 44, 'Suzuki': 45, 'Tesla': 46,
'Toyota': 47, 'Volkswagen': 48, 'Volvo': 49
}
idx_to_class = {v: k for k, v in class_to_idx.items()}
# --------------------
# Image Transform
# --------------------
transform = transforms.Compose([
transforms.Lambda(lambda x: x.convert("RGB")),
transforms.Resize((224,224)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.2),
transforms.RandomRotation(20),
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),
transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
transforms.ToTensor(),
transforms.Normalize([0.5]*3, [0.5]*3)
])
# --------------------
# Load Model
# --------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load pretrained ResNet50 correctly
base_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
# Replace final fully connected layer (your head)
in_features = base_model.fc.in_features # 2048 for ResNet50
base_model.fc = nn.Sequential(
nn.Linear(in_features, 512),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(512, 50) # 50 classes
)
# Load state dict
state_dict = torch.load("best_model.pth", map_location=device)
base_model.load_state_dict(state_dict)
base_model = base_model.to(device)
base_model.eval()
# --------------------
# Prediction Function
# --------------------
def predict(img):
img = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
outputs = base_model(img)
probs = torch.softmax(outputs, dim=1)[0]
top5_prob, top5_idx = torch.topk(probs, 5)
result = {idx_to_class[idx.item()]: float(top5_prob[i]) for i, idx in enumerate(top5_idx)}
return result
# --------------------
# Gradio UI
# --------------------
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=5),
title="Car Brand Classifier",
description="Upload a car image to predict its brand."
)
if __name__ == "__main__":
demo.launch()
|