first commmit

app.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import pytorch_lightning as pl
+
+import torchvision.models as models
+import gradio as gr
+
+
+class Net(pl.LightningModule):
+
+    def __init__(self):
+        super().__init__()
+
+        # init a pretrained resnet
+        backbone = models.resnet50()
+        num_filters = backbone.fc.in_features
+        layers = list(backbone.children())[:-1]
+        self.feature_extractor = nn.Sequential(*layers)
+
+        # use the pretrained model to classify cifar-10 (10 image classes)
+        num_target_classes = 2
+        self.classifier = nn.Linear(num_filters, num_target_classes)
+
+    def forward(self, x):
+        self.feature_extractor.eval()
+        with torch.no_grad():
+            representations = self.feature_extractor(x).flatten(1)
+        x = self.classifier(representations)
+        return x
+
+
+net = Net()
+net = net.load_from_checkpoint(
+    './resnet50_transfer.ckpt', map_location=torch.device('cpu'))
+net.eval()
+
+
+labels = ['Cat', 'Dog']
+val_transform = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+    ]
+)
+
+
+def inference(img):
+    img = val_transform(img)
+    img = torch.unsqueeze(img, 0)
+    with torch.no_grad():
+        pred = net(img)
+    pred = torch.softmax(pred, dim=1)
+    scores = pred.detach().numpy()[0]
+    confidences = {labels[i]: float(scores[i]) for i in range(len(labels))}
+    return confidences
+
+
+iface = gr.Interface(
+    fn=inference,
+    inputs=gr.components.Image(type="numpy"),
+    outputs=gr.Label(num_top_classes=2)
+)
+
+iface.launch()

requirements.txt

@@ -0,0 +1,3 @@
+pytorch
+pytorch-lightning
+torchvision