import gradio as gr
import torchvision.transforms as transforms
from PIL import Image
import torch
def predict(input_img):
    input_img =  Image.fromarray(np.uint8(input_img))
    model1 = models.__dict__['resnet50'](num_classes=1)
    model2 = models.__dict__['resnet50'](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()