import os
os.system("pip install git+https://www.github.com/hukkelas/DeepPrivacy")
import gradio
import numpy as np
import torch
import hashlib
from PIL import Image
import gradio.inputs
from deep_privacy.build import build_anonymizer
from deep_privacy.detection import ImageAnnotation
from typing import List
anonymizer = build_anonymizer()

cached_detections = {}
def anonymize(im: Image, truncation_value: float):
    anonymizer.truncation_level = truncation_value
    im = np.array(im.convert("RGB"))
    md5_ = hashlib.md5(im.tobytes()).hexdigest()
    if md5_ in cached_detections:
        detections = cached_detections[md5_]
    else:
        detections: List[ImageAnnotation]  = anonymizer.detector.get_detections([im])
        cached_detections[md5_] = detections
    if len(detections) == 0:
        return Image.fromarray(im)
    im = anonymizer.anonymize_images([im], detections)[0]
    im = Image.fromarray(im)
    return im


iface = gradio.Interface(
    anonymize, [gradio.inputs.Image(type="pil", label="Upload your image or try the example below!"), gradio.inputs.Slider(minimum=0, maximum=8, step=0.01, default=0.5, label="Truncation value (set to >0 to generate different bodies between runs)")],
    examples=[["coco_val2017_000000001000.jpg", 0], ["turing-2018-bengio-hinton-lecun.jpg", 0]],
    outputs="image",
    title="DeepPrivacy: A Generative Adversarial Network for Face Anonymization",
    description="A live demo of face anonymization with generative adversarial networks. See paper/code at: github.com/hukkelas/DeepPrivacy",
    live=True)
iface.launch()