File size: 1,131 Bytes
bad6934 bbe5e0c bad6934 daa5afe bad6934 |
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 |
import streamlit as st
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration
import torch
# Load model and processor
@st.cache_resource
def load_model():
model_name = "Nishthaaa/image_captioning"
processor = BlipProcessor.from_pretrained(model_name)
model = BlipForConditionalGeneration.from_pretrained(model_name)
return processor, model
processor, model = load_model()
# Streamlit UI
st.title("Cartoon Caption Generator 🖼️📜")
st.write("Upload a cartoon image and get a funny caption!")
uploaded_file = st.file_uploader("Upload a Cartoon Image", type=["jpg", "png", "jpeg"])
if uploaded_file is not None:
image = Image.open(uploaded_file).convert("RGB")
st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
# Preprocess and generate caption
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
generated_ids = model.generate(**inputs)
caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
st.subheader("Generated Caption:")
st.write(f"💬 {caption}")
|