def main():
dataset_root = 'datasets/ycb/YCB_Video_Dataset'
with open('datasets/ycb/dataset_config/classes.txt') as f:
classes = f.read().split()
classes.insert(0, 'background')
# with open('datasets/ycb/dataset_config/rendered_data_list.txt') as f:
# file_list = f.read().split()
for obj in tqdm(range(2, 22)):
file_list = [
'depth_renders_offset/{}/{:04d}'.format(classes[obj], i)
for i in range(3885)
]
for fn in tqdm(file_list):
img = Image.open('{0}/{1}-color.png'.format(dataset_root, fn))
obj_label = classes.index(fn.split('/')[-2])
quat = np.load('{0}/{1}-trans.npy'.format(dataset_root, fn))
label = np.where(np.array(img.split()[-1]) == 255, obj_label, 0)
cv2.imwrite('{0}/{1}-label.png'.format(dataset_root, fn), label)
poses = np.zeros([3, 4, 1])
poses[:3, :3, 0] = tfs.quaternion_matrix(quat)[:3, :3]
poses[:3, 3, 0] = [0., 0., 1.]
scio.savemat(
'{0}/{1}-meta.mat'.format(dataset_root, fn), {
'cls_indexes': np.array([[obj_label]]),
'factor_depth': np.array([[10000]]),
'poses': poses
})
def createScene(obj_filename, meta, scale=True):
scene = pyrender.Scene(ambient_light=np.ones(3))
obj_trimesh = trimesh.load(obj_filename)
if scale:
obj_trimesh.units = 'mm'
obj_trimesh.convert_units('m')
obj_mesh = pyrender.Mesh.from_trimesh(obj_trimesh)
obj_pose = np.eye(4)
obj_node = scene.add(obj_mesh, pose=obj_pose)
camera = pyrender.IntrinsicsCamera(fx=meta['camera_fx'],
fy=meta['camera_fy'],
cx=meta['camera_cx'],
cy=meta['camera_cy'])
camera_pose = quaternion_matrix([1, 0, 0, 0])
scene.add(camera, pose=camera_pose)
return scene, obj_node
parser.add_argument('--mtl_Ke', nargs=3, type=float, default=[0,0,0])
parser.add_argument('--mtl_Ns', type=float, default=18)
parser.add_argument('--mtl_Ni', type=float, default=1)
parser.add_argument('--mtl_d', type=float, default=1)
parser.add_argument('--mtl_illum', type=float, default=2)
args = parser.parse_args()
transform_verts = not (args.scale is None and args.quat is None and args.trans is None)
if(args.scale is not None):
s = np.array(args.scale)
else:
s = 1.0
if(args.quat is not None):
R = qm.quaternion_matrix(args.quat)[:3,:3]
else:
R = np.eye(3)
if(args.trans is not None):
t = np.array(args.trans)
else:
t = np.zeros(3)
if(args.mtl_Ka is not None):
with open(args.output_filename[:-3]+'mtl', 'w') as f:
f.write('newmtl obj_mtl\n')
f.write('Ns {}\n'.format(args.mtl_Ns))
f.write('Ka {} {} {}\n'.format(*args.mtl_Ka))
f.write('Kd {} {} {}\n'.format(*args.mtl_Kd))
f.write('Ks {} {} {}\n'.format(*args.mtl_Ks))
f.write('Ke {} {} {}\n'.format(*args.mtl_Ke))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--output_folder',
type=str,
help='Data render output location')
parser.add_argument('--model_list',
type=str,
default='datasets/ycb/YCB_Video_Dataset',
help='Dataset root dir ('
'YCB_Video_Dataset'
')')
parser.add_argument(
'--models',
type=str,
nargs='+',
help='List of model files or text file containing model file names')
parser.add_argument(
'--quats',
type=str,
default=None,
help='\'grid\' for 3885 grid or N for random quaternions')
parser.add_argument('--name_start_idx',
type=int,
default=-2,
help='Path depth for start of model name')
parser.add_argument('--name_end_idx',
type=int,
default=-1,
help='Path depth for end of model name')
args = parser.parse_args()
if (len(args.models) == 1 and args.models[0][-4:] == '.txt'):
with open(args.models[0]) as f:
args.models = f.readlines()
args.models = [x.strip() for x in args.models]
if (args.quats == 'grid'):
grid = S3Grid(2)
grid.Simplify()
quats = grid.vertices
elif (is_int(args.quats)):
quats = [random_quaternion() for _ in range(int(args.quats))]
elif (args.quats[-4:] == '.txt'):
with open(args.quats) as f:
quats = f.readlines()
for j, q in enumerate(quats):
quats[j] = np.array(q, dtype=float)
quats[j] /= np.linalg.norm(quats[j])
else:
raise ValueError(
'Bad quaternion format. Valid formats are \'grid\', N, or [file].txt'
)
digits = len(str(len(quats)))
fx = 1066.778
fy = 1067.487
px = 312.9869
py = 241.3109
pbar_model = tqdm(args.models)
for model_filename in pbar_model:
model = '/'.join(
model_filename.split('/')[args.name_start_idx:args.name_end_idx])
pbar_model.set_description('Rendering {}'.format(model))
renderer = BpyRenderer(transform_func=ycbRenderTransform)
renderer.loadModel(model_filename, emit=0.5)
renderer.setCameraMatrix(fx, fy, px, py, 640, 480)
renderer.setDepth()
render_dir = '{}/{}/'.format(args.output_folder, model)
pathlib.Path(render_dir).mkdir(parents=True, exist_ok=True)
filename_template = render_dir + '{:0' + str(digits) + 'd}-{}.{}'
for j, q in tqdm(enumerate(quats), total=len(quats)):
# change if there are different model ids
obj_label = 1
trans_mat = getYCBTransform(q, [0, 0, 1])
img, depth = renderer.renderTrans(trans_mat)
depth[depth > 1000] = 0
depth = depth * 10000
cv2.imwrite(filename_template.format(j, 'color', 'png'), img)
Image.fromarray(depth.astype(np.int32), "I").save(
filename_template.format(j, 'depth', 'png'))
np.save(filename_template.format(j, 'trans', 'npy'), q)
label = np.where(np.array(img[:, :, -1]) == 255, obj_label, 0)
cv2.imwrite(filename_template.format(j, 'label', 'png'), label)
poses = np.zeros([3, 4, 1])
poses[:3, :3, 0] = quaternion_matrix(q)[:3, :3]
poses[:3, 3, 0] = [0., 0., 1.]
scio.savemat(
filename_template.format(j, 'meta', 'mat'), {
'cls_indexes': np.array([[obj_label]]),
'factor_depth': np.array([[10000]]),
'poses': poses
})
def getYCBTransform(q, t=[0, 0, .5]):
trans_mat = quaternion_matrix(q)
ycb_mat = euler_matrix(-np.pi / 2, 0, 0)
trans_mat = trans_mat.dot(ycb_mat)
trans_mat[:3, 3] = t
return trans_mat