import os

import gradio as gr

from censor import _CENSOR_MODELS, _DEFAULT_CENSOR_MODEL, _gr_detect_censors
from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
from head import _gr_detect_heads, _HEAD_MODELS, _DEFAULT_HEAD_MODEL
from manbits import _MANBIT_MODELS, _DEFAULT_MANBIT_MODEL, _gr_detect_manbits
from person import _PERSON_MODELS, _DEFAULT_PERSON_MODEL, _gr_detect_person

if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Tabs():
            with gr.Tab('Face Detection'):
                with gr.Row():
                    with gr.Column():
                        gr_face_input_image = gr.Image(type='pil', label='Original Image')
                        gr_face_model = gr.Dropdown(_FACE_MODELS, value=_DEFAULT_FACE_MODEL, label='Model')
                        gr_face_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
                        with gr.Row():
                            gr_face_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
                            gr_face_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')

                        gr_face_submit = gr.Button(value='Submit', variant='primary')

                    with gr.Column():
                        gr_face_output_image = gr.Image(type='pil', label="Labeled")

                    gr_face_submit.click(
                        _gr_detect_faces,
                        inputs=[
                            gr_face_input_image, gr_face_model,
                            gr_face_infer_size, gr_face_score_threshold, gr_face_iou_threshold,
                        ],
                        outputs=[gr_face_output_image],
                    )

            with gr.Tab('Head Detection'):
                with gr.Row():
                    with gr.Column():
                        gr_head_input_image = gr.Image(type='pil', label='Original Image')
                        gr_head_model = gr.Dropdown(_HEAD_MODELS, value=_DEFAULT_HEAD_MODEL, label='Model')
                        gr_head_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
                        with gr.Row():
                            gr_head_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
                            gr_head_score_threshold = gr.Slider(0.0, 1.0, 0.3, label='Score Threshold')

                        gr_head_submit = gr.Button(value='Submit', variant='primary')

                    with gr.Column():
                        gr_head_output_image = gr.Image(type='pil', label="Labeled")

                    gr_head_submit.click(
                        _gr_detect_heads,
                        inputs=[
                            gr_head_input_image, gr_head_model,
                            gr_head_infer_size, gr_head_score_threshold, gr_head_iou_threshold,
                        ],
                        outputs=[gr_head_output_image],
                    )

            with gr.Tab('Person Detection'):
                with gr.Row():
                    with gr.Column():
                        gr_person_input_image = gr.Image(type='pil', label='Original Image')
                        gr_person_model = gr.Dropdown(_PERSON_MODELS, value=_DEFAULT_PERSON_MODEL, label='Model')
                        gr_person_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
                        with gr.Row():
                            gr_person_iou_threshold = gr.Slider(0.0, 1.0, 0.5, label='IOU Threshold')
                            gr_person_score_threshold = gr.Slider(0.0, 1.0, 0.3, label='Score Threshold')

                        gr_person_submit = gr.Button(value='Submit', variant='primary')

                    with gr.Column():
                        gr_person_output_image = gr.Image(type='pil', label="Labeled")

                    gr_person_submit.click(
                        _gr_detect_person,
                        inputs=[
                            gr_person_input_image, gr_person_model,
                            gr_person_infer_size, gr_person_score_threshold, gr_person_iou_threshold,
                        ],
                        outputs=[gr_person_output_image],
                    )

            with gr.Tab('Censor Point Detection'):
                with gr.Row():
                    with gr.Column():
                        gr_censor_input_image = gr.Image(type='pil', label='Original Image')
                        gr_censor_model = gr.Dropdown(_CENSOR_MODELS, value=_DEFAULT_CENSOR_MODEL, label='Model')
                        gr_censor_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
                        with gr.Row():
                            gr_censor_iou_threshold = gr.Slider(0.0, 1.0, 0.5, label='IOU Threshold')
                            gr_censor_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')

                        gr_censor_submit = gr.Button(value='Submit', variant='primary')

                    with gr.Column():
                        gr_censor_output_image = gr.Image(type='pil', label="Labeled")

                    gr_censor_submit.click(
                        _gr_detect_censors,
                        inputs=[
                            gr_censor_input_image, gr_censor_model,
                            gr_censor_infer_size, gr_censor_score_threshold, gr_censor_iou_threshold,
                        ],
                        outputs=[gr_censor_output_image],
                    )

            with gr.Tab('Manbits Detection'):
                with gr.Row():
                    with gr.Column():
                        gr_manbit_input_image = gr.Image(type='pil', label='Original Image')
                        gr_manbit_model = gr.Dropdown(_MANBIT_MODELS, value=_DEFAULT_MANBIT_MODEL, label='Model')
                        gr_manbit_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
                        with gr.Row():
                            gr_manbit_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
                            gr_manbit_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')

                        gr_manbit_submit = gr.Button(value='Submit', variant='primary')

                    with gr.Column():
                        gr_manbit_output_image = gr.Image(type='pil', label="Labeled")

                    gr_manbit_submit.click(
                        _gr_detect_manbits,
                        inputs=[
                            gr_manbit_input_image, gr_manbit_model,
                            gr_manbit_infer_size, gr_manbit_score_threshold, gr_manbit_iou_threshold,
                        ],
                        outputs=[gr_manbit_output_image],
                    )

    demo.queue(os.cpu_count()).launch()