prepare_data.py

import argparse
import math
import os
from datetime import datetime

import h5py
import numpy as np
import plyfile
from matplotlib import cm


def prepare_label(data_dir, output_dir):
  object_dict = {
    "clutter": 0,
    "ceiling": 1,
    "floor": 2,
    "wall": 3,
    "beam": 4,
    "column": 5,
    "door": 6,
    "window": 7,
    "table": 8,
    "chair": 9,
    "sofa": 10,
    "bookcase": 11,
    "board": 12,
  }

  for area in os.listdir(data_dir):
    path_area = os.path.join(data_dir, area)
    if not os.path.isdir(path_area):
      continue

    path_dir_rooms = os.listdir(path_area)
    for room in path_dir_rooms:
      path_annotations = os.path.join(data_dir, area, room, "Annotations")
      if not os.path.isdir(path_annotations):
        continue

      print(path_annotations)
      path_prepare_label = os.path.join(output_dir, area, room)
      if os.path.exists(os.path.join(path_prepare_label, ".labels")):
        print(f"{path_prepare_label} already processed, skipping")
        continue

      xyz_room = np.zeros((1, 6))
      label_room = np.zeros((1, 1))
      # make store directories
      if not os.path.exists(path_prepare_label):
        os.makedirs(path_prepare_label)

      path_objects = os.listdir(path_annotations)
      for obj in path_objects:
        object_key = obj.split("_", 1)[0]
        try:
          val = object_dict[object_key]
        except KeyError:
          continue
        print(f"{room}/{obj[:-4]}")
        xyz_object_path = os.path.join(path_annotations, obj)
        try:
          xyz_object = np.loadtxt(xyz_object_path)[:, :]  # (N,6)
        except ValueError as e:
          print(f"ERROR: cannot load {xyz_object_path}: {e}")
          continue
        label_object = np.tile(val, (xyz_object.shape[0], 1))  # (N,1)
        xyz_room = np.vstack((xyz_room, xyz_object))
        label_room = np.vstack((label_room, label_object))

      xyz_room = np.delete(xyz_room, [0], 0)
      label_room = np.delete(label_room, [0], 0)

      np.save(path_prepare_label + "/xyzrgb.npy", xyz_room)
      np.save(path_prepare_label + "/label.npy", label_room)

      # Marker indicating we've processed this room
      open(os.path.join(path_prepare_label, ".labels"), "w").close()


def main():
  default_data_dir = "data/s3dis/Stanford3dDataset_v1.2_Aligned_Version"
  default_output_dir = "data/s3dis/pointcnn"
  parser = argparse.ArgumentParser()
  parser.add_argument(
    "-d",
    "--data",
    dest="data_dir",
    default=default_data_dir,
    help=f"Path to S3DIS data (default is {default_data_dir})",
  )
  parser.add_argument(
    "-f",
    "--folder",
    dest="output_dir",
    default=default_output_dir,
    help=f"Folder to write labels (default is {default_output_dir})",
  )
  parser.add_argument(
    "--max_num_points",
    "-m",
    help="Max point number of each sample",
    type=int,
    default=8192,
  )
  parser.add_argument(
    "--block_size", "-b", help="Block size", type=float, default=1.5
  )
  parser.add_argument(
    "--grid_size", "-g", help="Grid size", type=float, default=0.03
  )
  parser.add_argument(
    "--save_ply", "-s", help="Convert .pts to .ply", action="store_true"
  )

  args = parser.parse_args()
  print(args)

  prepare_label(data_dir=args.data_dir, output_dir=args.output_dir)

  root = args.output_dir
  max_num_points = args.max_num_points

  batch_size = 2048
  data = np.zeros((batch_size, max_num_points, 9))
  data_num = np.zeros(batch_size, dtype=np.int32)
  label = np.zeros(batch_size, dtype=np.int32)
  label_seg = np.zeros((batch_size, max_num_points), dtype=np.int32)
  indices_split_to_full = np.zeros((batch_size, max_num_points), dtype=np.int32)

  for area_idx in range(1, 7):
    folder = os.path.join(root, "Area_%d" % area_idx)
    datasets = [dataset for dataset in os.listdir(folder)]
    for dataset_idx, dataset in enumerate(datasets):
      dataset_marker = os.path.join(folder, dataset, ".dataset")
      if os.path.exists(dataset_marker):
        print(
          f"{datetime.now()}-{folder}/{dataset} already processed, skipping"
        )
        continue
      filename_data = os.path.join(folder, dataset, "xyzrgb.npy")
      print(f"{datetime.now()}-Loading {filename_data}...")
      # Modified according to PointNet convensions.
      xyzrgb = np.load(filename_data)
      xyzrgb[:, 0:3] -= np.amin(xyzrgb, axis=0)[0:3]

      filename_labels = os.path.join(folder, dataset, "label.npy")
      print(f"{datetime.now()}-Loading {filename_labels}...")
      labels = np.load(filename_labels).astype(int).flatten()

      xyz, rgb = np.split(xyzrgb, [3], axis=-1)
      xyz_min = np.amin(xyz, axis=0, keepdims=True)
      xyz_max = np.amax(xyz, axis=0, keepdims=True)
      xyz_center = (xyz_min + xyz_max) / 2
      xyz_center[0][-1] = xyz_min[0][-1]
      # Remark: Don't do global alignment.
      # xyz = xyz - xyz_center
      rgb = rgb / 255.0
      max_room_x = np.max(xyz[:, 0])
      max_room_y = np.max(xyz[:, 1])
      max_room_z = np.max(xyz[:, 2])

      offsets = [("zero", 0.0), ("half", args.block_size / 2)]
      for offset_name, offset in offsets:
        idx_h5 = 0
        idx = 0

        print(
          f"{datetime.now()}-Computing block id of {xyzrgb.shape[0]} points..."
        )
        xyz_min = np.amin(xyz, axis=0, keepdims=True) - offset
        xyz_max = np.amax(xyz, axis=0, keepdims=True)
        block_size = (
          args.block_size,
          args.block_size,
          2 * (xyz_max[0, -1] - xyz_min[0, -1]),
        )
        # Note: Don't split over z axis.
        xyz_blocks = np.floor((xyz - xyz_min) / block_size).astype(np.int)

        print(f"{datetime.now()}-Collecting points belong to each block...")
        blocks, point_block_indices, block_point_counts = np.unique(
          xyz_blocks, return_inverse=True, return_counts=True, axis=0
        )
        block_point_indices = np.split(
          np.argsort(point_block_indices), np.cumsum(block_point_counts[:-1])
        )
        print(
          f"{datetime.now()}-{dataset} is split into {blocks.shape[0]} blocks."
        )

        block_to_block_idx_map = dict()
        for block_idx in range(blocks.shape[0]):
          block = (blocks[block_idx][0], blocks[block_idx][1])
          block_to_block_idx_map[(block[0], block[1])] = block_idx

        # merge small blocks into one of their big neighbors
        block_point_count_threshold = max_num_points / 10
        nbr_block_offsets = [
          (0, 1),
          (1, 0),
          (0, -1),
          (-1, 0),
          (-1, 1),
          (1, 1),
          (1, -1),
          (-1, -1),
        ]
        block_merge_count = 0
        for block_idx in range(blocks.shape[0]):
          if block_point_counts[block_idx] >= block_point_count_threshold:
            continue

          block = (blocks[block_idx][0], blocks[block_idx][1])
          for x, y in nbr_block_offsets:
            nbr_block = (block[0] + x, block[1] + y)
            if nbr_block not in block_to_block_idx_map:
              continue

            nbr_block_idx = block_to_block_idx_map[nbr_block]
            if block_point_counts[nbr_block_idx] < block_point_count_threshold:
              continue

            block_point_indices[nbr_block_idx] = np.concatenate(
              [
                block_point_indices[nbr_block_idx],
                block_point_indices[block_idx],
              ],
              axis=-1,
            )
            block_point_indices[block_idx] = np.array([], dtype=np.int)
            block_merge_count = block_merge_count + 1
            break
        print(
          f"{datetime.now()}-{block_merge_count} of {blocks.shape[0]} blocks are merged."
        )

        idx_last_non_empty_block = 0
        for block_idx in reversed(range(blocks.shape[0])):
          if block_point_indices[block_idx].shape[0] != 0:
            idx_last_non_empty_block = block_idx
            break

        # uniformly sample each block
        for block_idx in range(idx_last_non_empty_block + 1):
          point_indices = block_point_indices[block_idx]
          if point_indices.shape[0] == 0:
            continue
          block_points = xyz[point_indices]
          block_min = np.amin(block_points, axis=0, keepdims=True)
          xyz_grids = np.floor(
            (block_points - block_min) / args.grid_size
          ).astype(np.int)
          grids, point_grid_indices, grid_point_counts = np.unique(
            xyz_grids, return_inverse=True, return_counts=True, axis=0
          )
          grid_point_indices = np.split(
            np.argsort(point_grid_indices), np.cumsum(grid_point_counts[:-1])
          )
          grid_point_count_avg = int(np.average(grid_point_counts))
          point_indices_repeated = []
          for grid_idx in range(grids.shape[0]):
            point_indices_in_block = grid_point_indices[grid_idx]
            repeat_num = math.ceil(
              grid_point_count_avg / point_indices_in_block.shape[0]
            )
            if repeat_num > 1:
              point_indices_in_block = np.repeat(
                point_indices_in_block, repeat_num
              )
              np.random.shuffle(point_indices_in_block)
              point_indices_in_block = point_indices_in_block[
                :grid_point_count_avg
              ]
            point_indices_repeated.extend(
              list(point_indices[point_indices_in_block])
            )
          block_point_indices[block_idx] = np.array(point_indices_repeated)
          block_point_counts[block_idx] = len(point_indices_repeated)

        for block_idx in range(idx_last_non_empty_block + 1):
          point_indices = block_point_indices[block_idx]
          if point_indices.shape[0] == 0:
            continue

          block_point_num = point_indices.shape[0]
          block_split_num = int(
            math.ceil(block_point_num * 1.0 / max_num_points)
          )
          point_num_avg = int(
            math.ceil(block_point_num * 1.0 / block_split_num)
          )
          point_nums = [point_num_avg] * block_split_num
          point_nums[-1] = block_point_num - (
            point_num_avg * (block_split_num - 1)
          )
          starts = [0] + list(np.cumsum(point_nums))

          # Modified following convensions of PointNet.
          np.random.shuffle(point_indices)
          block_points = xyz[point_indices]
          block_rgb = rgb[point_indices]
          block_labels = labels[point_indices]
          x, y, z = np.split(block_points, (1, 2), axis=-1)
          norm_x = x / max_room_x
          norm_y = y / max_room_y
          norm_z = z / max_room_z

          minx = np.min(x)
          miny = np.min(y)
          x = x - (minx + args.block_size / 2)
          y = y - (miny + args.block_size / 2)

          block_xyzrgb = np.concatenate(
            [x, y, z, block_rgb, norm_x, norm_y, norm_z], axis=-1
          )
          for block_split_idx in range(block_split_num):
            start = starts[block_split_idx]
            point_num = point_nums[block_split_idx]
            end = start + point_num
            idx_in_batch = idx % batch_size
            data[idx_in_batch, 0:point_num, ...] = block_xyzrgb[start:end, :]
            data_num[idx_in_batch] = point_num
            label[idx_in_batch] = dataset_idx  # won't be used...
            label_seg[idx_in_batch, 0:point_num] = block_labels[start:end]
            indices_split_to_full[idx_in_batch, 0:point_num] = point_indices[
              start:end
            ]

            if ((idx + 1) % batch_size == 0) or (
              block_idx == idx_last_non_empty_block
              and block_split_idx == block_split_num - 1
            ):
              item_num = idx_in_batch + 1
              filename_h5 = os.path.join(
                folder, dataset, f"{offset_name}_{idx_h5:d}.h5"
              )
              print(f"{datetime.now()}-Saving {filename_h5}...")

              file = h5py.File(filename_h5, "w")
              file.create_dataset("data", data=data[0:item_num, ...])
              file.create_dataset("data_num", data=data_num[0:item_num, ...])
              file.create_dataset("label", data=label[0:item_num, ...])
              file.create_dataset("label_seg", data=label_seg[0:item_num, ...])
              file.create_dataset(
                "indices_split_to_full",
                data=indices_split_to_full[0:item_num, ...],
              )
              file.close()

              if args.save_ply:
                print(f"{datetime.now()}-Saving ply of {filename_h5}...")
                filepath_label_ply = os.path.join(
                  folder,
                  dataset,
                  "ply_label",
                  f"label_{offset_name}_{idx_h5:d}",
                )
                save_ply_property_batch(
                  data[0:item_num, :, 0:3],
                  label_seg[0:item_num, ...],
                  filepath_label_ply,
                  data_num[0:item_num, ...],
                  14,
                )

                filepath_rgb_ply = os.path.join(
                  folder, dataset, "ply_rgb", f"rgb_{offset_name}_{idx_h5:d}"
                )
                save_ply_color_batch(
                  data[0:item_num, :, 0:3],
                  data[0:item_num, :, 3:6] * 255,
                  filepath_rgb_ply,
                  data_num[0:item_num, ...],
                )

              idx_h5 = idx_h5 + 1
            idx = idx + 1

      # Marker indicating we've processed this dataset
      open(dataset_marker, "w").close()
  print(f"{datetime.now()}-Done.")


def save_ply(points, filename, colors=None, normals=None):
  vertex = np.core.records.fromarrays(
    points.transpose(), names="x, y, z", formats="f4, f4, f4"
  )
  n = len(vertex)
  desc = vertex.dtype.descr

  if normals is not None:
    vertex_normal = np.core.records.fromarrays(
      normals.transpose(), names="nx, ny, nz", formats="f4, f4, f4"
    )
    assert len(vertex_normal) == n
    desc = desc + vertex_normal.dtype.descr

  if colors is not None:
    vertex_color = np.core.records.fromarrays(
      colors.transpose() * 255, names="red, green, blue", formats="u1, u1, u1"
    )
    assert len(vertex_color) == n
    desc = desc + vertex_color.dtype.descr

  vertex_all = np.empty(n, dtype=desc)

  for prop in vertex.dtype.names:
    vertex_all[prop] = vertex[prop]

  if normals is not None:
    for prop in vertex_normal.dtype.names:
      vertex_all[prop] = vertex_normal[prop]

  if colors is not None:
    for prop in vertex_color.dtype.names:
      vertex_all[prop] = vertex_color[prop]

  ply = plyfile.PlyData(
    [plyfile.PlyElement.describe(vertex_all, "vertex")], text=False
  )
  if not os.path.exists(os.path.dirname(filename)):
    os.makedirs(os.path.dirname(filename))
  ply.write(filename)


def save_ply_property(
  points, property, property_max, filename, cmap_name="tab20"
):
  point_num = points.shape[0]
  colors = np.full(points.shape, 0.5)
  cmap = cm.get_cmap(cmap_name)
  for point_idx in range(point_num):
    if property[point_idx] == 0:
      colors[point_idx] = np.array([0, 0, 0])
    else:
      colors[point_idx] = cmap(property[point_idx] / property_max)[:3]
  save_ply(points, filename, colors)


def save_ply_color_batch(
  points_batch, colors_batch, file_path, points_num=None
):
  batch_size = points_batch.shape[0]
  if not isinstance(file_path, (list, tuple)):
    basename = os.path.splitext(file_path)[0]
    ext = ".ply"
  for batch_idx in range(batch_size):
    point_num = (
      points_batch.shape[1] if points_num is None else points_num[batch_idx]
    )
    if isinstance(file_path, (list, tuple)):
      save_ply(
        points_batch[batch_idx][:point_num],
        file_path[batch_idx],
        colors_batch[batch_idx][:point_num],
      )
    else:
      save_ply(
        points_batch[batch_idx][:point_num],
        f"{basename}_{batch_idx:04d}{ext}",
        colors_batch[batch_idx][:point_num],
      )


def save_ply_property_batch(
  points_batch,
  property_batch,
  file_path,
  points_num=None,
  property_max=None,
  cmap_name="tab20",
):
  batch_size = points_batch.shape[0]
  if not isinstance(file_path, (list, tuple)):
    basename = os.path.splitext(file_path)[0]
    ext = ".ply"
  property_max = (
    np.max(property_batch) if property_max is None else property_max
  )
  for batch_idx in range(batch_size):
    point_num = (
      points_batch.shape[1] if points_num is None else points_num[batch_idx]
    )
    if isinstance(file_path, (list, tuple)):
      save_ply_property(
        points_batch[batch_idx][:point_num],
        property_batch[batch_idx][:point_num],
        property_max,
        file_path[batch_idx],
        cmap_name,
      )
    else:
      save_ply_property(
        points_batch[batch_idx][:point_num],
        property_batch[batch_idx][:point_num],
        property_max,
        f"{basename}_{batch_idx:04d}{ext}",
        cmap_name,
      )


if __name__ == "__main__":
  main()