remove subprocess, U2Net for bg removal

.gitignore

@@ -1,3 +1,4 @@
+results/
 # Python build
 .eggs/
 gradio.egg-info/*

PIFu/apps/eval.py

@@ -1,28 +1,27 @@
+import tqdm
+import glob
+import torchvision.transforms as transforms
+from PIL import Image
+from lib.model import *
+from lib.train_util import *
+from lib.sample_util import *
+from lib.mesh_util import *
+# from lib.options import BaseOptions
+from torch.utils.data import DataLoader
+import torch
+import numpy as np
+import json
+import time
 import sys
 import os
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+sys.path.insert(0, os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '..')))
 ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-import time
-import json
-import numpy as np
-import torch
-from torch.utils.data import DataLoader
 
-from lib.options import BaseOptions
-from lib.mesh_util import *
-from lib.sample_util import *
-from lib.train_util import *
-from lib.model import *
-
-from PIL import Image
-import torchvision.transforms as transforms
-import glob
-import tqdm
-
-# get options
-opt = BaseOptions().parse()
+# # get options
+# opt = BaseOptions().parse()
 
 class Evaluator:
     def __init__(self, opt, projection_mode='orthogonal'):
@@ -34,19 +33,22 @@ class Evaluator:
             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
         ])
         # set cuda
-        cuda = torch.device('cuda:%d' % opt.gpu_id) if torch.cuda.is_available() else torch.device('cpu')
+        cuda = torch.device(
+            'cuda:%d' % opt.gpu_id) if torch.cuda.is_available() else torch.device('cpu')
 
         # create net
         netG = HGPIFuNet(opt, projection_mode).to(device=cuda)
         print('Using Network: ', netG.name)
 
         if opt.load_netG_checkpoint_path:
-            netG.load_state_dict(torch.load(opt.load_netG_checkpoint_path, map_location=cuda))
+            netG.load_state_dict(torch.load(
+                opt.load_netG_checkpoint_path, map_location=cuda))
 
         if opt.load_netC_checkpoint_path is not None:
             print('loading for net C ...', opt.load_netC_checkpoint_path)
             netC = ResBlkPIFuNet(opt).to(device=cuda)
-            netC.load_state_dict(torch.load(opt.load_netC_checkpoint_path, map_location=cuda))
+            netC.load_state_dict(torch.load(
+                opt.load_netC_checkpoint_path, map_location=cuda))
         else:
             netC = None
 
@@ -87,6 +89,30 @@ class Evaluator:
             'b_max': B_MAX,
         }
 
+    def load_image_from_memory(self, image_path, mask_path, img_name):
+        # Calib
+        B_MIN = np.array([-1, -1, -1])
+        B_MAX = np.array([1, 1, 1])
+        projection_matrix = np.identity(4)
+        projection_matrix[1, 1] = -1
+        calib = torch.Tensor(projection_matrix).float()
+        # Mask
+        mask = Image.fromarray(mask_path).convert('L')
+        mask = transforms.Resize(self.load_size)(mask)
+        mask = transforms.ToTensor()(mask).float()
+        # image
+        image = Image.fromarray(image_path).convert('RGB')
+        image = self.to_tensor(image)
+        image = mask.expand_as(image) * image
+        return {
+            'name': img_name,
+            'img': image.unsqueeze(0),
+            'calib': calib.unsqueeze(0),
+            'mask': mask.unsqueeze(0),
+            'b_min': B_MIN,
+            'b_max': B_MAX,
+        }
+
     def eval(self, data, use_octree=False):
         '''
         Evaluate a data point
@@ -98,18 +124,22 @@ class Evaluator:
             self.netG.eval()
             if self.netC:
                 self.netC.eval()
-            save_path = '%s/%s/result_%s.obj' % (opt.results_path, opt.name, data['name'])
+            save_path = '%s/%s/result_%s.obj' % (
+                opt.results_path, opt.name, data['name'])
             if self.netC:
-                gen_mesh_color(opt, self.netG, self.netC, self.cuda, data, save_path, use_octree=use_octree)
+                gen_mesh_color(opt, self.netG, self.netC, self.cuda,
+                               data, save_path, use_octree=use_octree)
             else:
-                gen_mesh(opt, self.netG, self.cuda, data, save_path, use_octree=use_octree)
+                gen_mesh(opt, self.netG, self.cuda, data,
+                         save_path, use_octree=use_octree)
 
 
 if __name__ == '__main__':
     evaluator = Evaluator(opt)
 
     test_images = glob.glob(os.path.join(opt.test_folder_path, '*'))
-    test_images = [f for f in test_images if ('png' in f or 'jpg' in f) and (not 'mask' in f)]
+    test_images = [f for f in test_images if (
+        'png' in f or 'jpg' in f) and (not 'mask' in f)]
     test_masks = [f[:-4]+'_mask.png' for f in test_images]
 
     print("num; ", len(test_masks))
@@ -120,4 +150,4 @@ if __name__ == '__main__':
             data = evaluator.load_image(image_path, mask_path)
             evaluator.eval(data, True)
         except Exception as e:
-           print("error:", e.args)
+            print("error:", e.args)

PIFu/inputs/.gitignore

@@ -1,2 +0,0 @@
-*
-!.gitignore

PIFu/lib/options.py

@@ -5,7 +5,7 @@ import os
 class BaseOptions():
     def __init__(self):
         self.initialized = False
-
+        argparse
     def initialize(self, parser):
         # Datasets related
         g_data = parser.add_argument_group('Data')
@@ -155,3 +155,7 @@ class BaseOptions():
     def parse(self):
         opt = self.gather_options()
         return opt
+
+    def parse_to_dict(self):
+        opt = self.gather_options()
+        return opt.__dict__

PIFu/results/spaces_demo/.gitignore

@@ -1,2 +0,0 @@
-*
-!.gitignore

app.py → PIFu/spaces.py

@@ -1,4 +1,5 @@
 import os
+from xml.etree.ElementPath import ops
 try:
     os.system("pip install --upgrade  torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html")
 except Exception as e:
@@ -7,67 +8,110 @@ except Exception as e:
 from pydoc import describe
 from huggingface_hub import hf_hub_download
 import gradio as gr
-import subprocess
 import os
-import datetime
+from datetime import datetime
 from PIL import Image
-from remove_bg import RemoveBackground
-import torch, torchvision, skimage 
+import torch
+import torchvision
+import skimage
+import paddlehub
+import numpy as np
+from lib.options import BaseOptions
+from apps.crop_img import process_img
+from apps.eval import Evaluator
+from types import SimpleNamespace
+import trimesh
 
 print(
     "torch: ", torch.__version__,
-    "\ntorchvision: ",torchvision.__version__,
+    "\ntorchvision: ", torchvision.__version__,
     "\nskimage:", skimage.__version__
 )
 
 net_C = hf_hub_download("radames/PIFu-upright-standing", filename="net_C")
 net_G = hf_hub_download("radames/PIFu-upright-standing", filename="net_G")
-torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet101', pretrained=True)
 
-remove_bg = RemoveBackground()
 
-env = {
-    **os.environ,
-    "CHECKPOINTS_NETG_PATH": net_G,
-    "CHECKPOINTS_NETC_PATH": net_C,
-    "RESULTS_PATH": './results',
-}
+opt = BaseOptions()
+opts = opt.parse_to_dict()
+opts['batch_size'] = 1
+opts['mlp_dim'] = [257, 1024, 512, 256, 128, 1]
+opts['mlp_dim_color'] = [513, 1024, 512, 256, 128, 3]
+opts['num_stack'] = 4
+opts['num_hourglass'] = 2
+opts['resolution'] = 128
+opts['hg_down'] = 'ave_pool'
+opts['norm'] = 'group'
+opts['norm_color'] = 'group'
+opts['load_netG_checkpoint_path'] = net_G
+opts['load_netC_checkpoint_path'] = net_C
+opts['results_path'] = "./results"
+opts['name'] = "spaces_demo"
+opts = SimpleNamespace(**opts)
+evaluator = Evaluator(opts)
+bg_remover_model = paddlehub.Module(name="U2Net")
 
 
 def process(img_path):
     base = os.path.basename(img_path)
     img_name = os.path.splitext(base)[0]
+    print("\n\n\nStarting Process", datetime.now())
     print("image name", img_name)
-    img_raw = Image.open(img_path)
+    img_raw = Image.open(img_path).convert('RGB')
+
     img = img_raw.resize(
         (800, int(800 * img_raw.size[1] / img_raw.size[0])),
         Image.Resampling.LANCZOS)
 
-    # remove background
-    print("Removeing background")
-    try: 
-        foreground = Image.fromarray(remove_bg.inference(img), 'RGBA')
-        foreground.save("./PIFu/inputs/" + img_name + ".png")
+    try:
+        # remove background
+        print("Removing Background")
+        masks = bg_remover_model.Segmentation(
+            images=[np.array(img)],
+            paths=None,
+            batch_size=1,
+            input_size=320,
+            output_dir='./PIFu/inputs',
+            visualization=False)
+        mask = masks[0]["mask"]
+        front = masks[0]["front"]
     except Exception as e:
         print(e)
+
     print("Aliging mask with input training image")
-    subprocess.Popen(["python", "./apps/crop_img.py", "--input_image",
-                      f'./inputs/{img_name}.png', "--out_path", "./inputs"], cwd="PIFu").communicate()
+    print("Not aligned", front.shape, mask.shape)
+    img_new, msk_new = process_img(front, mask)
+    print("Aligned", img_new.shape, msk_new.shape)
+
+    try:
+        time = datetime.now()
+        data = evaluator.load_image_from_memory(img_new, msk_new, img_name)
+        print("Evaluating via PIFu", time)
+        evaluator.eval(data, True)
+        print("Success Evaluating via PIFu", datetime.now() - time)
+        result_path = f'{opts.results_path}/{opts.name}/result_{img_name}'
+    except Exception as e:
+        print("Error evaluating via PIFu", e)
+
+    try:
+        mesh = trimesh.load(result_path + '.obj')\
+        # flip mesh
+        mesh.apply_transform([[1, 0, 0, 0],
+                              [0, 1, 0, 0],
+                              [0, 0, -1, 0],
+                              [0, 0, 0, 1]])
+        mesh.export(file_obj=result_path + '.glb')
+        result_gltf = result_path + '.glb'
+        return result_gltf
 
-    print("Generating 3D model")
-    subprocess.Popen("./scripts/test.sh", env={
-                     **env,
-                     "INPUT_IMAGE_PATH": f'./inputs/{img_name}.png',
-                     "VOL_RES": "128"},
-                     cwd="PIFu").communicate()
-    print("DONE 3D model")
-    return f'./PIFu/results/spaces_demo/result_{img_name}.glb'
+    except Exception as e:
+        print("error generating MESH", e)
 
 
 examples = [["./examples/" + img] for img in sorted(os.listdir("./examples/"))]
 description = '''
 # PIFu Clothed Human Digitization
-#### PIFu: Pixel-Aligned Implicit Function for High-Resolution Clothed Human Digitization
+# PIFu: Pixel-Aligned Implicit Function for High-Resolution Clothed Human Digitization
 <base target="_blank">
 
 This is a demo for <a href="https://github.com/shunsukesaito/PIFu" target="_blank"> PIFu model </a>.
@@ -76,17 +120,17 @@ The pre-trained model has the following warning:
 
 **The inference takes about 180seconds for a new image.**
 
-<details> 
+<details>
 <summary>More</summary>
 
-#### Image Credits
+# Image Credits
 
 * Julien and Clem
 * [StyleGAN Humans](https://huggingface.co/spaces/hysts/StyleGAN-Human)
 * [Renderpeople: Dennis](https://renderpeople.com)
 
 
-#### More
+# More
 * https://phorhum.github.io/
 * https://github.com/yuliangxiu/icon
 * https://shunsukesaito.github.io/PIFuHD/
@@ -102,6 +146,8 @@ iface = gr.Interface(
     examples=examples,
     allow_flagging="never",
     cache_examples=True
+
+
 )
 
 if __name__ == "__main__":

README.md

@@ -5,7 +5,7 @@ colorFrom: pink
 colorTo: green
 sdk: gradio
 sdk_version: 2.9.0b8
-app_file: app.py
+app_file: ./PIFu/spaces.py
 pinned: false
 python_version: 3.7.13
 ---

remove_bg.py

@@ -1,58 +0,0 @@
-# from https://huggingface.co/spaces/eugenesiow/remove-bg/blob/main/app.py
-import cv2
-import torch
-import numpy as np
-from torchvision import transforms
-
-class RemoveBackground(object):
-    def __init__(self):
-        self.model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet101', pretrained=True)
-        self.model.eval()
-
-    def make_transparent_foreground(self, pic, mask):
-        # split the image into channels
-        b, g, r = cv2.split(np.array(pic).astype('uint8'))
-        # add an alpha channel with and fill all with transparent pixels (max 255)
-        a = np.ones(mask.shape, dtype='uint8') * 255
-        # merge the alpha channel back
-        alpha_im = cv2.merge([b, g, r, a], 4)
-        # create a transparent background
-        bg = np.zeros(alpha_im.shape)
-        # setup the new mask
-        new_mask = np.stack([mask, mask, mask, mask], axis=2)
-        # copy only the foreground color pixels from the original image where mask is set
-        foreground = np.where(new_mask, alpha_im, bg).astype(np.uint8)
-
-        return foreground
-
-
-    def remove_background(self, input_image):
-        preprocess = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-        ])
-
-        input_tensor = preprocess(input_image)
-        input_batch = input_tensor.unsqueeze(0)  # create a mini-batch as expected by the model
-
-        # move the input and model to GPU for speed if available
-        if torch.cuda.is_available():
-          input_batch = input_batch.to('cuda')
-          self.model.to('cuda')
-
-        with torch.no_grad():
-            output = self.model(input_batch)['out'][0]
-        output_predictions = output.argmax(0)
-
-        # create a binary (black and white) mask of the profile foreground
-        mask = output_predictions.byte().cpu().numpy()
-        background = np.zeros(mask.shape)
-        bin_mask = np.where(mask, 255, background).astype(np.uint8)
-
-        foreground = self.make_transparent_foreground(input_image, bin_mask)
-
-        return foreground, bin_mask
-
-    def inference(self, img):
-        foreground, _ = self.remove_background(img)
-        return foreground

requirements.txt

@@ -19,4 +19,6 @@ six==1.14.0
 torch==1.4.0
 torchvision==0.5.0
 trimesh==3.5.23
-tqdm==4.64.0
+tqdm==4.64.0
+paddlehub
+paddlepaddle