Convert Project to a Pip Package

.gitignore

@@ -0,0 +1,11 @@
+.idea/
+__pycache__/
+checkpoints/
+sample_images/
+results/
+build/
+pifuhd.egg-info/
+dist/
+pifu_hd.egg-info/
+task_build
+task_upload

lib/render/__init__.py → PIFuHD/__init__.py

@@ -1 +1,3 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+__appname__ = "PIFuHD"
+__version__ = "1.0.1"

PIFuHD/__main__.py

@@ -0,0 +1,51 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+import argparse
+
+from PIFuHD.data import EvalWRectDataset, EvalWPoseDataset
+from PIFuHD.options import BaseOptions
+from PIFuHD.recontructor import Reconstructor
+
+###############################################################################################
+#                   Setting
+###############################################################################################
+parser = argparse.ArgumentParser()
+parser.add_argument('-i', '--input_path', type=str, default='./sample_images')
+parser.add_argument('-o', '--out_path', type=str, default='./results')
+parser.add_argument('-c', '--ckpt_path', type=str, default='./checkpoints/PIFuHD.pt')
+parser.add_argument('-r', '--resolution', type=int, default=512)
+parser.add_argument('--use_rect', action='store_true', help='use rectangle for cropping')
+args = parser.parse_args()
+
+
+###############################################################################################
+#                   Upper PIFu
+###############################################################################################
+
+def recon(opts: BaseOptions, use_rect=False):
+    if use_rect:
+        dataset = EvalWRectDataset(opts)
+    else:
+        dataset = EvalWPoseDataset(opts)
+
+    reconstructor = Reconstructor(opts)
+    reconstructor.evaluate(dataset)
+
+
+def main():
+    cmd = ['--dataroot', args.input_path,
+           '--results_path', args.out_path,
+           '--loadSize', '1024',
+           '--resolution', str(args.resolution),
+           '--load_netMR_checkpoint_path', args.ckpt_path,
+           '--start_id', '-1',
+           '--end_id', '-1'
+           ]
+
+    options_parser = BaseOptions()
+    opt = options_parser.parse(cmd)
+    recon(opt, args.use_rect)
+
+
+if __name__ == '__main__':
+    main()

PIFuHD/data/EvalDataset.py

@@ -0,0 +1,107 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+import os
+from abc import abstractmethod
+
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+class EvalDataset(Dataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def __init__(self, opt, items, projection='orthogonal'):
+        self.opt = opt
+        self.projection_mode = projection
+
+        self.root = self.opt.dataroot
+        self.items = items
+        self.IMG = os.path.join(self.root)
+
+        self.phase = 'val'
+        self.load_size = self.opt.loadSize
+
+        # PIL to tensor
+        self.to_tensor = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+
+        # only used in case of multi-person processing
+        self.person_id = 0
+
+    def __len__(self):
+        return len(self.items)
+
+    def get_item(self, index):
+        im = self.items[index].img
+        img_name = self.items[index].name
+
+        if im.shape[2] == 4:
+            im = im / 255.0
+            im[:, :, :3] /= im[:, :, 3:] + 1e-8
+            im = im[:, :, 3:] * im[:, :, :3] + 0.5 * (1.0 - im[:, :, 3:])
+            im = (255.0 * im).astype(np.uint8)
+        h, w = im.shape[:2]
+
+        intrinsic = np.identity(4)
+        trans_mat = np.identity(4)
+
+        rect = self.get_human_box(index)
+        im = self.crop_human_box(im, rect)
+
+        scale_im2ndc = 1.0 / float(w // 2)
+        scale = w / rect[2]
+        trans_mat *= scale
+        trans_mat[3, 3] = 1.0
+        trans_mat[0, 3] = -scale * (rect[0] + rect[2] // 2 - w // 2) * scale_im2ndc
+        trans_mat[1, 3] = scale * (rect[1] + rect[3] // 2 - h // 2) * scale_im2ndc
+
+        intrinsic = np.matmul(trans_mat, intrinsic)
+        im_512 = cv2.resize(im, (512, 512))
+        im = cv2.resize(im, (self.load_size, self.load_size))
+
+        image_512 = Image.fromarray(im_512[:, :, ::-1]).convert('RGB')
+        image = Image.fromarray(im[:, :, ::-1]).convert('RGB')
+
+        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()
+
+        calib_world = torch.Tensor(intrinsic).float()
+
+        # image
+        image_512 = self.to_tensor(image_512)
+        image = self.to_tensor(image)
+        return {
+            'name': img_name,
+            'img': image.unsqueeze(0),
+            'img_512': image_512.unsqueeze(0),
+            'calib': calib.unsqueeze(0),
+            'calib_world': calib_world.unsqueeze(0),
+            'b_min': B_MIN,
+            'b_max': B_MAX,
+        }
+
+    def __getitem__(self, index):
+        return self.get_item(index)
+
+    @abstractmethod
+    def get_n_person(self, index):
+        pass
+
+    @abstractmethod
+    def get_human_box(self, index):
+        pass
+
+    @abstractmethod
+    def crop_human_box(self, image, rect):
+        pass

PIFuHD/data/EvalWMetaDataset.py

@@ -0,0 +1,18 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+from PIFuHD.data.EvalDataset import EvalDataset
+from PIFuHD.data.helper_image_crop import crop_image
+
+
+class EvalWMetaDataset(EvalDataset):
+
+    def __init__(self, opt, items):
+        super().__init__(opt, items)
+
+    def get_n_person(self, index):
+        return len(self.items[index].meta)
+
+    def get_human_box(self, index):
+        return self.items[index].meta[0]
+
+    def crop_human_box(self, image, rect):
+        return crop_image(image, rect)

PIFuHD/data/EvalWPoseDataset.py

@@ -0,0 +1,82 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+import json
+import numpy as np
+from PIFuHD.data.EvalDataset import EvalDataset
+from .helper_dataset import make_bundles
+from .helper_image_crop import crop_image, face_crop, upperbody_crop, fullbody_crop
+
+crop_callbacks = {'face': face_crop, 'upperbody': upperbody_crop, 'fullbody': fullbody_crop}
+
+
+class EvalWPoseDataset(EvalDataset):
+
+    def __init__(self, opt):
+        items = make_bundles(opt.dataroot, '_keypoints.json')
+        super().__init__(opt, items)
+        if self.opt.crop_type == 'face':
+            self.crop_func = face_crop
+        elif self.opt.crop_type == 'upperbody':
+            self.crop_func = upperbody_crop
+        else:
+            self.crop_func = fullbody_crop
+
+    def get_n_person(self, index):
+        joint_path = self.items[index].meta
+        # Calib
+        with open(joint_path) as json_file:
+            data = json.load(json_file)
+            return len(data['people'])
+
+    def get_human_box(self, index):
+        joint_path = self.items[index].meta
+        with open(joint_path) as json_file:
+            data = json.load(json_file)
+            if len(data['people']) == 0:
+                raise IOError('non human found!!')
+
+            # if True, the person with the largest height will be chosen.
+            # set to False for multi-person processing
+            if True:
+                selected_data = data['people'][0]
+                height = 0
+                if len(data['people']) != 1:
+                    for i in range(len(data['people'])):
+                        tmp = data['people'][i]
+                        keypoints = np.array(tmp['pose_keypoints_2d']).reshape(-1, 3)
+
+                        flags = keypoints[:, 2] > 0.5  # openpose
+                        # flags = keypoints[:,2] > 0.2  #detectron
+                        if sum(flags) == 0:
+                            continue
+                        bbox = keypoints[flags]
+                        bbox_max = bbox.max(0)
+                        bbox_min = bbox.min(0)
+
+                        if height < bbox_max[1] - bbox_min[1]:
+                            height = bbox_max[1] - bbox_min[1]
+                            selected_data = tmp
+            else:
+                pid = min(len(data['people']) - 1, self.person_id)
+                selected_data = data['people'][pid]
+
+            keypoints = np.array(selected_data['pose_keypoints_2d']).reshape(-1, 3)
+
+            flags = keypoints[:, 2] > 0.5  # openpose
+            # flags = keypoints[:,2] > 0.2    #detectron
+
+            nflag = flags[0]
+            mflag = flags[1]
+
+            check_id = [2, 5, 15, 16, 17, 18]
+            cnt = sum(flags[check_id])
+            if self.opt.crop_type == 'face' and (not (nflag and cnt > 3)):
+                print('Waring: face should not be backfacing.')
+            if self.opt.crop_type == 'upperbody' and (not (mflag and nflag and cnt > 3)):
+                print('Waring: upperbody should not be backfacing.')
+            if self.opt.crop_type == 'fullbody' and sum(flags) < 15:
+                print('Waring: not sufficient keypoints.')
+
+        return self.crop_func(keypoints)
+
+    def crop_human_box(self, image, rect):
+        return crop_image(image, rect)

PIFuHD/data/EvalWRectDataset.py

@@ -0,0 +1,32 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+import numpy as np
+
+from PIFuHD.data.EvalDataset import EvalDataset
+from .helper_dataset import make_bundles
+from .helper_image_crop import crop_image
+
+
+class EvalWRectDataset(EvalDataset):
+
+    def __init__(self, opt):
+        items = make_bundles(opt.dataroot, '_rect.txt')
+        super().__init__(opt, items)
+
+    def get_rect(self, index):
+        rect_path = self.items[index].meta
+        return np.loadtxt(rect_path, dtype=np.int32)
+
+    def get_n_person(self, index):
+        rects = self.get_rect(index)
+        return rects.shape[0] if len(rects.shape) == 2 else 1
+
+    def get_human_box(self, index):
+        return self.get_rect(index)
+
+    def crop_human_box(self, image, rects):
+        if len(rects.shape) == 1:
+            rects = rects[None]
+        pid = min(rects.shape[0] - 1, self.person_id)
+
+        rect = rects[pid].tolist()
+        return crop_image(image, rect)

PIFuHD/data/ImageBundle.py

@@ -0,0 +1,11 @@
+class ImageBundle:
+    def __init__(self, img, name, meta):
+        self.img = img
+        self.name = name
+        self.meta = meta
+        self.img_512 = None
+        self.calib = None
+        self.calib_world = None
+        self.b_min = None
+        self.b_max = None
+        self.is_multi = False

lib/data/__init__.py → PIFuHD/data/__init__.py

@@ -1,4 +1,5 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
 from .EvalWPoseDataset import EvalWPoseDataset
-from .EvalDataset import EvalDataset
+from .EvalWRectDataset import EvalWRectDataset
+from .EvalWMetaDataset import EvalWMetaDataset

PIFuHD/data/helper_dataset.py

@@ -0,0 +1,40 @@
+import os
+import sys
+
+import cv2
+from PIFuHD.data.ImageBundle import ImageBundle
+
+
+def load_image(path):
+    return cv2.imread(path, cv2.IMREAD_UNCHANGED)
+
+
+def pair_ext(f, extension):
+    return f.replace('.%s' % (f.split('.')[-1]), extension)
+
+
+def is_image(file):
+    if file.split('.')[-1] in ['png', 'jpeg', 'jpg', 'PNG', 'JPG', 'JPEG']:
+        return True
+    return False
+
+
+def scan_image_folder(path, extension):
+    try:
+        return sorted([os.path.join(path, f) for f in os.listdir(path) if
+                       is_image(f) and os.path.exists(os.path.join(path, pair_ext(f, extension)))])
+    except FileNotFoundError:
+        print(f'error cannot found dataset path:{path}')
+        sys.exit(1)
+
+
+def make_bundles(path, ext):
+    path = os.path.abspath(path)
+    return [
+        ImageBundle(
+            img=load_image(file),
+            name=os.path.splitext(os.path.basename(file))[0],
+            meta=os.path.join(path, pair_ext(file, ext)),
+        )
+        for file in scan_image_folder(path, ext)
+    ]