Fer14/coffee_machine_captioning

1import streamlit as st2from PIL import Image3from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor4 5st.title("Coffe machine captioning app")6 7@st.cache_resource()8def load_model():9    with st.spinner('Loading model and tokenizer...'):10 11        model_id = "Fer14/paligemma_coffe_machine_caption"12 13        model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)14        processor = PaliGemmaProcessor.from_pretrained(model_id)15 16    st.success('Model loaded!')17    return model, processor18 19 20model, processor = load_model()21 22 23st.sidebar.title("Instructions")24st.sidebar.write(25    """26    1. Upload an image using the file uploader.27    2. Wait for the app to process and generate the caption.28    3. The caption will be displayed in the text area.29    4. Enjoy your caption!30    """31)32 33uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])34 35 36 37prompt  = (38            f"Generate a caption for the following coffee maker image. The caption has to be of the following structure:\n"39            "\"A <color> <type>, <accessories>, <shape> shaped, with <screen> and <number> <b_color> butons\"\n\n"40            "in which:\n"41            "- color: red, black, blue...\n"42            "- type: coffee machine, coffee maker, espresso coffee machine...\n"43            "- accessories: a list of accessories like the ones described above\n"44            "- shape: cubed, round...\n"45            "- screen: screen, no screen.\n"46            "- number: amount of buttons to add\n"47            "- b_color: color of the buttons"48        )49 50if uploaded_image is not None:51    # Display the uploaded image52    image = Image.open(uploaded_image).convert("RGB")53    st.image(image, caption='Uploaded Image.', use_column_width=True)54 55    inputs = processor(56            text=prompt,57            images=image,58            return_tensors="pt",59            padding="longest",60        )61    62 63    with st.spinner('Generating caption...'):64        output = model.generate(**inputs, max_length=1000)65 66    out = processor.decode(output[0], skip_special_tokens=True)[len(prompt) :]67 68    # Display the extracted text69    st.text_area("Coffe machine caption", out, height=100)70 71 72 73