コード例 #1
0
ファイル: can_smpl.py プロジェクト: bruinxiong/neuralbody
    def __getitem__(self, index):
        img_path = os.path.join(self.data_root, self.ims[index])
        img = imageio.imread(img_path).astype(np.float32) / 255.
        msk = self.get_mask(index, img)

        img = cv2.resize(img, (1000, 1002), interpolation=cv2.INTER_AREA)
        msk = cv2.resize(msk, (1000, 1002), interpolation=cv2.INTER_NEAREST)

        cam_ind = self.cam_inds[index]
        K = np.array(self.cams['K'][cam_ind])
        D = np.array(self.cams['D'][cam_ind])
        img = cv2.undistort(img, K, D)
        msk = cv2.undistort(msk, K, D)

        R = np.array(self.cams['R'][cam_ind])
        T = np.array(self.cams['T'][cam_ind]) / 1000.

        # reduce the image resolution by ratio
        H, W = int(img.shape[0] * cfg.ratio), int(img.shape[1] * cfg.ratio)
        img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
        msk = cv2.resize(msk, (W, H), interpolation=cv2.INTER_NEAREST)
        if cfg.mask_bkgd:
            img[msk == 0] = 0
            if cfg.white_bkgd:
                img[msk == 0] = 1
        K[:2] = K[:2] * cfg.ratio

        i = int(os.path.basename(img_path)[:-4])
        feature, coord, out_sh, can_bounds, bounds, Rh, Th = self.prepare_input(
            i)

        rgb, ray_o, ray_d, near, far, coord_, mask_at_box = if_nerf_dutils.sample_ray_h36m(
            img, msk, K, R, T, can_bounds, self.nrays, self.split)
        acc = if_nerf_dutils.get_acc(coord_, msk)

        ret = {
            'feature': feature,
            'coord': coord,
            'out_sh': out_sh,
            'rgb': rgb,
            'ray_o': ray_o,
            'ray_d': ray_d,
            'near': near,
            'far': far,
            'acc': acc,
            'mask_at_box': mask_at_box
        }

        R = cv2.Rodrigues(Rh)[0].astype(np.float32)
        i = index // self.num_cams
        meta = {'bounds': bounds, 'R': R, 'Th': Th, 'i': i, 'cam_ind': cam_ind}
        ret.update(meta)

        return ret
コード例 #2
0
    def __getitem__(self, index):
        img_path = os.path.join(self.data_root, self.ims[index])
        img = imageio.imread(img_path).astype(np.float32)[..., :3] / 255.
        msk = self.get_mask(index)

        cam_ind = self.cam_inds[index]
        K = np.array(self.cams['K'][cam_ind])

        R = np.array(self.cams['R'][cam_ind])
        T = np.array(self.cams['T'][cam_ind])

        # reduce the image resolution by ratio
        H, W = int(img.shape[0] * cfg.ratio), int(img.shape[1] * cfg.ratio)
        img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
        msk = cv2.resize(msk, (W, H), interpolation=cv2.INTER_NEAREST)
        if cfg.mask_bkgd:
            img[msk == 0] = 0
        K[:2] = K[:2] * cfg.ratio

        feature, coord, out_sh, can_bounds, bounds, Rh, Th, center, rot, trans = self.prepare_input(
            index)

        if cfg.sample_smpl:
            depth_path = os.path.join(self.data_root, 'depth',
                                      self.ims[index])[:-4] + '.npy'
            depth = np.load(depth_path)
            rgb, ray_o, ray_d, near, far, coord_, mask_at_box = if_nerf_dutils.sample_smpl_ray(
                img, msk, depth, K, R, T, self.nrays, self.split)
        elif cfg.sample_grid:
            # print('sample_grid')
            rgb, ray_o, ray_d, near, far, coord_, mask_at_box = if_nerf_dutils.sample_ray_grid(
                img, msk, K, R, T, can_bounds, self.nrays, self.split)
        else:
            rgb, ray_o, ray_d, near, far, coord_, mask_at_box = if_nerf_dutils.sample_ray_h36m(
                img, msk, K, R, T, can_bounds, self.nrays, self.split)
        acc = if_nerf_dutils.get_acc(coord_, msk)

        ret = {
            'feature': feature,
            'coord': coord,
            'out_sh': out_sh,
            'rgb': rgb,
            'ray_o': ray_o,
            'ray_d': ray_d,
            'near': near,
            'far': far,
            'acc': acc,
            'mask_at_box': mask_at_box,
            'index': index,
        }

        R = cv2.Rodrigues(Rh)[0].astype(np.float32)
        i = index  #// self.num_cams
        meta = {
            'bounds': bounds,
            'R': R,
            'Th': Th,
            'center': center,
            'rot': rot,
            'trans': trans,
            'i': i,
            'cam_ind': cam_ind
        }
        ret.update(meta)

        return ret
コード例 #3
0
    def __getitem__(self, index):
        # get one image from the ims array
        img_path = os.path.join(self.data_root, self.ims[index])
        # normalize the values to 0~1
        img = imageio.imread(img_path).astype(np.float32) / 255.0
        # resize the image to 1024, 1024
        img = cv2.resize(img, (1024, 1024))
        # get the image mask according to the index
        msk = self.get_mask(index)

        # get the camera index of this image
        cam_ind = self.cam_inds[index]
        # get the camera intrinsic parameter K (focal length and optical center)
        K = np.array(self.cams["K"][cam_ind])
        # get the camera intrinsic parameter D
        D = np.array(self.cams["D"][cam_ind])

        img = cv2.undistort(img, K, D)
        msk = cv2.undistort(msk, K, D)

        R = np.array(self.cams["R"][cam_ind])
        T = np.array(self.cams["T"][cam_ind]) / 1000.0

        # reduce the image resolution by ratio
        H, W = int(img.shape[0] * cfg.ratio), int(img.shape[1] * cfg.ratio)
        img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
        msk = cv2.resize(msk, (W, H), interpolation=cv2.INTER_NEAREST)
        if cfg.mask_bkgd:
            img[msk == 0] = 0
            if cfg.white_bkgd:
                img[msk == 0] = 1
        K[:2] = K[:2] * cfg.ratio

        if self.human in ["CoreView_313", "CoreView_315"]:
            i = int(os.path.basename(img_path).split("_")[4])
            frame_index = i - 1
        else:
            i = int(os.path.basename(img_path)[:-4])
            frame_index = i
        (
            feature,  # (6890, 6)
            coord,  #
            out_sh,
            can_bounds,
            bounds,
            Rh,
            Th,
            center,
            rot,
            trans,
        ) = self.prepare_input(i)

        (
            rgb,
            ray_o,
            ray_d,
            near,
            far,
            coord_,
            mask_at_box,
        ) = if_nerf_dutils.sample_ray_h36m(img, msk, K, R, T, can_bounds,
                                           self.nrays, self.split)
        acc = if_nerf_dutils.get_acc(coord_, msk)

        ret = {
            "feature": feature,
            "coord": coord,
            "out_sh": out_sh,
            "rgb": rgb,
            "ray_o": ray_o,
            "ray_d": ray_d,
            "near": near,
            "far": far,
            "acc": acc,
            "mask_at_box": mask_at_box,
        }

        R = cv2.Rodrigues(Rh)[0].astype(np.float32)
        i = index // self.num_cams
        if cfg.test_novel_pose:
            i = 0
        meta = {
            "bounds": bounds,
            "R": R,
            "Th": Th,
            "center": center,
            "rot": rot,
            "trans": trans,
            "i": i,
            "frame_index": frame_index,
            "cam_ind": cam_ind,
        }
        ret.update(meta)

        return ret