Spaces:
Running
Running
File size: 1,737 Bytes
74a9cf9 4df3bee 74a9cf9 4df3bee 74a9cf9 4df3bee 74a9cf9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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
|