test token

app.py

@@ -18,6 +18,10 @@ options = vision.ImageSegmenterOptions(base_options=base_options,output_category
 segmenter = vision.ImageSegmenter.create_from_options(options)
 labels = segmenter.labels
 
+def get_session_token(request: gr.Request):
+    x_ip_token = request.headers['x-ip-token']
+    return x_ip_token
+
 def segment(input_image, category):
     image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
     segmentation_result = segmenter.segment(image)
@@ -32,7 +36,11 @@ def segment(input_image, category):
         target_mask = category_mask_np == 0
     croper = Croper(input_image, target_mask)
 
-    return croper.corp_mask_image()
+    croper.corp_mask_image()
+    restore_image = croper.restore_result(croper.resized_square_image)
+    mask_image = croper.resized_square_mask_image
+    
+    return mask_image, restore_image
 
 def get_clothes_mask(category_mask_np):
     body_skin_mask = category_mask_np == 2
@@ -101,6 +109,7 @@ with gr.Blocks() as app:
             input_image = gr.Image(type='pil', label='Upload image')
             category = gr.Dropdown(label='Category', choices=category_options, value=category_options[0])
             submit_btn = gr.Button(value='Submit', variant='primary')
+            session_token = gr.Textbox(label='Session token', value='')
         with gr.Column():
             mask_image = gr.Image(type='pil', label='Segmentation mask')
             output_image = gr.Image(type='pil', label='Segmented image')
@@ -114,4 +123,6 @@ with gr.Blocks() as app:
         outputs=[mask_image, output_image]
     )
 
+    app.load(get_session_token, None, session_token)
+
 app.launch(debug=False, show_error=True)

croper.py

@@ -8,27 +8,30 @@ class Croper:
         self,
         input_image: PIL.Image,
         target_mask: np.ndarray,
+        mask_size: int = 256,
+        mask_expansion: int = 20
     ):
         self.input_image = input_image
         self.target_mask = target_mask
+        self.mask_size = mask_size
+        self.mask_expansion = mask_expansion
     
     def corp_mask_image(self):
         target_mask = self.target_mask
         input_image = self.input_image
-        crop_length = 512
-        expand_size = 40
+        mask_expansion = self.mask_expansion
         original_width, original_height = input_image.size
         mask_indices = np.where(target_mask)
-        start_y = np.min(mask_indices[0]) - expand_size
+        start_y = np.min(mask_indices[0]) - mask_expansion
         if start_y < 0:
             start_y = 0
-        end_y = np.max(mask_indices[0]) + expand_size
+        end_y = np.max(mask_indices[0]) + mask_expansion
         if end_y > original_height:
             end_y = original_height
-        start_x = np.min(mask_indices[1]) - expand_size
+        start_x = np.min(mask_indices[1]) - mask_expansion
         if start_x < 0:
             start_x = 0
-        end_x = np.max(mask_indices[1]) + expand_size
+        end_x = np.max(mask_indices[1]) + mask_expansion
         if end_x > original_width:
             end_x = original_width
         mask_height = end_y - start_y
@@ -65,7 +68,21 @@ class Croper:
         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))
+        mask_size = self.mask_size
+        self.resized_square_mask_image = square_mask_image.resize((mask_size, mask_size))
+        self.resized_square_image = square_image.resize((mask_size, mask_size))
 
-        return self.square_image, self.resized_square_image
+        return self.resized_square_mask_image
+    
+    def restore_result(self, generated_image):
+        square_length = self.square_length
+        generated_image = generated_image.resize((square_length, square_length))
+        square_mask_image = self.square_mask_image
+        cropped_generated_image = generated_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+        cropped_square_mask_image = square_mask_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+
+        restored_image = self.input_image.copy()
+        restored_image.paste(cropped_generated_image, (self.origin_start_x, self.origin_start_y), cropped_square_mask_image)
+        
+        return restored_image
+