import os
os.system("pip install --upgrade pip")
os.system("pip install ftfy regex tqdm")
os.system("pip install git+https://github.com/openai/CLIP.git")
os.system("pip install git+https://github.com/facebookresearch/detectron2@96c752ce821a3340e27edd51c28a00665dd32a30#subdirectory=projects/DensePose")
os.system("pip install git+https://github.com/hukkelas/DSFD-Pytorch-Inference")
import gradio
import numpy as np
import torch
from PIL import Image
from dp2 import utils
from tops.config import instantiate
import tops
import gradio.inputs


cfg_body = utils.load_config("configs/anonymizers/FB_cse.py")
anonymizer_body = instantiate(cfg_body.anonymizer, load_cache=False)
anonymizer_body.initialize_tracker(fps=1)
cfg_face = utils.load_config("configs/anonymizers/face.py")
anonymizer_face = instantiate(cfg_face.anonymizer, load_cache=False)
anonymizer_face.initialize_tracker(fps=1)


class ExampleDemo:

    def __init__(self, anonymizer, multi_modal_truncation=False) -> None:
        self.multi_modal_truncation = multi_modal_truncation
        self.anonymizer = anonymizer
        with gradio.Row():
            input_image = gradio.Image(type="pil", label="Upload your image or try the example below!")
            output_image = gradio.Image(type="numpy", label="Output")
        with gradio.Row():
            update_btn = gradio.Button("Update Anonymization").style(full_width=True)
        visualize_det = gradio.Checkbox(value=False, label="Show Detections")
        visualize_det.change(self.anonymize, inputs=[input_image, visualize_det], outputs=[output_image])
        gradio.Examples(
            ["erling.jpg", "g7-summit-leaders-distraction.jpg"], inputs=[input_image]
        )
        update_btn.click(self.anonymize, inputs=[input_image, visualize_det], outputs=[output_image])
        input_image.change(self.anonymize, inputs=[input_image, visualize_det], outputs=[output_image])
        self.track = False

    def anonymize(self, img: Image, visualize_detection: bool):
        img, cache_id = pil2torch(img)
        img = tops.to_cuda(img)
        if visualize_detection:
            img = self.anonymizer.visualize_detection(img, cache_id=cache_id)
        else:
            img = self.anonymizer(
                img, truncation_value=0 if self.multi_modal_truncation else 1, multi_modal_truncation=self.multi_modal_truncation, amp=True,
                cache_id=cache_id, track=self.track)
        img = utils.im2numpy(img)
        return img


def pil2torch(img: Image.Image):
    img = img.convert("RGB")
    img = np.array(img)
    img = np.rollaxis(img, 2)
    return torch.from_numpy(img), None


with gradio.Blocks() as demo:
    gradio.Markdown("# <center> DeepPrivacy2 - Realistic Image Anonymization </center>")
    gradio.Markdown("### <center> Håkon Hukkelås, Rudolf Mester, Frank Lindseth </center>")
    gradio.Markdown("<center> DeepPrivacy2 is a toolbox for realistic anonymization of humans, including a face and a full-body anonymizer. </center>")
    gradio.Markdown("<center> See more information at: <a href='https://github.com/hukkelas/deep_privacy2'> https://github.com/hukkelas/deep_privacy2 </a> </center>")

    
    
    with gradio.Tab("Full-Body Anonymization"):
        ExampleDemo(anonymizer_body, multi_modal_truncation=True)
    with gradio.Tab("Face Anonymization"): 
        ExampleDemo(anonymizer_face, multi_modal_truncation=False)


demo.launch()