krypticmouse's picture
Update app.py
86b2682
raw
history blame
No virus
1.82 kB
import torch
import re
import gradio as gr
from pathlib import Path
from transformers import AutoTokenizer, AutoFeatureExtractor, VisionEncoderDecoderModel
def predict(image, max_length=64, num_beams=4):
image = image.convert('RGB')
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
pixel_values = pixel_values.to(device)
with torch.no_grad():
text = tokenizer.decode(model.generate(pixel_values.cpu())[0])
text = text.replace('<|endoftext|>', '').split('\n')
return text[0]
model_path = "team-indain-image-caption/hindi-image-captioning"
device = "cpu"
# Load model.
model = VisionEncoderDecoderModel.from_pretrained(model_path)
model.to(device)
print("Loaded model")
feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
print("Loaded feature_extractor")
tokenizer = AutoTokenizer.from_pretrained(model_path)
print("Loaded tokenizer")
title = "Hindi Image Captioning"
description = ""
input = gr.inputs.Image(type="pil")
#example_images = sorted([f.as_posix() for f in Path("examples").glob("*.jpg")])
#print(f"Loaded {len(example_images)} example images")
article = "This huggingface presents a demo for Image captioning in Hindi built with VIT Encoder and GPT2 Decoder"
'''interface = gr.Interface(
fn=predict,
inputs=input,
outputs="textbox",
title=title,
description=description,
#examples=example_images,
live=True,
theme="darkpeach"
)'''
#inp=gr.inputs.Textbox(lines=1, placeholder=None, default="", label="search you query here")
output = gr.outputs.Textbox(type="auto",label="Captions")
interface = gr.Interface(fn=predict, inputs=input, outputs=output,examples=exp,article=article,title=title,theme="huggingface",layout='vertical')
interface.launch(share=True)