Upload 9 files

app.py

@@ -0,0 +1,43 @@
+import torch
+from torch import nn
+import torchvision
+import gradio as gr
+
+# Define and load my resnet50 model
+model = torchvision.models.resnet50()
+num_ftrs = model.fc.in_features
+model.fc = nn.Sequential(
+    # Add dropout layer with 50% probability
+    nn.Dropout(0.5),
+    # Add a linear layer in order to deal with 5 classes
+    nn.Linear(num_ftrs, 5),
+)
+
+model.load_state_dict(
+    torch.load("model/final_model_state_dict.pth", map_location=torch.device("cpu"))
+)
+model.eval()
+
+# Define the labels
+labels = ["bird", "cat", "dog", "horse", "sheep"]
+
+
+# Define the predict function
+def predict(inp):
+    inp = torchvision.transforms.ToTensor()(inp).unsqueeze(0)
+    with torch.no_grad():
+        prediction = model(inp)
+        # Map prediction to label
+        prediction = labels[prediction.argmax()]
+    return prediction
+
+
+# Define the gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=5),
+    examples=[["demo/input_imgs/cat.jpeg"], ["demo/input_imgs/dog.jpeg"]],
+)
+
+demo.launch()

model/final_model_state_dict.pth

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

requirements.txt

@@ -0,0 +1,8 @@
+pillow==10.3.0
+pycocotools==2.0.7
+scikit-learn==1.4.2
+tensorboard==2.16.2
+torch==2.3.0
+torch-tb-profiler==0.4.3
+torchvision==0.18.0
+gradio==4.32.2