Image_Caption_Generator / Image Caption Generator.py
Tevfik istanbullu
Update Image Caption Generator.py
7b4c32c verified
raw
history blame
1.63 kB
import gradio as gr
import numpy as np
from PIL import Image
from transformers import AutoProcessor, BlipForConditionalGeneration
import os
# Load the pretrained processor and model
processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
def caption_image(input_image: np.ndarray):
# Convert numpy array to PIL Image and convert to RGB
raw_image = Image.fromarray(input_image).convert('RGB')
# Process the image
inputs = processor(raw_image, return_tensors="pt")
# Generate a caption for the image
out = model.generate(**inputs,max_length=50)
# Decode the generated tokens to text
caption = processor.decode(out[0], skip_special_tokens=True)
return caption
# Save the data to the Hugging Face dataset
HF_TOKEN = os.getenv("HF_TOKEN")
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-images-data")
# Define examples
examples = [
["1.jpg"],
["2.jpg"],
["3.jpg"],
["4.jpg"],
]
# Create a Gradio interface
iface = gr.Interface(
fn=caption_image,
inputs=gr.Image(),
outputs=gr.Textbox(label="Generated Caption", lines=2),
title="πŸ” Image Caption Generator πŸ–ΌοΈ ",
description = "Generate stunning captions for your images with our AI-powered model! 🌟\n\nπŸš«πŸ“š Note: Please avoid entering any sensitive or personal information, as inputs may be reviewed or used for training purposes.",
allow_flagging="auto",
flagging_callback=hf_writer,
examples=examples,
)
iface.launch()