diff --git a/cli/modules/process.py b/cli/modules/process.py index 924250ae2..00f92f072 100755 --- a/cli/modules/process.py +++ b/cli/modules/process.py @@ -20,6 +20,8 @@ from sdapi import postsync params = Map({ 'src': '', # source folder 'dst': '', # destination folder + 'extract_face': True, # extract face from image + 'extract_body': True, # extract face from image 'clear_dst': True, # remove all files from destination at the start 'target_size': 512, # target resolution 'square_images': True, # should output images be squared @@ -34,6 +36,10 @@ params = Map({ 'body_pad': 0.2, # pad body image percentage 'body_model': 2, # body model to use 0/low 1/medium 2/high 'body_blur_score': 1.6, # max score for body blur detection + 'segmentation_face': True, # segmentation enabled + 'segmentation_body': False, # segmentation enabled + 'segmentation_model': 0, # segmentation model 0/general 1/landscape + 'segmentation_background': (192, 192, 192), # segmentation background color 'interrogate_model': 'clip' # interrogate model }) @@ -52,7 +58,21 @@ def detect_blur(image): return mean +def segmentation(image): + with mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=params.segmentation_model) as selfie_segmentation: + data = np.array(image) + results = selfie_segmentation.process(data) + condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1 + background = np.zeros(data.shape, dtype=np.uint8) + background[:] = params.segmentation_background + data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine + segmented = Image.fromarray(data) + return segmented + + def extract_face(img): + if not params.extract_face: + return None, True if img.mode == 'RGBA': img = img.convert('RGB') scale = max(img.size[0], img.size[1]) / params.target_size @@ -80,6 +100,8 @@ def extract_face(img): if params.square_images: squared = Image.new('RGB', (params.target_size, params.target_size)) squared.paste(cropped, (0, 0)) + if params.segmentation_face: + squared = segmentation(squared) else: squared = cropped blur = detect_blur(squared) @@ -93,6 +115,8 @@ def extract_face(img): def extract_body(img): + if not params.extract_body: + return None, True if img.mode == 'RGBA': img = img.convert('RGB') scale = max(img.size[0], img.size[1]) / params.target_size @@ -122,6 +146,8 @@ def extract_body(img): if params.square_images: squared = Image.new('RGB', (params.target_size, params.target_size)) squared.paste(cropped, (0, 0)) + if params.segmentation_body: + squared = segmentation(squared) else: squared = cropped blur = detect_blur(squared)