app.py

import gradio as gr
import torchvision.transforms as transforms
from PIL import Image
import torch
from timm.models import create_model
import numpy as np
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,)
    
    checkpoint1 = torch.load("./machine_full_best.tar",map_location=torch.device('cpu'))
    model1.load_state_dict(checkpoint1['state_dict'])
    checkpoint2 = torch.load("./human_full_best.tar",map_location=torch.device('cpu'))
    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)