318 lines
8.7 KiB
Python
318 lines
8.7 KiB
Python
#!/usr/bin/python3
|
|
import cv2
|
|
import numpy as np
|
|
|
|
import torch
|
|
import einops
|
|
|
|
import facer
|
|
|
|
|
|
def load_image(image_path):
|
|
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
|
|
image = cv2.cvtColor(image, code=cv2.COLOR_BGR2RGB)
|
|
image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
|
|
return image
|
|
|
|
|
|
def do_recolor(vis_seg_probs, n_classes):
|
|
val = int(255 / n_classes)
|
|
vis_seg_probs = vis_seg_probs.cpu().detach().numpy()
|
|
not_visible = (vis_seg_probs == 0).astype(dtype=np.uint8)
|
|
not_visible = 1 - not_visible
|
|
not_visible *= 255
|
|
vis_seg_probs *= val
|
|
ret = np.array((vis_seg_probs, not_visible, not_visible), np.uint8)
|
|
ret = einops.rearrange(ret, 'c h w -> h w c')
|
|
ret = cv2.cvtColor(ret, cv2.COLOR_HSV2BGR_FULL)
|
|
return ret
|
|
|
|
|
|
def detect_face_from_tensor(image):
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
image *= 255
|
|
image = image.to(dtype=torch.uint8)
|
|
image = facer.hwc2bchw(image).to(device=device)
|
|
face_detector = facer.face_detector('retinaface/mobilenet', device=device)
|
|
|
|
with torch.inference_mode():
|
|
faces = face_detector(image)
|
|
|
|
face_parser = facer.face_parser(
|
|
'farl/lapa/448', device=device) # optional "farl/celebm/448"
|
|
|
|
with torch.inference_mode():
|
|
faces = face_parser(image, faces)
|
|
|
|
seg_logits = faces['seg']['logits']
|
|
num_faces = seg_logits.shape[0]
|
|
print(num_faces)
|
|
|
|
if num_faces >= 1:
|
|
|
|
seg_probs = seg_logits.softmax(dim=1) # nfaces x nclasses x h x w
|
|
n_classes = seg_probs.size(1)
|
|
|
|
vis_seg_probs = seg_probs.argmax(dim=1)
|
|
vis_seg_probs = einops.einsum(vis_seg_probs, 'b h w -> h w')
|
|
|
|
return (vis_seg_probs, n_classes, num_faces)
|
|
|
|
else:
|
|
|
|
vis_seg_probs = torch.zeros((image.shape[0], image.shape[1]),
|
|
dtype=torch.int64)
|
|
|
|
n_classes = 11
|
|
|
|
return (vis_seg_probs, n_classes, num_faces)
|
|
|
|
|
|
def full_work_wrapper(image):
|
|
|
|
try:
|
|
res, n_classes, num_faces = detect_face_from_tensor(image)
|
|
except:
|
|
res = torch.zeros((image.shape[0], image.shape[1]), dtype=torch.int64)
|
|
n_classes = 11
|
|
tup = do_recolor(res, n_classes)
|
|
return tup
|
|
|
|
|
|
def run_slave(input_image_path, output_image_path, tmp_file_path):
|
|
import os
|
|
|
|
EXEC_STRING = '''
|
|
import os
|
|
|
|
try:
|
|
del os.environ['AUX_ANNOTATOR_CKPTS_PATH']
|
|
os.unsetenv('AUX_ANNOTATOR_CKPTS_PATH')
|
|
except:
|
|
print('Failed to unset AUX_ANNOTATOR_CKPTS_PATH')
|
|
try:
|
|
del os.environ['AUX_ORT_PROVIDERS']
|
|
os.unsetenv('AUX_ORT_PROVIDERS')
|
|
except:
|
|
print('Failed to unset AUX_ORT_PROVIDERS')
|
|
try:
|
|
del os.environ['AUX_TEMP_DIR']
|
|
os.unsetenv('AUX_TEMP_DIR')
|
|
except:
|
|
print('Failed to unset AUX_TEMP_DIR')
|
|
try:
|
|
del os.environ['AUX_USE_SYMLINKS']
|
|
os.unsetenv('AUX_USE_SYMLINKS')
|
|
except:
|
|
print('Failed to unset AUX_USE_SYMLINKS')
|
|
try:
|
|
del os.environ['CUBLAS_WORKSPACE_CONFIG']
|
|
os.unsetenv('CUBLAS_WORKSPACE_CONFIG')
|
|
except:
|
|
print('Failed to unset CUBLAS_WORKSPACE_CONFIG')
|
|
try:
|
|
del os.environ['CUDA_MODULE_LOADING']
|
|
os.unsetenv('CUDA_MODULE_LOADING')
|
|
except:
|
|
print('Failed to unset CUDA_MODULE_LOADING')
|
|
try:
|
|
del os.environ['DWPOSE_ONNXRT_CHECKED']
|
|
os.unsetenv('DWPOSE_ONNXRT_CHECKED')
|
|
except:
|
|
print('Failed to unset DWPOSE_ONNXRT_CHECKED')
|
|
try:
|
|
del os.environ['KINETO_LOG_LEVEL']
|
|
os.unsetenv('KINETO_LOG_LEVEL')
|
|
except:
|
|
print('Failed to unset KINETO_LOG_LEVEL')
|
|
try:
|
|
del os.environ['KMP_DUPLICATE_LIB_OK']
|
|
os.unsetenv('KMP_DUPLICATE_LIB_OK')
|
|
except:
|
|
print('Failed to unset KMP_DUPLICATE_LIB_OK')
|
|
try:
|
|
del os.environ['KMP_INIT_AT_FORK']
|
|
os.unsetenv('KMP_INIT_AT_FORK')
|
|
except:
|
|
print('Failed to unset KMP_INIT_AT_FORK')
|
|
try:
|
|
del os.environ['PYTORCH_CUDA_ALLOC_CONF']
|
|
os.unsetenv('PYTORCH_CUDA_ALLOC_CONF')
|
|
except:
|
|
print('Failed to unset PYTORCH_CUDA_ALLOC_CONF')
|
|
try:
|
|
del os.environ['PYTORCH_ENABLE_MPS_FALLBACK']
|
|
os.unsetenv('PYTORCH_ENABLE_MPS_FALLBACK')
|
|
except:
|
|
print('Failed to unset PYTORCH_ENABLE_MPS_FALLBACK')
|
|
try:
|
|
del os.environ['PYTORCH_NVML_BASED_CUDA_CHECK']
|
|
os.unsetenv('PYTORCH_NVML_BASED_CUDA_CHECK')
|
|
except:
|
|
print('Failed to unset PYTORCH_NVML_BASED_CUDA_CHECK')
|
|
try:
|
|
del os.environ['TF_CPP_MIN_LOG_LEVEL']
|
|
os.unsetenv('TF_CPP_MIN_LOG_LEVEL')
|
|
except:
|
|
print('Failed to unset TF_CPP_MIN_LOG_LEVEL')
|
|
try:
|
|
del os.environ['TOKENIZERS_PARALLELISM']
|
|
os.unsetenv('TOKENIZERS_PARALLELISM')
|
|
except:
|
|
print('Failed to unset TOKENIZERS_PARALLELISM')
|
|
try:
|
|
del os.environ['TORCH_CPP_LOG_LEVEL']
|
|
os.unsetenv('TORCH_CPP_LOG_LEVEL')
|
|
except:
|
|
print('Failed to unset TORCH_CPP_LOG_LEVEL')
|
|
|
|
import torch
|
|
import facer
|
|
import cv2
|
|
import einops
|
|
import numpy as np
|
|
import sys
|
|
|
|
|
|
def load_image(image_path):
|
|
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
|
|
image = cv2.cvtColor(image, code=cv2.COLOR_BGR2RGB)
|
|
image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
|
|
return image
|
|
|
|
|
|
def do_recolor(vis_seg_probs, n_classes):
|
|
val = int(255 / n_classes)
|
|
vis_seg_probs = vis_seg_probs.cpu().detach().numpy()
|
|
not_visible = (vis_seg_probs == 0).astype(dtype=np.uint8)
|
|
not_visible = 1 - not_visible
|
|
not_visible *= 255
|
|
vis_seg_probs *= val
|
|
ret = np.array((vis_seg_probs, not_visible, not_visible), np.uint8)
|
|
ret = einops.rearrange(ret, 'c h w -> h w c')
|
|
ret = cv2.cvtColor(ret, cv2.COLOR_HSV2BGR_FULL)
|
|
return ret
|
|
|
|
|
|
def detect_face_from_tensor(image):
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
image *= 255
|
|
image = image.to(dtype=torch.uint8)
|
|
image = facer.hwc2bchw(image).to(device=device)
|
|
face_detector = facer.face_detector('retinaface/mobilenet', device=device)
|
|
|
|
with torch.inference_mode():
|
|
faces = face_detector(image)
|
|
|
|
face_parser = facer.face_parser(
|
|
'farl/lapa/448', device=device) # optional "farl/celebm/448"
|
|
|
|
with torch.inference_mode():
|
|
faces = face_parser(image, faces)
|
|
|
|
seg_logits = faces['seg']['logits']
|
|
seg_probs = seg_logits.softmax(dim=1) # nfaces x nclasses x h x w
|
|
n_classes = seg_probs.size(1)
|
|
|
|
vis_seg_probs = seg_probs.argmax(dim=1)
|
|
vis_seg_probs = einops.einsum(vis_seg_probs, 'b h w -> h w')
|
|
return (vis_seg_probs, n_classes)
|
|
|
|
|
|
def full_work_wrapper(image):
|
|
try:
|
|
res, n_classes = detect_face_from_tensor(image)
|
|
tup = do_recolor(res, n_classes)
|
|
except:
|
|
print('Warning: Failed to find a face.')
|
|
tup = np.zeros(image.shape, dtype=np.uint8)
|
|
return tup
|
|
|
|
tup = full_work_wrapper(image=load_image(image_path=sys.argv[1]))
|
|
cv2.imwrite(sys.argv[2], tup)
|
|
'''
|
|
|
|
with open(tmp_file_path, 'w', encoding='utf-8') as f:
|
|
f.write(EXEC_STRING)
|
|
|
|
CMD = 'env > ~/env.txt ; python3 ' + tmp_file_path + ' ' + input_image_path + ' ' + output_image_path
|
|
|
|
print(CMD)
|
|
os.system(CMD)
|
|
|
|
|
|
def run_slave_tensor(image):
|
|
|
|
import tempfile
|
|
import cv2
|
|
import os
|
|
|
|
device = image.device
|
|
outtype = image.dtype
|
|
|
|
path_dir = tempfile.TemporaryDirectory(
|
|
suffix='.dir',
|
|
prefix='facer.',
|
|
dir=None,
|
|
ignore_cleanup_errors=False,
|
|
)
|
|
|
|
path_input = path_dir.name + '/input.png'
|
|
path_output = path_dir.name + '/output.png'
|
|
path_source = path_dir.name + '/exec.py'
|
|
|
|
image = image.detach().cpu().numpy() * 255.0
|
|
image = image.astype(dtype=np.uint8)
|
|
image = cv2.cvtColor(src=image, code=cv2.COLOR_RGB2BGR)
|
|
cv2.imwrite(path_input, image)
|
|
|
|
run_slave(input_image_path=path_input,
|
|
output_image_path=path_output,
|
|
tmp_file_path=path_source)
|
|
|
|
os.unlink(path_input)
|
|
os.unlink(path_source)
|
|
image = cv2.imread(path_output, cv2.IMREAD_COLOR)
|
|
os.unlink(path_output)
|
|
os.rmdir(path_dir.name)
|
|
# image = cv2.cvtColor(src=image, code=cv2.COLOR_BGR2RGB)
|
|
image = image.astype(np.float32) / 255.0
|
|
# image = torch.from_numpy(image).to(dtype=outtype, device=device) / 255.0
|
|
return image
|
|
|
|
|
|
class main_face_segment():
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"to_run": ("BOOLEAN", )
|
|
},
|
|
}
|
|
|
|
FUNCTION = "run"
|
|
RETURN_TYPES = ("IMAGE", )
|
|
CATEGORY = "TRI3D"
|
|
|
|
def run(self, image, to_run):
|
|
if to_run:
|
|
batch_size = image.shape[0]
|
|
ret = []
|
|
for i in range(batch_size):
|
|
ret.append(run_slave_tensor(image[i].clone()))
|
|
# ret.append(full_work_wrapper(image[i].clone()))
|
|
|
|
ret = np.array(ret)
|
|
|
|
ret = torch.from_numpy(ret).to(dtype=image.dtype,
|
|
device=image.device)
|
|
|
|
return (ret, )
|
|
else:
|
|
return (torch.zeros_like(image), )
|