preprocess_data.py

import os
import json
import argparse
import numpy as np
import numpy as np
from tqdm import tqdm
from copy import deepcopy
from easydict import EasyDict as edict
from dataset.riseVAE_2cam import RiseVAEDataset
from process_pointcloud import *
from PIL import Image
import h5py
from dataset.projector import Projector
from utils.constants import INTRINSICS, IMAGE_SIZE
default_args = edict({
    "num_action": 20,
    "voxel_size": 0.005,
    "vis":False
})

def save_to_hdf5(data_dict, filename):
    """Save preprocessed data to HDF5 with compression"""
    with h5py.File(filename, 'w') as f:
        # Save regular arrays
        f.create_dataset('handPose_data', data=data_dict['handPose_data'], compression='gzip')
        f.create_dataset('tcp_data', data=data_dict['tcp_data'], compression='gzip')
        
def load_from_hdf5(filename):
    """Load preprocessed data from HDF5"""
    with h5py.File(filename, 'r') as f:
        handPose_data = np.array(f['handPose_data'])
        tcp_data = np.array(f['tcp_data'])
            
    return {
        "handPose_data": handPose_data,
        "tcp_data": tcp_data,
    }



def preprocess(args_override):
    # load default arguments
    args = deepcopy(default_args)
    for key, value in args_override.items():
        args[key] = value
    calib_path=os.path.expanduser(args.calib_path)
    projector=Projector(calib_path=calib_path)
    print("Loading Dataset")
    dataset = RiseVAEDataset(
        path = args.data_path,
        split = 'train',
        num_action = args.num_action,
        voxel_size = args.voxel_size,
        calib_path=args.calib_path,
        preprocess=True
    ) 
    print("Finish Loading Dataset")

    # Saving Actions
    handPose_data=[]
    tcp_data=[]
    pc_data=[]

    print("Total Length for preprocessing: ",len(dataset))
    # for index in range(len(dataset)):
    os.makedirs(args.preprocessed_pc_dir, exist_ok=True)
    for index in range(5):
        print("Processing data id: ",index)   
        data_path = dataset.data_paths[index]
        obs_frame_ids = dataset.obs_frame_ids[index]
        color_dir = os.path.join(data_path, "cam_{}".format(dataset.cams[0]), 'color')
        depth_dir = os.path.join(data_path, "cam_{}".format(dataset.cams[0]), 'depth')
        second_cam_color_dir=os.path.join(data_path, "cam_{}".format(dataset.cams[1]), 'color')
        second_cam_depth_dir = os.path.join(data_path, "cam_{}".format(dataset.cams[1]), 'depth')
        colors1 = []
        depth1 = []
        for frame_id in obs_frame_ids:
            color = Image.open(os.path.join(color_dir, "{}.png".format(frame_id)))
            colors1.append(color)
            depth1.append(
                np.array(Image.open(os.path.join(depth_dir, "{}.png".format(frame_id))), dtype = np.float32)
            )
        colors2 = []
        depth2 = []
        matching_pairs=torch.load(f"{color_dir}/pairs.pth")
        for frame_id in obs_frame_ids:
            closest_id=matching_pairs[str(frame_id)]
            color = Image.open(os.path.join(second_cam_color_dir, "{}.png".format(closest_id)))
            colors2.append(color)
            depth2.append(
                np.array(Image.open(os.path.join(second_cam_depth_dir, "{}.png".format(closest_id))), dtype = np.float32)
            )
        colors1=np.stack(colors1,axis=0)
        depth1=np.stack(depth1,axis=0)
        colors2=np.stack(colors2,axis=0)
        depth2=np.stack(depth2,axis=0)
        pc1 = process([colors1], [depth1], dataset.cams[0], INTRINSICS=INTRINSICS,IMAGE_SIZE=IMAGE_SIZE,projector=projector, voxel_size=args.voxel_size)[0]
        pc2 = process([colors2], [depth2], dataset.cams[1], INTRINSICS=INTRINSICS,IMAGE_SIZE=IMAGE_SIZE,projector=projector, voxel_size=args.voxel_size)[0]
        pc1_points, pc1_colors = pc1[:, :3], pc1[:, 3:]
        pc2_points, pc2_colors = pc2[:, :3], pc2[:, 3:]

        combined_points = np.concatenate([pc1_points, pc2_points], axis=0)
        combined_colors = np.concatenate([pc1_colors, pc2_colors], axis=0)
        pointcloud=np.concatenate([combined_points, combined_colors], axis=1)
        np.save(f"{args.preprocessed_pc_dir}/processed_pc_{index}.npy",pointcloud)
        if args.vis:
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(pointcloud[:, :3])
            pcd.colors = o3d.utility.Vector3dVector(pointcloud[:, 3:])
            filename="vis_output/point.ply"
            o3d.io.write_point_cloud(filename, pcd)
            print(f"Saved point cloud to {filename}")
            print("Number of Points: ",pointcloud.shape[0])
            print(data_path,obs_frame_ids)
        tcp_dir = os.path.join(data_path, "lowdim")
        hand_dir = os.path.join(data_path, "lowdim")
        action_seq=dataset.action_frame_ids[index]
        handPose_list=[]
        tcp_list=[]
        for action_id in action_seq:
            tcp=np.load(os.path.join(tcp_dir,'tcp.npz'))[str(action_id)][:7].astype(np.float32)
            handPose=np.load(os.path.join(hand_dir,"pos.npz"))[str(action_id)]
            handPose_list.append(handPose)
            tcp_list.append(tcp)
        handPose_data.append(np.stack(handPose_list)) 
        tcp_data.append(np.stack(tcp_list))
        pc_data.append(pointcloud)
    tcp_data=np.stack(tcp_data)
    
    handPose_data=np.stack(handPose_data)
    save_to_hdf5({
        "handPose_data": handPose_data,
        "tcp_data": tcp_data,
    }, args.saved_h5_path)
    print("Finished Preprocessing Data")



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_path', action = 'store', type = str, help = 'data path', required = False, default="")
    parser.add_argument('--calib_path', action = 'store', type = str, help = 'calib path', required = False, default="")
    parser.add_argument('--saved_h5_path', action = 'store', type = str, help = 'the path of your saved h5 data file', required = True, default="")
    parser.add_argument('--preprocessed_pc_dir', action = 'store', type = str, help = 'the directory to save your point clouds', required = True, default="")
    parser.add_argument('--num_action', action = 'store', type = int, help = 'number of action steps', required = False, default = 20)
    parser.add_argument('--voxel_size', action = 'store', type = float, help = 'voxel size', required = False, default = 0.005)

    preprocess(vars(parser.parse_args()))