import streamlit as st
from camera_input_live import camera_input_live
from utils import object_detection_brainai as odb
object_detection = odb.ObjectDetectionModel()

st.set_page_config(
    page_title = "객체 인식", 
    page_icon = ":black_cat:",
    layout = "wide" )

if "random_class" not in st.session_state:
    st.session_state.random_class = object_detection.call_class()

st.sidebar.header("메뉴")
source_radio = st.sidebar.radio("선택하세요", ["IMAGE", "VIDEO", "GAME", "WEBCAM"])

col1, col2 = st.columns([3,2])
with col1:
    st.title(":blue[Object Detection] :black_cat:")
with col2:
    if source_radio == "GAME":
        if st.button('새로운 객체'):
            st.session_state.random_class = object_detection.call_class()

if source_radio == "IMAGE":
    st.write(":green[왼쪽 메뉴 'Browse files' 버튼을 클릭하여 이미지 파일을 선택하면 AI 추론이 시작됩니다.]")
    st.sidebar.header("이미지 파일 업로드")
    input_img = st.sidebar.file_uploader("이미지 파일을 선택하세요.", type=("jpg", "png"))

    if input_img is not None:
        result_img, detected_object = object_detection.process(input_img)
        col1, col2 = st.columns(2)
        with col1:
            st.image(result_img)
        with col2:
            st.header(detected_object)

    else:
        col1, col2 = st.columns(2)
        with col1:
            st.image("data/table.jpg")
        with col2:
            st.header("Objected detected: chair, potted plant, vase, dining table")

if source_radio == "VIDEO":
    st.write(":green[왼쪽 메뉴 'Browse files' 버튼을 클릭하여 비디오 파일을 선택하면 AI 추론이 시작됩니다.]")
    st.sidebar.header("비디오 파일 업로드")
    input_video = st.sidebar.file_uploader("비디오 파일을 선택하세요.", type=("mp4"))
    if input_video is not None:
        output_video_path = object_detection.play_video(input_video)
        st.video(output_video_path)
    else:
        st.video("data/breakfast.mp4")
  
if source_radio == "GAME":
    st.sidebar.header("이미지 파일 업로드")
    input_img = st.sidebar.file_uploader("이미지 파일을 선택하세요.", type=("jpg", "png")) 

    if input_img is not None:
        result_img, detected_object = object_detection.process(input_img)
        if st.session_state.random_class in detected_object:
            st.title(f':{st.session_state.random_class}: 찾습니다!')

        else: 
            st.title(f':{st.session_state.random_class}: 찾을 수 없습니다!')
        st.image(result_img)
    else:
        st.title(f':{st.session_state.random_class}: 있는 이미지를 업로드')

if source_radio == "WEBCAM":
    input_img = camera_input_live()
    if input_img:
        result_img, detected_object = object_detection.process(input_img)
        st.image(result_img)