image-caption / app.py
ChandraP12330's picture
Update app.py
cb5683a verified
raw
history blame
1.33 kB
import streamlit as st
from transformers import BlipForConditionalGeneration, BlipProcessor
# Load the BLIP model and processor
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
def generate_caption(image):
# Preprocess the image
pixel_values = processor(images=image, return_tensors="pt").pixel_values
# Generate caption using the BLIP model
output_ids = model.generate(pixel_values, max_length=50, num_beams=4, early_stopping=True)
# Decode the caption
caption = processor.decode(output_ids[0], skip_special_tokens=True)
return caption
def main():
st.title("Image Caption Generator")
# Upload image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Display the uploaded image
image = st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
# Generate caption
if st.button("Generate Caption"):
with st.spinner("Generating caption..."):
caption = generate_caption(uploaded_file.getvalue())
st.success(f"Caption: {caption}")
if __name__ == "__main__":
main()