diff --git a/ffn/inference/align.py b/ffn/inference/align.py index af57e9a..e5d8499 100644 --- a/ffn/inference/align.py +++ b/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 diff --git a/ffn/inference/executor.py b/ffn/inference/executor.py index df78185..5907f6e 100644 --- a/ffn/inference/executor.py +++ b/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.') diff --git a/ffn/inference/inference.py b/ffn/inference/inference.py index 54c47b0..ce9a2d4 100644 --- a/ffn/inference/inference.py +++ b/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), diff --git a/ffn/inference/runner.py b/ffn/inference/runner.py index f7cb53f..ac10d7b 100644 --- a/ffn/inference/runner.py +++ b/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] diff --git a/ffn/inference/seed.py b/ffn/inference/seed.py index 1334e54..fb5349e 100644 --- a/ffn/inference/seed.py +++ b/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 diff --git a/ffn/inference/segmentation.py b/ffn/inference/segmentation.py index b43cd28..710ca2b 100644 --- a/ffn/inference/segmentation.py +++ b/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)) diff --git a/ffn/inference/storage.py b/ffn/inference/storage.py index bdfe413..bfe2865 100644 --- a/ffn/inference/storage.py +++ b/ffn/inference/storage.py @@ -60,7 +60,7 @@ 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): @@ -68,7 +68,7 @@ def __init__(self, *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) diff --git a/ffn/jax/input_pipeline.py b/ffn/jax/input_pipeline.py index 7d08e78..f4fc4c9 100644 --- a/ffn/jax/input_pipeline.py +++ b/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,13 +346,13 @@ 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] @@ -360,7 +360,7 @@ def _update( 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][ :, # @@ -368,13 +368,13 @@ def _update( 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. diff --git a/ffn/jax/train.py b/ffn/jax/train.py index e84276c..788b5e7 100644 --- a/ffn/jax/train.py +++ b/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) diff --git a/ffn/utils/bounding_box.py b/ffn/utils/bounding_box.py index a3bd270..c2c2fc9 100644 --- a/ffn/utils/bounding_box.py +++ b/ffn/utils/bounding_box.py @@ -47,7 +47,7 @@ def __init__(self, start=None, size=None, end=None): ValueError: on bad inputs. """ if start is not None: - if (isinstance(start, bounding_box_pb2.BoundingBox) or + if (isinstance(start, bounding_box_pb2.BoundingBox) or # pyrefly: ignore[invalid-argument] isinstance(start, BoundingBox)): if size is not None or end is not None: raise ValueError('a BoundingBox object/proto must be specified alone') @@ -170,7 +170,7 @@ def __repr__(self): tuple(self.size)) def __eq__(self, other): - if isinstance(other, bounding_box_pb2.BoundingBox): + if isinstance(other, bounding_box_pb2.BoundingBox): # pyrefly: ignore[invalid-argument] other = BoundingBox(other) elif not isinstance(other, BoundingBox): return False @@ -191,9 +191,9 @@ def _required(bbox: Optional[BoundingBox]) -> BoundingBox: def intersection(box0, box1): """Get intersection between two bounding boxes, or None.""" - if isinstance(box0, bounding_box_pb2.BoundingBox): + if isinstance(box0, bounding_box_pb2.BoundingBox): # pyrefly: ignore[invalid-argument] box0 = BoundingBox(box0.start, box0.size) - if isinstance(box1, bounding_box_pb2.BoundingBox): + if isinstance(box1, bounding_box_pb2.BoundingBox): # pyrefly: ignore[invalid-argument] box1 = BoundingBox(box1.start, box1.size) if not isinstance(box0, BoundingBox): raise ValueError('box0 must be a BoundingBox') diff --git a/ffn/utils/geom_utils.py b/ffn/utils/geom_utils.py index cb84d47..de8fec2 100644 --- a/ffn/utils/geom_utils.py +++ b/ffn/utils/geom_utils.py @@ -38,7 +38,7 @@ def ToVector3j(*args): else: raise ValueError('Expected three ints, a 3-sequence of ints, or a Vector3j') - if isinstance(seq, vector_pb2.Vector3j): return seq + if isinstance(seq, vector_pb2.Vector3j): return seq # pyrefly: ignore[invalid-argument] if isinstance(seq, numpy.ndarray) and seq.dtype.kind in 'iu': seq = [int(s) for s in seq] if len(seq) != 3: @@ -61,7 +61,7 @@ def To3Tuple(vector): Raises: ValueError: Unsupported argument type. """ - if isinstance(vector, (vector_pb2.Vector3j, vector_pb2.Vector3f)): + if isinstance(vector, (vector_pb2.Vector3j, vector_pb2.Vector3f)): # pyrefly: ignore[invalid-argument] return (vector.x, vector.y, vector.z) if isinstance(vector, numpy.ndarray): if vector.shape != (3,): diff --git a/ffn/utils/ortho_plane_visualization.py b/ffn/utils/ortho_plane_visualization.py index 5e4246a..2e75081 100644 --- a/ffn/utils/ortho_plane_visualization.py +++ b/ffn/utils/ortho_plane_visualization.py @@ -45,19 +45,19 @@ def cut_ortho_planes( yx, zx, zy. """ if center is None: - center = np.array(vol.shape[:3]) // 2 + center = np.array(vol.shape[:3]) // 2 # pyrefly: ignore[bad-assignment] planes = [] full_slice = [slice(None)] * 3 - for axis, ix in enumerate(center): + for axis, ix in enumerate(center): # pyrefly: ignore[bad-argument-type] cut_slice = list(full_slice) - cut_slice[axis] = ix + cut_slice[axis] = ix # pyrefly: ignore[unsupported-operation] planes.append(vol[tuple(cut_slice)]) if cross_hair: # Copy because cross hair is written into array data. plane = planes[-1].copy() i = 0 - for ax, c in enumerate(center): + for ax, c in enumerate(center): # pyrefly: ignore[bad-argument-type] if ax != axis: # Make axis i the 0-axis an work in-place. view = np.rollaxis(plane, i)