Update lib/pymaf/utils/imutils.py

Fix bug for bbox selection and update rembg

app.py

@@ -10,8 +10,6 @@ import subprocess
 
 if os.getenv('SYSTEM') == 'spaces':
     subprocess.run('pip install pyembree'.split())
-    subprocess.run(
-        'pip install git+https://github.com/YuliangXiu/rembg.git@hf'.split())
     subprocess.run(
         'pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html'.split())
     subprocess.run(

lib/pymaf/utils/imutils.py

@@ -6,7 +6,8 @@ import io
 import torch
 import numpy as np
 from PIL import Image
-from rembg.bg import remove
+from rembg import remove
+from rembg.session_factory import new_session
 from torchvision.models import detection
 
 from lib.pymaf.core import constants
@@ -43,6 +44,7 @@ def get_bbox(img, det):
     bbox = bboxes[0, 0, 0].cpu().numpy()
 
     return bbox
+# Michael Black is
 
 
 def get_transformer(input_res):
@@ -107,9 +109,8 @@ def process_image(img_file, hps_type, input_res=512, device=None, seg_path=None)
     detector.eval()
     predictions = detector(
         [torch.from_numpy(img_for_crop).permute(2, 0, 1) / 255.])[0]
-    human_ids = torch.logical_and(
-        predictions["labels"] == 1,
-        predictions["scores"] == predictions["scores"].max()).nonzero().squeeze(1)
+    human_ids = torch.where(
+        predictions["scores"] == predictions["scores"][predictions['labels'] == 1].max())
     bbox = predictions["boxes"][human_ids, :].flatten().detach().cpu().numpy()
     
     width = bbox[2] - bbox[0]
@@ -126,12 +127,8 @@ def process_image(img_file, hps_type, input_res=512, device=None, seg_path=None)
         img_np, cropping_parameters = crop(
             img_for_crop, center, scale, (input_res, input_res))
 
-    with torch.no_grad():
-        buf = io.BytesIO()
-        Image.fromarray(img_np).save(buf, format='png')
-        img_pil = Image.open(
-            io.BytesIO(remove(buf.getvalue()))).convert("RGBA")
-
+    img_pil = Image.fromarray(remove(img_np, post_process_mask=True, session=new_session("u2net")))
+    
     # for icon
     img_rgb = image_to_tensor(img_pil.convert("RGB"))
     img_mask = torch.tensor(1.0) - (mask_to_tensor(img_pil.split()[-1]) <

requirements.txt

@@ -1,7 +1,7 @@
 pip>=21.2.4
 numpy==1.22.4
 matplotlib==3.5.0
-Pillow==9.0.0
+Pillow==9.2.0
 PyOpenGL
 PyOpenGL_accelerate
 PyYAML>=6.0
@@ -16,11 +16,12 @@ shapely==1.7.1
 rtree==0.9.7
 pytorch_lightning==1.2.5
 PyMCubes
-chumpy
 opencv-python
 opencv_contrib_python
 scikit-learn
 protobuf==3.20.0
 pymeshlab
 iopath
-fvcore
+fvcore
+chumpy
+git+https://github.com/YuliangXiu/rembg.git