Update app.py

app.py

@@ -1,29 +1,40 @@
-import os
 import subprocess
 import sys
 
+# Ensure required libraries are installed
+def install(package):
+    subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+# Install torch and torchvision
 try:
     import torch
 except ImportError:
-    subprocess.check_call([sys.executable, "-m", "pip", "install", 
-                           "torch==2.0.1+cpu", 
-                           "torchvision==0.15.2+cpu", 
-                           "-f", "https://download.pytorch.org/whl/torch_stable.html"])
+    install("torch==2.0.1+cpu")
+    install("torchvision==0.15.2+cpu")
+    install("-f https://download.pytorch.org/whl/torch_stable.html")
 
+# Install numpy
 try:
     import numpy as np
 except ImportError:
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy<2"])
+    install("numpy<2")
+
+# Install Pillow
+try:
+    from PIL import Image
+except ImportError:
+    install("Pillow==9.5.0")
 
+# Imports
 import torch
-import numpy
 import torch.nn as nn
+import torch.nn.functional as F
 import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
 
 
-
+# Define the model
 class ModifiedLargeNet(nn.Module):
     def __init__(self):
         super(ModifiedLargeNet, self).__init__()
@@ -32,7 +43,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)  # classify into "Rope"/"Hammer"/"others"
 
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
@@ -40,38 +51,44 @@ class ModifiedLargeNet(nn.Module):
         x = x.view(-1, 10 * 29 * 29)
         x = F.relu(self.fc1(x))
         x = self.fc2(x)
-        x = x.squeeze(1) # Flatten to [batch_size]
         return x
 
 
+# Load the trained model
 model = ModifiedLargeNet()
 model.load_state_dict(torch.load("modified_large_net.pt", map_location=torch.device("cpu")))
 model.eval()
 
-
+# Define image transformation pipeline
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
 ])
 
-
+# Prediction function
 def predict(image):
+    # Verify input image is a PIL image
+    if not isinstance(image, Image.Image):
+        raise ValueError("Input must be a PIL Image.")
 
-    image = transform(image).unsqueeze(0) 
+    # 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)}
 
+# Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),  
-    outputs=gr.Label(num_top_classes=3),  
+    inputs=gr.Image(type="pil"),  # Ensure input is a PIL image
+    outputs=gr.Label(num_top_classes=3),  # Display top 3 class probabilities
     title="Mechanical Tools Classifier",
     description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",
 )
 
+# Launch the interface
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(share=True)  # Add `share=True` for a public link