import gradio as gr
from torch.nn import functional as F

from model_loader import ModelType, type_to_transforms, type_to_loaded_model

def predict_from_model(model_type, raw_image):
    tfm = type_to_transforms[model_type]
    model = type_to_loaded_model[model_type]
    img = tfm(raw_image)
    y = None
    if model_type == ModelType.SYNTHETIC_DETECTOR_V2:
        y = model.forward(img.unsqueeze(0).to("cpu"))
    else:
        y = model.forward(img[None, ...])
    y_1 = F.softmax(y, dim=1)[:, 1].cpu().detach().numpy()
    y_2 = F.softmax(y, dim=1)[:, 0].cpu().detach().numpy()
    return {'created by AI': y_1.tolist(),
            'created by human': y_2.tolist()}

def get_y(model_type, model, image):
    if model_type == ModelType.SYNTHETIC_DETECTOR_V2:
        return model.forward(image.unsqueeze(0).to("cpu"))
    return model.forward(image[None, ...])

def predict(raw_image, model_name):
    if model_name not in ModelType.get_list():
        return {'error': [0.]}

    model_type = ModelType[str(model_name).upper()].value
    model = type_to_loaded_model[model_type]
    tfm = type_to_transforms[model_type]
    image = tfm(raw_image)
    y = get_y(model_type, model, image)
    y_1 = F.softmax(y, dim=1)[:, 1].cpu().detach().numpy()
    y_2 = F.softmax(y, dim=1)[:, 0].cpu().detach().numpy()
    return {'created by AI': y_1.tolist(),
            'created by human': y_2.tolist()}

general_examples = [
    ["images/general/img_1.jpg"],
    ["images/general/img_2.jpg"],
    ["images/general/img_3.jpg"],
    ["images/general/img_4.jpg"],
    ["images/general/img_5.jpg"],
    ["images/general/img_6.jpg"],
    ["images/general/img_7.jpg"],
    ["images/general/img_8.jpg"],
    ["images/general/img_9.jpg"],
    ["images/general/img_10.jpg"],
]

optic_examples = [
    ["images/optic/img_1.jpg"],
    ["images/optic/img_2.jpg"],
    ["images/optic/img_3.jpg"],
    ["images/optic/img_4.jpg"],
    ["images/optic/img_5.jpg"],
]

famous_deepfake_examples = [
    ["images/famous_deepfakes/img_1.jpg"],
    ["images/famous_deepfakes/img_2.jpg"],
    ["images/famous_deepfakes/img_3.jpg"],
    ["images/famous_deepfakes/img_4.webp"],
]

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        <h1 style="text-align: center;">For Fake's Sake: a set of models for detecting generated and synthetic images</h3>
        This is a demo space for synthetic image detectors:
        <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_mj_200'>midjourney200M</a>,
        <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_mj_5'>midjourney5M</a>,
        <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_sd_200'>diffusions200M</a>,
        <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_sd_5'>diffusions5M</a>,
        <a href=''>synthetic_detector_v2</a>.
        <br>We provide several detectors for images generated by popular tools, such as Midjourney and Stable Diffusion.<br>
        Please refer to model cards for evaluation metrics and limitations.
        """
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil")
            drop_down = gr.Dropdown(ModelType.get_list(), type="value", label="Model", value=ModelType.SYNTHETIC_DETECTOR_V2)
            with gr.Row():
                gr.ClearButton(components=[image_input])
                submit_button = gr.Button("Submit", variant="primary")
        with gr.Column():
            result_score = gr.Label(label='result', num_top_classes=2)
    with gr.Tab("Examples"):
        gr.Examples(examples=general_examples, inputs=image_input)
    # with gr.Tab("More examples"):
    #     gr.Examples(examples=optic_examples, inputs=image_input)
    with gr.Tab("Widely known deepfakes"):
        gr.Examples(examples=famous_deepfake_examples, inputs=image_input)

    submit_button.click(predict, inputs=[image_input, drop_down], outputs=result_score)

    gr.Markdown(
        """
        <h3>Models</h3>
        <p><code>*_200M</code> models are based on <code>convnext_large_mlp.clip_laion2b_soup_ft_in12k_in1k_384</code> with image size <code>640x640</code></p>
        <p><code>*_5M</code> models are based on <code>tf_mobilenetv3_large_100.in1k</code> with image size <code>224x224</code></p>
        <p><code>synthetic_detector_2.0</code> models are based on <code>convnext_large_mlp.clip_laion2b_soup_ft_in12k_in1k_384</code> with image size <code>384x384</code></p>

        <h3>Details</h3>
        <li>Model cards: <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_mj_200'>midjourney200M</a>,
                         <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_mj_5'>midjourney5M</a>,
                         <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_sd_200'>diffusions200M</a>,
                         <a href='https://huggingface.co/Sumsub/Sumsub-ffs-synthetic-1.0_sd_5'>diffusions5M</a>,
                         <a href=''>synthetic_detector_v2</a>.
        </li>
        <li>License: CC-By-SA-3.0</li>
        """
    )

demo.launch()