import os
from models import Noise2Same
import gradio as gr

os.system("mkdir trained_models/denoising_ImageNet")
os.system("cd trained_models/denoising_ImageNet; gdown https://drive.google.com/uc?id=1asrwULW1lDFasystBc3UfShh5EeTHpkW; gdown https://drive.google.com/uc?id=1Re1ER7KtujBunN0-74QmYrrOx77WpVXK; gdown https://drive.google.com/uc?id=1QdlyUPUKyyGtqD0zBrj5F7qQZtmUELSu; gdown https://drive.google.com/uc?id=1LQsYR26ldHebcdQtP2zt4Mh-ZH9vXQ2S; gdown https://drive.google.com/uc?id=1AxTDD4dS0DtzmBywjGyeJYgDrw-XjYbc; gdown https://drive.google.com/uc?id=1w4UdNAbOjvWSL0Jgbq8_hCniaxqsbLaQ; cd ../..")
os.system("wget -O arch.png https://i.imgur.com/NruRABn.png")
os.system("wget -O parrot.png https://i.imgur.com/zdji3xv.png")
os.system("wget -O lion.png https://i.imgur.com/qNT0lJJ.png")

model = Noise2Same('trained_models/', 'denoising_ImageNet', dim=2, in_channels=3)

def norm(x):
    x = (x-x.min())/(x.max()-x.min())
    return x
    
def predict(img):
  pred = model.predict(img.astype('float32'))
  return norm(pred)

img = gr.inputs.Image()

title = "Noise2Same: Optimizing A Self-Supervised Bound for Image Denoising"
description = "Interactive demo of Noise2Same, an image denoising method developed by Yaochen Xie"

denoise = gr.Interface(fn=predict,
                       inputs=gr.Image(placeholder="Drag image here.", label='Input Image'),
                       outputs=gr.Image(placeholder="Output image will appear here.", label='Input Image'),
                       examples=[["lion.png"], ["arch.png"], ["parrot.png"]],
                       title=title, description=description)

#launching the app
if __name__ == "__main__":
    denoise.launch(debug=True)