<NME> resegmentation.py
<BEF> # Copyright 2017 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.
# ==============================================================================
"""Helper functions for resegmentation.

Resegmentation is local segmentation targeted to specific points in an already
segmented volume. The results of resegmentation can be compared to the original
segments in order to perform object agglomeration.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import hashlib
import logging
import os

import numpy as np

from scipy import ndimage
from scipy.special import expit

from tensorflow import gfile

from . import storage
from .inference_utils import timer_counter
from ..utils import bounding_box


def get_starting_location(dists, exclusion_radius):
  z, y, x = np.unravel_index(np.argmax(dists), tuple(dists.shape))
  # Mark area around the new point as 'excluded' by clearing the distance
  # map around it.
  er = exclusion_radius
  dists[max(z - er.z, 0):z + er.z + 1,
        max(y - er.y, 0):y + er.y + 1,
        max(x - er.x, 0):x + er.x + 1] = 0
  return z, y, x


def get_target_path(request, point_num):
  """Computes the output path for a specific point.

  Args:
    request: ResegmentationRequest proto
    point_num: index of the point of interest within the proto

  Returns:
    path to the output file where resegmentation results will be saved
  """
  # Prepare the output directory.
  output_dir = request.output_directory

  id_a = request.points[point_num].id_a
  id_b = request.points[point_num].id_b

  if request.subdir_digits > 1:
    m = hashlib.md5()
    m.update(str(id_a))
    m.update(str(id_b))
    output_dir = os.path.join(output_dir, m.hexdigest()[:request.subdir_digits])
  gfile.MakeDirs(output_dir)

  # Terminate early if the output already exists.
  dp = request.points[point_num].point
  target_path = os.path.join(output_dir, '%d-%d_at_%d_%d_%d.npz' % (
      id_a, id_b, dp.x, dp.y, dp.z))
  if gfile.Exists(target_path):
    logging.info('Output already exists: %s', target_path)
    return

  return target_path


def get_canvas(point, radius, runner):
  """Creates an FFN Canvas.

  Args:
    point: decision point as (z, y, x)
    radius: radius around decision point as (z, y, x)
    runner: inference Runner object

  Returns:
    inference Canvas object
  """

  origin = np.array(point)
  radius = np.array(radius)
  corner = origin - radius
  subvol_size = radius * 2 + 1
  end = subvol_size + corner

  if (np.any(corner < 0) or
      runner.init_seg_volume.shape[1] <= end[0] or
      runner.init_seg_volume.shape[2] <= end[1] or
      runner.init_seg_volume.shape[3] <= end[2]):
                  point[2], point[1], point[0], corner, end)
    return None, None

  return runner.make_canvas(corner, subvol_size, keep_history=True,
                            relabel_init_ids=False)


def process_point(request, runner, point_num):
def process_point(request, runner, point_num, voxel_size):
  """Runs resegmentation for a specific point.

  Args:
    runner: inference Runner object
    point_num: index of the point of interest within the proto
  """
  runner.counters.reset()
  with timer_counter(runner.counters, 'resegmentation'):
    target_path = get_target_path(request, point_num)
    if target_path is None:
    target_path = get_target_path(request, point_num)
    if target_path is None:
      return

    curr = request.points[point_num]
    point = curr.point
    point = point.z, point.y, point.x
    radius = (request.radius.z, request.radius.y, request.radius.x)
      logging.warning('Could not get a canvas object.')
      return

    def unalign_image(img):
      return alignment.align_and_crop(canvas.corner_zyx,
                                      img,
                                      alignment.corner,
                                      alignment.size,
                                      forward=False)

    is_shift = (canvas.restrictor is not None and
                np.any(canvas.restrictor.shift_mask))
    is_endpoint = not curr.HasField('id_b')

    seg_a = canvas.segmentation == curr.id_a
    size_a = np.sum(seg_a)

    if is_endpoint:
      size_b = -1
      todo = [(seg_a, size_a)]
    else:
      seg_b = canvas.segmentation == curr.id_b
      size_b = np.sum(seg_b)
      todo = [(seg_a, size_a), (seg_b, size_b)]

    if size_a == 0 or size_b == 0:
      logging.warning('Segments not found in input. Present values are: %r.',
                      np.unique(canvas.segmentation))
      return

    if is_endpoint:

    if is_endpoint:
      canvas.seg_prob[:] = 0.0
      canvas.segmentation[:] = 0
    else:
      # Clear the two segments in question, but keep everything else as
      # context.
      canvas.segmentation[seg_a] = 0
      canvas.segmentation[seg_b] = 0
      canvas.seg_prob[seg_a] = 0.0
      canvas.seg_prob[seg_b] = 0.0

    transformed_point = alignment.transform(np.array([point]).T)
    tz, ty, tx = transformed_point[:, 0]
    oz, oy, ox = canvas.corner_zyx
    tz -= oz
    ty -= oy
    tx -= ox

    # First index enumerates the original segments. Second index,
    # when present, enumerates segmentation attempts.
    raw_probs = []
    histories = []
    start_points = [[], []]

    options = request.inference.inference_options
    for i, (seg, start_size) in enumerate(todo):
      logging.info('processing object %d', i)

      with timer_counter(canvas.counters, 'edt'):

      with timer_counter(canvas.counters, 'edt'):
        dists = ndimage.distance_transform_edt(seg, sampling=voxel_size)
        # Do not seed where not enough context is available.
        dists[:canvas.margin[0], :, :] = 0
        dists[:, :canvas.margin[1], :] = 0
        dists[-canvas.margin[0]:, :, :] = 0
        dists[:, -canvas.margin[1]:, :] = 0
        dists[:, :, -canvas.margin[2]:] = 0
        logging.info('EDT computation done')

      # Optionally exclude a region around the decision point from seeding.
      if request.HasField('init_exclusion_radius'):
      if request.HasField('init_exclusion_radius'):
        ier = request.init_exclusion_radius
        dists[tz - ier.z:tz + ier.z + 1,
              ty - ier.y:ty + ier.y + 1,
              tx - ier.x:tx + ier.x + 1] = 0

      seg_prob = None
      recovered = False

      for _ in range(request.max_retry_iters):
        if not seg[z0, y0, x0]:
          continue

        logging.info('.. starting segmentation at (xyz): %d %d %d', x0, y0, z0)
        canvas.segment_at((z0, y0, x0))
        seg_prob = expit(canvas.seed)
        start_points[i].append((x0, y0, z0))
        start_points[i].append((x0, y0, z0))
        # Check if we recovered an acceptable fraction of the initial segment
        # in which the seed was located.
        recovered = True
        segmented_voxels = np.sum((seg_prob >= options.segment_threshold) & seg)
        if request.segment_recovery_fraction > 0:
          if segmented_voxels / start_size >= request.segment_recovery_fraction:
            break
            break
        elif segmented_voxels >= options.min_segment_size:
          break

        recovered = False

      if seg_prob is not None:
        qprob = storage.quantize_probability(seg_prob)
        raw_probs.append(qprob)
        probs.append(unalign_image(qprob))
        deletes.append(np.array(canvas.history_deleted))
        histories.append(np.array(canvas.history))


      if request.terminate_early:
          break
        if (request.segment_recovery_fraction > 0 and i == 0 and
            len(todo) > 1):
          seg2, size2 = todo[1]
          segmented_voxels2 = np.sum(
              (seg_prob >= options.segment_threshold) & seg2)
          if segmented_voxels2 / size2 < request.segment_recovery_fraction:
            break

  logging.info('saving results to %s', target_path)
  with storage.atomic_file(target_path) as fd:
    np.savez_compressed(fd,
                        probs=np.array(probs),
                        probs=np.array(probs),
                        raw_probs=np.array(raw_probs),
                        deletes=np.array(deletes),
                        histories=np.array(histories),
                        start_points=start_points,
                        counters=canvas.counters.dumps(),
                        corner_zyx=canvas.corner_zyx,
                        is_shift=is_shift)
  logging.info('.. save complete')


def process(request, runner):
def process(request, runner):
  num_points = len(request.points)
  for i in range(num_points):
    logging.info('processing %d/%d', i, num_points)
    process_point(request, runner, i)


<MSG> Add support for the executor interface in inference.

<DFF> @@ -108,8 +108,7 @@ def get_canvas(point, radius, runner):
                   point[2], point[1], point[0], corner, end)
     return None, None
 
-  return runner.make_canvas(corner, subvol_size, keep_history=True,
-                            relabel_init_ids=False)
+  return runner.make_canvas(corner, subvol_size, keep_history=True)
 
 
 def process_point(request, runner, point_num):
@@ -120,7 +119,6 @@ def process_point(request, runner, point_num):
     runner: inference Runner object
     point_num: index of the point of interest within the proto
   """
-  runner.counters.reset()
   with timer_counter(runner.counters, 'resegmentation'):
     target_path = get_target_path(request, point_num)
     if target_path is None:
@@ -135,31 +133,36 @@ def process_point(request, runner, point_num):
       logging.warning('Could not get a canvas object.')
       return
 
-    def unalign_image(img):
-      return alignment.align_and_crop(canvas.corner_zyx,
-                                      img,
-                                      alignment.corner,
-                                      alignment.size,
-                                      forward=False)
+    def unalign_prob(prob):
+      return alignment.align_and_crop(
+          canvas.corner_zyx,
+          prob,
+          alignment.corner,
+          alignment.size,
+          forward=False)
 
     is_shift = (canvas.restrictor is not None and
                 np.any(canvas.restrictor.shift_mask))
     is_endpoint = not curr.HasField('id_b')
 
-    seg_a = canvas.segmentation == curr.id_a
+    seg_a = canvas.segmentation == canvas.local_id(curr.id_a)
     size_a = np.sum(seg_a)
 
     if is_endpoint:
       size_b = -1
-      todo = [(seg_a, size_a)]
+      todo = [seg_a]
     else:
-      seg_b = canvas.segmentation == curr.id_b
+      seg_b = canvas.segmentation == canvas.local_id(curr.id_b)
       size_b = np.sum(seg_b)
-      todo = [(seg_a, size_a), (seg_b, size_b)]
+      todo = [seg_a, seg_b]
 
     if size_a == 0 or size_b == 0:
-      logging.warning('Segments not found in input. Present values are: %r.',
+      logging.warning('Segments (%d, %d) local ids (%d, %d) not found in input '
+                      'at %r.  Current values are: %r.',
+                      curr.id_a, curr.id_b, canvas.local_id(curr.id_a),
+                      canvas.local_id(curr.id_b), point,
                       np.unique(canvas.segmentation))
+      canvas._deregister_client()  # pylint:disable=protected-access
       return
 
     if is_endpoint:
@@ -188,8 +191,19 @@ def process_point(request, runner, point_num):
     histories = []
     start_points = [[], []]
 
+    if request.HasField('analysis_radius'):
+      ar = request.analysis_radius
+      analysis_box = bounding_box.BoundingBox(
+          start=(radius[2] - ar.x,
+                 radius[1] - ar.y,
+                 radius[0] - ar.z),
+          size=(2 * ar.x + 1, 2 * ar.y + 1, 2 * ar.z + 1))
+    else:
+      analysis_box = bounding_box.BoundingBox(
+          (0, 0, 0), canvas.image.shape[::-1])
+
     options = request.inference.inference_options
-    for i, (seg, start_size) in enumerate(todo):
+    for i, seg in enumerate(todo):
       logging.info('processing object %d', i)
 
       with timer_counter(canvas.counters, 'edt'):
@@ -202,7 +216,7 @@ def process_point(request, runner, point_num):
         dists[-canvas.margin[0]:, :, :] = 0
         dists[:, -canvas.margin[1]:, :] = 0
         dists[:, :, -canvas.margin[2]:] = 0
-        logging.info('EDT computation done')
+        canvas.log_info('EDT computation done')
 
       # Optionally exclude a region around the decision point from seeding.
       if request.HasField('init_exclusion_radius'):
@@ -219,7 +233,8 @@ def process_point(request, runner, point_num):
         if not seg[z0, y0, x0]:
           continue
 
-        logging.info('.. starting segmentation at (xyz): %d %d %d', x0, y0, z0)
+        canvas.log_info('.. starting segmentation at (xyz): %d %d %d',
+                        x0, y0, z0)
         canvas.segment_at((z0, y0, x0))
         seg_prob = expit(canvas.seed)
         start_points[i].append((x0, y0, z0))
@@ -227,7 +242,12 @@ def process_point(request, runner, point_num):
         # Check if we recovered an acceptable fraction of the initial segment
         # in which the seed was located.
         recovered = True
-        segmented_voxels = np.sum((seg_prob >= options.segment_threshold) & seg)
+
+        crop_seg = seg[analysis_box.to_slice()]
+        crop_prob = seg_prob[analysis_box.to_slice()]
+        start_size = np.sum(crop_seg)
+        segmented_voxels = np.sum((crop_prob >= options.segment_threshold) &
+                                  crop_seg)
         if request.segment_recovery_fraction > 0:
           if segmented_voxels / start_size >= request.segment_recovery_fraction:
             break
@@ -240,7 +260,7 @@ def process_point(request, runner, point_num):
       if seg_prob is not None:
         qprob = storage.quantize_probability(seg_prob)
         raw_probs.append(qprob)
-        probs.append(unalign_image(qprob))
+        probs.append(unalign_prob(qprob))
         deletes.append(np.array(canvas.history_deleted))
         histories.append(np.array(canvas.history))
 
@@ -249,13 +269,16 @@ def process_point(request, runner, point_num):
           break
         if (request.segment_recovery_fraction > 0 and i == 0 and
             len(todo) > 1):
-          seg2, size2 = todo[1]
+          seg2 = todo[1]
+          crop_seg = seg2[analysis_box.to_slice()]
+          size2 = np.sum(crop_seg)
           segmented_voxels2 = np.sum(
-              (seg_prob >= options.segment_threshold) & seg2)
+              (crop_prob >= options.segment_threshold) & crop_seg)
+
           if segmented_voxels2 / size2 < request.segment_recovery_fraction:
             break
 
-  logging.info('saving results to %s', target_path)
+  canvas.log_info('saving results to %s', target_path)
   with storage.atomic_file(target_path) as fd:
     np.savez_compressed(fd,
                         probs=np.array(probs),
@@ -267,7 +290,10 @@ def process_point(request, runner, point_num):
                         counters=canvas.counters.dumps(),
                         corner_zyx=canvas.corner_zyx,
                         is_shift=is_shift)
-  logging.info('.. save complete')
+  canvas.log_info('.. save complete')
+  # Cannot `del canvas` here in Python 2 -- deleting an object referenced
+  # in a nested scope is a syntax error.
+  canvas._deregister_client()  # pylint:disable=protected-access
 
 
 def process(request, runner):

{'addition_count': 49, 'commit_subject': 'Add support for the executor interface in inference.', 'deletion_count': 23, 'file_extension': '.py', 'lang': 'py', 'license': 'apache-2.0', 'repo_name': 'google/ffn'}