Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+import yaml
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+import os
+
+CONFIG_PATH = 'staging_config.yaml'
+CHECKPOINT_FILENAME = 'model.pt'
+
+
+def get_model(model_name, num_classes):
+    """Factory function to create a model shell for loading weights."""
+    model = None
+    if model_name == "efficientnet_b0":
+        model = models.efficientnet_b0(weights=None)
+        num_ftrs = model.classifier[1].in_features
+        model.classifier[1] = nn.Linear(num_ftrs, num_classes)
+    elif model_name == "resnet50":
+        model = models.resnet50(weights=None)
+        num_ftrs = model.fc.in_features
+        model.fc = nn.Linear(num_ftrs, num_classes)
+    elif model_name == "vit_b_16":
+        model = models.vit_b_16(weights=None)
+        num_ftrs = model.heads.head.in_features
+        model.heads.head = nn.Linear(num_ftrs, num_classes)
+    else:
+        raise ValueError(f"Model '{model_name}' not supported.")
+    return model
+
+def load_checkpoint(checkpoint_path, device):
+    """Loads a checkpoint and returns the model and class mapping."""
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Checkpoint file not found at: {checkpoint_path}")
+    
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model_name = checkpoint['model_name']
+    class_to_idx = checkpoint['class_to_idx']
+    
+    model = get_model(model_name, num_classes=1) 
+    model.load_state_dict(checkpoint['state_dict'])
+    model.to(device)
+    model.eval()
+    
+    idx_to_class = {v: k for k, v in class_to_idx.items()}
+    return model, idx_to_class
+
+
+try:
+    with open(CONFIG_PATH, 'r') as f:
+        config = yaml.safe_load(f)
+except FileNotFoundError:
+    raise RuntimeError(f"ERROR: Config file not found at '{CONFIG_PATH}'. Make sure it's uploaded to the Space.")
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL, IDX_TO_CLASS = load_checkpoint(CHECKPOINT_FILENAME, DEVICE)
+print(f"Model loaded successfully on {DEVICE}.")
+print(f"Class mapping: {IDX_TO_CLASS}")
+
+
+IMG_SIZE = config['data_params']['image_size']
+inference_transform = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def predict(pil_image):
+    """
+    Runs prediction on a single PIL image and returns a dictionary of class probabilities.
+    Gradio's `Label` component expects a dictionary format for its output.
+    """
+    if pil_image is None:
+        return None
+        
+    pil_image = pil_image.convert("RGB")
+    
+    image_tensor = inference_transform(pil_image).unsqueeze(0).to(DEVICE)
+    
+    with torch.no_grad():
+        output = MODEL(image_tensor)
+        prob = torch.sigmoid(output).item()
+        
+    class_0_name = IDX_TO_CLASS.get(0, "Class 0")
+    class_1_name = IDX_TO_CLASS.get(1, "Class 1")
+    
+    confidences = {
+        class_0_name: 1 - prob,
+        class_1_name: prob
+    }
+    return confidences
+
+
+
+title = "Image Classifier API"
+description = """
+Upload an image and the model will predict its class. 
+This model was trained to distinguish between two classes.
+The API returns the probabilities for each class.
+"""
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Label(num_top_classes=2, label="Predictions"),
+    title=title,
+    description=description,
+)
+
+iface.launch()