import streamlit as st
from PIL import Image

# Custom CSS for enhanced styling
st.markdown("""
    <style>
        body {
            background-color: #f7f9fc;
            font-family: 'Arial', sans-serif;
        }
        .title {
            color: #2c3e50;
            font-size: 40px;
            font-weight: 700;
            text-align: center;
            margin-bottom: 40px;
        }
        .upload-area {
            background-color: #e8f4f8;
            padding: 20px;
            border-radius: 10px;
            border: 2px dashed #1abc9c;
            text-align: center;
            margin-bottom: 20px;
        }
        .upload-area:hover {
            background-color: #d1ecf1;
        }
        .generate-button {
            background-color: #1abc9c;
            color: white;
            border-radius: 5px;
            padding: 10px;
            font-weight: bold;
            cursor: pointer;
            text-align: center;
            margin-top: 20px;
        }
        .generate-button:hover {
            background-color: #16a085;
        }
        .result-section {
            margin-top: 40px;
        }
        .result-title {
            color: #34495e;
            font-size: 24px;
            font-weight: 600;
            margin-bottom: 20px;
        }
        .success {
            background-color: #dff0d8;
            color: #3c763d;
            padding: 15px;
            border-radius: 5px;
        }
    </style>
""", unsafe_allow_html=True)

# Initialize session state to block re-running
if 'has_run' not in st.session_state:
    st.session_state.has_run = False

# Main UI container
st.markdown('<div class="title">Florence-2 Image Captioning Demo</div>', unsafe_allow_html=True)



# Image upload area
st.sidebar.markdown('<div class="upload-area">', unsafe_allow_html=True)
uploaded_image = st.sidebar.file_uploader("Upload your image here", type=["jpg", "jpeg", "png"])
st.sidebar.markdown('</div>', unsafe_allow_html=True)


# Display the uploaded image and process it if available
if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    # Task prompt input
    task_prompt = st.sidebar.text_input("Task Prompt", value="Describe the image in detail:")
    # Additional text input (optional)
    text_input = st.text_area("Input Questions", height=50)
    # Generate Caption button
    if st.button("Generate Caption", key="Generate") and not st.session_state.has_run:
        # Mark that the script has been run
        st.session_state.has_run = True
        st.write(task_prompt,"\n",text_input)

# Clear session state button
if st.session_state.has_run:
    if st.button("Clear and Upload New Image"):
        st.session_state.has_run = False
        st.experimental_rerun()

st.markdown('</div>', unsafe_allow_html=True)