convert.py

#!/usr/bin/env python

from __future__ import print_function

import argparse

import os
import os.path as osp
import re

import numpy as np
import scipy.io


parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("data_dir", help="e.g., /home/wkentaro/data/ycb_video/YCB_Video_Dataset/data")
parser.add_argument("out_dir", help="e.g., ycb_video_data_converted")
args = parser.parse_args()

data_dir = args.data_dir
out_dir = args.out_dir

for trajectory_dir_basename in os.listdir(data_dir):
    trajectory_dir = osp.join(data_dir, trajectory_dir_basename)
    for frame_file_basename in sorted(os.listdir(trajectory_dir)):
        m = re.match(r"([0-9]+8)-meta.mat", frame_file_basename)
        if not m:
            continue
        frame_id = m.groups()[0]
        meta_file = osp.join(trajectory_dir, frame_file_basename)
        meta = scipy.io.loadmat(meta_file, squeeze_me=True)
        out_file_stem = osp.join(out_dir, trajectory_dir_basename, frame_id)

        # transform
        out_transform_file = out_file_stem + "-transform.txt"
        T_world2cam = np.r_[meta["rotation_translation_matrix"], [[0, 0, 0, 1]]]
        if not osp.exists(osp.dirname(out_transform_file)):
            os.makedirs(osp.dirname(out_transform_file))
        np.savetxt(out_transform_file, T_world2cam, fmt="%.18f")

        # class ids
        out_class_file = out_file_stem + "-class.txt"
        np.savetxt(out_class_file, meta["cls_indexes"], fmt="%d")

        print(out_transform_file)
        print(out_class_file)