Create app.py

app.py

@@ -0,0 +1,38 @@
+import gradio as gr
+import torchvision.transforms as transforms
+from PIL import Image
+import torch
+from timm.models import create_model
+def predict(input_img):
+    input_img =  Image.fromarray(np.uint8(input_img))
+    model1 = create_model(
+        'resnet50',
+        drop_rate=0.5,
+        num_classes=1,)
+    model2 = create_model(
+        'resnet50',
+        drop_rate=0.5,
+        num_classes=1,)
+    
+    loc = 'cuda:{}'.format(0)
+    checkpoint1 = torch.load("./machine_full_best.tar", map_location=loc)
+    model1.load_state_dict(checkpoint1['state_dict'])
+    checkpoint2 = torch.load("./human_full_best.tar", map_location=loc)
+    model2.load_state_dict(checkpoint2['state_dict'])
+    
+    my_transform = transforms.Compose([
+        transforms.RandomResizedCrop(224, (1, 1)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                            std=[0.229, 0.224, 0.225]),])
+    
+    input_img = my_transform(input_img).view(1,3,224,224)
+    model1.eval()
+    model2.eval()
+    result1 = round(model1(input_img).item(), 3)
+    result2 = round(model2(input_img).item(), 3)
+    result = 'MachineMem score = ' + str(result1) + ', HumanMem score = ' + str(result2) +'.'
+    return result
+
+demo = gr.Interface(predict, gr.Image(), "text")
+demo.launch(debug = True)