Update app.py

app.py

@@ -1,54 +1,67 @@
-import torch
-import torch.nn as nn
-import torchvision.transforms as transforms
-from PIL import Image
-import gradio as gr
-
-
-class ModifiedLargeNet(nn.Module):
-    def __init__(self):
-        super(ModifiedLargeNet, self).__init__()
-        self.name = "modified_large"
-        self.fc1 = nn.Linear(128 * 128 * 3, 256)
-        self.fc2 = nn.Linear(256, 128)
-        self.fc3 = nn.Linear(128, 3)  # 3 classes: Rope, Hammer, Other
-
-    def forward(self, x):
-        x = x.view(-1, 128 * 128 * 3)
-        x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = self.fc3(x)
-        return x
-
-
-model = ModifiedLargeNet()
-model.load_state_dict(torch.load("modified_large_net.pt", map_location=torch.device("cpu")))
-model.eval()
-
-
-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]),
-])
-
-
-def predict(image):
-
-    image = transform(image).unsqueeze(0) 
-    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)}
-
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),  
-    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'.",
-)
-
-if __name__ == "__main__":
-    interface.launch()
+import os
+import subprocess
+import sys
+
+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"])
+
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+import gradio as gr
+
+
+
+class ModifiedLargeNet(nn.Module):
+    def __init__(self):
+        super(ModifiedLargeNet, self).__init__()
+        self.name = "modified_large"
+        self.fc1 = nn.Linear(128 * 128 * 3, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, 3)  # 3 classes: Rope, Hammer, Other
+
+    def forward(self, x):
+        x = x.view(-1, 128 * 128 * 3)
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+
+model = ModifiedLargeNet()
+model.load_state_dict(torch.load("modified_large_net.pt", map_location=torch.device("cpu")))
+model.eval()
+
+
+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]),
+])
+
+
+def predict(image):
+
+    image = transform(image).unsqueeze(0) 
+    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)}
+
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),  
+    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'.",
+)
+
+if __name__ == "__main__":
+    interface.launch()