Draft of using new face parsing

app.py

@@ -4,7 +4,7 @@ from huggingface_hub import hf_hub_download
 import onnxruntime as ort
 import numpy as np
 from PIL import Image
-from faceparsing import get_face_mask
+from faceparsing2 import get_face_mask
 
 # import torch
 # from your_pix2pixhd_code import YourPix2PixHDModel, load_image, tensor2im # Adapt these imports
@@ -34,7 +34,8 @@ ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider
 
 #     # return output_image
 
-def predict(input_image, mode):    
+def predict(input_image, mode):
+    input_image = input_image.convert("RGB")   
     if mode == "Classic":
         # Use the transition_onnx function for side-by-side comparison
         zombie_image = zombie.transition_onnx(input_image, ort_session)
@@ -42,6 +43,7 @@ def predict(input_image, mode):
             return "No face found"
         return zombie_image
     elif mode == "In Place":
+        im_array = np.array(input_image)
         # Use the make_faces_zombie_from_array function for in-place transformation
         #zombie_image = zombie.make_faces_zombie_from_array(im_array, None, ort_session)
         #if zombie_image is None:

faceparsing2.py

@@ -0,0 +1,68 @@
+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
+import matplotlib.pyplot as plt
+from huggingface_hub import hf_hub_download
+import cv2
+
+
+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:
+    # Resize the image
+    resized_image = cv2.resize(image, input_size, interpolation=cv2.INTER_LINEAR)
+    
+    # # Define transformation pipeline
+    # transform = transforms.Compose([
+    #     transforms.ToTensor(),
+    #     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    # ])
+
+    # # Apply transformations
+    # image_tensor = transform(resized_image)
+    # image_batch = image_tensor.unsqueeze(0)
+
+    image_batch = np.array(resized_image)
+
+
+
+    
+    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
+    h,w = image.shape[:2]
+    
+    # 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
+

zombie.py

@@ -8,6 +8,8 @@ import align
 # import time
 import cv2
 
+from faceparsing2 import get_face_mask
+
 # sys.path.append(r"..\face-parsing.PyTorch")
 # import inferface
 
@@ -118,7 +120,7 @@ def ImageOpen(filepath):
         # cases: image don't have getexif
         return Image.open(filepath)
 
-def do_face(facenet, aligned, box, im_array, ort_session):
+def do_face(aligned, box, im_array, ort_session):
         z = square_human_2_zombie_onnx(aligned,ort_session)
         t1 = time.time()
 
@@ -127,14 +129,9 @@ def do_face(facenet, aligned, box, im_array, ort_session):
         h,w,c = im_array.shape
         face_in_place = cv2.warpAffine(np.array(z),t,(w,h))
         
-        face_classes = inferface.run_net(facenet,aligned)
+        face_classes = get_face_mask(aligned)
         face_mask = np.logical_and(face_classes < 14, face_classes > 0)
         
-        # face_classes_z = inferface.run_net(facenet,z)
-        # face_mask_z = np.logical_and(face_classes_z < 14, face_classes > 0)
-        
-        # face_mask = np.logical_or(face_mask,face_mask_z)
-
         imagemask = cv2.warpAffine(face_mask.astype("uint8"),t,(w,h))
 
         imagemask = imagemask.astype("uint8")
@@ -158,6 +155,6 @@ def make_faces_zombie_from_array(im_array_rgb, facenet, ort_session):
         faces = []
     
     for aligned,box in faces:
-        do_face(facenet,aligned,box,im_array_rgb,ort_session)
+        do_face(aligned,box,im_array_rgb,ort_session)
 
     return Image.fromarray(im_array_rgb)