File size: 2,696 Bytes
b750cc9 f70ae67 b750cc9 d6b84d7 f70ae67 b750cc9 f70ae67 b750cc9 f70ae67 b750cc9 f70ae67 b750cc9 f70ae67 d6b84d7 b750cc9 f70ae67 b750cc9 f70ae67 d6b84d7 b750cc9 d6b84d7 b750cc9 f70ae67 b750cc9 f70ae67 b750cc9 f70ae67 d6b84d7 0bae8cc f70ae67 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download
# 1. Define the architecture EXACTLY as it was in training notebook
class VehicleClassifier(nn.Module):
def __init__(self):
super(VehicleClassifier, self).__init__()
# Block 1: 3 -> 16 channels. Input 224x224 -> Output 112x112 (after pool)
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(16)
# Block 2: 16 -> 32 channels. Output 56x56 (after pool)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(32)
# Block 3: 32 -> 64 channels. Output 28x28 (after pool)
self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.bn3 = nn.BatchNorm2d(64)
self.pool = nn.MaxPool2d(2, 2)
# Fully Connected Layers
self.fc1 = nn.Linear(64 * 28 * 28, 512)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, 8) # 8 Classes
def forward(self, x):
x = self.pool(F.relu(self.bn1(self.conv1(x))))
x = self.pool(F.relu(self.bn2(self.conv2(x))))
x = self.pool(F.relu(self.bn3(self.conv3(x))))
x = x.view(-1, 64 * 28 * 28) # Flatten
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)
return x
# 2. Setup Device and Load Model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = VehicleClassifier().to(device)
# Download and Load Weights
repo_id = "abhiprd20/vehicle_classification_model-utd"
checkpoint_path = hf_hub_download(repo_id=repo_id, filename="model.pth")
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
model.eval()
# 3. Prediction Logic
classes = ['Bicycle', 'Bus', 'Car', 'Motorcycle', 'NonVehicles', 'Taxi', 'Truck', 'Van']
data_transforms = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def predict(img):
if img is None: return None
img_tensor = data_transforms(img).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(img_tensor)
probs = torch.nn.functional.softmax(outputs[0], dim=0)
confidences = {classes[i]: float(probs[i]) for i in range(8)}
return confidences
# 4. Interface
gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=3),
title="🏎️ Autonomous Vehicle Perception Demo",
description="Custom CNN Model classifying type of vehicles."
).launch() |