text: #  Copyright (c) 2020. Author: Hanchen Wang, hw501@cam.ac.uk
  
  import os, argparse, numpy as np
  from tensorpack import DataFlow, dataflow
  from open3d.open3d.io import read_triangle_mesh, read_point_cloud
  
  
  def sample_from_mesh(filename, num_samples=16384):
      pcd = read_triangle_mesh(filename).sample_points_uniformly(number_of_points=num_samples)
      return np.array(pcd.points)
  
  
  class pcd_df(DataFlow):
      def __init__(self, model_list, num_scans, partial_dir, complete_dir, num_partial_points=1024):
          self.model_list = [_file for _file in model_list if 'train' in _file]
          self.num_scans = num_scans
          self.partial_dir = partial_dir
          self.complete_dir = complete_dir
          self.num_ppoints = num_partial_points
  
      def size(self):
          return len(self.model_list) * self.num_scans
  
      @staticmethod
      def read_pcd(filename):
          pcd = read_point_cloud(filename)
          return np.array(pcd.points)
  
      def get_data(self):
          for model_id in self.model_list:
              complete = sample_from_mesh(os.path.join(self.complete_dir, '%s.obj' % model_id))
              for i in range(self.num_scans):
                  partial = self.read_pcd(os.path.join(self.partial_dir, model_id + '_%d.pcd' % i))
                  partial = partial[np.random.choice(len(partial), self.num_ppoints)]
                  yield model_id.replace('/', '_'), partial, complete
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('--list_path', default=r'../render/ModelNet_flist_normalised.txt')
      parser.add_argument('--num_scans', type=int, default=10)
      parser.add_argument('--partial_dir', default=r'../render/dump_modelnet_normalised_supercoarse/pcd')
      parser.add_argument('--complete_dir', default=r'../data/ModelNet40')
      parser.add_argument('--output_file', default=r'../data/ModelNet40_train_1024_supercoarse.lmdb')
      args = parser.parse_args()
  
      with open(args.list_path) as file:
          model_list = file.read().splitlines()
      df = pcd_df(model_list, args.num_scans, args.partial_dir, args.complete_dir)
      if os.path.exists(args.output_file):
          os.system('rm %s' % args.output_file)
      dataflow.LMDBSerializer.save(df, args.output_file)
  