Spaces:
Sleeping
Sleeping
File size: 2,393 Bytes
19e3c32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
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()
|