fix sam mask

app.py

@@ -80,8 +80,18 @@ def sam_segment(predictor, input_image, *bbox_coords):
     out_image_bbox = out_image.copy()
     out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
     torch.cuda.empty_cache()
-    return Image.fromarray(out_image_bbox, mode='RGB')
-    # return Image.fromarray(out_image_bbox, mode='RGBA')
+    # return Image.fromarray(out_image_bbox, mode='RGB')
+
+    x_nonzero = np.nonzero(masks_bbox[-1].astype(np.uint8).sum(axis=0))
+    y_nonzero = np.nonzero(masks_bbox[-1].astype(np.uint8).sum(axis=1))
+    x_min = int(x_nonzero[0].min())
+    y_min = int(y_nonzero[0].min())
+    x_max = int(x_nonzero[0].max())
+    y_max = int(y_nonzero[0].max())
+
+    out_image_bbox = out_image_bbox[y_min:y_max, x_min:x_max]
+
+    return Image.fromarray(out_image_bbox, mode='RGBA')
 
 
 def expand2square(pil_img, background_color):
@@ -113,28 +123,8 @@ def preprocess(predictor, input_image, chk_group=None, segment=False):
         x_max = int(x_nonzero[0].max())
         y_max = int(y_nonzero[0].max())
         input_image = sam_segment(predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
-    # Rescale and recenter
-    # if rescale:
-    #     image_arr = np.array(input_image)
-    #     in_w, in_h = image_arr.shape[:2]
-    #     out_res = min(RES, max(in_w, in_h))
-    #     ret, mask = cv2.threshold(np.array(input_image.split()[-1]), 0, 255, cv2.THRESH_BINARY)
-    #     x, y, w, h = cv2.boundingRect(mask)
-    #     max_size = max(w, h)
-    #     ratio = 0.75
-    #     side_len = int(max_size / ratio)
-    #     padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8)
-    #     center = side_len//2
-    #     padded_image[center-h//2:center-h//2+h, center-w//2:center-w//2+w] = image_arr[y:y+h, x:x+w]
-    #     rgba = Image.fromarray(padded_image).resize((out_res, out_res), Image.LANCZOS)
-
-    #     rgba_arr = np.array(rgba) / 255.0
-    #     rgb = rgba_arr[...,:3] * rgba_arr[...,-1:] + (1 - rgba_arr[...,-1:])
-    #     input_image = Image.fromarray((rgb * 255).astype(np.uint8))
-    # else:
-    #     input_image = expand2square(input_image, (127, 127, 127, 0))
     
-    input_image = expand2square(input_image, (0, 0, 0))
+    input_image = expand2square(input_image, (0, 0, 0, 0))
     return input_image, input_image.resize((256, 256), Image.Resampling.LANCZOS)
 
 
@@ -545,7 +535,7 @@ def run_pipeline(model_items, cfgs, input_img):
         mesh_image = save_images(shading, mask_pred)
         mesh_bones_image = save_images(image_with_bones, mask_final)
 
-        shape_glb, shape_obj = process_mesh(shape, 'reconstruced_shape')
+        shape_glb, shape_obj = process_mesh(shape, 'reconstructed_shape')
         base_shape_glb, base_shape_obj = process_mesh(prior_shape, 'reconstructed_base_shape')
 
         return mesh_image, mesh_bones_image, shape_glb, shape_obj, base_shape_glb, base_shape_obj