Update app.py

app.py

@@ -1,5 +1,6 @@
 import subprocess
 import sys
+import os
 
 # Function to install or reinstall specific packages
 def install(package):
@@ -75,25 +76,38 @@ transform = transforms.Compose([
 
 # Prediction function
 def predict(image):
-    # Verify input image is a PIL image
+    if image is None:
+        raise ValueError("Please provide an image")
+        
+    # Convert to PIL Image if necessary
     if not isinstance(image, Image.Image):
-        raise ValueError("Input must be a PIL Image.")
+        try:
+            image = Image.fromarray(image)
+        except Exception as e:
+            raise ValueError(f"Failed to convert input to PIL Image: {str(e)}")
 
     # Transform and predict
-    image = transform(image).unsqueeze(0)  # Add batch dimension
-    with torch.no_grad():
-        outputs = model(image)
-        probabilities = torch.softmax(outputs, dim=1).numpy()[0]
-        classes = ["Rope", "Hammer", "Other"]
-        return {cls: float(prob) for cls, prob in zip(classes, probabilities)}
+    try:
+        image = transform(image).unsqueeze(0)  # Add batch dimension
+        with torch.no_grad():
+            outputs = model(image)
+            probabilities = torch.softmax(outputs, dim=1).numpy()[0]
+            classes = ["Rope", "Hammer", "Other"]
+            return {cls: float(prob) for cls, prob in zip(classes, probabilities)}
+    except Exception as e:
+        raise ValueError(f"Error during prediction: {str(e)}")
 
 # Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),  # Ensure input is a PIL image
-    outputs=gr.Label(num_top_classes=3),  # Display top 3 class probabilities
+    inputs=gr.Image(),  # Remove type="pil" constraint
+    outputs=gr.Label(num_top_classes=3),
     title="Mechanical Tools Classifier",
     description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",
+    examples=[
+        ["example_rope.jpg"],
+        ["example_hammer.jpg"],
+    ] if os.path.exists("example_rope.jpg") else None  # Optional examples
 )
 
 # Launch the interface