first commit

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+vitb16_v1.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,58 @@
+import model
+import os
+import torch
+import gradio as gr
+from timeit import default_timer as timer
+from torch import nn
+from torchvision import transforms
+
+class_names = ['chicken', 'elephant' ,'sheep']
+
+DEVICE = 'cpu'
+
+vit_model, vit_transform = model.create_vitb16_model(
+  num_classes=len(class_names)
+)
+
+vit_model.load_state_dict(
+  torch.load(
+    f='vitb16_v1.pth',
+    map_location=torch.device(DEVICE),
+  )
+)
+
+def predict_single_image(image):
+  start_time = timer()
+  image = vit_transform(image).unsqueeze(0).to(DEVICE)
+  vit_model.eval()
+  logits = vit_model(image)
+  with torch.inference_mode():
+    probs = torch.softmax(logits, dim=1)
+
+  classes_and_probs = {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
+  inference_time = round(timer() - start_time, 5)
+
+  return classes_and_probs, inference_time
+
+title = 'AnimalsVision \U0001F413\U0001F418\U0001F411'
+description = 'A ViT computer vision model to classify images of animals as chicken, elephant or sheep.'
+article = 'GitHub Repo: https://github.com/oschan77/AnimalsVision-App'
+
+examples = [['examples/' + example] for example in os.listdir('examples/')]
+
+app = gr.Interface(
+  fn=predict_single_image,
+  inputs=gr.Image(type='pil'),
+  outputs=[
+    gr.Label(num_top_classes=len(class_names), label='Predictions'),
+    gr.Number(label='Prediction time (sec)'),
+  ],
+  examples=examples,
+  title=title,
+  description=description,
+  article=article,
+)
+
+app.launch(
+  share=True,
+)

model.py

@@ -0,0 +1,18 @@
+import torchvision
+import torch.nn as nn
+
+def create_vitb16_model(
+  num_classes: int,
+):
+  vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT
+  vit_model = torchvision.models.vit_b_16(weights=vit_weights)
+  vit_transform = vit_weights.transforms()
+
+  for param in vit_model.parameters():
+    param.requires_grad = False
+
+  vit_model.heads = nn.Sequential(
+    nn.Linear(in_features=768, out_features=num_classes, bias=True),
+  )
+
+  return vit_model, vit_transform

requirements.txt

@@ -0,0 +1,3 @@
+torch==1.13.1
+torchvision==0.14.1
+gradio==3.17.1

vitb16_v1.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a9ee27519da16ff37b67fe57efb0383d9633de799f588590355e43abb636241
+size 343264069