import os
import onnxruntime as rt
import sys
from PIL import Image, ImageOps
import numpy as np
from pathlib import Path
import collections
from typing import Union, List
import scipy.ndimage
import requests
import gradio as gr
import face_detection  # Ensure this is the adjusted face_detection.py

MODEL_FILE = "ffhqu2vintage512_pix2pixHD_v1E11-inp2inst-simp.onnx"
so = rt.SessionOptions()
so.inter_op_num_threads = 4
so.intra_op_num_threads = 4
session = rt.InferenceSession(MODEL_FILE, sess_options=so)
input_name = session.get_inputs()[0].name
print("input_name = " + str(input_name))
output_name = session.get_outputs()[0].name
print("output_name = " + str(output_name))

def array_to_image(array_in):
    array_in = np.squeeze(255 * (array_in + 1) / 2)
    array_in = np.transpose(array_in, (1, 2, 0))
    im = Image.fromarray(array_in.astype(np.uint8))
    return im

def image_as_array(image_in):
    im_array = np.array(image_in, np.float32)
    im_array = (im_array / 255) * 2 - 1
    im_array = np.transpose(im_array, (2, 0, 1))
    im_array = np.expand_dims(im_array, 0)
    return im_array

def find_aligned_face(image_in, size=512):
    aligned_image, n_faces, quad = face_detection.align(image_in, face_index=0, output_size=size)
    return aligned_image, n_faces, quad

def align_first_face(image_in, size=512):
    aligned_image, n_faces, quad = find_aligned_face(image_in, size=size)
    if n_faces == 0:
        try:
            image_in = ImageOps.exif_transpose(image_in)
        except:
            print("exif problem, not rotating")
        image_in = image_in.resize((size, size))
        im_array = image_as_array(image_in)
    else:
        im_array = image_as_array(aligned_image)

    return im_array

def img_concat_h(im1, im2):
    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (im1.width, 0))
    return dst

def face2vintage(img: Image.Image, size: int) -> Image.Image:
    aligned_img = align_first_face(img)
    if aligned_img is None:
        output = None
    else:
        output = session.run([output_name], {input_name: aligned_img})[0]
        output = array_to_image(output)
        aligned_img = array_to_image(aligned_img).resize((output.width, output.height))
        output = img_concat_h(aligned_img, output)

    return output

def inference(img):
    out = face2vintage(img, 512)
    return out

title = "Vintage style Pix2PixHD"
description = "Style a face to look more \"Vintage\". Upload an image with a face, or click on one of the examples below."
article = "<hr><p style='text-align: center'>See the <a href='https://github.com/justinpinkney/pix2pixHD' target='_blank'>Github Repo</a></p>"

examples = [['Example00001.jpg'], ['Example00002.jpg'], ['Example00003.jpg'], ['Example00004.jpg'], ['Example00005.jpg']]

demo = gr.Interface(
    inference, 
    inputs=[gr.Image(type="pil", label="Input")],
    outputs=[gr.Image(type="pil", label="Output")],
    title=title,
    description=description,
    article=article,
    
)

demo.queue()
demo.launch()