Spaces:
Build error
Build error
File size: 1,317 Bytes
86c62f2 c4b9028 86c62f2 c4b9028 86c62f2 1f1350f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 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) |