Upload 3 files

app.py

@@ -0,0 +1,54 @@
+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()

modified_large_net.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10ca86bf582a3820406346d57fa09083279fbb4aa1862d44b966144a556d228c
+size 1086788

requriements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+gradio
+Pillow