import gradio as gr

from nudenetupdated import NudeDetector

detector = NudeDetector()


def predict(image, checkbox):
    if checkbox:
        mode = "fast"
    else:
        mode = "default"

    try:
        result = detector.detect(image, mode)

        sections = []
        for i in range(len(result)):
            print()
            sections.append((
                    result[i]['box'],
                    result[i]['label']
            ))

        output = (image, sections)
        return (output, result)
    except Exception as e:
        print(f"Got uncaught exception {type(e)}: {e}")


gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(),
        gr.Checkbox(label="Fast"),
    ],
    outputs=[
        gr.AnnotatedImage(label="Annotated Image"),
        gr.JSON(label="JSON")
    ]
).launch(server_name="0.0.0.0")