Create app.py

app.py

@@ -0,0 +1,44 @@
+import gradio as gr
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
+from torchvision import transforms
+
+# Load your model and processor from Hugging Face Hub
+model_name = "1ancelot/base_rn"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+
+# Define the same preprocessing transformations as for validation
+val_test_resnet_combined_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),  # Convert image to tensor
+    transforms.Lambda(lambda img: processor(images=img, return_tensors="pt")['pixel_values'].squeeze(0))
+])
+
+# Define the prediction function
+def predict(img):
+    # Apply the same preprocessing transformations
+    img = val_test_resnet_combined_transforms(img)
+    
+    # Unsqueeze the image to add batch dimension, as the model expects batch input
+    img = img.unsqueeze(0)
+
+    # Perform inference
+    outputs = model(img)
+    
+    # Assuming the model returns logits
+    logits = outputs.logits
+    predicted_class = logits.argmax(-1).item()  # Get the predicted class index
+    
+    # Return the class as a text output
+    return f"Predicted class: {predicted_class}"
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict, 
+    inputs=gr.Image(type="pil"),  # Input is a PIL image
+    outputs="text",  # Output is the predicted class as text
+    title="Image Classification with base_rn"
+)
+
+# Launch the app
+iface.launch()