app.py

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()