File size: 1,867 Bytes
7a9fafa 20dc7c2 7a9fafa 9e0c87c 7a9fafa 9e0c87c 9e29eb1 13250bc 7a9fafa |
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 |
import torch
import gradio as gr
from transformers import AutoTokenizer, ViTImageProcessor, VisionEncoderDecoderModel
device = 'cpu'
# Load the pretrained model, feature extractor, and tokenizer
model = VisionEncoderDecoderModel.from_pretrained("premanthcharan/Image_Captioning_Model").to(device)
feature_extractor = ViTImageProcessor.from_pretrained("premanthcharan/Image_Captioning_Model")
tokenizer = AutoTokenizer.from_pretrained("premanthcharan/Image_Captioning_Model")
def predict(image, max_length=64, num_beams=4):
# Process the input image
image = image.convert('RGB')
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
# Generate the caption
caption_ids = model.generate(pixel_values, max_length=max_length, num_beams=num_beams)[0]
# Decode and clean the generated caption
caption = tokenizer.decode(caption_ids, skip_special_tokens=True)
return caption
css = '''
h1#title {
text-align: center;
}
h3#header {
text-align: center;
}
img#overview {
max-width: 800px;
max-height: 600px;
}
img#style-image {
max-width: 1000px;
max-height: 600px;
}
'''
demo = gr.Blocks(css=css)
with demo:
gr.Markdown('''<h1 id="title">Automated Image Captioning Using Generative AI: A Transformer based approach 🖼️</h1>''')
gr.Markdown('Contributed by : Charan Gudivada, Premanth Alahari')
with gr.Column():
input_image = gr.Image(label="Upload your Image", type='pil')
output_caption = gr.Textbox(label="Generated Caption")
btn = gr.Button("Generate Caption")
btn.click(fn=predict, inputs=input_image, outputs=output_caption)
with demo:
gr.Markdown('''<h1 id="title">Features:</h1>''')
gr.Markdown("1. Drop the Image Here or Click on Upload\n2. Click to Access Webcam\n3. Paste from Clipboard")
demo.launch() |