No public description

ffn/inference/align.py

@@ -109,10 +109,10 @@ def align_and_crop(self,
 
     zyx_offset = np.array(src_corner) - np.array(dst_corner)
     src_size = np.array(source.shape)
-    dst_beg = np.clip(zyx_offset, 0, dst_size).astype(np.int)
-    dst_end = np.clip(dst_size, 0, src_size + zyx_offset).astype(np.int)
-    src_beg = np.clip(-zyx_offset, 0, src_size).astype(np.int)
-    src_end = np.clip(src_size, 0, dst_size - zyx_offset).astype(np.int)
+    dst_beg = np.clip(zyx_offset, 0, dst_size).astype(np.int)  # pyrefly: ignore[missing-attribute]
+    dst_end = np.clip(dst_size, 0, src_size + zyx_offset).astype(np.int)  # pyrefly: ignore[missing-attribute]
+    src_beg = np.clip(-zyx_offset, 0, src_size).astype(np.int)  # pyrefly: ignore[missing-attribute]
+    src_end = np.clip(src_size, 0, dst_size - zyx_offset).astype(np.int)  # pyrefly: ignore[missing-attribute]
 
     if np.any(dst_end - dst_beg == 0) or np.any(src_end - src_beg == 0):
       return destination

ffn/inference/executor.py

@@ -227,7 +227,7 @@ def __init__(self,
                counters: inference_utils.Counters,
                batch_size: int,
                expected_clients: int = 1):
-    super(ThreadingBatchExecutor, self).__init__(interface, model, model_info,
+    super(ThreadingBatchExecutor, self).__init__(interface, model, model_info,  # pyrefly: ignore[bad-argument-type]
                                                  session, counters, batch_size)
 
     # Total clients seen during the lifetime of the executor.
@@ -310,7 +310,7 @@ def _run_executor(self):
             ready.append(client_id)
 
       if ready:
-        self._schedule_batch(ready, fetches)
+        self._schedule_batch(ready, fetches)  # pyrefly: ignore[unbound-name]
 
     logging.info('Executor terminating.')
 

ffn/inference/inference.py

@@ -653,7 +653,7 @@ def segment_all(self, seed_policy=seed.PolicyPeaks, partial_segment_iters=0):
         sid = self.get_next_segment_id()
         self.segmentation[sel][mask] = sid
         if self.keep_probability_maps:
-          self.seg_prob[sel][mask] = storage.quantize_probability(
+          self.seg_prob[sel][mask] = storage.quantize_probability(  # pyrefly: ignore[unsupported-operation]
               expit(self.seed[sel][mask])
           )
 
@@ -806,8 +806,8 @@ def save_checkpoint(self, path: str, partial_segment_iters: int):
             ),
             segmentation=self.segmentation,
             seed=self.seed,
-            origins=self.origins,
-            overlaps=self.overlaps,
+            origins=self.origins,  # pyrefly: ignore[bad-argument-type]
+            overlaps=self.overlaps,  # pyrefly: ignore[bad-argument-type]
             min_pos=self._min_pos,
             max_pos=self._max_pos,
             history=np.array(self.history),

ffn/inference/runner.py

@@ -64,7 +64,7 @@ class Runner:
 
   ALL_MASKED = 1
 
-  request: inference_pb2.InferenceRequest
+  request: inference_pb2.InferenceRequest  # pyrefly: ignore[not-a-type]
   executor: executor.BatchExecutor
   init_seg_volume: storage.Volume | None
   _image_volume: storage.Volume | None
@@ -76,7 +76,7 @@ class Runner:
 
   def __init__(self):
     self.counters = inference_utils.Counters()
-    self.executor = None
+    self.executor = None  # pyrefly: ignore[bad-assignment]
     self._exec_interface = executor.ExecutorInterface()
     self.canvases = {}
 
@@ -93,7 +93,7 @@ def stop_executor(self):
         self.executor.stop_server()
       except executor.TerminationException:
         pass
-      self.executor = None
+      self.executor = None  # pyrefly: ignore[bad-assignment]
 
   def _load_tf_model_checkpoint(
       self, session: tf.Session, checkpoint_path: str
@@ -115,7 +115,7 @@ def _get_tf_model(self, model_name: str):
 
   def _init_tf_model(
       self,
-      request: inference_pb2.InferenceRequest,
+      request: inference_pb2.InferenceRequest,  # pyrefly: ignore[not-a-type]
       batch_size: int,
       session: Optional[tf.Session] = None,
   ) -> tf.Session:
@@ -201,11 +201,11 @@ def _open_or_none(settings):
       self._shift_mask_volume = _open_or_none(request.shift_mask)
 
       alignment_options = request.alignment_options
-      null_alignment = inference_pb2.AlignmentOptions.NO_ALIGNMENT
+      null_alignment = inference_pb2.AlignmentOptions.NO_ALIGNMENT  # pyrefly: ignore[missing-attribute]
       if not alignment_options or alignment_options.type == null_alignment:
         self._aligner = align.Aligner()
       else:
-        type_name = inference_pb2.AlignmentOptions.AlignType.Name(
+        type_name = inference_pb2.AlignmentOptions.AlignType.Name(  # pyrefly: ignore[missing-attribute]
             alignment_options.type
         )
         error_string = 'Alignment for type %s is not implemented' % type_name
@@ -292,7 +292,7 @@ def make_restrictor(self, corner, subvol_size, image, alignment):
         else:
           shift_mask_diameter = np.array(self._model_info.input_image_size)
           shift_mask_fov = bounding_box.BoundingBox(
-              start=-(shift_mask_diameter // 2), size=shift_mask_diameter
+              start=-(shift_mask_diameter // 2), size=shift_mask_diameter  # pyrefly: ignore[bad-argument-type]
           )
 
         kwargs.update({
@@ -530,7 +530,7 @@ def run(self, corner: Tuple3i, subvol_size: Tuple3i, reset_counters=True):
     self.canvases[corner] = canvas
     canvas.segment_all(
         seed_policy=self.get_seed_policy(corner, subvol_size),
-        partial_segment_iters=partial_segment_iters,
+        partial_segment_iters=partial_segment_iters,  # pyrefly: ignore[unbound-name]
     )
     self.save_segmentation(canvas, alignment, seg_path, prob_path)
     del self.canvases[corner]

ffn/inference/seed.py

@@ -51,7 +51,7 @@ def __init__(self, canvas, **kwargs):
 
     # TODO(mjanusz): Remove circular reference between Canvas and seed policies.
     self.canvas = weakref.proxy(canvas)
-    self.coords: np.ndarray = None  # shape: [N, 3] (zyx order in the last dim)
+    self.coords: np.ndarray = None  # shape: [N, 3] (zyx order in the last dim)  # pyrefly: ignore[bad-assignment]
     self.idx = 0
 
   def init_coords(self):
@@ -180,7 +180,7 @@ def init_coords(self):
 
     with PolicyPeaks._sem:
       logging.info('peaks: filtering done')
-      dt = edt.edt(
+      dt = edt.edt(  # pyrefly: ignore[not-callable]
           1 - filt_edges,
           anisotropy=self.canvas.voxel_size_zyx).astype(np.float32)
       logging.info('peaks: edt done')
@@ -250,7 +250,7 @@ def init_coords(self):
         filt_edges[self.canvas.restrictor.mask[z, :, :]] = 1
 
       # Distance transform
-      dt = edt.edt(1 - filt_edges).astype(np.float32)
+      dt = edt.edt(1 - filt_edges).astype(np.float32)  # pyrefly: ignore[not-callable]
 
       idxs = _find_peaks(
           dt,
@@ -293,7 +293,7 @@ class PolicyFillEmptySpace(BaseSeedPolicy):
   def init_coords(self):
     logging.info('fill_empty: starting')
 
-    dt = edt.edt(self.canvas.segmentation == 0).astype(np.float32)
+    dt = edt.edt(self.canvas.segmentation == 0).astype(np.float32)  # pyrefly: ignore[not-callable]
 
     # Set absolute threshold to <1 to avoid generating seeds in areas that are
     # already segmented, where dt >= 1. This also helps to avoid slow execution

ffn/inference/segmentation.py

@@ -128,13 +128,13 @@ def clean_up_and_count(seg: np.ndarray,
   cc_to_orig, cc_to_count = None, None
 
   if compute_id_map or compute_counts:
-    unique_result_tuple = np.unique(
+    unique_result_tuple = np.unique(  # pyrefly: ignore[no-matching-overload]
         seg.ravel(), return_index=compute_id_map, return_counts=compute_counts)
     cc_ids = unique_result_tuple[0]
   if compute_id_map:
-    cc_idx = unique_result_tuple[1]
-    orig_ids = seg_orig.ravel()[cc_idx]
-    cc_to_orig = dict(zip(cc_ids, orig_ids))
+    cc_idx = unique_result_tuple[1]  # pyrefly: ignore[unbound-name]
+    orig_ids = seg_orig.ravel()[cc_idx]  # pyrefly: ignore[unbound-name]
+    cc_to_orig = dict(zip(cc_ids, orig_ids))  # pyrefly: ignore[unbound-name]
   if compute_counts:
     cc_counts = unique_result_tuple[-1]
     cc_to_count = dict(zip(cc_ids, cc_counts))

ffn/inference/storage.py

@@ -60,15 +60,15 @@ class NumpyArray(np.ndarray):
 
   def __new__(cls, default_value=0, **kwargs):
     ret = super(NumpyArray, cls).__new__(cls, **kwargs)
-    ret.default_value = default_value
+    ret.default_value = default_value  # pyrefly: ignore[missing-attribute]
     return ret
 
   def __init__(self, *args, **kwargs):
     del args, kwargs
     self.clear()
 
   def clear(self):
-    self[...] = self.default_value
+    self[...] = self.default_value  # pyrefly: ignore[missing-attribute]
 
 
 def decorated_volume(settings, **kwargs) -> Volume:
@@ -317,7 +317,7 @@ def clip_subvolume_to_bounds(corner, size, volume):
   volume_bounds = bounding_box.BoundingBox(start=(0, 0, 0), size=volume_size)
   subvolume_bounds = bounding_box.BoundingBox(start=corner, size=size)
   clipped_bounds = volume_bounds.intersection(subvolume_bounds)
-  return clipped_bounds.start, clipped_bounds.size
+  return clipped_bounds.start, clipped_bounds.size  # pyrefly: ignore[missing-attribute]
 
 
 def build_mask(masks, corner, subvol_size, mask_volume_map=None,
@@ -385,8 +385,8 @@ def build_mask(masks, corner, subvol_size, mask_volume_map=None,
       else:
         logging.fatal('Unsupported mask source: %s', source_type)
 
-      for chan_config in channels:
-        channel_mask = mask[chan_config.channel, ...]
+      for chan_config in channels:  # pyrefly: ignore[unbound-name]
+        channel_mask = mask[chan_config.channel, ...]  # pyrefly: ignore[unbound-name]
         channel_mask = alignment.align_and_crop(
             src_corner, channel_mask, corner, subvol_size)
 

ffn/jax/input_pipeline.py

@@ -124,7 +124,7 @@ def _add_ffn_data(ex: volume.Example) -> volume.Example:
       if config.loss_mask_invert:
         loss_mask = tf.equal(ex['loss_mask'], 0)
       else:
-        loss_mask = ex['loss_mask'] > 0
+        loss_mask = ex['loss_mask'] > 0  # pyrefly: ignore[unsupported-operation]
 
       weights *= 1.0 - tf.cast(loss_mask, tf.float32)
 
@@ -136,7 +136,7 @@ def _add_ffn_data(ex: volume.Example) -> volume.Example:
         seg.shape[3] // 2,  #
         0,
     ]
-    lom = tf.logical_and(seg > 0, tf.equal(seg, center_val))
+    lom = tf.logical_and(seg > 0, tf.equal(seg, center_val))  # pyrefly: ignore[unsupported-operation]
     labels = inputs.soften_labels(lom)
 
     lx, ly, lz = load_shape
@@ -266,7 +266,7 @@ def __init__(
     self._fs_lock = threading.Lock()
     self._fs = set()
     for i, gen in enumerate(self._generators):
-      self._fs.add(self._tpe.submit(lambda gen=gen, i=i: (i, next(gen))))
+      self._fs.add(self._tpe.submit(lambda gen=gen, i=i: (i, next(gen))))  # pyrefly: ignore[missing-argument]
 
     # Prefetching of complete batches.
     self._batch_tpe = futures.ThreadPoolExecutor(max_workers=batch_prefetch)
@@ -337,7 +337,7 @@ def __next__(self):
         'weight': batched_weights,
     }
 
-  def update_seeds(self, batched_seeds: list[jax.Array]):
+  def update_seeds(self, batched_seeds: list[jax.Array]):  # pyrefly: ignore[bad-override]
     """Propagates data from `batched_seeds` back to the example generators."""
 
     def _update(
@@ -346,35 +346,35 @@ def _update(
         current: list[int],
     ):
       # Transfer data from device to host.
-      batched_seeds = np.array(batched_seeds)
+      batched_seeds = np.array(batched_seeds)  # pyrefly: ignore[bad-assignment]
       # Fold batch dimensions back to a single one.
-      batched_seeds = np.reshape(
-          batched_seeds, [-1] + list(batched_seeds.shape[-4:])
+      batched_seeds = np.reshape(  # pyrefly: ignore[bad-assignment]
+          batched_seeds, [-1] + list(batched_seeds.shape[-4:])  # pyrefly: ignore[missing-attribute]
       )
 
-      assert batched_seeds.shape[0] == len(seeds)
+      assert batched_seeds.shape[0] == len(seeds)  # pyrefly: ignore[missing-attribute]
 
       dx = self._info.input_seed_size[0] - self._info.pred_mask_size[0]
       dy = self._info.input_seed_size[1] - self._info.pred_mask_size[1]
       dz = self._info.input_seed_size[2] - self._info.pred_mask_size[2]
 
       for i, _ in enumerate(current):
         if dz == 0 and dy == 0 and dx == 0:
-          seeds[i][:] = batched_seeds[i, ...]
+          seeds[i][:] = batched_seeds[i, ...]  # pyrefly: ignore[bad-index]
         else:
           seeds[i][
               :,  #
               dz // 2 : -(dz - dz // 2),  #
               dy // 2 : -(dy - dy // 2),  #
               dx // 2 : -(dx - dx // 2),  #
               :,
-          ] = batched_seeds[i, ...]
+          ] = batched_seeds[i, ...]  # pyrefly: ignore[bad-index]
 
       with self._fs_lock:
         for gen_idx in current:
           gen = self._generators[gen_idx]
           self._fs.add(
-              self._tpe.submit(lambda gen=gen, i=gen_idx: (i, next(gen)))
+              self._tpe.submit(lambda gen=gen, i=gen_idx: (i, next(gen)))  # pyrefly: ignore[missing-argument]
           )
 
     # Distribute data asynchronously.

ffn/jax/train.py

@@ -83,7 +83,7 @@ def create_train_state(
     The initialized TrainState with the optimizer.
   """
   model = model_util.model_from_config(config)
-  rng = {'params': rng, 'dropout': jax.random.PRNGKey(1)}
+  rng = {'params': rng, 'dropout': jax.random.PRNGKey(1)}  # pyrefly: ignore[bad-assignment]
   variables = model.init(rng, jnp.ones(input_shape))
   params = variables['params']
 
@@ -106,9 +106,9 @@ def create_train_state(
 
 @flax.struct.dataclass
 class TrainMetrics(metrics.Collection):
-  loss: metrics.Average.from_output('loss')
-  loss_std: metrics.Std.from_output('loss')
-  learning_rate: metrics.LastValue.from_output('learning_rate')
+  loss: metrics.Average.from_output('loss')  # pyrefly: ignore[invalid-annotation]
+  loss_std: metrics.Std.from_output('loss')  # pyrefly: ignore[invalid-annotation]
+  learning_rate: metrics.LastValue.from_output('learning_rate')  # pyrefly: ignore[invalid-annotation]
 
 
 def _updated_seed(seed: jnp.ndarray, update: jnp.ndarray) -> jnp.ndarray:
@@ -366,7 +366,7 @@ def _make_ckpt_args(
 ) -> ocp.args.CheckpointArgs:
   args = {'train_state': ocp.args.StandardSave(state)}
   if 'train_iter' in checkpoint_items:
-    args['train_iter'] = _get_ocp_args(train_iter, restore=False)
+    args['train_iter'] = _get_ocp_args(train_iter, restore=False)  # pyrefly: ignore[bad-assignment]
   return ocp.args.Composite(**args)
 
 
@@ -377,8 +377,8 @@ def train_and_evaluate(
     checkpoint_items: Sequence[str] = ('train_state', 'train_iter'),
 ):
   """Main training loop."""
-  workdir = epath.Path(workdir)
-  workdir.mkdir(parents=True, exist_ok=True)
+  workdir = epath.Path(workdir)  # pyrefly: ignore[bad-assignment]
+  workdir.mkdir(parents=True, exist_ok=True)  # pyrefly: ignore[missing-attribute]
 
   rng = training.get_rng(config.seed)
 
@@ -468,7 +468,7 @@ def train_and_evaluate(
   )
   checkpointed_state = {'train_state': state}
   if 'train_iter' in checkpoint_items:
-    checkpointed_state['train_iter'] = train_iter
+    checkpointed_state['train_iter'] = train_iter  # pyrefly: ignore[bad-assignment]
   latest_step = checkpoint_manager.latest_step()
   # If an initial checkpoint is provided and the checkpointing library does not
   # report a 'latest' checkpoint, then we are starting a new experiment.
@@ -487,8 +487,8 @@ def train_and_evaluate(
       if isinstance(train_iter, tf.data.Iterator):
         iter_handler = item_handlers['train_iter']
         args = DatasetArgs(train_iter)
-      checkpointed_state['train_iter'] = iter_handler.restore(
-          train_iter_path, args=args
+      checkpointed_state['train_iter'] = iter_handler.restore(  # pyrefly: ignore[bad-assignment]
+          train_iter_path, args=args  # pyrefly: ignore[bad-argument-type]
       )
 
     logging.info('Initializing training from %r', config.init_from_cpoint)
@@ -497,7 +497,7 @@ def train_and_evaluate(
         'train_state': ocp.args.StandardRestore(state),
     }
     if 'train_iter' in checkpoint_items:
-      restore_args['train_iter'] = _get_ocp_args(train_iter)
+      restore_args['train_iter'] = _get_ocp_args(train_iter)  # pyrefly: ignore[bad-assignment, bad-specialization]
     checkpointed_state = checkpoint_manager.restore(
         latest_step,
         args=ocp.args.Composite(**restore_args),
@@ -602,7 +602,7 @@ def train_fn(state, batch, loss_scale):
   eval_tracker = tracker.EvalTracker(eval_shape_zyx, fov_shifts)
 
   batch_iter = input_pipeline.get_batch_iter(
-      train_iter,
+      train_iter,  # pyrefly: ignore[bad-argument-type]
       eval_tracker,
       policy_fn,
       info,
@@ -691,7 +691,7 @@ def _reshape(x):
           train_state = jax.tree.map(np.array, state)
           checkpoint_manager.save(
               step,
-              args=_make_ckpt_args(train_state, train_iter, checkpoint_items),
+              args=_make_ckpt_args(train_state, train_iter, checkpoint_items),  # pyrefly: ignore[bad-specialization]
           )
 
         if checkpoint_manager.reached_preemption(step):
@@ -727,7 +727,7 @@ def _reshape(x):
             assert tfw is not None
             # TODO(mjanusz): Find a cleaner and less brittle way of saving
             # raw summaries.
-            with tfw._summary_writer.as_default():
+            with tfw._summary_writer.as_default():  # pyrefly: ignore[missing-attribute]
               for s in raws:
                 tf.summary.experimental.write_raw_pb(s, step=step)