Add application file

app.py

@@ -0,0 +1,62 @@
+"""
+This module provides an interface for classifying images using the ResNet-18 model.
+The interface allows users to upload an image and receive the top 3 predicted labels.
+"""
+
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from datasets import load_dataset
+from PIL import Image
+
+# Load dataset and get test image
+dataset = load_dataset("huggingface/cats-image", trust_remote_code=True)
+test_image = dataset["test"]["image"][0]
+
+# Initialize the image processor and model
+image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-18")
+model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18").to("cuda")
+
+@spaces.GPU
+def predict(image: Image, top_k: int = 3) -> dict:
+    """
+    Predicts the top 'top_k' labels for an image using the ResNet-18 model.
+
+    Args:
+        image (Image): The input image as a PIL Image object.
+        top_k (int): The number of top predictions to return.
+
+    Returns:
+        dict: A dictionary with the top 'top_k' labels and their probabilities.
+    """
+    inputs = image_processor(image, return_tensors="pt").to("cuda")
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    # Apply softmax to logits to get probabilities
+    probabilities = torch.softmax(logits, dim=-1)
+
+    # Get the top 'top_k' probabilities and their corresponding indices
+    top_k_probs, top_k_indices = torch.topk(input=probabilities, k=top_k, dim=-1)
+
+    # Map the indices to labels and probabilities
+    predicted_labels = {
+        model.config.id2label[idx.item()]: prob.item()
+        for idx, prob in zip(top_k_indices[0], top_k_probs[0])
+    }
+    
+    return predicted_labels
+
+# Define the Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Classifying Images with ResNet-18",
+    description="Upload an image to predict the top 3 labels.",
+    examples=[test_image]
+)
+
+# Launch the Gradio interface
+demo.launch()