Update app.py

app.py

@@ -12,9 +12,9 @@ def dummy_warmup():
 
 import numpy as np
 import gradio as gr
-# import rembg
+import rembg
 # import trimesh
-# from moge.model.v1 import MoGeModel
+from moge.model.v1 import MoGeModel
 # from utils.geometry import compute_pointmap
 # import os, shutil
 import cv2
@@ -177,24 +177,24 @@ import cv2
 #     output = moge_model.infer(input_img_torch).cpu()
 #     return output 
 
-# @GPU(duration = 180)
-# def mask_rembg(input_img):
-#     #masked_img = rembg.remove(input_img,)
-#     output_img = rembg.remove(input_img, alpha_matting=False, post_process_mask=True)
-
-#     # Convert to NumPy array
-#     output_np = np.array(output_img)
-#     alpha = output_np[..., 3]
-
-#     # Step 2: Erode the alpha mask to shrink object slightly
-#     kernel = np.ones((3, 3), np.uint8)  # Adjust size for aggressiveness
-#     eroded_alpha = cv2.erode(alpha, kernel, iterations=1)
-#     # Step 3: Replace alpha channel
-#     output_np[..., 3] = eroded_alpha
+@GPU(duration = 180)
+def mask_rembg(input_img):
+    #masked_img = rembg.remove(input_img,)
+    output_img = rembg.remove(input_img, alpha_matting=False, post_process_mask=True)
+
+    # Convert to NumPy array
+    output_np = np.array(output_img)
+    alpha = output_np[..., 3]
+
+    # Step 2: Erode the alpha mask to shrink object slightly
+    kernel = np.ones((3, 3), np.uint8)  # Adjust size for aggressiveness
+    eroded_alpha = cv2.erode(alpha, kernel, iterations=1)
+    # Step 3: Replace alpha channel
+    output_np[..., 3] = eroded_alpha
   
-#     mask = output_np[:,:,-1] >= 128
-#     rgb = output_np[:,:,:3]
-#     return mask, rgb 
+    mask = output_np[:,:,-1] >= 128
+    rgb = output_np[:,:,:3]
+    return mask, rgb 
 
 @GPU(duration = 180)
 def process_image(input_img):
@@ -205,7 +205,7 @@ def process_image(input_img):
         rotated = True
     input_img = cv2.resize(input_img, (640, 480))
     mask, rgb = mask_rembg(input_img)
-    depth_dict = depth_moge(input_img)
+    # depth_dict = depth_moge(input_img)
     if os.path.exists(outdir):
         shutil.rmtree(outdir)
     os.makedirs(outdir)