import os
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_url, cached_download


os.system("git clone https://github.com/AK391/stylegan2-ada-pytorch")


os.chdir("stylegan2-ada-pytorch")

os.mkdir("outputs")
os.mkdir("outputs/images")

config_file_url = hf_hub_url("AUBMC-AIM/MammoGANesis", filename="mammoGANesis.pkl")
cached_file = cached_download(config_file_url)

def inference(truncation,seeds):
  os.system("python generate.py --outdir=./outputs/images/ --trunc="+str(truncation)+" --seeds="+str(int(seeds))+" --network="+cached_file)
  seeds = int(seeds)
  image = Image.open(f"./outputs/images/seed{seeds:04d}.png")
  return image

title = "MammoGANesis"
description = "Gradio demo for MammoGANesis: Controlled Generation of High-Resolution Mammograms for Radiology Education. This paper demonstrates the model’s ability to generate anatomically and medically relevant mammograms by achieving an average AUC of 0.54 in a double-blind study on four expert mammography radiologists to distinguish between generated and real images, ascribing to the high visual quality of the synthesized and edited mammograms, and to their potential use in advancing and facilitating medical education. To use it, add seed and truncation, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center'><a href='https://cyrilzakka.github.io/radiology/2020/10/13/mammogenesis.html' target='_blank'>MammoGANesis: Controlled Generation of High-Resolution Mammograms for Radiology Education</a><center></a></center></p><center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_mammogan' alt='visitor badge'></center>"

gr.Interface(inference,[gr.inputs.Slider(label="truncation",minimum=0, maximum=5, step=0.1, default=0.8),gr.inputs.Slider(label="Seed",minimum=0, maximum=1000, step=1, default=0)],"pil",title=title,description=description,article=article, examples=[
    [0.8,0]
    ]).launch(enable_queue=True,cache_examples=True)