code: import numpy as np
  import torch
  from torch.utils.data import Dataset, DataLoader
  from torchvision import transforms as tfs
  from PIL import Image
  import os, cv2, copy, time
  from config import *
  
  
  # args.
  image_height, image_width = opt.image_height, opt.image_width
  intrinsics = opt.intrinsics
  close_radius, far_radiuses = 0, opt.far_radiuses
  n_neighb_pts = opt.n_neighb_pts
  
  
  def isSon(son, fa):
      for i in range(len(fa)):
          if son[i] != fa[i]:
              return False
      return True
  
  
  # todo: to be migrated...
  def depth2local(depth):  # depth: float32, meter.
      cx, cy, fx, fy = intrinsics[0, 2], intrinsics[1, 2], intrinsics[0, 0], intrinsics[1, 1]
      u_base = np.tile(np.arange(image_width), (image_height, 1))
      v_base = np.tile(np.arange(image_height)[:, np.newaxis], (1, image_width))
      X = (u_base - cx) * depth / fx
      Y = (v_base - cy) * depth / fy
      coord_local = np.stack((X, Y, depth), axis=2)
      return coord_local
  def partial_pts(pts_all_in, p, r_min, r_max):  # pts_all_in.shape (#points, #channel)
      pts_all = copy.deepcopy(pts_all_in)
      p_mat = p[np.newaxis, 0:3].repeat(pts_all.shape[0], axis=0)
      norms = np.linalg.norm((p_mat - pts_all[:, 0:3]), axis=1)
      return pts_all[np.logical_and(norms >= r_min, norms <= r_max)]
  def sample_pts(pts_in, num):  # pts_in.shape (#points, #channel)
      pts = copy.deepcopy(pts_in)
      while pts.shape[0] < num:
          pts = np.concatenate((pts, pts), axis=0)
      rand_ids = np.arange(pts.shape[0])
      np.random.shuffle(rand_ids)
      return pts[rand_ids[0:num], :]
  def sample_pts_rc(pts_in, rcs_in, num):  # pts_in.shape (#points, #channel)
      pts = copy.deepcopy(pts_in)
      rcs = copy.deepcopy(rcs_in)
      while pts.shape[0] < num:
          pts = np.concatenate((pts, pts), axis=0)
      rand_ids = np.arange(pts.shape[0])
      np.random.shuffle(rand_ids)
      return pts[rand_ids[0:num], :], rcs_in[rand_ids[0:num], :]
  def sample_pts9d_r3d(pts_in, num, radius):  # pts_in.shape (#points, #channel)
      pts = copy.deepcopy(pts_in)
      thresh = 500
      # remove background by 3d radius
      xyz = pts[:, 0:3]
      pts = pts[np.linalg.norm(xyz, axis=1) <= radius]
      # print('pt num after r3d {}'.format(pts.shape[0]))
      if pts.shape[0] < thresh:  # avoid infinite loop.
          return None
      while pts.shape[0] < num:
          pts = np.concatenate((pts, pts), axis=0)
      rand_ids = np.arange(pts.shape[0])
      np.random.shuffle(rand_ids)
      return pts[rand_ids[0:num], :]
  def shift_pts(pts_in, cen):  # pts_in.shape (#points, #channel)
      pts = copy.deepcopy(pts_in)
      cen_mat = cen[np.newaxis, :].repeat(pts.shape[0], axis=0)
      pts[:, 0:3] = pts[:, 0:3] - cen_mat
      return pts
  def shift_pts6d(pts_in, cen):  # pts_in.shape (#points, #channel)
      pts = copy.deepcopy(pts_in)
      cen_mat = cen[np.newaxis, :].repeat(pts.shape[0], axis=0)
      pts[:, :] = pts[:, :] - cen_mat
      return pts
  def shift_pts9d(pts_in, cen):  # pts_in.shape (#points, #channel)
      cpt = copy.deepcopy(cen)
      cpt[3:6] = np.zeros(3)  # remove shift of normal
      pts = copy.deepcopy(pts_in)
      cpt_mat = cpt[np.newaxis, :].repeat(pts.shape[0], axis=0)
      pts[:, :] = pts[:, :] - cpt_mat
      return pts
  
  
  def make_ppf(pts9d, cen9d):  # (N,9), (9,)
      # prepare
      n_pts = pts9d.shape[0]
      d = pts9d[:, 0:3]
      n2 = pts9d[:, 3:6]
      n1 = np.repeat(cen9d[3:6].reshape(1, 3), n_pts, axis=0)
      # ppf
      dim1 = np.linalg.norm(d, axis=1).reshape(n_pts, 1)
      d = d / (dim1.reshape(n_pts, 1))
      dim2 = np.sum(n1 * d, axis=1).reshape(n_pts, 1)
      dim3 = np.sum(n2 * d, axis=1).reshape(n_pts, 1)
      dim4 = np.sum(n1 * n2, axis=1).reshape(n_pts, 1)
      ppf = np.concatenate((dim1, dim2, dim3, dim4), axis=1)
      ppf7d = np.concatenate((ppf, pts9d[:, 6:9]), axis=1)
      return ppf7d
  
  def compute_points_normal(pts):
      raw_shape = pts.shape
      normal = np.zeros((raw_shape))  # (r,c,3)
      t0 = time.time()
      for r in range(2, raw_shape[0] - 2):
          for c in range(2, raw_shape[1] - 2):
              pts_local = pts[r - 2:r + 3, c - 2:c + 3, :]  # (5,5,3)
              pts_local = pts_local.reshape(-1, 3)  # (N,3)
              pts_local = pts_local[np.linalg.norm(pts_local - pts[r, c, :], axis=1) < 0.1]  # remove outliers.
              if pts_local.shape[0] < 4:
                  continue
              pts_local = pts_local - np.mean(pts_local, axis=0)
              C = pts_local.T @ pts_local / pts_local.shape[0]
              e, v = np.linalg.eig(C)
              d = v[:, np.where(e == np.min(e))[0][0]]
              n = d / np.linalg.norm(d)
              if np.dot(n, np.array([0, 0, 1])) > 0:
                  n = -n
              normal[r, c, :] = n
      t1 = time.time()
      print('preprocess data: compute normal cost {:.2f}s'.format(t1 - t0))
      return normal
  
  
  # for depth adaptive 2d
  def partial_pts_2d(pts_rc, cen_rc, list_drdc):
      result = None
      r_max, c_max = int(image_height / 4 - 1), int(image_width / 4 - 1)
      mat_drdc = (np.array(list_drdc) / 4).astype(int)
      mat_cen_rc = np.array(cen_rc)
      mat_targ_rc = cen_rc + mat_drdc
      mat_targ_rc[mat_targ_rc < 0] = 0
      targ_r = mat_targ_rc[:, 0]
      targ_r[targ_r > r_max] = r_max
      targ_c = mat_targ_rc[:, 1]
      targ_c[targ_c > c_max] = c_max
      result = pts_rc[targ_r, targ_c]
      return copy.deepcopy(result)
  
  
  # for depth adaptive 2d
  def partial_pts_2d_rc(pts_rc, cen_rc, list_drdc):
      result = None
      r_max, c_max = int(image_height / 4 - 1), int(image_width / 4 - 1)
      mat_drdc = (np.array(list_drdc) / 4).astype(int)
      mat_cen_rc = np.array(cen_rc)
      mat_targ_rc = cen_rc + mat_drdc
      mat_targ_rc[mat_targ_rc < 0] = 0
      targ_r = mat_targ_rc[:, 0]
      targ_r[targ_r > r_max] = r_max
      targ_c = mat_targ_rc[:, 1]
      targ_c[targ_c > c_max] = c_max
      result = pts_rc[targ_r, targ_c]
      return copy.deepcopy(result), copy.deepcopy(
          np.concatenate((targ_r.reshape(targ_r.shape[0], 1), targ_c.reshape(targ_c.shape[0], 1)), axis=1))
  
  
  # for depth adaptive 2d with dynamics label
  def partial_pts_2d_with_label(pts_rc, cen_rc, list_drdc, mask):  # mask: 0 for static pixel, 255 for dynamic pixel.
      result = None
      r_max, c_max = int(image_height / 4 - 1), int(image_width / 4 - 1)
      mat_drdc = (np.array(list_drdc) / 4).astype(int)
      mat_cen_rc = np.array(cen_rc)
      mat_targ_rc = cen_rc + mat_drdc
      mat_targ_rc[mat_targ_rc < 0] = 0
      targ_r = mat_targ_rc[:, 0]
      targ_r[targ_r > r_max] = r_max
      targ_c = mat_targ_rc[:, 1]
      targ_c[targ_c > c_max] = c_max
      m1 = np.zeros((mask.shape[0], mask.shape[1]))
      m1[mask == 0] = 1
      m2 = np.zeros((mask.shape[0], mask.shape[1]))
      m2[targ_r, targ_c] = 1
      m3 = np.logical_and(m1, m2)
      result = pts_rc[m3]
      return copy.deepcopy(result)
  
  
  class LevelDataset_PPF(Dataset):
  
      def __init__(self, data_dir, the_list, n_pts_per_frame=opt.n_pts_per_frame, neighbor_da2d=None, far_radius=None,
                   enable_color_aug=True, specified_node=None):
          super().__init__()
          self.data_dir, self.the_list = data_dir, the_list
          self.n_pts_per_frame = n_pts_per_frame
          self.neighbor_da2d = neighbor_da2d  # (n_pts, dim_pt).
          self.far_radius = far_radius  # scalar.
          self.enable_color_aug = enable_color_aug
          self.specified_node = specified_node
  
      def __len__(self):
          return len(self.the_list)
  
      def __getitem__(self, idx):
          fid, rc_route = self.the_list[idx]
          # load 
          depth = cv2.imread('{}/{}'.format(self.data_dir, opt.depth_format.format(fid)), cv2.IMREAD_UNCHANGED) / 1000.0
          color = cv2.imread('{}/{}'.format(self.data_dir, opt.color_format.format(fid)), cv2.IMREAD_UNCHANGED)[:, :, 0:3]
          # color jitter
          if self.enable_color_aug:
              img = Image.fromarray(color)
              if np.random.rand() < 0.5:
                  img = tfs.ColorJitter(brightness=1.)(img)
              if np.random.rand() < 0.5:
                  img = tfs.ColorJitter(contrast=1.)(img)
              if np.random.rand() < 0.5:
                  img = tfs.ColorJitter(saturation=1.)(img)
              color = np.array(img)
          if np.max(color) > 1:
              color = color / 255. - 0.5
          local = depth2local(depth)
          r_ids, c_ids = list(range(0, image_height, 4)), list(range(0, image_width, 4))
          depth, color, local = depth[r_ids, :][:, c_ids], color[r_ids, :, :][:, c_ids, :], local[r_ids, :, :][:, c_ids, :]
          # normal by 3d neighbor plane fitting.
          normal_path = '{}/frame-{:06d}.scaled.normal.npy'.format(self.data_dir, fid)
          if os.path.exists(normal_path):
              # print('fid {}'.format(fid)) # to debug rio10 scene09 10
              # normal = np.load(normal_path)
              if os.path.getsize(normal_path) > 1:
                  normal = np.load(normal_path, encoding='bytes', allow_pickle=True)
              else:
                  normal = compute_points_normal(local)
                  np.save(normal_path, normal)
          else:
              normal = compute_points_normal(local)
              np.save(normal_path, normal)
          lclnmlclr = np.concatenate((np.concatenate((local, normal), axis=2), color), axis=2)
          # build a patch
          rand_ids = np.arange(len(rc_route))
          np.random.shuffle(rand_ids)
          selected_ids = rand_ids[0:self.n_pts_per_frame * 2]  # more candidates
          pt_in = torch.zeros((self.n_pts_per_frame, 7, 1))
          nb_in = torch.zeros((self.n_pts_per_frame, 7, opt.n_neighb_pts))
          route_labs = torch.zeros((self.n_pts_per_frame, opt.tree_height - 1)).fill_(ary)
          rc_list = []
          # da2d+3d neighbor
          if not self.neighbor_da2d is None:
              sid = 0
              for tmp_idx in range(len(selected_ids)):
                  r, c = rc_route[selected_ids[tmp_idx]][0], rc_route[selected_ids[tmp_idx]][1]
                  if np.isnan(lclnmlclr[r, c, 3]):
                      continue
                  if self.specified_node:
                      if not isSon(rc_route[selected_ids[tmp_idx]][2], self.specified_node):
                          continue
                  route_labs[sid] = torch.Tensor(rc_route[selected_ids[tmp_idx]][2])
                  rc_list.append([r, c])
                  pt_in[sid] = torch.Tensor(
                      np.concatenate((np.array([[0.], [0.], [0.], [0.]]), color[r, c, 0:3][:, np.newaxis]), axis=0))
                  da2d_list = (np.array(self.neighbor_da2d) / depth[r, c]).astype(int)
                  # ppf
                  pts9d = shift_pts9d(sample_pts(partial_pts_2d(lclnmlclr, (r, c), da2d_list), opt.n_neighb_pts),
                                      lclnmlclr[r, c, :])
                  cen9d = copy.deepcopy(lclnmlclr[r, c, :])
                  cen9d[0:3] = np.zeros(3)
                  ppf7d = make_ppf(pts9d, cen9d)  # (N,9), (9,)
                  ppf7d[np.isnan(ppf7d)] = 0.0
                  nb_in[sid] = torch.Tensor(ppf7d).transpose(1, 0)
                  # remove background by 3d radius
                  xyz = pts9d[:, 0:3]
                  ids_out_of_bound = np.linalg.norm(xyz, axis=1) > self.far_radius
                  nb_in[sid, :, ids_out_of_bound] = 0.
                  # count
                  sid += 1
                  if sid >= self.n_pts_per_frame:
                      break
          pt_in = pt_in[:sid]
          nb_in = nb_in[:sid]
          route_labs = route_labs[:sid]
          return pt_in, nb_in, route_labs, fid, torch.Tensor(np.array(rc_list))
  
  
  class TestDataset_PPF(Dataset):
  
      def __init__(self, data_dir, the_list, n_pts_per_frame=opt.n_pts_per_frame, neighbor_da2d=None):
          super().__init__()
          self.data_dir, self.the_list = data_dir, the_list
          self.n_pts_per_frame = n_pts_per_frame
          self.neighbor_da2d = neighbor_da2d  # list of (n_pts, dim_pt)
  
      def __len__(self):
          return len(self.the_list)
  
      def __getitem__(self, idx):
          fid = self.the_list[idx]
          # load 
          depth = cv2.imread('{}/{}'.format(self.data_dir, opt.depth_format.format(fid)), cv2.IMREAD_UNCHANGED) / 1000.0
          color = cv2.imread('{}/{}'.format(self.data_dir, opt.color_format.format(fid)), cv2.IMREAD_UNCHANGED)[:, :, 0:3]
          if np.max(color) > 1:
              color = color / 255. - 0.5
          local = depth2local(depth)
          r_ids, c_ids = list(range(0, image_height, 4)), list(range(0, image_width, 4))
          depth, color, local = depth[r_ids, :][:, c_ids], color[r_ids, :, :][:, c_ids, :], local[r_ids, :, :][:, c_ids,
                                                                                            :]
          # normal by 3d neighbor plane fitting.
          normal_path = '{}/frame-{:06d}.scaled.normal.npy'.format(self.data_dir, fid)
          if os.path.exists(normal_path):
              # normal = np.load(normal_path)
              if os.path.getsize(normal_path) > 1:
                  normal = np.load(normal_path, encoding='bytes', allow_pickle=True)
              else:
                  normal = compute_points_normal(local)
                  np.save(normal_path, normal)
          else:
              normal = compute_points_normal(local)
              np.save(normal_path, normal)
          lclnmlclr = np.concatenate((np.concatenate((local, normal), axis=2), color), axis=2)
          # build a patch
          pt_in = torch.zeros((self.n_pts_per_frame, 7, 1))
          nb_ms_in = torch.zeros((self.n_pts_per_frame, opt.tree_height - 1, 7, opt.n_neighb_pts))
          route_labs = torch.zeros((self.n_pts_per_frame, opt.tree_height - 1))
          r_max, c_max = int(image_height / 4 - 1), int(image_width / 4 - 1)
          rc_list = []
          # da2d+3d neighbor
          if not self.neighbor_da2d is None:
              sid, count_crt, count_max = 0, 0, 9999
              mask = np.zeros((r_max, c_max))
              while len(rc_list) < self.n_pts_per_frame:
                  # avoid infinite loop
                  count_crt += 1
                  if count_crt > count_max:
                      break
                  r, c = np.random.randint(0, r_max), np.random.randint(0, c_max)
                  if depth[r, c] == 0. or mask[r, c] == 1.:
                      continue
                  if np.isnan(lclnmlclr[r, c, 3]):
                      continue
                  mask[r, c] = 1.
                  rc_list.append([r, c])
                  pt_in[sid] = torch.Tensor(
                      np.concatenate((np.array([[0.], [0.], [0.], [0.]]), color[r, c, 0:3][:, np.newaxis]), axis=0))
                  for lid in range(opt.tree_height - 1):
                      da2d_list = (np.array(self.neighbor_da2d[lid]) / depth[r, c]).astype(int)
                      # ppf
                      pts9d = shift_pts9d(
                          sample_pts(partial_pts_2d(lclnmlclr, (r, c), da2d_list), opt.n_neighb_pts),
                          lclnmlclr[r, c, :])
                      cen9d = copy.deepcopy(lclnmlclr[r, c, :])
                      cen9d[0:3] = np.zeros(3)
                      ppf7d = make_ppf(pts9d, cen9d)  # (N,9), (9,)
                      ppf7d[np.isnan(ppf7d)] = 0.0
                      nb_ms_in[sid, lid, :, :] = torch.Tensor(ppf7d).transpose(1, 0)
                      # remove background by 3d radius
                      xyz = pts9d[:, 0:3]
                      ids_out_of_bound = np.linalg.norm(xyz, axis=1) > opt.far_radiuses[lid]
                      nb_ms_in[sid, lid, :, ids_out_of_bound] = 0.
                  # count
                  sid += 1
          return pt_in, nb_ms_in, -1, fid, torch.Tensor(np.array(rc_list))
  
  
  # # debug
  # if __name__ == '__main__':
  # 	print('done.')
  
  
apis: numpy.sum numpy.save numpy.stack numpy.logical_and numpy.random.rand numpy.isnan numpy.linalg.eig torch.Tensor numpy.mean numpy.load numpy.zeros torch.zeros numpy.arange numpy.max numpy.min numpy.linalg.norm numpy.random.shuffle numpy.array numpy.concatenate numpy.random.randint