init

app.py

@@ -0,0 +1,76 @@
+import torch
+import torchvision.transforms as transforms
+import gradio as gr
+from PIL import Image
+import torch.nn as nn
+import torch.nn.functional as F
+
+def get_model_name(name, batch_size, learning_rate, epoch):
+    """ Generate a name for the model consisting of all the hyperparameter values
+
+    Args:
+        config: Configuration object containing the hyperparameters
+    Returns:
+        path: A string with the hyperparameter name and value concatenated
+    """
+    path = "model_{0}_bs{1}_lr{2}_epoch{3}".format(name,
+                                                   batch_size,
+                                                   learning_rate,
+                                                   epoch)
+    return path
+
+class LargeNet(nn.Module):
+    def __init__(self):
+        super(LargeNet, self).__init__()
+        self.name = "large"
+        self.conv1 = nn.Conv2d(3, 5, 5)
+        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, 8)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        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
+
+transform = transforms.Compose([
+    transforms.Resize((128, 128)),  # Resize to 128x128
+    transforms.ToTensor(),         # Convert to Tensor
+    transforms.Normalize((0.5,), (0.5,))  # Normalize to [-1, 1]
+])
+
+def load_model():
+    net = LargeNet() #small or large network
+    model_path = get_model_name(net.name, batch_size=128, learning_rate=0.001, epoch=29)
+    state = torch.load(model_path)
+    net.load_state_dict(state)
+    
+    net.eval()
+    return net
+
+class_names = ["Gasoline_Can", "Pebbels", "pliers", "Screw_Driver", "Toolbox", "Wrench", "other"]
+
+
+def predict(image):
+    model = load_model()
+    image = transform(image).unsqueeze(0)  
+    with torch.no_grad():
+        output = model(image)
+        _, pred = torch.max(output, 1)
+    return class_names[pred.item()]
+
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="label",
+    title="Mechanical Tools Classifier",
+    description="Upload an image to classify it as one of the mechanical tools."
+)
+
+if __name__ == "__main__":
+    interface.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+gradio
+Pillow
+numpy
+pandas