import os from os.path import join import json import cv2 import copy from ...mytools.reader import read_smpl, read_json from ...mytools.camera_utils import Undistort, read_cameras from ...mytools.debug_utils import myerror, mywarn, oncewarn from ...mytools.timer import Timer from .utils_sample import get_bounds, get_rays, sample_rays_rate, sample_rays import numpy as np from .utils_reader import img_to_numpy, numpy_to_img, read_json_with_cache, parse_semantic class BaseBase: def __init__(self, split): self.split = split self.infos = [] self.file_cache = {} self.back_mask_cache = {} self.cache_ray_o_d = {} self.timer = False def __len__(self): return len(self.infos) def __getitem__(self, index): raise NotImplementedError def scale_and_undistort(self, img, info, undis=True): img = img_to_numpy(img) if self.image_args.scale != 1: H, W = int(img.shape[0] * self.image_args['scale']), int(img.shape[1] * self.image_args['scale']) img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST) K, D = info['camera']['K'], info['camera']['dist'] sub = info['sub'] if self.image_args.undis and np.linalg.norm(D) > 0. and undis: img = Undistort.image(img, K, D, sub=sub) return img def parse_object_args(self, object_keys, object_args, ignore_keys): for _class in ['human', 'ball']: if 'all'+_class in object_args.keys(): object_args = object_args.copy() _allargs = object_args.pop('all'+_class) pids = _allargs.pop('pids') for pid in pids: _args = copy.deepcopy(_allargs) _args['args'].pid = pid object_args[_class+'_{}'.format(pid)] = _args self.ignore_keys = ignore_keys self.object_args = object_args self.object_keys = object_keys def object_factory(self, root, obj_type, obj_args, info, params, file_cache): # 通用的一些参数：例如SMPL pid = obj_args.get('pid', -1) # pid only available in human or ball sub = info['sub'] if 'frame' in params: frame = params['frame'] else: frame = info['frame'] feat = {} reader = obj_args.get('reader', {}) # reader only available in human use_param_foreachview = obj_args.get('use_param_foreachview', False) for key, args in reader.items(): if key == 'smpl': if use_param_foreachview and sub.startswith('novel'): sub_ = sorted(os.listdir(join(root, args.root)))[0] smplname = join(root, args.root, sub_, '{:06d}.json'.format(frame)) elif use_param_foreachview and not sub.startswith('novel'): smplname = join(root, args.root, sub, '{:06d}.json'.format(frame)) else: smplname = join(root, args.root, '{:06d}.json'.format(frame)) data = read_json_with_cache(smplname, file_cache) if isinstance(data, dict): data = data['annots'] data = [d for d in data if d['id'] == pid] assert len(data) > 0, "{} don't have {}".format(smplname, pid) data = data[0] for key in ['poses', 'shapes', 'Rh', 'Th']: feat[key] = np.array(data[key], dtype=np.float32) feat['R'] = cv2.Rodrigues(feat['Rh'])[0] elif key == 'vertices': if use_param_foreachview and sub.startswith('novel'): sub_ = sorted(os.listdir(join(root, args.root)))[0] smplname = join(root, args.root, sub_, '{:06d}.json'.format(frame)) elif use_param_foreachview and not sub.startswith('novel'): smplname = join(root, args.root, sub, '{:06d}.json'.format(frame)) else: smplname = join(root, args.root, '{:06d}.json'.format(frame)) data = read_json_with_cache(smplname, file_cache) if isinstance(data, dict): data = data['annots'] data = [d for d in data if d['id'] == pid][0] if 'keypoints3d' in data.keys() and 'vertices' not in data.keys(): # use keypoints3d instead k3d = np.array(data['keypoints3d'], dtype=np.float32) conf = k3d[:, 3] k3d = k3d[:, :3] data['vertices'] = k3d[conf>0.1] for key in ['vertices']: feat[key] = np.array(data[key], dtype=np.float32) feat['bounds'] = get_bounds(feat['vertices'], delta=args.padding) elif key in ['depth']: depthname = join(root, args.root, sub, '{:06d}.png'.format(frame)) if not os.path.exists(depthname): depth = None else: depth = cv2.imread(depthname, cv2.IMREAD_ANYDEPTH) depth = depth.astype(np.float32)/1000 feat['depth'] = depth elif key in ['mask', 'instance', 'label', 'semantic']: if args.root == 'none': continue mskname = join(root, args.root, sub, '{:06d}_{}.png'.format(frame, pid)) if not os.path.exists(mskname): mskname0 = join(root, args.root, sub, '{:06d}.png'.format(frame)) if pid == 0 and os.path.exists(mskname0): mskname = mskname0 oncewarn('Using null pid to read mask') else: if not 'novel_' in mskname: print('!!!{} not exists'.format(mskname)) feat[key] = None continue assert os.path.exists(mskname), mskname if key == 'semantic': msk = cv2.imread(mskname) else: msk = cv2.imread(mskname, 0) msk = self.scale_and_undistort(msk, info, undis=args.undis) if key == 'mask': feat[key] = msk>0 elif key == 'label': feat[key] = msk elif key == 'semantic': feat[key] = parse_semantic(msk) feat['mask'] = feat[key] > 0 else: raise NotImplementedError if obj_type in ['nearfar', 'nearfardepth']: from .utils_sample import NearFarSampler obj = NearFarSampler(split=self.split, near=obj_args.near, far=obj_args.far, depth=feat.get('depth', None)) elif obj_type == 'bbox': from .utils_sample import AABBSampler obj = AABBSampler(split=self.split, bounds=obj_args.bounds) elif obj_type == 'twobbox': from .utils_sample import TwoAABBSampler obj = TwoAABBSampler(split=self.split, bbox_inter=obj_args.bbox_inter, bbox_outer=obj_args.bbox_outer) elif obj_type == 'compose': objlist = [] for key, obj_args_ in obj_args.items(): obj_ = self.object_factory(root, obj_args_.model, obj_args_, info, params, file_cache) objlist.append(obj_) from .utils_sample import ComposeSampler obj = ComposeSampler(split=self.split, objlist=objlist) elif obj_type == 'bodybbox': from .utils_sample import AABBwMask, AABBSampler if 'vertices' in feat.keys(): vertices = feat['vertices'] if 'label' in feat.keys() and feat['label'] is not None: obj = AABBwMask.from_vertices( label=feat['label'].astype(np.float32), mask=feat['label'], # use label to represents mask rate_body=obj_args.rate_body, dilate=True, split=self.split, vertices=vertices, delta=obj_args.reader.vertices.padding) else: obj = AABBSampler.from_vertices( split=self.split, vertices=vertices, delta=obj_args.reader.vertices.padding) else: center = feat['Th'] obj = AABBSampler(self.split, center=center, scale=obj_args.scale) elif obj_type == 'maskbbox': from .utils_sample import AABBwMask center = feat['Th'] obj = AABBwMask(split=self.split, center=center, scale=obj_args.scale, mask=feat['mask']) elif obj_type == 'neuralbody' or obj_type == 'neuralbody-smplmask': from .utils_sample import AABBwMask # ATTN dilate = obj_type == 'neuralbody' if 'rotate_axis' in params.keys(): axis = { 'y': np.array([0., 1., 0.], dtype=np.float32) }[params['rotate_axis']] Rrel = cv2.Rodrigues(params['rotate_angle']*axis)[0] feat['vertices'] = (feat['vertices'] - feat['Th']) @ Rrel.T + feat['Th'] feat['R'] = Rrel @ feat['R'] feat['Rh'] = cv2.Rodrigues(feat['R'])[0].reshape(1, 3) feat['bounds'] = get_bounds(feat['vertices'], delta=obj_args.reader.vertices.padding) if 'rotation' in params.keys() or 'translation' in params.keys(): # support rotation and translation rotation = params.get('rotation', [0., 0., 0.]) translation = params.get('translation', [0., 0., 0.]) Rrel = cv2.Rodrigues(np.array(rotation))[0].astype(np.float32) Trel = np.array(translation).reshape(1, 3).astype(np.float32) # update the original Rh, Th feat['R'] = Rrel @ feat['R'] feat['Rh'] = cv2.Rodrigues(feat['R'])[0].reshape(1, 3) feat['Th'] = (Rrel @ feat['Th'].T + Trel.T).T feat['vertices'] = feat['vertices'] @ Rrel.T + Trel feat['bounds'] = get_bounds(feat['vertices'], delta=obj_args.reader.vertices.padding) # calculate canonical vertices and bounds feat['vertices_canonical'] = (feat['vertices'] - feat['Th']) @ feat['R'].T.T feat['bounds_canonical'] = get_bounds(feat['vertices_canonical'], delta=obj_args.reader.vertices.padding) obj = AABBwMask(split=self.split, bounds=feat['bounds'], mask=feat.get('mask', None), label=feat.get('label', None), dilate=dilate, rate_body=obj_args.rate_body) for key in ['R', 'Rh', 'Th', 'vertices', 'poses', 'shapes', 'vertices_canonical', 'bounds_canonical']: obj.feature[key] = feat[key] # extra keys for key in ['semantic']: if key in feat.keys(): obj.feature_input[key] = feat[key] elif obj_type == 'neuralbody-womask': from .utils_sample import AABBSampler obj = AABBSampler(split=self.split, bounds=feat['bounds']) for key in ['R', 'Rh', 'Th', 'vertices', 'poses', 'shapes']: obj.feature[key] = feat[key] elif obj_type == 'trajbbox': from .utils_sample import AABBSampler annname = join(self.root, obj_args.root, '{:06d}.json'.format(frame)) annots = read_json_with_cache(annname, file_cache) annots = [a for a in annots if a['id'] == pid] assert len(annots) == 1, annname annots = annots[0] center = np.array(annots['keypoints3d'], dtype=np.float32)[:, :3] obj = AABBSampler(self.split, center=center, scale=obj_args.scale) obj.feature['center'] = center elif obj_type == 'cylinder': from .utils_sample import CylinderSampler obj = CylinderSampler(split=self.split, **obj_args) elif obj_type == 'plane': from .utils_sample import PlaneSampler obj = PlaneSampler(split=self.split, **obj_args) else: myerror('[Error] Unknown object type: {}'.format(obj_type)) raise NotImplementedError for key, val in params.items(): obj.feature[key] = val return obj def get_objects(self, root, info, object_keys, object_args): objects = {} current_keys = object_keys.copy() file_cache = {} if len(current_keys) == 0: current_keys = list(object_args.keys()) for oid, wrapkey in enumerate(current_keys): key, params = wrapkey, {} if key in self.ignore_keys: continue if '@' in wrapkey: key_ = wrapkey.split('_@')[0] params = json.loads(wrapkey.split('_@')[1].replace("'", '"')) val = object_args[key_] else: val = object_args[key] with Timer(wrapkey, not self.timer): model = self.object_factory(root, val['model'], val['args'], info, params, file_cache) objects[key] = model return objects def sample_ray(self, img, back_mask, info, objects, debug=False): H, W = img.shape[:2] K, R, T = info['camera']['K'].copy(), info['camera']['R'], info['camera']['T'] sub = info['sub'] if sub not in self.cache_ray_o_d: # cache this ray as it takes 0.2s ray_o, ray_d = get_rays(H, W, K, R, T) self.cache_ray_o_d[sub] = (ray_o, ray_d) ray_o, ray_d = self.cache_ray_o_d[sub] # sample rays according to the rate bounds, rates = {}, {} for key, obj in objects.items(): with Timer('mask '+key, not self.timer): ret = obj.mask(K, R, T, H, W, ray_o=ray_o, ray_d=ray_d) if isinstance(ret, np.ndarray): bounds[key] = ret if self.split == 'train': rates[key] = self.object_args[key]['rate'] elif isinstance(ret, dict): for key_, val_ in ret.items(): bounds[key+key_] = val_['mask'] if self.split == 'train': rates[key+key_] = self.object_args[key]['rate'] * val_['rate'] else: raise NotImplementedError if self.split == 'train': # if self.sample_args['method'] == 'rate' and self.split == 'train': with Timer('sample ray', not self.timer): coord = sample_rays_rate(bounds, rates, back_mask, **self.sample_args) else: bounds = np.dstack(list(bounds.values())) bound_sum = np.sum(bounds, axis=-1) coord = sample_rays(bound_sum, back_mask, self.split, **self.sample_args) if debug: import matplotlib.pyplot as plt plt.imshow(img) plt.scatter(coord[:, 1], coord[:, 0]) plt.figure() bounds = np.dstack(bounds.values()).copy() for i in range(bounds.shape[-1]): bounds[..., i] *= 2**i bound_sum = np.sum(bounds, axis=-1) plt.imshow(bound_sum) plt.show() import ipdb;ipdb.set_trace() ray_o = ray_o[coord[:, 0], coord[:, 1]] ray_d = ray_d[coord[:, 0], coord[:, 1]] rgb = img[coord[:, 0], coord[:, 1]] return ray_o, ray_d, rgb, coord def sample_near_far(self, rgb, ray_o, ray_d, coord, objects): ret = { 'rgb': rgb, 'coord': coord, 'ray_o': ray_o, 'ray_d': ray_d, 'viewdirs': ray_d/np.linalg.norm(ray_d, axis=-1, keepdims=True), } # TODO:这里是每个物体分别进行采样，没有统一调度 # sample the ray in the fore for key, model in objects.items(): with Timer('sample '+key, not self.timer): near, far, mask = model(ray_o, ray_d, coord) # 这里把单个物体的相关信息直接返回了，用于兼容单人的情况，单人就不用考虑如何取mask了 for k in ['rgb', 'coord', 'ray_o', 'ray_d', 'viewdirs']: ret[key+'_'+k] = ret[k][mask] ret.update({ key+'_near': near, key+'_far': far, key+'_mask': mask, key+'_bounds': model.bounds}) # update other features for k, v in model.feature.items(): if 'coord' in k: ret[key+'_'+k] = v[coord[:, 0], coord[:, 1]] else: ret[key+'_'+k] = v return ret def create_cameras(self, camera_args): from .utils_sample import create_center_radius RT = create_center_radius( center=camera_args.center, radius=camera_args.radius, up=camera_args.up, angle_x=camera_args.angle_x, ranges=camera_args.ranges) focal = camera_args.focal K = np.array([ focal, 0, camera_args.H/2, 0, focal, camera_args.W/2, 0, 0, 1 ]).reshape(3, 3) cameras = [{'K': K, 'R': RT[i, :3, :3], 'T': RT[i, :3, 3:]} for i in range(RT.shape[0])] return cameras class BaseDataset(BaseBase): # This class is mainly for multiview or single view training dataset # contains the operation for image def __init__(self, root, subs, ranges, split, image_args, object_keys, object_args, ignore_keys, sample_args) -> None: super().__init__(split=split) self.root = root self.cameras = read_cameras(root) self.subs = self.check_subs(root, subs, image_args.root) self.ranges = ranges self.infos = self.get_allnames(root, self.subs, ranges, image_args) self.split = split self.image_args = image_args self.sample_args = sample_args self.parse_object_args(object_keys, object_args, ignore_keys) self.debug = False self.check = False self.check_data() def check_data(self): from tqdm import tqdm visited = set() for i in tqdm(range(len(self)), desc='check all the data'): info = self.infos[i] sub = info['sub'] if sub in visited:continue visited.add(sub) data = self[i] def get_allnames(self, root, subs, ranges, image_args): infos = [] index = 0 unsync = {} if image_args.get('unsync', 'none') != 'none': unsync = read_json(join(root, image_args['unsync'])) for nv, sub in enumerate(subs): camera = self.cameras[sub].copy() K = camera['K'].copy() K[:2] *= image_args.scale camera['K'] = K for nnf, nf in enumerate(range(*ranges)): imgname = join(root, image_args.root, sub, '{:06d}{}'.format(nf, image_args.ext)) info = { 'imgname': imgname, 'sub': sub, 'frame': nf, 'nf': nnf, 'nv': nv, 'index': index, 'camera': camera } if sub in unsync.keys(): info['time'] = nf + unsync[sub] infos.append(info) return infos @staticmethod def check_subs(root, subs, image_root): if len(subs) == 0: subs = sorted(os.listdir(join(root, image_root))) if subs[0].isdigit(): subs.sort(key=lambda x:int(x)) return subs def read_image(self, imgname, image_args, info, isgray=False, skip_mask=False, mask_global='_0.png'): if isgray: img = cv2.imread(imgname, 0) else: img = cv2.imread(imgname) img = img_to_numpy(img) if image_args.mask_bkgd and not skip_mask: # TODO: polish mask name mskname = imgname.replace(image_args.root, image_args.mask).replace(image_args.ext, mask_global) if not os.path.exists(mskname): mskname = mskname.replace(mask_global, '.png') assert os.path.exists(mskname), mskname msk = cv2.imread(mskname, 0) msk_rand = np.random.rand(msk.shape[0], msk.shape[1], 3) msk_rand = (msk_rand * 255).astype(np.uint8) img[msk==0] = -1. # img[msk==0] = msk_rand[msk==0] # 这里选择了先进行畸变矫正，再进行scale，这样相机参数就不需要修改 # 由于预先进行了map，所以畸变矫正不会很慢 if self.image_args.scale != 1: H, W = int(img.shape[0] * self.image_args['scale']), int(img.shape[1] * self.image_args['scale']) img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST) K, D = info['camera']['K'], info['camera']['dist'] sub = info['sub'] if self.image_args.undis and np.linalg.norm(D) > 0.: img = Undistort.image(img, K, D, sub=sub) if info.get('flip_lr', False): img = cv2.flip(img, 1) return img def read_bkgd(self, imgname, image_args, info): backname = join(self.root, 'background', info['sub'], '000000.jpg') if not os.path.exists(backname) or not image_args.get('read_background'): return None back = self.read_image(backname, image_args, info, skip_mask=True) return back def read_backmask(self, imgname, image_args, info, blankshape): sub = info['sub'] if self.split != 'demo' and sub in self.back_mask_cache.keys(): back_mask = self.back_mask_cache[sub].copy() elif self.split != 'demo': mskname = join(self.root, 'mask-background', sub, '000000.png') # if not exists, should read and undistort back_mask = np.ones(blankshape, dtype=np.float32) K, D = info['camera']['K'], info['camera']['dist'] if self.image_args.undis and np.linalg.norm(D) > 0.: back_mask = Undistort.image(back_mask, K, D, sub=sub) if os.path.exists(mskname) and self.split == 'train' and self.image_args.ignore_back: oncewarn('using mask of background') _back_mask = self.read_image(mskname, image_args, info, isgray=True) # 畸变矫正后边缘上填充的值是0，这里取反之后边缘上的填充值又变成了1 _back_mask[_back_mask>0] = 1. back_mask = (1. - _back_mask) * back_mask self.back_mask_cache[sub] = back_mask back_mask = self.back_mask_cache[sub].copy() else: # in demo mode sample all pixels back_mask = np.ones(blankshape, dtype=np.float32) # get the framewise mask mskname = imgname.replace('images', 'mask-ignore').replace('.jpg', '.png') if os.path.exists(mskname): # oncewarn('using mask of frame {}'.format(mskname)) mask_frame = self.read_image(mskname, image_args, info, isgray=True) mask_frame[mask_frame>0] = 1. mask_frame = 1. - mask_frame back_mask = back_mask * mask_frame return back_mask def __getitem__(self, index): info = self.infos[index] imgname = info['imgname'] with Timer('read image', not self.timer): img = self.read_image(imgname, self.image_args, info) with Timer('read background', not self.timer): back = self.read_bkgd(imgname, self.image_args, info) with Timer('read back', not self.timer): back_mask = self.read_backmask(imgname, self.image_args, info, blankshape=img.shape[:2]) object_keys = info.get('object_keys', self.object_keys) objects = self.get_objects(self.root, info, object_keys, self.object_args) # sample the ray from image ray_o, ray_d, rgb, coord = self.sample_ray(img, back_mask, info, objects, debug=self.debug) ret = self.sample_near_far(rgb, ray_o, ray_d, coord, objects) # append the background if back is not None: ret['background'] = back[coord[:, 0], coord[:, 1]] meta = { 'split': self.split, 'H': img.shape[0], 'W': img.shape[1], 'index': index, 'nframe': info['nf'], 'nview': info['nv'], 'time': info.get('time', info['nf']), 'sub': info['sub'], 'keys': list(objects.keys()), } if self.check: meta['img'] = img meta['object_keys'] = meta['keys'] ret['meta'] = meta return ret class BaseDatasetDemo(BaseDataset): def __init__(self, camera_args, demo_args=None, **kwargs): self.camera_args = camera_args super().__init__(**kwargs) scale = self.image_args.scale self.blank = np.zeros((int(camera_args.H*scale), int(camera_args.W*scale), 3), dtype=np.uint8) self.infos = self.create_demo_cameras(self.image_args.scale, camera_args, demo_args) # self.camera_args = camera_args # K, RTs, D = self.create_demo_camera(camera_method, camera_args, self.cameras) @staticmethod def _demo_play_frames(ranges, K, R, T): infos = [] index = 0 frames = [i for i in range(*ranges)] for nv in range(K.shape[0]): nnf = nv % (len(frames)) nf = frames[nnf] info = { 'imgname': 'none', 'sub': 'novel_'+str(nv), 'frame': nf, 'nf': nnf, 'nv': nv, 'index': index, 'camera': { 'K': K[nv], 'dist': np.zeros((1, 5)), 'R': R[nv], 'T': T[nv] } } infos.append(info) index += 1 return infos @staticmethod def _demo_keyframe(ranges, K, R, T, frame, nFrames): infos = [] index = 0 frames = [i for i in range(*ranges)] for nnf, nf in enumerate(frames): info = { 'imgname': 'none', 'sub': 'novel_'+str(0), 'frame': nf, 'nf': nnf, 'nv': 0, 'index': index, 'camera': { 'K': K[0], 'dist': np.zeros((1, 5)), 'R': R[0], 'T': T[0] } } if nf == frame: for i in range(nFrames): angle = i/nFrames * 2 * np.pi object_keys = ["human_0_@{{'rotate_angle': {}, 'rotate_axis': 'y'}}".format(angle)] _info = info.copy() _info['object_keys'] = object_keys infos.append(_info) info['index'] += 1 index = info['index'] + 1 continue infos.append(info) index += 1 return infos def _demo_script(self, ranges, K, R, T, stages): infos = [] index = 0 frames = [i for i in range(*ranges)] for name, stage in stages.items(): _infos = [] _frames = list(range(*stage.frame)) _views = list(range(*stage.view)) if len(_frames) == 1 and len(_views) != 1: _frames = _frames * len(_views) elif len(_views) == 1 and len(_frames) != 1: _views = _views * len(_frames) elif len(_views) == 1 and len(_frames) == 1 and 'steps' in stage.keys(): _views = _views * stage.steps _frames = _frames * stage.steps elif len(_views) != 1 and len(_frames) != 1 and len(_views) != len(_frames): raise NotImplementedError _index = [i for i in range(len(_frames))] for _i in _index: nv, nf = _views[_i], _frames[_i] nv = nv % (K.shape[0]) info = { 'imgname': 'none', 'sub': 'novel_'+str(nv), 'frame': nf, 'nf': frames.index(nf), 'nv': nv, 'index': _i + index, 'camera': { 'K': K[nv], 'dist': np.zeros((1, 5)), 'R': R[nv], 'T': T[nv] } } # create object float_i = _i*1./(len(_index) - 1) object_keys = stage.object_keys.copy() if len(object_keys) == 0: object_keys = list(self.object_args.keys()).copy() if 'effect' in stage.keys(): if stage.effect in ['disappear', 'appear']: for _obj in stage.effect_args.key: object_keys.remove(_obj) if stage.effect == 'disappear': occ = (1 - float_i)**3 elif stage.effect == 'appear': occ = float_i**3 object_keys.append(_obj+"_@{{'scale_occ': {}, 'min_acc': 0.5}}".format(occ)) if stage.effect in ['zoom']: scale = float_i * stage.effect_args.scale[1] + (1-float_i) * stage.effect_args.scale[0] cx = float_i * stage.effect_args.cx[1] + (1-float_i) * stage.effect_args.cx[0] cy = float_i * stage.effect_args.cy[1] + (1-float_i) * stage.effect_args.cy[0] _K = info['camera']['K'].copy() _K[:2, :2] *= scale _K[0, 2] *= cx _K[1, 2] *= cy info['camera']['K'] = _K info['camera']['K'] = _K info['sub'] = info['sub'] + '_scale_{}'.format(scale) if stage.effect_args.get('use_previous_K', False): info['camera']['K'] = infos[-1]['camera']['K'] info['object_keys'] = object_keys _infos.append(info) index += len(_index) infos.extend(_infos) return infos def create_demo_cameras(self, scale, camera_args, demo_args=None): if camera_args.method == 'none': from .utils_sample import create_center_radius RTs = create_center_radius(**camera_args) K = np.array([ camera_args.focal, 0, camera_args.W/2, 0, camera_args.focal, camera_args.H/2, 0, 0, 1], dtype=np.float32).reshape(3, 3)[None].repeat(RTs.shape[0], 0) R = RTs[:, :3, :3] T= RTs[:, :3, 3:] elif camera_args.method == 'mean': from .utils_sample import create_cameras_mean K, R, T = create_cameras_mean(list(self.cameras.values()), camera_args) K[:, 0, 2] = camera_args.W / 2 K[:, 1, 2] = camera_args.H / 2 elif camera_args.method == 'static': assert len(self.subs) == 1, "Only support monocular videos" camera = self.cameras[self.subs[0]] K = camera['K'][None] R = camera['R'][None] T = camera['T'][None] elif camera_args.method == 'line': for key, camera in self.cameras.items(): R = camera['R'] T = camera['T'] center_old = - R.T @ T print(key, center_old.T[0]) camera = self.cameras[str(camera_args.ref_sub)] K = camera['K'][None] R = camera['R'][None] T = camera['T'][None] t = np.linspace(0., 1., camera_args.allstep).reshape(-1, 1) t = t - 0.33 t[t<0.] = 0. t = t/t.max() start = np.array(camera_args.center_start).reshape(1, 3) end = np.array(camera_args.center_end).reshape(1, 3) center = end * t + start * (1-t) K = K.repeat(camera_args.allstep, 0) R = R.repeat(camera_args.allstep, 0) T = - np.einsum('fab,fb->fa', R, center) T = T.reshape(-1, 3, 1) K[:, :2] *= scale if demo_args is None: infos = self._demo_play_frames(self.ranges, K, R, T) return infos # create scripts if demo_args.mode == 'scripts': infos = self._demo_script(self.ranges, K, R, T, demo_args.stages) elif demo_args.mode == 'keyframe+rotate': infos = self._demo_keyframe(self.ranges, K, R, T, demo_args.frame, demo_args.nFrames) else: raise NotImplementedError return infos def get_allnames(self, root, subs, ranges, image_args): return [] def read_image(self, imgname, image_args, info): return self.blank class BaseNovelPose(BaseBase): def __init__(self, root, object_keys, ignore_keys, object_args, sample_args, camera_args): super().__init__(split='demo') self.root = root self.parse_object_args(object_keys, object_args, ignore_keys) self.params = self.load_all_smpl(root) self.cameras = self.create_cameras(camera_args) self.infos = self.combine_frame_and_camera(self.params, self.cameras) self.H = camera_args.H self.W = camera_args.W self.sample_args = sample_args self.check = False def load_all_smpl(self, root): from glob import glob from tqdm import tqdm smplnames = sorted(glob(join(root, 'smpl', '*.json'))) params = [] for nf, smplname in enumerate(tqdm(smplnames)): frame = int(os.path.basename(smplname).split('.')[0]) nv, sub = 0, 'novelpose0' # nf, frame = 0, 0 imgname = 'smplname' params.append({ 'imgname': imgname, 'sub': sub, 'frame': frame, 'nf': nf, 'nv': nv, 'index': nf, }) return params def combine_frame_and_camera(self, params, cameras): infos = [] for i in range(len(params)): info = params[i].copy() info['camera'] = cameras[0] infos.append(info) return infos def read_blank_image(self, *args): img = np.zeros([self.H, self.W, 3], dtype=np.float32) return img def __len__(self): return len(self.infos) def __getitem__(self, index): info = self.infos[index] objects = self.get_objects(self.root, info, self.object_keys, self.object_args) info = self.infos[index] imgname = info['imgname'] img = self.read_blank_image(imgname, info) back_mask = np.ones_like(img[:, :, 0]) # sample the ray from image ray_o, ray_d, rgb, coord = self.sample_ray(img, back_mask, info, objects) ret = self.sample_near_far(rgb, ray_o, ray_d, coord, objects) info['nf'] = 0 meta = { 'H': img.shape[0], 'W': img.shape[1], 'index': index, 'nframe': info['nf'], 'nview': info['nv'], 'sub': info['sub'], 'keys': list(objects.keys()), } if self.check: meta['img'] = img meta['object_keys'] = meta['keys'] ret['meta'] = meta return ret class BaseCanonical(BaseNovelPose): def __init__(self, nFrames, **kwargs): super().__init__(**kwargs) self.nFrames = nFrames def load_all_smpl(self, root): smpl = read_smpl(root)[0] for key in ['poses', 'Rh', 'Th']: smpl[key] = np.zeros_like(smpl[key]) return smpl def combine_frame_and_camera(self, params, cameras): infos = [] for i in range(len(cameras)): info = params.copy() info['camera'] = cameras[i] info.update({ 'imgname': f'view_{i}_frame_0', 'sub': 'novel_' + str(i), 'frame': 0, 'nf': 0, 'nv': i, 'index': i, }) infos.append(info) return infos def get_objects(self, root, info, object_keys, object_args): objects = {'human_0': AABBSampler(split='test', bounds=[[-1, -1.3, -0.3], [1, 0.7, 0.3]])} objects['human_0'].feature['shapes'] = self.params['shapes'] return objects if __name__ == '__main__': from ...config import Config, load_object from copy import deepcopy # config = Config.load('config/neuralbody/dataset/multiview_custom.yml') # config = Config.load('config/neuralbody/dataset/neuralbody_lightstage.yml') config = Config.load('config/neuralbody/dataset/neuralbody_soccer.yml') data_share = config.pop('data_share_args') data_share.root = os.environ['data'] for split in ['train', 'val', 'demo']: data = deepcopy(data_share) data.merge_from_other_cfg(config['data_{}_args'.format(split)]) config['data_{}_args'.format(split)] = data print(config) dataset = load_object(config.data_train_module, config.data_train_args) dataset.debug = True for data in dataset: print(data.keys())