Spaces:
Runtime error
Runtime error
| import torch | |
| from torch import nn | |
| from torchvision import transforms | |
| import gradio as gr | |
| title = "PyTorch Cat vs Dog" | |
| description = "Classifying cats and dogs with Pytorch" | |
| article = "<p style='text-align: center'><a href='https://github.com/TirendazAcademy'>Github Repo</a></p>" | |
| # The model architecture | |
| class ImageClassifier(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.conv_layer_1 = nn.Sequential( | |
| nn.Conv2d(3, 64, 3, padding=1), | |
| nn.ReLU(), | |
| nn.BatchNorm2d(64), | |
| nn.MaxPool2d(2)) | |
| self.conv_layer_2 = nn.Sequential( | |
| nn.Conv2d(64, 512, 3, padding=1), | |
| nn.ReLU(), | |
| nn.BatchNorm2d(512), | |
| nn.MaxPool2d(2)) | |
| self.conv_layer_3 = nn.Sequential( | |
| nn.Conv2d(512, 512, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.BatchNorm2d(512), | |
| nn.MaxPool2d(2)) | |
| self.classifier = nn.Sequential( | |
| nn.Flatten(), | |
| nn.Linear(in_features=512*3*3, out_features=2) | |
| ) | |
| def forward(self, x: torch.Tensor): | |
| x = self.conv_layer_1(x) | |
| x = self.conv_layer_2(x) | |
| x = self.conv_layer_3(x) | |
| x = self.conv_layer_3(x) | |
| x = self.conv_layer_3(x) | |
| x = self.conv_layer_3(x) | |
| x = self.classifier(x) | |
| return x | |
| model = ImageClassifier() | |
| model.load_state_dict(torch.load('image_classifier.pth')) | |
| def predict(inp): | |
| image_transform = transforms.Compose([ transforms.Resize(size=(224,224)), transforms.ToTensor()]) | |
| labels = ['cat', 'dog'] | |
| inp = image_transform(inp).unsqueeze(dim=0) | |
| with torch.no_grad(): | |
| prediction = torch.nn.functional.softmax(model(inp)) | |
| confidences = {labels[i]: float(prediction.squeeze()[i]) for i in range(len(labels))} | |
| return confidences | |
| gr.Interface(fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Label(num_top_classes=2), | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=['cat.jpg', 'dog.jpg']).launch() |