import gradio as gr
from fastai.vision.all import *
from os.path import exists
import requests

model_fn = 'quick_224px'
url = 'https://huggingface.co/johnowhitaker/sketchy_unet_rn34/resolve/main/quick_224px'

if not exists(model_fn):
  print('starting download')
  with requests.get(url, stream=True) as r:
    r.raise_for_status()
    with open(model_fn, 'wb') as f:
      for chunk in r.iter_content(chunk_size=8192):
        f.write(chunk)
  print('done')
else:
  print('file exists')

# Load the model (requires dummy itemgetters)
def get_x(item):return None
def get_y(item):return None
sketch_model = load_learner(model_fn)


def sketchify(image_path):
   pred = sketch_model.predict(image_path)
   np_im = pred[0].permute(1, 2, 0).numpy()
   return np_im
    
title = "Sketchy Unet Demo"
description = """
<center>
A resnet34-based unet model trained (briefly) to sketchify faces.
</center>
"""

article = "Blog post: https://datasciencecastnet.home.blog/2022/03/29/sketchy-unet/ \n Model training (colab): https://colab.research.google.com/drive/1ydcC4Gs2sLOelj0YqwJfRqDPU2sjQunb?usp=sharing \n My Twitter (questions and feedback welcome) https://twitter.com/johnowhitaker"
    

iface = gr.Interface(fn=sketchify, 
  inputs=[gr.inputs.Image(label="Input Image", shape=(512, 512), type="filepath")], 
  outputs=[gr.outputs.Image(type="numpy", label="Model Output")],
  title = title, description = description, article = article
)
iface.launch()