Update app.py

app.py

@@ -1,44 +1,47 @@
 import gradio as gr
-from transformers import ViTImageProcessor, ViTForImageClassification
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
 import torch
 
-# Load the pre-trained ViT model and processor
-processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')  # Using the in21k pre-trained model
-model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k')
+# Load the ResNet-50 model and feature extractor
+model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
+feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
 
-# Inference function for predicting with ViT model
+# Define the prediction function
 def predict(image):
     try:
         # Ensure the image is in PIL format
-        if isinstance(image, str):
-            image = Image.open(image)
-        
-        # Preprocess the input image using the processor, with padding enabled
-        inputs = processor(images=image, return_tensors="pt", padding=True)
+        if not isinstance(image, Image.Image):
+            return "Invalid image format. Please upload a valid image."
+
+        # Preprocess the image using the feature extractor
+        inputs = feature_extractor(images=image, return_tensors="pt")
         
-        # Get the model's predictions
+        # Perform inference
         with torch.no_grad():
             outputs = model(**inputs)
             logits = outputs.logits
-        
-        # Get the predicted class index (class with the highest logit value)
+
+        # Get the class with the highest score
         predicted_class_idx = logits.argmax(-1).item()
 
-        # Get the human-readable label for the predicted class
+        # Map the predicted index to its human-readable label
         predicted_class_label = model.config.id2label[predicted_class_idx]
-        
+
         return f"Predicted class: {predicted_class_label}"
 
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Create Gradio Interface (Note the change here: `gr.Image` and `gr.Text`)
-interface = gr.Interface(fn=predict, 
-                         inputs=gr.Image(type="pil", label="Upload Image"),
-                         outputs=gr.Text(), 
-                         title="ViT Image Classification (ImageNet)", 
-                         description="Upload an image to classify it into one of the 1000 ImageNet classes.")
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Text(label="Prediction"),
+    title="ResNet-50 Image Classification",
+    description="Upload an image to classify it into one of the ImageNet classes using the ResNet-50 model."
+)
 
-# Launch the interface
-interface.launch()
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()