Update app.py

app.py

@@ -1,13 +1,15 @@
 import gradio as gr
+import torch
 from torchvision import models, transforms
 from PIL import Image
-import torch
+import requests
+from io import BytesIO
 
-# Load the pre-trained MobileNetV2 model
+# Load pre-trained MobileNetV2 model (you can choose another model as needed)
 model = models.mobilenet_v2(pretrained=True)
 model.eval()
 
-# Image transformation for input
+# Define the image transformation (resize, normalization)
 transform = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
@@ -15,17 +17,27 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-def classify_image(image):
-    # Apply transformations
-    img_tensor = transform(image).unsqueeze(0)
+# Download the ImageNet class labels (you can replace this with your own if needed)
+LABELS_URL = "https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json"
+class_idx = requests.get(LABELS_URL).json()
+idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]
+
+# Function to perform image inference
+def predict_image(image):
+    image = Image.open(BytesIO(image)).convert("RGB")
+    image = transform(image).unsqueeze(0)
     
     # Perform inference
     with torch.no_grad():
-        outputs = model(img_tensor)
-        _, predicted_class = torch.max(outputs, 1)
-        
-    return predicted_class.item()
+        output = model(image)
+    
+    # Get the predicted label
+    _, predicted_class = torch.max(output, 1)
+    label = idx2label[predicted_class.item()]
+    return label
 
 # Gradio interface
-interface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil"), outputs="text")
-interface.launch()
+with gr.Interface(fn=predict_image, 
+                  inputs=gr.inputs.Image(type="bytes"), 
+                  outputs=gr.outputs.Textbox()) as demo:
+    demo.launch(debug=True)