Spaces:
Build error
Build error
File size: 3,068 Bytes
b1bd157 9cd489b b1bd157 9cd489b b1bd157 9cd489b b1bd157 9cd489b 3c8c97b b1bd157 | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import numpy as np
import random,json
import torch
from PIL import Image
from transformer.model import*
import torchvision.transforms as transforms
import torch
import gradio as gr
import argparse
import torchvision.transforms as T
parser = argparse.ArgumentParser(description='Image Captioning inference script')
# Data args
parser.add_argument('--max_seq_len', default=60, type=int, help='max sequence length')
# Model parameters
parser.add_argument('--height', default=224, type=int, metavar='N', help='image height')
parser.add_argument('--width', default=224, type=int, metavar='N', help='image width')
parser.add_argument('--channel', default=3, type=int, help='disable cuda')
parser.add_argument('--enc_heads', default=8, type=int, help='number of encoder heads')
parser.add_argument('--enc_depth', default=4, type=int, help='number of encoder blocks')
parser.add_argument('--dec_heads', default=8, type=int, help='number of decoder heads')
parser.add_argument('--dec_depth', default=1, type=int, help='number of decoder blocks')
parser.add_argument('--patch_size', default=16, type=int, help='patch size')
parser.add_argument('--dim', default=512, type=int, help='embedding dim of patch')
parser.add_argument('--enc_mlp_dim', default=1024, type=int, help='feed forward hidden_dim for an encoder block')
parser.add_argument('--dec_mlp_dim', default=1024, type=int, help='feed forward hidden_dim for a decoder block')
args = parser.parse_args()
height, width, n_channels = args.height, args.width, args.channel
patch_size, dim, enc_head = args.patch_size, args.dim, args.enc_heads
enc_feed_forward, enc_depth = args.enc_mlp_dim, args.enc_depth
dec_feed_forward, dec_depth = args.dec_mlp_dim, args.dec_depth
dec_head = args.dec_heads
# to load the dictionary from the JSON file
with open('vocab.json', 'r') as f:
vocab = json.load(f)
device = ""
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
transforms = T.Compose([
T.Resize((height,width)),
T.ToTensor()
])
vocab_size = len(vocab)
max_seq_len = 60
padding_idx = vocab['<PAD>']
model = Transformer(height,width,n_channels,patch_size,dim,enc_head,enc_feed_forward,enc_depth,
dec_head,dec_feed_forward,dec_depth,max_seq_len,vocab_size,padding_idx)
state = torch.load("captioning.ptrom",map_location="cpu")
model.load_state_dict(state['model_state_dict'])
model = model.to(device)
itos = {v: k for k, v in vocab.items()}
def predict(image):
#img_tensor = torch.from_numpy(image)
image = Image.fromarray(image)
img_tensor = transforms(image).unsqueeze(0).to(device)
preds = greedy_decoding(model, img_tensor, max_seq_len,1,2)[0].detach().cpu().numpy()
output = [itos[token] for token in preds.tolist()]
return " ".join(output)
if __name__ == '__main__':
demo = gr.Interface(
fn=predict,
inputs=["image"],
outputs=["text"],
title='Image captioning'
)
demo.launch()
|