AshProg/AppliedMachineLearning_BirdClassifierInterface

1"""2Gradio App for Bird Species Classification3Deployed on Hugging Face Spaces4"""5 6import gradio as gr7import torch8import torch.nn as nn9from torchvision import transforms10from torchvision.models import convnext_base11from PIL import Image12import json13 14# Load class names15with open('class_names.json', 'r') as f:16    class_names = json.load(f)17 18# Device configuration19device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')20 21# Create model architecture (same as training)22def create_model(num_classes=200):23    """Create ConvNeXt model with same architecture as training"""24    model = convnext_base(weights=None)25    26    # Same classifier architecture as training27    num_ftrs = model.classifier[2].in_features28    model.classifier = nn.Sequential(29        nn.Flatten(1),30        nn.LayerNorm((num_ftrs,)),31        nn.Dropout(0.6),32        nn.Linear(num_ftrs, 512),33        nn.GELU(),34        nn.Dropout(0.5),35        nn.Linear(512, num_classes)36    )37    38    return model39 40# Load the trained model41print("Loading model...")42model = create_model(num_classes=200)43 44# Load weights45import gradio as gr46import torch47import torch.nn as nn48from torchvision import transforms49from torchvision.models import convnext_base50from PIL import Image51import json52from huggingface_hub import hf_hub_download53 54# Load class names55with open('class_names.json', 'r') as f:56    class_names = json.load(f)57 58# Device configuration59device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')60 61# Create model architecture (same as training)62def create_model(num_classes=200):63    """Create ConvNeXt model with same architecture as training"""64    model = convnext_base(weights=None)65    66    # Same classifier architecture as training67    num_ftrs = model.classifier[2].in_features68    model.classifier = nn.Sequential(69        nn.Flatten(1),70        nn.LayerNorm((num_ftrs,)),71        nn.Dropout(0.6),72        nn.Linear(num_ftrs, 512),73        nn.GELU(),74        nn.Dropout(0.5),75        nn.Linear(512, num_classes)76    )77    78    return model79 80# Download model from Hugging Face Model Hub81print("Downloading model from Hugging Face Model Hub...")82model_path = hf_hub_download(83    repo_id="AshProg/bird-classifier-convnext",84    filename="final_model.pth"85)86 87# Load the trained model88model = create_model(num_classes=200)89checkpoint = torch.load(model_path, map_location=device)90if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:91    model.load_state_dict(checkpoint['model_state_dict'])92    if 'val_acc' in checkpoint:93        val_acc = checkpoint['val_acc']94        print(f"Model loaded! Validation accuracy: {val_acc:.2f}%")95else:96    model.load_state_dict(checkpoint)97    print("Model loaded!")98 99model = model.to(device)100model.eval()101 102# Image preprocessing (same as validation transforms)103transform = transforms.Compose([104    transforms.Resize((224, 224)),105    transforms.ToTensor(),106    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])107])108 109def predict(image):110    """111    Make prediction on uploaded image112    113    Args:114        image: PIL Image115        116    Returns:117        dict: Top 5 predictions with confidence scores118    """119    # Preprocess image120    img_tensor = transform(image).unsqueeze(0).to(device)121    122    # Make prediction123    with torch.no_grad():124        outputs = model(img_tensor)125        probabilities = torch.nn.functional.softmax(outputs, dim=1)126        127        # Get top 5 predictions128        top5_prob, top5_idx = torch.topk(probabilities, 5)129        130    # Format results131    results = {}132    for i in range(5):133        class_id = top5_idx[0][i].item()134        prob = top5_prob[0][i].item()135        species_name = class_names.get(str(class_id), f"Class {class_id}")136        results[species_name] = float(prob)137    138    return results139 140# Create Gradio interface141title = "🐦 Bird Species Classification"142description = """143Upload an image of a bird and the model will predict the species!144 145**Model Details:**146- Architecture: ConvNeXt-Base (87M parameters)147- Dataset: CUB-200-2011 (200 bird species)148- Test Accuracy: 83.64%149- Average Per-Class Accuracy: 83.29%150 151Upload a clear image of a bird to get started!152"""153 154article = """155### About This Model156 157This bird classifier was trained on the CUB-200-2011 dataset containing 200 North American bird species.158 159**Key Features:**160- ✅ 200 bird species classification161- ✅ State-of-the-art ConvNeXt architecture162- ✅ 83.64% test accuracy163- ✅ Real-time inference164 165"""166 167examples = [168    # You can add example images here if you have them169    # ["examples/bird1.jpg"],170    # ["examples/bird2.jpg"],171]172 173# Create interface174iface = gr.Interface(175    fn=predict,176    inputs=gr.Image(type="pil", label="Upload Bird Image"),177    outputs=gr.Label(num_top_classes=5, label="Top 5 Predictions"),178    title=title,179    description=description,180    article=article,181    examples=examples if examples else None,182    theme=gr.themes.Soft(),183    allow_flagging="never",184)185 186# Launch the app187if __name__ == "__main__":188    iface.launch()189