add crop

app.py

@@ -15,8 +15,10 @@ base_options = python.BaseOptions(model_asset_path=MODEL_PATH)
 options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True)
 segmenter = vision.ImageSegmenter.create_from_options(options)
 labels = segmenter.labels
+expand_size = 40
 
 def segment(input_image, category):
+    original_height, original_width = input_image.size
     image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
     segmentation_result = segmenter.segment(image)
     category_mask = segmentation_result.category_mask
@@ -29,6 +31,41 @@ def segment(input_image, category):
     else:
         target_mask = category_mask_np == 0
 
+    target_indices = np.where(target_mask)
+    start_y = np.min(target_indices[0]) - expand_size
+    if start_y < 0:
+        start_y = 0
+    end_y = np.max(target_indices[0]) + expand_size
+    if end_y > original_height:
+        end_y = original_height
+    start_x = np.min(target_indices[1]) - expand_size
+    if start_x < 0:
+        start_x = 0
+    end_x = np.max(target_indices[1]) + expand_size
+    if end_x > original_width:
+        end_x = original_width
+    target_height = end_y - start_y
+    target_width = end_x - start_x
+
+    # choose the max side length
+    max_side_length = max(target_height, target_width)
+    # calculate the crop area
+    crop_mask = target_mask[start_y:end_y, start_x:end_x]
+    crop_mask_height, crop_mask_width = crop_mask.shape
+    crop_mask_start_y = (max_side_length - crop_mask_height) // 2
+    crop_mask_end_y = crop_mask_start_y + crop_mask_height
+    crop_mask_start_x = (max_side_length - crop_mask_width) // 2
+    crop_mask_end_x = crop_mask_start_x + crop_mask_width
+    # create a square mask
+    crop_mask_square = np.zeros((max_side_length, max_side_length), dtype=target_mask.dtype)
+    crop_mask_square[crop_mask_start_y:crop_mask_end_y, crop_mask_start_x:crop_mask_end_x] = crop_mask
+    # create a square image
+    crop_mask_image = Image.fromarray((crop_mask_square * 255).astype(np.uint8))
+
+    crop_image = input_image.crop((start_x, start_y, end_x, end_y))
+    crop_image_square = Image.new("RGB", (max_side_length, max_side_length))
+    crop_image_square.paste(crop_image, (crop_mask_start_x, crop_mask_start_y))
+
     # Generate solid color images for showing the output segmentation mask.
     image_data = image.numpy_view()
     fg_image = np.zeros(image_data.shape, dtype=np.uint8)
@@ -40,7 +77,7 @@ def segment(input_image, category):
 
     output_image = np.where(condition, fg_image, bg_image)
     output_image = Image.fromarray(output_image)
-    return output_image
+    return crop_mask_image, crop_image_square
 
 def get_clothes_mask(category_mask_np):
     body_skin_mask = category_mask_np == 2
@@ -82,10 +119,10 @@ def get_hair_mask(category_mask_np, should_dilate=False):
         expanded_hair_mask = binary_dilation(expanded_hair_mask, iterations=10)
     
     # Trim the expanded_hair_mask
-    # 1. Remove the area above hair_mask by 20 pixels
+    # 1. Remove the area above hair_mask by 10 pixels
     hair_indices = np.where(hair_mask)
     min_hair_y = np.min(hair_indices[0])
-    expanded_hair_mask[:min_hair_y - 20, :] = 0
+    expanded_hair_mask[:min_hair_y - 10, :] = 0
 
     # 2. Remove the areas on both sides that exceed the clothing coordinates
     clothes_indices = np.where(clothes_mask)
@@ -110,7 +147,8 @@ with gr.Blocks() as app:
             category = gr.Dropdown(label='Category', choices=category_options, value=category_options[0])
             submit_btn = gr.Button(value='Submit', variant='primary')
         with gr.Column():
-            output_image = gr.Image(type='pil', label='Image Output')
+            mask_image = gr.Image(type='pil', label='Segmentation mask')
+            output_image = gr.Image(type='pil', label='Segmented image')
 
     submit_btn.click(
         fn=segment,
@@ -118,7 +156,7 @@ with gr.Blocks() as app:
             input_image,
             category,
         ],
-        outputs=[output_image]
+        outputs=[mask_image, output_image]
     )
 
 app.launch(debug=False, show_error=True)