import streamlit as st
from PIL import Image
import kosmos

st.title('KOSMOS2 Chat box')

uploaded_file = st.file_uploader("Pick an image...", type="jpg")

if uploaded_file is not None:
    try:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image.', use_column_width=True)
        st.write("")
        num_tokens = st.text_input("Number of new tokens", "20")
        prompt = st.chat_input("Detect the main object in the image. The image is a")
        num_tokens = int(num_tokens)
        if prompt:
            st.write(f"User: {prompt}")
            with st.spinner('Processing...'):
                label = kosmos.single_image_classification(image, prompt, num_tokens)
            st.write(f"Model: {label}")
    except Exception as e:
        st.error(f"An error occurred: {e}")