Dex_YCB dataset organization

[zerchen]

Dex_YCB dataset organization

[zerchen]

Hi @Tonyzhang2000,

Thanks for your interests in this work.
I attach the script to generate meta files below. You could make some modifications to make it compatiable with the code. You need to use it together with the official dexycb toolkit and put it under the examples folder.
For the image preprocessing, I project the hand wrist location to the image plane and use it as a guide to crop a 480x480 image around the projected hand wrist out of the original 640x480 image and then resize it to 256x256.

Best,
Zerui

import json
import pickle
from tqdm import tqdm
import os
import cv2
from manopth.manolayer import ManoLayer
import trimesh
import torch
import shutil
import numpy as np
from scipy.spatial import cKDTree as KDTree

from dex_ycb_toolkit.factory import get_dataset


def main():
#   for setup in ('s0', 's1', 's2', 's3'):
  output_rgb_dir = '/data/zerui/dexycb_processed_v2/{}/rgb/'
  output_seg_dir = '/data/zerui/dexycb_processed_v2/{}/segm/'
  output_meta_dir = '/data/zerui/dexycb_processed_v2/{}/meta/'
  output_hmesh_dir = '/data/zerui/dexycb_processed_v2/{}/mesh_hand/'
  output_omesh_dir = '/data/zerui/dexycb_processed_v2/{}/mesh_obj/'
  for setup in ('s0',):
    for split in ('test',):
    # for split in ('train', 'val', 'test'):
      name = '{}_{}'.format(setup, split)
      print('Dataset name: {}'.format(name))

      output_rgb_dir  = output_rgb_dir.format(split)
      output_seg_dir  = output_seg_dir.format(split)
      output_meta_dir  = output_meta_dir.format(split)
      output_hmesh_dir  = output_hmesh_dir.format(split)
      output_omesh_dir  = output_omesh_dir.format(split)
      os.makedirs(output_rgb_dir, exist_ok=True)
      os.makedirs(output_seg_dir, exist_ok=True)
      os.makedirs(output_meta_dir, exist_ok=True)
      os.makedirs(output_hmesh_dir, exist_ok=True)
      os.makedirs(output_omesh_dir, exist_ok=True)

      dataset = get_dataset(name)

      print('Dataset size: {}'.format(len(dataset)))

      for i in tqdm(range(len(dataset))):
        sample = dataset[i]
        filename = str(i).rjust(8, "0") + '.pkl'
        label = np.load(sample['label_file'])

        obj_id = sample['ycb_ids'][sample['ycb_grasp_ind']]
        obj_mesh = trimesh.load(dataset.obj_file[obj_id])
        obj_verts = obj_mesh.vertices
        obj_faces = obj_mesh.faces
        homo_obj_verts = np.ones((obj_verts.shape[0], 4))
        homo_obj_verts[:, :3] = obj_verts
        pose_y = label['pose_y'][sample['ycb_grasp_ind']]
        if not np.all(pose_y == 0):
            obj_verts = np.dot(pose_y, homo_obj_verts.transpose(1, 0)).transpose(1, 0)
            obj_mesh = trimesh.Trimesh(vertices=obj_verts, faces=obj_faces)
        else:
            continue

        pose_m = label['pose_m']
        mano_layer = ManoLayer(flat_hand_mean=False, ncomps=45, side=sample['mano_side'], mano_root='/home2/zerui/code/dex-ycb-toolkit/manopth/mano/models', use_pca=True)
        hand_faces = np.load('closed_fmano.npy')
        betas = torch.tensor(sample['mano_betas'], dtype=torch.float32).unsqueeze(0)

        if sample['mano_side'] != 'right':
            continue

        # Add MANO meshes.
        if not np.all(pose_m == 0.0):
          hand_poses = torch.from_numpy(pose_m)
          hand_verts, _ = mano_layer(hand_poses[:, 0:48], betas, hand_poses[:, 48:51])
          hand_verts /= 1000
          hand_verts = hand_verts.view(778, 3)
          hand_verts = hand_verts.numpy()
          hand_mesh = trimesh.Trimesh(vertices=hand_verts, faces=hand_faces)
        else:
            continue

        hand_points_kd_tree = KDTree(hand_verts)
        obj2hand_distances, _ = hand_points_kd_tree.query(obj_verts)
        if obj2hand_distances.min() <= 0.005:
            shutil.copy2(sample['color_file'], os.path.join(output_rgb_dir, filename.replace('pkl', 'jpg')))
            cv2.imwrite(os.path.join(output_seg_dir, filename.replace('pkl', 'png')), label['seg'])
            obj_mesh.export(os.path.join(output_omesh_dir, filename.replace('pkl', 'obj')))
            hand_mesh.export(os.path.join(output_hmesh_dir, filename.replace('pkl', 'obj')))
            meta_data = {}

            if np.any(label['joint_2d'][0] < 0) or np.any(label['joint_2d'][0][:, 0] > 640) or np.any(label['joint_2d'][0][:, 1] > 480):
                continue

            meta_data['coords_2d'] = label['joint_2d'][0]
            meta_data['coords_3d'] = label['joint_3d'][0]
            meta_data['verts_3d'] = hand_verts
            meta_data['hand_pose'] = label['pose_m'][0][3:48]
            meta_data['trans'] = label['pose_m'][0][48:51]
            affine_matrix = np.zeros((4, 4))
            affine_matrix[3, 3] = 1.
            affine_matrix[:3, :4] = label['pose_y'][sample['ycb_grasp_ind']]
            meta_data['affine_transform'] = affine_matrix
            meta_data['side'] = sample['mano_side']
            cam_intr = np.zeros((3, 4))
            fx = sample['intrinsics']['fx']
            fy = sample['intrinsics']['fy']
            cx = sample['intrinsics']['ppx']
            cy = sample['intrinsics']['ppy']
            cam_intr[:3, :3] = np.array([[fx, 0., cx], [0., fy, cy], [0., 0., 1.]])
            meta_data['cam_intr'] = cam_intr
            with open(output_meta_dir + filename, 'wb') as f:
                pickle.dump(meta_data, f)


if __name__ == '__main__':
  main()

[Tonyzhang2000]

Thank you for your reply!