File size: 1,608 Bytes
4aa568e |
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 |
import gradio as gr
from PIL import Image
import torch
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
# Load model and processor
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Captioning function
def generate_caption(upload_img, webcam_img):
# Choose image from upload or webcam
image = webcam_img if webcam_img is not None else upload_img
if image is None:
return "No image provided."
# Preprocess
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
# Generate
output_ids = model.generate(pixel_values, max_length=16, num_beams=4)
caption = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
return caption
# Build Gradio UI
with gr.Blocks() as demo:
gr.Markdown("# Image Captioning with Gradio")
with gr.Row():
upload_input = gr.Image(source="upload", type="pil", label="Upload Image")
webcam_input = gr.Image(source="webcam", type="pil", label="Use Camera")
output_text = gr.Textbox(label="Caption", interactive=False)
generate_btn = gr.Button("Generate Caption")
generate_btn.click(
fn=generate_caption,
inputs=[upload_input, webcam_input],
outputs=output_text
)
demo.launch()
|