import gradio as gr
from preprocessor import Preprocessor

# Initialize Preprocessor
preprocessor = Preprocessor()

# Define processing function with extended options
def process_image(image, preprocessor_name, preprocess_resolution=None, mlsd_value_threshold=None, mlsd_distance_threshold=None, canny_low_threshold=None, canny_high_threshold=None):
    preprocessor.load(preprocessor_name)
    kwargs = {}
    
    if preprocess_resolution:
        kwargs['preprocess_resolution'] = preprocess_resolution
    if mlsd_value_threshold and preprocessor_name == "MLSD":
        kwargs['mlsd_value_threshold'] = mlsd_value_threshold
    if mlsd_distance_threshold and preprocessor_name == "MLSD":
        kwargs['mlsd_distance_threshold'] = mlsd_distance_threshold
    if canny_low_threshold and preprocessor_name == "Canny":
        kwargs['canny_low_threshold'] = canny_low_threshold
    if canny_high_threshold and preprocessor_name == "Canny":
        kwargs['canny_high_threshold'] = canny_high_threshold

    return preprocessor(image, **kwargs)

# UI creation with segmentation options
def create_ui():
    with gr.Blocks() as demo:
        with gr.Row():
            image_input = gr.Image(type="pil")
            preprocessor_dropdown = gr.Dropdown(choices=["ContentShuffle", "Openpose", "Midas", "MLSD", "Canny", "Lineart", "DPT", "UPerNet", "HED", "PidiNet"], label="Preprocessor")
            preprocess_resolution = gr.Slider(128, 512, step=1, label="Preprocess Resolution", visible=False)
            # Additional options for MLSD and Canny
            mlsd_value_threshold = gr.Slider(0.01, 2.0, step=0.01, label="MLSD Value Threshold", visible=False)
            mlsd_distance_threshold = gr.Slider(0.01, 20.0, step=0.01, label="MLSD Distance Threshold", visible=False)
            canny_low_threshold = gr.Slider(1, 255, step=1, label="Canny Low Threshold", visible=False)
            canny_high_threshold = gr.Slider(1, 255, step=1, label="Canny High Threshold", visible=False)
            submit_button = gr.Button("Process")
            result_image = gr.Image(label="Processed Image")

        def update_options(preprocessor_name):
            # Update visibility based on preprocessor choice
            options_visibility = {
                'preprocess_resolution': preprocessor_name in ["Openpose", "Midas", "MLSD", "Lineart", "DPT", "UPerNet", "HED", "PidiNet"],
                'mlsd_value_threshold': preprocessor_name == "MLSD",
                'mlsd_distance_threshold': preprocessor_name == "MLSD",
                'canny_low_threshold': preprocessor_name == "Canny",
                'canny_high_threshold': preprocessor_name == "Canny",
            }
            return list(options_visibility.values())

        preprocessor_dropdown.change(fn=update_options, inputs=[preprocessor_dropdown], outputs=[preprocess_resolution, mlsd_value_threshold, mlsd_distance_threshold, canny_low_threshold, canny_high_threshold])
        submit_button.click(fn=process_image, inputs=[image_input, preprocessor_dropdown, preprocess_resolution, mlsd_value_threshold, mlsd_distance_threshold, canny_low_threshold, canny_high_threshold], outputs=[result_image])

    return demo

if __name__ == "__main__":
    create_ui().launch()