Spaces:
Configuration error
Configuration error
File size: 1,347 Bytes
2a3a041 |
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 39 40 41 42 43 44 |
import torch
import torch.nn as nn
import numpy as np
import os
from src.args import get_parser
import pickle
from src.model import get_model
from torchvision import transforms
from src.utils.output_ing import prepare_output
from PIL import Image
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
map_loc = None if torch.cuda.is_available() else 'cpu'
def im2ingr(image, ingrs_vocab, model):
transf_list_batch = []
transf_list_batch.append(transforms.ToTensor())
transf_list_batch.append(transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225)))
to_input_transf = transforms.Compose(transf_list_batch)
greedy = True
beam = -1
temperature = 1.0
transf_list = []
transf_list.append(transforms.Resize(256))
transf_list.append(transforms.CenterCrop(224))
transform = transforms.Compose(transf_list)
image_transf = transform(image)
image_tensor = to_input_transf(image_transf).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model.sample(image_tensor, greedy=greedy,
temperature=temperature, beam=beam, true_ingrs=None)
ingr_ids = outputs['ingr_ids'].cpu().numpy()
outs = prepare_output(ingr_ids[0], ingrs_vocab)
return outs['ingrs']
|