only convert mask area to gray

app_haircolor.py

@@ -5,7 +5,7 @@ import torch
 
 from PIL import Image
 from segment_utils import(
-    segment_image,
+    segment_image_with_gray,
     restore_result,
 )
 from enhance_utils import enhance_image
@@ -111,7 +111,7 @@ def create_demo() -> gr.Blocks:
                 generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
         
         g_btn.click(
-            fn=segment_image,
+            fn=segment_image_with_gray,
             inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
             outputs=[origin_area_image, croper],
         ).success(

inversion_run_adapter.py

@@ -177,9 +177,6 @@ def run(
         adapter_conditioning_scale=conditioning_scale,
     )
 
-    # convert to grayscale
-    input_image = convert_to_grayscale(input_image)
-    
     x_0_image = input_image
     x_0 = encode_image(x_0_image, pipeline)
     x_ts = create_xts(1, None, 0, generator, pipeline.scheduler, timesteps, x_0, no_add_noise=False)
@@ -283,7 +280,3 @@ def get_timesteps(pipe, num_inference_steps, strength, device, denoising_start=N
         return timesteps, num_inference_steps
 
     return timesteps, num_inference_steps - t_start
-
-def convert_to_grayscale(pil_image: Image):
-    gray_image = pil_image.convert('L')
-    return Image.merge('RGB', (gray_image, gray_image, gray_image))

segment_utils.py

@@ -58,6 +58,33 @@ def segment_image(input_image, category, input_size, mask_expansion, mask_dilati
 
     return origin_area_image, croper
 
+def segment_image_with_gray(input_image, category, input_size, mask_expansion, mask_dilation):
+    mask_size = int(input_size)
+    mask_expansion = int(mask_expansion)
+
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
+    segmentation_result = segmenter.segment(image)
+    category_mask = segmentation_result.category_mask
+    category_mask_np = category_mask.numpy_view()
+
+    if category == "hair":
+        target_mask = get_hair_mask(category_mask_np, mask_dilation)
+    elif category == "clothes":
+        target_mask = get_clothes_mask(category_mask_np, mask_dilation)
+    elif category == "face":
+        target_mask = get_face_mask(category_mask_np, mask_dilation)
+    else:
+        target_mask = get_face_mask(category_mask_np, mask_dilation)
+    
+    croper = Croper(input_image, target_mask, mask_size, mask_expansion)
+    croper.corp_mask_image()
+    origin_area_image = croper.resized_square_image
+    mask_image = croper.resized_square_mask_image
+    gray_area_image = origin_area_image.convert('L')
+    origin_area_image.paste(gray_area_image, (0, 0), mask_image)
+
+    return origin_area_image, croper
+
 def get_face_mask(category_mask_np, dilation=1):
     face_skin_mask = category_mask_np == 3
     if dilation > 0: