File size: 1,112 Bytes
b9e51b5 f3bbc2a |
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 |
import gradio as gr
import requests
from io import BytesIO
from PIL import Image
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
# Load GPT-2 model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
def generate_caption(image):
# Preprocess image
response = requests.get(image)
img = Image.open(BytesIO(response.content)).convert('RGB')
img = img.resize((224, 224))
# Generate caption using GPT-2
input_text = "This is an image of " + tokenizer.decode(tokenizer.encode(image)) + ". "
input_ids = tokenizer.encode(input_text, return_tensors='pt')
output = model.generate(input_ids=input_ids, max_length=200, do_sample=True)
caption = tokenizer.decode(output[0], skip_special_tokens=True)
return caption
# Create Gradio interface
inputs = gr.inputs.Image()
outputs = gr.outputs.Textbox()
gr.Interface(fn=generate_caption, inputs=inputs, outputs=outputs, title='Image Captioning with GPT-2', description='Upload an image and get a detailed caption generated by GPT-2.').launch()
|