import streamlit as st
import os

pwd = os.getcwd()
jpg_files = [f for f in os.listdir(os.path.join(pwd, 'images')) if f.endswith('.jpg')]
# Sort the files by name
jpg_files = sorted(jpg_files, key=lambda x: int(x.split('.')[0]))
txt_files = [f for f in os.listdir(os.path.join(pwd, 'images')) if f.endswith('.txt')]
# Sort the files by name
txt_files = sorted(txt_files, key=lambda x: int(x.split('.')[0]))

# Session state to keep track of the current file index
if 'file_index' not in st.session_state:
    st.session_state.file_index = 0

# Display "Previous" and "Next" buttons
col1, col2 = st.columns(2)
if col1.button("Previous"):
    st.session_state.file_index -= 1
    st.session_state.file_index = max(0, st.session_state.file_index)
if col2.button("Next"):
    st.session_state.file_index += 1
    st.session_state.file_index = min(len(jpg_files) - 1, st.session_state.file_index)
st.markdown(f"**File {st.session_state.file_index + 1} of {len(jpg_files)}**")
st.text("")
st.text("")

# Display the image and text
jpg_file = jpg_files[st.session_state.file_index]
txt_file = txt_files[st.session_state.file_index]

st.image(os.path.join(pwd, 'images', jpg_file))
with open(os.path.join(pwd, 'images', txt_file)) as f:
    text = f.read()
    generated_caption = text.split('\n')[0]
    laion_caption = text.split('\n')[1]
    
st.markdown(f"**Generated Caption:** {generated_caption}")
st.markdown(f"**Laion Caption:** {laion_caption}")