import os
import argparse
import logging
from typing import List, Tuple, Optional
from pathlib import Path

import numpy as np
from tqdm import tqdm
import onnxruntime as ort
from huggingface_hub import hf_hub_download
import cv2
from PIL import Image


model_path = hf_hub_download(repo_id="jbrownkramer/face-parsing", filename="resnet18.onnx")
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider']
session = ort.InferenceSession(model_path, providers=providers)

def prepare_image(image, input_size: Tuple[int, int] = (512, 512)) -> np.ndarray:
    image_batch = np.array(image)
        # Resize the image
    resized_image = cv2.resize(image_batch, input_size, interpolation=cv2.INTER_LINEAR)

    image_batch = image_batch / 255.0
    image_batch -= np.array([[[0.485, 0.456, 0.406]]])
    image_batch /= np.array([[[0.229, 0.224, 0.225]]])
    image_batch = image_batch.transpose(2, 0, 1)
    image_batch = image_batch.astype(np.float32)
    image_batch = image_batch.reshape(1, 3, input_size[1], input_size[0])

    return image_batch

def get_face_mask(image):
    # Store original image resolution
    w,h = image.size
    
    # Prepare image for inference
    image_batch = prepare_image(image)
    
    # Run ONNX inference
    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: image_batch})
    
    # Get the first output (assuming it's the segmentation map)
    output = outputs[0]
    
    # Convert to segmentation mask
    predicted_mask = output.squeeze(0).argmax(0)
    
    #resize to original size
    restored_mask = cv2.resize(predicted_mask, (w,h), interpolation=cv2.INTER_NEAREST)

    return restored_mask