build_coordinates.py

#!/usr/bin/env python
# Copyright 2017-2023 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Builds a TFRecord file of coordinates for training.

Use ./compute_partitions.py to generate data for --partition_volumes.
Note that the volume names you provide in --partition_volumes will
have to match the volume labels you pass to the training script.
"""
from collections import defaultdict

from absl import app
from absl import flags
from absl import logging

import h5py
import numpy as np
import tensorflow.compat.v1 as tf

FLAGS = flags.FLAGS

flags.DEFINE_list('partition_volumes', None,
                  'Partition volumes as '
                  '<volume_name>:<volume_path>:<dataset>, where volume_path '
                  'points to a HDF5 volume, and <volume_name> is an arbitrary '
                  'label that will have to also be used during training.')
flags.DEFINE_string('coordinate_output', None,
                    'Path to a TF Record file in which to save the '
                    'coordinates.')
flags.DEFINE_list('margin', None, '(z, y, x) tuple specifying the '
                  'number of voxels adjacent to the border of the volume to '
                  'exclude from sampling. This should normally be set to the '
                  'radius of the FFN training FoV (i.e. network FoV radius '
                  '+ deltas.')


IGNORE_PARTITION = 255


def _int64_feature(values):
  return tf.train.Feature(int64_list=tf.train.Int64List(value=values))


def _bytes_feature(value):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def main(argv):
  del argv  # Unused.

  totals = defaultdict(int)  # partition -> voxel count
  indices = defaultdict(list)  # partition -> [(vol_id, 1d index)]

  vol_labels = []
  vol_shapes = []
  mz, my, mx = [int(x) for x in FLAGS.margin]

  for i, partvol in enumerate(FLAGS.partition_volumes):
    name, path, dataset = partvol.split(':')
    with h5py.File(path, 'r') as f:
      partitions = f[dataset][mz:-mz, my:-my, mx:-mx]
      vol_shapes.append(partitions.shape)
      vol_labels.append(name)

      uniques, counts = np.unique(partitions, return_counts=True)
      for val, cnt in zip(uniques, counts):
        if val == IGNORE_PARTITION:
          continue

        totals[val] += cnt
        indices[val].extend(
            [(i, flat_index) for flat_index in
             np.flatnonzero(partitions == val)])

  logging.info('Partition counts:')
  for k, v in totals.items():
    logging.info(' %d: %d', k, v)

  logging.info('Resampling and shuffling coordinates.')

  max_count = max(totals.values())
  indices = np.concatenate(
      [np.resize(np.random.permutation(v), (max_count, 2)) for
       v in indices.values()], axis=0)
  np.random.shuffle(indices)

  logging.info('Saving coordinates.')
  record_options = tf.python_io.TFRecordOptions(
      tf.python_io.TFRecordCompressionType.GZIP)
  with tf.python_io.TFRecordWriter(FLAGS.coordinate_output,
                                   options=record_options) as writer:
    for i, coord_idx in indices:
      z, y, x = np.unravel_index(coord_idx, vol_shapes[i])

      coord = tf.train.Example(features=tf.train.Features(feature=dict(
          center=_int64_feature([mx + x, my + y, mz + z]),
          label_volume_name=_bytes_feature(vol_labels[i].encode('utf-8'))
      )))
      writer.write(coord.SerializeToString())


if __name__ == '__main__':
  flags.mark_flag_as_required('margin')
  flags.mark_flag_as_required('coordinate_output')
  flags.mark_flag_as_required('partition_volumes')

  app.run(main)