Create app.py

app.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import gradio as gr
+
+transformer = models.ResNet18_Weights.IMAGENET1K_V1.transforms()
+class_names = ['anger', 'disgust', 'fear', 'happy', 'pain', 'sad']
+classes_count = len(class_names)
+
+model = models.resnet18(weights='DEFAULT')
+model.fc = nn.Sequential(
+    nn.Linear(512, classes_count)
+)
+model.load_state_dict(torch.load('./model_param.pt'), strict=False)
+
+def predict(img):
+  img = transformer(img).unsqueeze(0)
+
+  model.eval()
+
+  with torch.inference_mode():
+    pred = torch.softmax(model(img), dim=1)
+
+  pred_and_labels = {class_names[i] : pred[0][i].item() for i in range(len(pred[0])) }
+  
+  return pred_and_labels
+
+app = gr.Interface(
+    predict,
+    gr.Image(type='pil'),
+    gr.Label(num_top_classes=classes_count)               
+)
+app.launch()