847 lines
37 KiB
Python
847 lines
37 KiB
Python
|
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())
|