add croper

.gitignore

@@ -1,2 +1,3 @@
 .vscode
-.DS_Store
+.DS_Store
+__pycache__

app.py

@@ -1,10 +1,12 @@
 import gradio as gr
 import mediapipe as mp
 import numpy as np
+
 from PIL import Image
 from mediapipe.tasks import python
 from mediapipe.tasks.python import vision
 from scipy.ndimage import binary_dilation, label
+from croper import Croper
 
 BG_COLOR = np.array([0, 0, 0], dtype=np.uint8) # black
 MASK_COLOR = np.array([255, 255, 255], dtype=np.uint8) # white
@@ -15,10 +17,8 @@ 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
@@ -30,54 +30,9 @@ def segment(input_image, category):
         target_mask = get_clothes_mask(category_mask_np)
     else:
         target_mask = category_mask_np == 0
+    croper = Croper(input_image, target_mask)
 
-    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)
-    fg_image[:] = MASK_COLOR
-    bg_image = np.zeros(image_data.shape, dtype=np.uint8)
-    bg_image[:] = BG_COLOR
-
-    condition = np.stack((target_mask,) * 3, axis=-1) > 0.2
-
-    output_image = np.where(condition, fg_image, bg_image)
-    output_image = Image.fromarray(output_image)
-    return crop_mask_image, crop_image_square
+    return croper.corp_mask_image()
 
 def get_clothes_mask(category_mask_np):
     body_skin_mask = category_mask_np == 2

croper.py

@@ -0,0 +1,71 @@
+import PIL
+import numpy as np
+
+from PIL import Image
+
+class Croper:
+    def __init__(
+        self,
+        input_image: PIL.Image,
+        target_mask: np.ndarray,
+    ):
+        self.input_image = input_image
+        self.target_mask = target_mask
+    
+    def corp_mask_image(self):
+        target_mask = self.target_mask
+        input_image = self.input_image
+        crop_length = 512
+        expand_size = 40
+        original_width, original_height = input_image.size
+        mask_indices = np.where(target_mask)
+        start_y = np.min(mask_indices[0]) - expand_size
+        if start_y < 0:
+            start_y = 0
+        end_y = np.max(mask_indices[0]) + expand_size
+        if end_y > original_height:
+            end_y = original_height
+        start_x = np.min(mask_indices[1]) - expand_size
+        if start_x < 0:
+            start_x = 0
+        end_x = np.max(mask_indices[1]) + expand_size
+        if end_x > original_width:
+            end_x = original_width
+        mask_height = end_y - start_y
+        mask_width = end_x - start_x
+
+        # choose the max side length
+        max_side_length = max(mask_height, mask_width)
+        # calculate the crop area
+        crop_mask = target_mask[start_y:end_y, start_x:end_x]
+        crop_mask_start_y = (max_side_length - mask_height) // 2
+        crop_mask_end_y = crop_mask_start_y + mask_height
+        crop_mask_start_x = (max_side_length - mask_width) // 2
+        crop_mask_end_x = crop_mask_start_x + mask_width
+        # create a square mask
+        square_mask = np.zeros((max_side_length, max_side_length), dtype=target_mask.dtype)
+        square_mask[crop_mask_start_y:crop_mask_end_y, crop_mask_start_x:crop_mask_end_x] = crop_mask
+        square_mask_image = Image.fromarray((square_mask * 255).astype(np.uint8))
+
+        crop_image = input_image.crop((start_x, start_y, end_x, end_y))
+        square_image = Image.new("RGB", (max_side_length, max_side_length))
+        square_image.paste(crop_image, (crop_mask_start_x, crop_mask_start_y))
+
+        self.origin_start_x = start_x
+        self.origin_start_y = start_y
+        self.origin_end_x = end_x
+        self.origin_end_y = end_y
+
+        self.square_start_x = crop_mask_start_x
+        self.square_start_y = crop_mask_start_y
+        self.square_end_x = crop_mask_end_x
+        self.square_end_y = crop_mask_end_y
+
+        self.square_length = max_side_length
+        self.square_mask_image = square_mask_image
+        self.square_image = square_image
+
+        self.resized_square_mask_image = square_mask_image.resize((crop_length, crop_length))
+        self.resized_square_image = square_image.resize((crop_length, crop_length))
+
+        return self.square_image, self.resized_square_image