Spaces:
Sleeping
Sleeping
import streamlit as st | |
from PIL import Image | |
import requests | |
from transformers import BlipProcessor, BlipForConditionalGeneration | |
# Load the model and processor outside the main function to avoid reloading on every run | |
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") | |
def generate_caption(img_url): | |
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB') | |
# Conditional image captioning | |
text = "a photography of" | |
inputs = processor(raw_image, text, return_tensors="pt") | |
out = model.generate(**inputs) | |
conditional_caption = processor.decode(out[0], skip_special_tokens=True) | |
# Unconditional image captioning | |
inputs = processor(raw_image, return_tensors="pt") | |
out = model.generate(**inputs) | |
unconditional_caption = processor.decode(out[0], skip_special_tokens=True) | |
return conditional_caption, unconditional_caption | |
def main(): | |
st.title("Image Captioning App") | |
img_url = st.text_input("Enter the image URL:") | |
if img_url: | |
try: | |
# Display the image | |
image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB') | |
st.image(image, caption='Input Image', use_column_width=True) | |
# Generate captions | |
conditional_caption, unconditional_caption = generate_caption(img_url) | |
# Display captions | |
st.subheader("Conditional Image Caption") | |
st.write(conditional_caption) | |
st.subheader("Unconditional Image Caption") | |
st.write(unconditional_caption) | |
except Exception as e: | |
st.error(f"Error processing the image: {e}") | |
if __name__ == "__main__": | |
main() | |