import streamlit as st
import os

from core.ui import show_choosing_pic_bar, show_right_sidebar
from core.model import load_model, run_yolo_inference
from core.helpers import get_image_list, show_note_for_classes, show_hyperparameters
        
def main():
    st.set_page_config(layout="wide")
    st.title("Vietnam Traffic Sign Deployment")
    
    model = load_model()
    st.success("Model load successfully!")
    
    # Left sidebar for image selection
    images = get_image_list()
    img_folder = "test_dataset/images"
    show_choosing_pic_bar(images, img_folder)
    
    if "selected_image" in st.session_state:
        selected_image = os.path.join(img_folder, st.session_state.selected_image)

        st.subheader("🖼️ Selected Image")
        st.image(selected_image, caption=f"Original Image {st.session_state.selected_image}", use_container_width=True)

        # Now create 2 columns below the Selected Image
        middle_part, right_bar = st.columns([4, 1], gap="large")
        
        # Right sidebar for thresholds
        conf_threshold, iou_threshold, conf_show = show_right_sidebar(right_bar)

        with middle_part:
            st.subheader("🔍 Detection Result")
            detected_img, classes_with_conf = run_yolo_inference(model, selected_image, conf_threshold, iou_threshold, conf_show)
            st.image(detected_img, caption="Detected Objects", use_container_width=True)

            show_hyperparameters(conf_threshold, iou_threshold)
            show_note_for_classes(classes_with_conf)

if __name__ == "__main__":
    main()