# CS5330 Lab 4
# This file builds the Gradio interface

import gradio as gr
from segmentation import extract_person
from insert import Background, Disparity, insert_person_in_pair
from anaglyph import create_anaglyph


# Gradio interface function
def gradio_interface(person_image, background, disparity):
    # Gradio dropdown menu does not accept enum as a type
    # So convert str to enum
    if background == 'Nature':
        background_arg = Background.NATURE
    elif background == 'NEU':
        background_arg = Background.NEU
    else:
        background_arg = Background.GASTOWN

    if disparity == 'Close':
        disparity_arg = Disparity.CLOSE
    elif disparity == 'Medium':
        disparity_arg = Disparity.MED
    else:
        disparity_arg = Disparity.FAR
    
    person_arg = extract_person(person_image)

    # If there's no person detected, raise an error message
    if not person_arg:
        return None, 'Error: No person detected in the image.'

    left_image, right_image = insert_person_in_pair(background_arg, extract_person(person_image), disparity_arg)
    anaglyph = create_anaglyph(left_image, right_image)
    return anaglyph, 'Successfully generated the anaglyph image.'

demo = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.Image(type='pil', label='Upload an image of a person'),
            gr.Dropdown(choices=['Nature', 'NEU', 'Gastown'], label='Choose background'),
            gr.Dropdown(choices=['Close', 'Medium', 'Far'], label='Choose disparity')],
    outputs=[gr.Image(type='pil', label='Anaglyph Image'),
             gr.Textbox(label='Message')],
    title='Anaglyph Generator',
    description='Upload an image of a person, then choose a background image and disparity to determine how the anaglyph is generated.'
)

demo.launch(share=True)