Update app.py

app.py

@@ -6,36 +6,33 @@ import os
 def install(package):
     subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
 
-# Ensure NumPy and pandas are compatible
+# First, ensure NumPy is installed with the correct version
 try:
     import numpy as np
-    import pandas as pd
-    if not (np.__version__.startswith("1.24")):
-        print(f"Detected incompatible versions. Reinstalling NumPy...")
+    if not np.__version__.startswith("1.24"):
+        print("Installing compatible NumPy version...")
         install("numpy==1.24.3")
 except ImportError:
-    print("NumPy or pandas not found. Installing compatible versions...")
+    print("NumPy not found. Installing...")
     install("numpy==1.24.3")
 
-# Ensure other dependencies are installed with specific versions
-try:
-    import torch
-    import torchvision
-except ImportError:
-    install("torch==2.0.1")
-    install("torchvision==0.15.2")
+# Then install other dependencies
+packages = {
+    "torch": "2.0.1",
+    "torchvision": "0.15.2",
+    "Pillow": "9.5.0",
+    "gradio": "3.50.2"
+}
 
-try:
-    from PIL import Image
-except ImportError:
-    install("Pillow==9.5.0")
-
-try:
-    import gradio as gr
-except ImportError:
-    install("gradio==3.50.2")
+for package, version in packages.items():
+    try:
+        __import__(package.lower())
+    except ImportError:
+        print(f"Installing {package}...")
+        install(f"{package}=={version}")
 
-# Import libraries after ensuring installations
+# Import all required libraries
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -52,7 +49,7 @@ class ModifiedLargeNet(nn.Module):
         self.pool = nn.MaxPool2d(2, 2)
         self.conv2 = nn.Conv2d(5, 10, 5)
         self.fc1 = nn.Linear(10 * 29 * 29, 32)
-        self.fc2 = nn.Linear(32, 3)  # classify into "Rope"/"Hammer"/"others"
+        self.fc2 = nn.Linear(32, 3)
 
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
@@ -74,42 +71,50 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
 ])
 
-# Prediction function
-def predict(image):
+def process_image(image):
     if image is None:
-        raise ValueError("Please provide an image")
-        
-    # Convert to PIL Image if necessary
-    if not isinstance(image, Image.Image):
-        try:
-            image = Image.fromarray(image)
-        except Exception as e:
-            raise ValueError(f"Failed to convert input to PIL Image: {str(e)}")
+        return None
+    
+    # Convert to RGB if necessary
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    return image
 
-    # Transform and predict
+def predict(image):
+    if image is None:
+        return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
+    
     try:
-        image = transform(image).unsqueeze(0)  # Add batch dimension
+        # Process the image
+        processed_image = process_image(Image.fromarray(image))
+        if processed_image is None:
+            return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
+        
+        # Transform for model
+        tensor_image = transform(processed_image).unsqueeze(0)
+        
+        # Make prediction
         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)}
+            outputs = model(tensor_image)
+            probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
+            
+        # Return results
+        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)}")
+        print(f"Prediction error: {str(e)}")
+        return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
 
 # Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(),  # Remove type="pil" constraint
+    inputs=gr.Image(type="numpy"),
     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
 if __name__ == "__main__":
-    interface.launch()  # Removed share=True for Hugging Face Spaces
+    interface.launch()