add segment_image

app_upscale.py

@@ -6,6 +6,11 @@ import torch
 import gradio as gr
 import spaces
 
+from segment_utils import(
+    segment_image,
+    restore_result,
+)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print(f'{device} is available')
@@ -24,13 +29,13 @@ def create_demo() -> gr.Blocks:
         input_image: Image,
         prompt: str,
     ):
+        time_cost_str = ''
+        run_task_time = 0
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
         upscaled_image = upscale_pipe(prompt=prompt, image=input_image).images[0]
-        extension = 'png'
-
-        path = f"output/{uuid.uuid4()}.{extension}"
-        upscaled_image.save(path, quality=100)
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
         
-        return upscaled_image, path, time_cost_str
+        return upscaled_image, time_cost_str
 
     def get_time_cost(run_task_time, time_cost_str):
         now_time = int(time.time()*1000)
@@ -55,14 +60,27 @@ def create_demo() -> gr.Blocks:
             with gr.Column():
                 input_image = gr.Image(label="Input Image", type="pil")
             with gr.Column():
+                origin_area_image = gr.Image(label="Origin Area Image", format="png", type="pil", interactive=False, visible=False)
                 upscaled_image = gr.Image(label="Upscaled Image", format="png", type="pil", interactive=False)
                 download_path = gr.File(label="Download the output image", interactive=False)
                 generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
-        
+                category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
+                generate_size = gr.Number(label="Generate Size", value=1024, visible=False)
+                mask_expansion = gr.Number(label="Mask Expansion", value=20, visible=False)
+                mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation", visible=False)
+
         g_btn.click(
+            fn=segment_image,
+            inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
+            outputs=[origin_area_image, croper],
+        ).success(
             fn=upscale_image,
-            inputs=[input_image, input_image_prompt],
-            outputs=[upscaled_image, download_path, generated_cost],
+            inputs=[origin_area_image, input_image_prompt],
+            outputs=[upscaled_image, generated_cost],
+        ).success(
+            fn=restore_result,
+            inputs=[croper, category, enhanced_image],
+            outputs=[upscaled_image, download_path],
         )
 
     return demo

checkpoints/selfie_multiclass_256x256.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6748b1253a99067ef71f7e26ca71096cd449baefa8f101900ea23016507e0e0
+size 16371837

enhance_utils.py

@@ -38,7 +38,7 @@ face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=1, arch='clean', c
 
 def enhance_image(
     pil_image: Image,
-    enhance_face: bool = True,
+    enhance_face: bool = False,
 ):
     img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
 

segment_utils.py

@@ -0,0 +1,98 @@
+import numpy as np
+import mediapipe as mp
+import uuid
+
+from PIL import Image
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+from scipy.ndimage import binary_dilation
+from croper import Croper
+
+segment_model = "checkpoints/selfie_multiclass_256x256.tflite"
+base_options = python.BaseOptions(model_asset_path=segment_model)
+options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True)
+segmenter = vision.ImageSegmenter.create_from_options(options)
+
+def restore_result(croper, category, generated_image):
+    square_length = croper.square_length
+    generated_image = generated_image.resize((square_length, square_length))
+
+    cropped_generated_image = generated_image.crop((croper.square_start_x, croper.square_start_y, croper.square_end_x, croper.square_end_y))
+    cropped_square_mask_image = get_restore_mask_image(croper, category, cropped_generated_image)
+
+    restored_image = croper.input_image.copy()
+    restored_image.paste(cropped_generated_image, (croper.origin_start_x, croper.origin_start_y), cropped_square_mask_image)
+
+    extension = 'png'
+    # if restored_image.mode == 'RGBA':
+    #     extension = 'png'
+    # else:
+    #     extension = 'jpg'
+
+    path = f"output/{uuid.uuid4()}.{extension}"
+    restored_image.save(path, quality=100)
+
+    return restored_image, path
+
+def segment_image(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
+
+    return origin_area_image, croper
+
+def get_face_mask(category_mask_np, dilation=1):
+    face_skin_mask = category_mask_np == 3
+    if dilation > 0:
+        face_skin_mask = binary_dilation(face_skin_mask, iterations=dilation)
+
+    return face_skin_mask
+
+def get_clothes_mask(category_mask_np, dilation=1):
+    body_skin_mask = category_mask_np == 2
+    clothes_mask = category_mask_np == 4
+    combined_mask = np.logical_or(body_skin_mask, clothes_mask)
+    combined_mask = binary_dilation(combined_mask, iterations=4)
+    if dilation > 0:
+        combined_mask = binary_dilation(combined_mask, iterations=dilation)
+    return combined_mask
+
+def get_hair_mask(category_mask_np, dilation=1):
+    hair_mask = category_mask_np == 1
+    if dilation > 0:
+        hair_mask = binary_dilation(hair_mask, iterations=dilation)
+    return hair_mask
+
+def get_restore_mask_image(croper, category, generated_image):
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(generated_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, 0)
+    elif category == "clothes":
+        target_mask = get_clothes_mask(category_mask_np, 0)
+    elif category == "face":
+        target_mask = get_face_mask(category_mask_np, 0)
+    
+    combined_mask = np.logical_or(target_mask, croper.corp_mask)
+    mask_image = Image.fromarray((combined_mask * 255).astype(np.uint8))
+    return mask_image