import os
import torch
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader
from torch.utils.data.dataset import random_split
from torch.utils.data import TensorDataset, Dataset
import h5py
import tables as pytables
import random
import numpy as np
import time
import matplotlib.pyplot as plt
import os.path as osp
import glob
import time
from torch_geometric.nn import fps
from tqdm import tqdm

xyz2 = h5py.File("./trainset.h5",'r')["xyz2"]

store = np.random.randn(230000,64,3)

for i in tqdm(range(len(xyz2))):
    data = torch.tensor(xyz2[i])
    ind = fps(a,ratio=64.0/1024.0)
    store[i] = np.array(data[ind])

with h5py.File('trainset_fps256.h5', 'w') as f:
    f.create_dataset('xyz2', data=store)