diff --git a/src/napari_convpaint/__init__.py b/src/napari_convpaint/__init__.py index 1cb0fce..84c4404 100644 --- a/src/napari_convpaint/__init__.py +++ b/src/napari_convpaint/__init__.py @@ -11,4 +11,5 @@ # Import the model, widget and feature extractor superclass to make them available at the package level from .convpaint_model import ConvpaintModel from .convpaint_widget import ConvpaintWidget -from .feature_extractor import FeatureExtractor \ No newline at end of file +from .feature_extractor import FeatureExtractor +from .utils import CancelToken, CancelledError \ No newline at end of file diff --git a/src/napari_convpaint/_tests/conftest.py b/src/napari_convpaint/_tests/conftest.py index 48dba37..db93c8a 100644 --- a/src/napari_convpaint/_tests/conftest.py +++ b/src/napari_convpaint/_tests/conftest.py @@ -7,15 +7,24 @@ import torch import pytest -# Fix for MPS memory leaks in tests: +# Run long-running widget operations (train / predict / predict_all) on the +# calling thread in tests. The test assertions check layer state immediately +# after _on_train() / _on_predict(), which assumes synchronous execution. +from napari_convpaint.convpaint_widget import ConvpaintWidget # noqa: E402 +ConvpaintWidget._sync_workers = True + + +# Belt-and-suspenders MPS cleanup between tests. The root-cause fix for the +# worker-pinning leak is in convpaint_widget (we use superqt.utils.thread_worker +# so napari's task_status manager doesn't retain workers); this fixture adds an +# extra gc.collect() + torch.mps.empty_cache() after each test so any lingering +# per-test allocations don't carry over into the next one. @pytest.fixture(autouse=True) def cleanup_mps_after_test(): yield - # delete any local refs in test code if possible, then: gc.collect() - # empty PyTorch MPS cache if available if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): try: torch.mps.empty_cache() except Exception: - pass \ No newline at end of file + pass diff --git a/src/napari_convpaint/_tests/test_cancellation.py b/src/napari_convpaint/_tests/test_cancellation.py new file mode 100644 index 0000000..249c210 --- /dev/null +++ b/src/napari_convpaint/_tests/test_cancellation.py @@ -0,0 +1,354 @@ +"""Tests for cooperative cancellation of training and prediction.""" +import threading +import warnings + +import numpy as np +import pytest + +from napari_convpaint.convpaint_model import ConvpaintModel +from napari_convpaint.feature_extractors.gaussian import GaussianFeatures +from napari_convpaint.utils import CancelToken, CancelledError + + +def _tiny_dataset(): + rng = np.random.default_rng(0) + image = rng.random((64, 64), dtype=np.float32) + annot = np.zeros((64, 64), dtype=np.uint8) + annot[10:20, 10:20] = 1 + annot[30:40, 30:40] = 2 + return image, annot + + +def test_cancel_token_starts_uncancelled(): + t = CancelToken() + assert t.cancelled is False + t.raise_if_cancelled() # should not raise + + +def test_cancel_token_raises_after_cancel(): + t = CancelToken() + t.cancel() + assert t.cancelled is True + with pytest.raises(CancelledError): + t.raise_if_cancelled() + + +def test_train_aborts_when_token_is_pre_cancelled(): + """A token cancelled before train() starts should abort at the first checkpoint.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + + token = CancelToken() + token.cancel() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, cancel_token=token) + + +def test_segment_aborts_when_token_is_pre_cancelled(): + """Same for segment() — once trained, a pre-cancelled token aborts prediction.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + model.train(image, annot) + + token = CancelToken() + token.cancel() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.segment(image, cancel_token=token) + + +def test_train_completes_with_fresh_token(): + """Passing an un-cancelled token must not affect normal completion.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + clf = model.train(image, annot, cancel_token=CancelToken()) + assert clf is not None + + +class _CancelOnNthCheck(CancelToken): + """Test double: auto-cancels itself after `n` checkpoints have been hit. + + Lets us deterministically verify that cancel-checks are reached mid-run, + without racing a wall-clock timer against a fast feature extractor.""" + + def __init__(self, n): + super().__init__() + self._n = n + self._checks = 0 + + def raise_if_cancelled(self): + self._checks += 1 + if self._checks >= self._n and not self.cancelled: + self.cancel() + super().raise_if_cancelled() + + +def test_cancel_mid_train_is_honored(): + """Proves that cancel-checkpoints are actually reached during training: + the token cancels itself after 2 checks, so train must raise CancelledError + (it would otherwise complete normally).""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + token = _CancelOnNthCheck(n=2) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, cancel_token=token) + assert token._checks >= 2, "cancel_token was never checked during train" + + +def test_cancel_mid_segment_is_honored(): + """Same idea for prediction.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + model.train(image, annot) + + token = _CancelOnNthCheck(n=2) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.segment(image, cancel_token=token) + assert token._checks >= 2, "cancel_token was never checked during segment" + + +def test_cancel_preserves_previous_classifier(): + """A cancelled training must not clobber a classifier from an earlier + successful training. Cancel fires before _clf_train (which is an + uninterruptible C call), so classifier state stays on the previous fit.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + model.train(image, annot) + first_clf = model.classifier + assert first_clf is not None + + token = _CancelOnNthCheck(n=2) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, cancel_token=token) + + assert model.classifier is first_clf, ( + "classifier object changed after a cancelled training — previous fit was clobbered" + ) + + +def test_cancel_in_memory_mode_leaves_state_retrainable(): + """In memory-mode training, _register_and_update_annots mutates self.annot_dict + and self.table before feature extraction. If we cancel without rolling those + back, the *next* train sees no new annotations and raises 'No features or + targets found'. Verify the rollback makes retraining work after a cancel.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + + # Cancel the first training mid-run + token = _CancelOnNthCheck(n=2) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, memory_mode=True, img_ids='img0', + cancel_token=token) + + # After cancel, the memory state should be back to where it started, so a + # fresh train with the same annotations can still find work to do. + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + clf = model.train(image, annot, memory_mode=True, img_ids='img0') + assert clf is not None, "retrain after cancel should succeed" + + +class _OldSignatureFE(GaussianFeatures): + """Simulates a third-party FE written before cancellation existed: its + overrides use the pre-cancellation signatures (no cancel_token, no + **kwargs). Cancellation is carried by an ambient ContextVar, so such FEs + must keep working unmodified AND be cancellable through the base-class + loop checkpoints.""" + + def extract_features_pyramid(self, data, param, patched=True, device=None): + return super().extract_features_pyramid(data, param, patched=patched, device=device) + + def extract_features_from_stack(self, image, device=None): + return super().extract_features_from_stack(image, device=device) + + +def test_custom_fe_with_old_signature_still_works(): + """A custom FE with pre-cancellation override signatures must neither crash + (TypeError from an unexpected cancel_token kwarg) nor lose cancellability.""" + model = ConvpaintModel(fe_name='gaussian_features') + model.fe_model = _OldSignatureFE() + image, annot = _tiny_dataset() + + # Trains fine without a token ... + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + clf = model.train(image, annot) + assert clf is not None + + # ... and with one ... + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + clf = model.train(image, annot, cancel_token=CancelToken()) + assert clf is not None + + # ... and is still cancellable mid-run through the ambient token, even + # though the FE itself never sees or forwards it. + token = _CancelOnNthCheck(n=2) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, cancel_token=token) + assert token._checks >= 2, "ambient cancel_token was never checked with a legacy FE" + + +def test_ambient_token_does_not_leak_out_of_the_call(): + """After a cancelled call returns, the ambient token must be uninstalled — + a later call without a token must not see the stale cancelled one.""" + from napari_convpaint.utils import check_cancel + + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + + token = CancelToken() + token.cancel() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, cancel_token=token) + + check_cancel() # must not raise: no ambient token installed anymore + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + clf = model.train(image, annot) # must not be affected by the old token + assert clf is not None + + +def test_cancel_during_catboost_fit_preserves_classifier(): + """The CatBoost fit itself is cancellable on CPU via a per-iteration + callback. A cancel that lands mid-fit must abort with CancelledError and + leave the previously trained classifier in place (the partial fit is + discarded, and self.classifier is only reassigned after a successful fit).""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + model.train(image, annot) + first_clf = model.classifier + assert first_clf is not None + + # Only a handful of checkpoints run before the fit (per-scale, per-Z, + # pre-fit), so n=20 lands inside the boosting loop's per-iteration + # callback checks (default 100 iterations). + token = _CancelOnNthCheck(n=20) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with pytest.raises(CancelledError): + model.train(image, annot, cancel_token=token) + assert token._checks >= 20, "cancel checkpoints were never reached" + assert model.classifier is first_clf, ( + "classifier changed after a cancelled fit — partial fit was adopted" + ) + + # And the model must still be retrainable afterwards + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + clf = model.train(image, annot) + assert clf is not None and clf is not first_clf + + +def test_cancel_from_another_thread_aborts_train(): + """Cross-thread cancel: main thread calls cancel() while a worker runs train(). + Uses the auto-cancel token so the outcome is deterministic regardless of how + fast the feature extractor happens to be on this machine.""" + model = ConvpaintModel(fe_name='gaussian_features') + image, annot = _tiny_dataset() + + # Give enough checkpoints that the main thread has time to call cancel(), + # but few enough that the worker will actually reach the cancel on the next check. + token = CancelToken() + result = {} + started = threading.Event() + + def worker(): + started.set() + try: + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + model.train(image, annot, cancel_token=token) + result['ok'] = True + except CancelledError: + result['cancelled'] = True + except Exception as e: # noqa: BLE001 + result['error'] = repr(e) + + t = threading.Thread(target=worker) + t.start() + started.wait(timeout=5.0) + token.cancel() + t.join(timeout=30.0) + + assert not t.is_alive(), "worker did not exit within the timeout after cancel" + # Either cancelled mid-run (ideal) or completed before cancel landed (acceptable + # on very fast hardware — the deterministic mid-run case is covered separately). + assert 'error' not in result, f"worker crashed: {result!r}" + assert 'cancelled' in result or 'ok' in result + + +def test_widget_async_worker_completes_without_main_thread_violation(make_napari_viewer): + """Regression test for the 'NSWindow should only be instantiated on the main + thread' crash on macOS. The rest of the widget test suite runs with + ConvpaintWidget._sync_workers = True, which executes the worker on the + calling thread and would hide any QWidget-constructed-from-worker-thread + bug. Flip that off and drive one full _on_train / _on_predict cycle end + to end in a real Qt worker thread.""" + import time + from qtpy.QtWidgets import QApplication + from napari_convpaint.convpaint_widget import ConvpaintWidget + from napari_convpaint.testing_data import ( + generate_synthetic_square, + generate_synthetic_circle_annotation, + ) + + im, _ = generate_synthetic_square(im_dims=(252, 252), square_dims=(70, 70)) + im_annot = generate_synthetic_circle_annotation( + im_dims=(252, 252), circle1_xy=(125, 70), circle2_xy=(125, 125) + ) + + orig_sync = ConvpaintWidget._sync_workers + ConvpaintWidget._sync_workers = False + try: + viewer = make_napari_viewer() + widget = ConvpaintWidget(viewer) + widget.ensure_init() + # Swap to a cheap CPU-only feature extractor so the test finishes + # quickly; the bug we're guarding against is about Qt thread affinity, + # not about the specific FE choice. + widget.cp_model = ConvpaintModel(fe_name='gaussian_features') + widget.auto_seg = False # keep the test to a single worker round + viewer.add_image(im) + widget._on_add_annot_layer() + viewer.layers['annotations'].data[...] = im_annot + widget.cp_model.set_params(channel_mode='rgb') + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + widget._on_train() + deadline = time.monotonic() + 30 + while widget._op is not None and time.monotonic() < deadline: + QApplication.processEvents() + time.sleep(0.01) + assert widget._op is None, "async worker did not finish within 30s" + # If QWidgets had been constructed on the worker thread, the + # process would have aborted before reaching this line on macOS. + assert widget.trained + finally: + ConvpaintWidget._sync_workers = orig_sync diff --git a/src/napari_convpaint/convpaint_model.py b/src/napari_convpaint/convpaint_model.py index 6c13f94..3ade08d 100644 --- a/src/napari_convpaint/convpaint_model.py +++ b/src/napari_convpaint/convpaint_model.py @@ -1,3 +1,4 @@ +import copy import pickle from pathlib import Path import importlib @@ -763,7 +764,7 @@ def get_fe_description(self): def train(self, image, annotations, memory_mode=False, img_ids=None, use_rf=False, allow_writing_files=False, in_channels=None, skip_norm=False, - fe_use_device=None, clf_use_device=None): + fe_use_device=None, clf_use_device=None, cancel_token=None): """ Trains the Convpaint model's classifier given images and annotations. @@ -798,18 +799,30 @@ def train(self, image, annotations, memory_mode=False, img_ids=None, use_rf=Fals Device policy for feature extractor ("auto", "gpu", "cpu"). clf_use_device : str, optional Device policy for classifier training ("auto", "gpu", "cpu"). + cancel_token : CancelToken, optional + Cooperative cancellation token. Call its `cancel()` method (typically + from another thread) to abort training at the next checkpoint, which + raises a `CancelledError`. On cancellation the classifier keeps its + previous state, and in memory mode the annotation bookkeeping is + rolled back so the same training can simply be retried. Returns ---------- clf : CatBoostClassifier or RandomForestClassifier Trained classifier (also saved inside the model instance) + + Raises + ---------- + CancelledError + If `cancel_token` is cancelled before training completes. """ clf, _, _ = self._train(image, annotations, memory_mode=memory_mode, img_ids=img_ids, use_rf=use_rf, allow_writing_files=allow_writing_files, in_channels=in_channels, skip_norm=skip_norm, - fe_use_device=fe_use_device, clf_use_device=clf_use_device) + fe_use_device=fe_use_device, clf_use_device=clf_use_device, + cancel_token=cancel_token) return clf - def segment(self, image, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None): + def segment(self, image, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None, cancel_token=None): """ Segments images by predicting the most probable class of each pixel using the trained classifier. @@ -829,18 +842,27 @@ def segment(self, image, in_channels=None, skip_norm=False, use_dask=False, fe_u Whether to use dask for parallel processing fe_use_device : str, optional Device policy for feature extractor ("auto", "gpu", "cpu"). + cancel_token : CancelToken, optional + Cooperative cancellation token. Call its `cancel()` method (typically + from another thread) to abort segmentation at the next checkpoint, + which raises a `CancelledError`. Returns ---------- seg : np.ndarray or list[np.ndarray] Segmented image or list of segmented images (according to the input) Dimensions are equal to the input image(s) without the channel dimension + + Raises + ---------- + CancelledError + If `cancel_token` is cancelled before segmentation completes. """ _, seg = self._predict(image, add_seg=True, in_channels=in_channels, skip_norm=skip_norm, - use_dask=use_dask, fe_use_device=fe_use_device) + use_dask=use_dask, fe_use_device=fe_use_device, cancel_token=cancel_token) return seg - def predict_probas(self, image, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None): + def predict_probas(self, image, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None, cancel_token=None): """ Predicts the probabilities of the classes of the pixels in an image using the trained classifier. @@ -860,6 +882,10 @@ def predict_probas(self, image, in_channels=None, skip_norm=False, use_dask=Fals Whether to use dask for parallel processing fe_use_device : str, optional Device policy for feature extractor ("auto", "gpu", "cpu"). + cancel_token : CancelToken, optional + Cooperative cancellation token. Call its `cancel()` method (typically + from another thread) to abort prediction at the next checkpoint, + which raises a `CancelledError`. Returns ---------- @@ -867,15 +893,21 @@ def predict_probas(self, image, in_channels=None, skip_norm=False, use_dask=Fals Predicted probabilities of the classes of the pixels in the image or list of images Dimensions are equal to the input image(s) without the channel dimension, with the class dimension added first + + Raises + ---------- + CancelledError + If `cancel_token` is cancelled before prediction completes. """ probas = self._predict(image, add_seg=False, in_channels=in_channels, - skip_norm=skip_norm, use_dask=use_dask, fe_use_device=fe_use_device) + skip_norm=skip_norm, use_dask=use_dask, fe_use_device=fe_use_device, + cancel_token=cancel_token) return probas def get_feature_image(self, data, in_channels=None, skip_norm=False, pca_components=0, kmeans_clusters=0, - use_device=None): + use_device=None, cancel_token=None): """ Returns the feature images extracted by the feature extractor model. For details, see the underlying `_get_features` method. @@ -891,26 +923,38 @@ def get_feature_image(self, data, If True, the images are not normalized according to the parameter `normalize` in the model parameters. pca_components : int, optional Number of PCA components to reduce the features to (0 for no PCA) + cancel_token : CancelToken, optional + Cooperative cancellation token. Call its `cancel()` method (typically + from another thread) to abort feature extraction at the next + checkpoint, which raises a `CancelledError`. Returns ---------- features : np.ndarray or list[np.ndarray] Extracted features of the image(s) or list of features for each image if input is a list. - Reshaped to the input imges' shapes. Features dimension is added first (FHW or FZHW). - """ - # Extract features - features = self._get_features( - data, - annotations=None, - restore_input_form=True, - memory_mode=False, # Only valid when using annotations - img_ids=None, # Only needed when using memory_mode - in_channels=in_channels, - skip_norm=skip_norm, - use_device=use_device, - pca_components=pca_components, - kmeans_clusters=kmeans_clusters - ) + Reshaped to the input imges' shapes. Features dimension is added first (FHW or FZHW). + + Raises + ---------- + CancelledError + If `cancel_token` is cancelled before feature extraction completes. + """ + # Extract features; the cancel token is installed as the ambient token + # (utils.cancel_scope), so all downstream check_cancel() calls see it + # without it being passed through every signature. + with utils.cancel_scope(cancel_token): + features = self._get_features( + data, + annotations=None, + restore_input_form=True, + memory_mode=False, # Only valid when using annotations + img_ids=None, # Only needed when using memory_mode + in_channels=in_channels, + skip_norm=skip_norm, + use_device=use_device, + pca_components=pca_components, + kmeans_clusters=kmeans_clusters, + ) return features @@ -1123,12 +1167,12 @@ def _get_features(self, data, annotations=None, restore_input_form=True, supported_devices=self.fe_model.supported_devices(), warn=True, ) - features = [self.fe_model.extract_features_pyramid( - d, - params_for_extract, - patched=keep_patched, + features = [ + self.fe_model.extract_features_pyramid( + d, params_for_extract, patched=keep_patched, device=fe_runtime_device) - for d in data] + for d in data + ] if pca_components: features = [utils.apply_pca_to_f_image(f, n_components=pca_components) @@ -1216,12 +1260,15 @@ def _clf_train(self, features, targets, use_rf=False, allow_writing_files=False, clf : CatBoostClassifier or RandomForestClassifier Trained classifier (also saved in the model instance) """ + # NOTE: The new classifier is fit into a local variable and adopted as + # self.classifier only after a successful fit — a failed or cancelled + # fit must not clobber the previously trained classifier. if not use_rf: use_device = self.check_locked_device(use_device, part='clf') task_type = utils.get_catboost_device(use_device, warn=True) # Fixed seed for reproducibility; can be set to None for random seed from catboost import CatBoostClassifier - self.classifier = CatBoostClassifier( + clf = CatBoostClassifier( iterations=self._param.clf_iterations, learning_rate=self._param.clf_learning_rate, depth=self._param.clf_depth, @@ -1229,16 +1276,35 @@ def _clf_train(self, features, targets, use_rf=False, allow_writing_files=False, task_type=task_type, random_seed=0, ) - self.classifier.fit(features, targets) + fit_kwargs = {} + if task_type == 'CPU': + # Make the boosting loop cancellable: the callback checks the + # ambient cancel token after every iteration and stops the fit + # early when cancelled; the check_cancel() after fit then turns + # the (partially trained, to-be-discarded) result into a proper + # CancelledError. CatBoost only supports callbacks on CPU; + # GPU fits remain uninterruptible. + class _CancelFitCallback: + def after_iteration(self, info): + try: + utils.check_cancel() + except utils.CancelledError: + return False # stop the fit at this iteration + return True + fit_kwargs['callbacks'] = [_CancelFitCallback()] + clf.fit(features, targets, **fit_kwargs) + utils.check_cancel() # Discard the partial fit if it was stopped by cancellation + self.classifier = clf self._param.classifier = 'CatBoost' else: # train a random forest classififer (does not support GPU) from sklearn.ensemble import RandomForestClassifier # Fix random_state for reproducibility; can be set to None for random seed - self.classifier = RandomForestClassifier(n_estimators=100, n_jobs=-1, random_state=0) - self.classifier.fit(features, targets) + # (sklearn has no fit callbacks, so an RF fit is uninterruptible) + clf = RandomForestClassifier(n_estimators=100, n_jobs=-1, random_state=0) + clf.fit(features, targets) + self.classifier = clf self._param.classifier = 'RandomForest' - clf = self.classifier self.num_features = features.shape[1] if isinstance(features, np.ndarray) else features[0].shape[1] if isinstance(features, list) else None return clf @@ -1268,12 +1334,23 @@ class dimension is added first (if return_proba is True). features = np.moveaxis(features, 0, -1) features = np.reshape(features, (-1, nb_features)) # flatten - # Predict + # Predict in row chunks: a single predict/predict_proba call over a full + # plane (H*W rows) is one uninterruptible C call that can take many + # seconds on large images; chunking bounds the cancel latency to one + # chunk while producing bit-identical results. + chunk_size = 1_000_000 + num_rows = features.shape[0] + predict_fn = self.classifier.predict_proba if return_proba else self.classifier.predict + if num_rows > chunk_size: + parts = [] + for i in range(0, num_rows, chunk_size): + utils.check_cancel() + parts.append(predict_fn(features[i:i+chunk_size])) + predictions = np.concatenate(parts, axis=0) + else: + predictions = predict_fn(features) if return_proba: - predictions = self.classifier.predict_proba(features) predictions = np.moveaxis(predictions, -1, 0) # [nb_classes, width*height] - else: - predictions = self.classifier.predict(features) return predictions @@ -1294,7 +1371,8 @@ def reset_classifier(self): def _train(self, data, annotations, memory_mode=False, img_ids=None, use_rf=False, allow_writing_files=False, in_channels=None, skip_norm=False, - fe_use_device=None, clf_use_device=None, sort_features=True): + fe_use_device=None, clf_use_device=None, sort_features=True, + cancel_token=None): """ Backend training method for the Convpaint model. """ @@ -1304,6 +1382,36 @@ def _train(self, data, annotations, memory_mode=False, img_ids=None, use_rf=Fals warnings.warn('No annotations provided for training.') return self.classifier, None, None + # Snapshot the memory-mode bookkeeping so we can roll it back on + # CancelledError. _register_and_update_annots mutates self.annot_dict + # and self.table *before* feature extraction runs; without rollback a + # cancelled train would leave the annotations registered, and the next + # Train click would see "no new annotations" and fail with + # "No features or targets found". + mem_backup = None + if memory_mode: + mem_backup = (copy.deepcopy(self.annot_dict), self.table.copy(deep=True)) + + # Install the cancel token as the ambient token (utils.cancel_scope), so + # all downstream check_cancel() calls see it without it being passed + # through every signature (including custom FE subclasses). + try: + with utils.cancel_scope(cancel_token): + return self._train_body(data, annotations, memory_mode=memory_mode, + img_ids=img_ids, use_rf=use_rf, + allow_writing_files=allow_writing_files, + in_channels=in_channels, skip_norm=skip_norm, + fe_use_device=fe_use_device, + clf_use_device=clf_use_device, + sort_features=sort_features) + except utils.CancelledError: + if mem_backup is not None: + self.annot_dict, self.table = mem_backup + raise + + def _train_body(self, data, annotations, memory_mode=False, img_ids=None, use_rf=False, + allow_writing_files=False, in_channels=None, skip_norm=False, + fe_use_device=None, clf_use_device=None, sort_features=True): if not memory_mode: # Use _get_features to extract features and the suiting annotations parts (returns lists if restore_input_form=False) feature_parts, annot_parts = self._get_features( @@ -1345,6 +1453,9 @@ def _train(self, data, annotations, memory_mode=False, img_ids=None, use_rf=Fals # print("Targets:", targets, "values:", np.unique(targets)) raise ValueError('Not enough classes found in the targets. At least two classes are required for training.') + # Last cancel checkpoint before the (uninterruptible) classifier fit + utils.check_cancel() + # Train the classifier self._clf_train(features, targets, use_rf=use_rf, allow_writing_files=allow_writing_files, @@ -1504,7 +1615,7 @@ def _register_and_get_all_features_annots(self, features, annotations, coords, i return features, annotations - def _predict(self, data, add_seg=False, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None): + def _predict(self, data, add_seg=False, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None, cancel_token=None): """ Backend method to predict images as a whole or tiling and parallelizing the prediction. @@ -1519,6 +1630,16 @@ def _predict(self, data, add_seg=False, in_channels=None, skip_norm=False, use_d if self.classifier is None: raise ValueError('No trained classifier found.') + # Install the cancel token as the ambient token (utils.cancel_scope), so + # all downstream check_cancel() calls see it without it being passed + # through every signature (including custom FE subclasses). + with utils.cancel_scope(cancel_token): + return self._predict_body(data, add_seg=add_seg, in_channels=in_channels, + skip_norm=skip_norm, use_dask=use_dask, + fe_use_device=fe_use_device) + + def _predict_body(self, data, add_seg=False, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None): + # Check if we have only a single image input single_input = hasattr(data, 'ndim') and data.ndim >= 2 and not isinstance(data, list) input_shapes = [data.shape] if single_input else [d.shape for d in data] @@ -1535,9 +1656,11 @@ def _predict(self, data, add_seg=False, in_channels=None, skip_norm=False, use_d if not skip_norm: data = [self._norm_single_image(d) for d in data] - # Get class probabilities, using tiling if enabled + # Get class probabilities, using tiling if enabled. The per-tile check + # inside _parallel_predict_image takes over cancellation responsiveness. if self._param.tile_image: - probas = [self._parallel_predict_image(d, return_proba=True, use_dask=use_dask, fe_use_device=fe_use_device) + probas = [self._parallel_predict_image( + d, return_proba=True, use_dask=use_dask, fe_use_device=fe_use_device) for d in data] else: probas = self._predict_image(data, return_proba=True, fe_use_device=fe_use_device) # Can handle lists directly @@ -1592,7 +1715,11 @@ def _predict_image(self, image, return_proba=True, feature_img=None, fe_use_devi # Predict pixels based on the features and classifier # NOTE: We always first predict probabilities and then take the argmax feature_img = feature_img if isinstance(feature_img, list) else [feature_img] - predictions = [self._clf_predict(f, return_proba=True) for f in feature_img] + predictions = [] + for f in feature_img: + utils.check_cancel() + predictions.append(self._clf_predict(f, return_proba=True)) + utils.check_cancel() # Checkpoint before the (potentially large) reshaping # Reshape the predictions to the original image shape padded_shapes = self.padded_shapes # Saved when extracting features pre_pad_shapes = self.pre_pad_shapes # Saved when extracting features @@ -1676,91 +1803,112 @@ def _parallel_predict_image(self, image, return_proba=True, use_dask=False, fe_u if plot_tiles: img_to_plot = image.copy() - for row in range(nblocks_rows+1): - for col in range(nblocks_cols+1): - min_row = np.max([0, row*maxblock-margin]) - min_col = np.max([0, col*maxblock-margin]) - max_row = np.min([image.shape[-2], (row+1)*maxblock+margin]) - max_col = np.min([image.shape[-1], (col+1)*maxblock+margin]) - - min_row_ind = 0 - new_min_row_ind = 0 - if min_row > 0: - min_row_ind = min_row + margin - new_min_row_ind = margin - min_col_ind = 0 - new_min_col_ind = 0 - if min_col > 0: - min_col_ind = min_col + margin - new_min_col_ind = margin - - max_col = (col+1)*maxblock+margin - max_col_ind = np.min([min_col_ind+maxblock,image.shape[-1]]) - new_max_col_ind = new_min_col_ind + (max_col_ind-min_col_ind) - if max_col > image.shape[-1]: - max_col = image.shape[-1] - max_row = (row+1)*maxblock+margin - max_row_ind = np.min([min_row_ind+maxblock,image.shape[-2]]) - new_max_row_ind = new_min_row_ind + (max_row_ind-min_row_ind) - if max_row > image.shape[-2]: - max_row = image.shape[-2] - - image_block = image[..., min_row:max_row, min_col:max_col] - - # For plotting: - # Take the entire image, add boarders for the block - if plot_tiles: - img_to_plot[..., min_row, min_col:max_col] = 1 - img_to_plot[..., max_row-1, min_col:max_col] = 1 - img_to_plot[..., min_row:max_row, min_col] = 1 - img_to_plot[..., min_row:max_row, max_col-1] = 1 - img_to_plot[..., min_row_ind, min_col_ind:max_col_ind] = 0.5 - img_to_plot[..., max_row_ind-1, min_col_ind:max_col_ind] = 0.5 - img_to_plot[..., min_row_ind:max_row_ind, min_col_ind] = 0.5 - img_to_plot[..., min_row_ind:max_row_ind, max_col_ind-1] = 0.5 - - # Predict the block using dask or directly (with no normalization, as it is done outside) - if use_dask: - processes.append(client.submit( - self._predict_image, image=image_block, return_proba=return_proba, fe_use_device=fe_use_device)) - - min_row_ind_collection.append(min_row_ind) - min_col_ind_collection.append(min_col_ind) - max_row_ind_collection.append(max_row_ind) - max_col_ind_collection.append(max_col_ind) - new_max_col_ind_collection.append(new_max_col_ind) - new_max_row_ind_collection.append(new_max_row_ind) - new_min_col_ind_collection.append(new_min_col_ind) - new_min_row_ind_collection.append(new_min_row_ind) + # The try/finally guarantees the dask cluster is torn down even when a + # cancellation (or any other error) aborts the loops below — otherwise + # a cancel mid-run would leak the client and its worker processes. + # NOTE: cancellation inside dask workers is not supported (the ambient + # token does not cross process boundaries); cancel takes effect at the + # per-tile submission and gathering checkpoints, and pending tiles are + # cancelled in the finally block. + try: + for row in range(nblocks_rows+1): + for col in range(nblocks_cols+1): + utils.check_cancel() + min_row = np.max([0, row*maxblock-margin]) + min_col = np.max([0, col*maxblock-margin]) + max_row = np.min([image.shape[-2], (row+1)*maxblock+margin]) + max_col = np.min([image.shape[-1], (col+1)*maxblock+margin]) + + min_row_ind = 0 + new_min_row_ind = 0 + if min_row > 0: + min_row_ind = min_row + margin + new_min_row_ind = margin + min_col_ind = 0 + new_min_col_ind = 0 + if min_col > 0: + min_col_ind = min_col + margin + new_min_col_ind = margin + + max_col = (col+1)*maxblock+margin + max_col_ind = np.min([min_col_ind+maxblock,image.shape[-1]]) + new_max_col_ind = new_min_col_ind + (max_col_ind-min_col_ind) + if max_col > image.shape[-1]: + max_col = image.shape[-1] + max_row = (row+1)*maxblock+margin + max_row_ind = np.min([min_row_ind+maxblock,image.shape[-2]]) + new_max_row_ind = new_min_row_ind + (max_row_ind-min_row_ind) + if max_row > image.shape[-2]: + max_row = image.shape[-2] + + image_block = image[..., min_row:max_row, min_col:max_col] + + # For plotting: + # Take the entire image, add boarders for the block + if plot_tiles: + img_to_plot[..., min_row, min_col:max_col] = 1 + img_to_plot[..., max_row-1, min_col:max_col] = 1 + img_to_plot[..., min_row:max_row, min_col] = 1 + img_to_plot[..., min_row:max_row, max_col-1] = 1 + img_to_plot[..., min_row_ind, min_col_ind:max_col_ind] = 0.5 + img_to_plot[..., max_row_ind-1, min_col_ind:max_col_ind] = 0.5 + img_to_plot[..., min_row_ind:max_row_ind, min_col_ind] = 0.5 + img_to_plot[..., min_row_ind:max_row_ind, max_col_ind-1] = 0.5 + + # Predict the block using dask or directly (with no normalization, as it is done outside) + if use_dask: + processes.append(client.submit( + self._predict_image, image=image_block, return_proba=return_proba, fe_use_device=fe_use_device)) + + min_row_ind_collection.append(min_row_ind) + min_col_ind_collection.append(min_col_ind) + max_row_ind_collection.append(max_row_ind) + max_col_ind_collection.append(max_col_ind) + new_max_col_ind_collection.append(new_max_col_ind) + new_max_row_ind_collection.append(new_max_row_ind) + new_min_col_ind_collection.append(new_min_col_ind) + new_min_row_ind_collection.append(new_min_row_ind) - else: - predicted_image = self._predict_image(image_block, return_proba=return_proba, fe_use_device=fe_use_device) - crop_pred = predicted_image[..., - new_min_row_ind: new_max_row_ind, - new_min_col_ind: new_max_col_ind] + else: + predicted_image = self._predict_image(image_block, return_proba=return_proba, fe_use_device=fe_use_device) + crop_pred = predicted_image[..., + new_min_row_ind: new_max_row_ind, + new_min_col_ind: new_max_col_ind] + if not return_proba: + crop_pred = crop_pred.astype(np.uint8) + predicted_image_complete[..., + min_row_ind:max_row_ind, + min_col_ind:max_col_ind] = crop_pred + + # Gather the results of the dask processes if enabled + if use_dask: + from dask.distributed import TimeoutError as DaskTimeoutError + for k in range(len(processes)): + future = processes[k] + # Poll instead of blocking indefinitely, so a cancel is + # honored while waiting for a tile that is still computing. + while True: + utils.check_cancel() + try: + out = future.result(timeout=1) + break + except DaskTimeoutError: + continue + crop_out = out[..., + new_min_row_ind_collection[k]:new_max_row_ind_collection[k], + new_min_col_ind_collection[k]:new_max_col_ind_collection[k]] if not return_proba: - crop_pred = crop_pred.astype(np.uint8) + crop_out = crop_out.astype(np.uint8) + # Release the future's result once it is written to the complete image + future.cancel() predicted_image_complete[..., - min_row_ind:max_row_ind, - min_col_ind:max_col_ind] = crop_pred - - # Terminate dask processes if enabled - if use_dask: - for k in range(len(processes)): - future = processes[k] - out = future.result() - crop_out = out[..., - new_min_row_ind_collection[k]:new_max_row_ind_collection[k], - new_min_col_ind_collection[k]:new_max_col_ind_collection[k]] - if not return_proba: - crop_out = crop_out.astype(np.uint8) - # Write the result to the complete image - future.cancel() - del future - predicted_image_complete[..., - min_row_ind_collection[k]:max_row_ind_collection[k], - min_col_ind_collection[k]:max_col_ind_collection[k]] = crop_out - client.close() + min_row_ind_collection[k]:max_row_ind_collection[k], + min_col_ind_collection[k]:max_col_ind_collection[k]] = crop_out + finally: + if use_dask: + for future in processes: + future.cancel() + client.close() if plot_tiles: from matplotlib import pyplot as plt diff --git a/src/napari_convpaint/convpaint_widget.py b/src/napari_convpaint/convpaint_widget.py index 649f069..b7ea0c6 100644 --- a/src/napari_convpaint/convpaint_widget.py +++ b/src/napari_convpaint/convpaint_widget.py @@ -1,5 +1,7 @@ +from dataclasses import dataclass, field +from typing import Optional from qtpy.QtWidgets import (QWidget, QPushButton,QVBoxLayout, - QLabel, QComboBox,QFileDialog, QListWidget, + QLabel, QComboBox,QFileDialog, QListWidget, QApplication, QCheckBox, QAbstractItemView, QGridLayout, QSpinBox, QButtonGroup, QRadioButton,QDoubleSpinBox, QTableWidget, QTableWidgetItem, QHeaderView, QMessageBox) @@ -9,6 +11,13 @@ import napari from napari.utils import progress from napari.utils.notifications import show_info +# Use superqt's thread_worker rather than napari.qt.threading.thread_worker: the +# latter registers every worker with window._task_status_manager (and never +# unregisters), so the worker's closure — including cp_model with its +# VGG16 MPS weights — is pinned for the lifetime of the viewer. In a test +# loop creating many widgets this accumulates and blows past the macOS +# runner's 7.93 GiB MPS cap. +from superqt.utils import thread_worker from napari_guitils.gui_structures import VHGroup, TabSet from pathlib import Path import numpy as np @@ -21,8 +30,29 @@ # import torch # from .utils import normalize_image, compute_image_stats, normalize_image_percentile, normalize_image_imagenet, get_fe_device # from .convpaint_model import ConvpaintModel +# CancelToken / CancelledError are also imported inline inside the slot +# methods that need them, to avoid pulling in .utils (and its torch import) +# at widget-module load time. + + +@dataclass +class _ActiveOp: + name: str # 'train' | 'train_multiple' | 'predict' | 'predict_all' | 'features' | 'features_all' | 'segment_files' + cancel_token: object # CancelToken — not annotated as a forward ref so @dataclass doesn't try to resolve it at decoration time + button: QPushButton + button_orig_text: str + disabled_buttons: list = field(default_factory=list) + cancel_was_requested: bool = False + pbar: object = None # napari progress bar, for switching to 'Cancelling…' + class ConvpaintWidget(QWidget): + + # When True, long-running operations (train/predict/predict_all) run on the + # calling thread instead of a worker thread. Test code sets this to keep the + # existing synchronous test assertions valid. + _sync_workers = False + """ Implementation of a napari widget for interactive segmentation performed via multiple means of feature extraction combined with a CatBoost Classifier @@ -1710,113 +1740,246 @@ def _on_add_annot_layer(self, event=None, force_add=True): # Train def _on_train(self, event=None): - """Given a set of new annotations, update the CatBoost classifier.""" + """Button slot: start training, or cancel the in-progress training.""" + if self._handle_cancel_click('train'): + return - # Get the data img = self._get_selected_img(check=True) annot = self.annotations_layer_selection_widget.value mem_mode = (self.cont_training == "Image" or self.cont_training == "Global") - # Check if annotations of at least 2 classes are present if annot is None: raise Exception('No annotations layer selected. Please create/select one.') unique_labels = np.unique(annot.data) unique_labels = unique_labels[unique_labels != 0] if len(unique_labels) < 2: if not mem_mode: - raise Exception('You need annotations for at least foreground and background') + raise Exception('Training requires annotations of at least 2 classes ' + '(e.g. foreground and background).') if self.cp_model.num_trainings == 0: - raise Exception('Model has not yet been trained. You need annotations for at least foreground and background') + raise Exception('Training requires annotations of at least 2 classes ' + '(e.g. foreground and background). With continuous training, ' + 'a single class is only allowed once the model has been trained before.') # Check if annotations layer has correct shape for the chosen data type if not self._approve_annotations_layer_shape(annot, img): raise Exception('annotations layer has wrong shape for the chosen data') - # Set the current model path to 'in training' and adjust the model description self.current_model_path = 'in training' self._set_model_description() - # Get the image data and normalize it; also get the annotations + # Snapshot inputs on the main thread so the worker never touches UI state. + # Copy the annotations: the UI stays responsive during training, so the + # user could otherwise keep painting into the very array the worker is + # reading, yielding inconsistent features/targets. image_stack_norm = self._get_data_channel_first_norm(img) - annot = annot.data - - # Start training + annot_data = annot.data.copy() + img_name = img.name + in_channels = self._parse_in_channels(self.input_channels) + fe_device = self.fe_device + clf_device = self.clf_device + cp_model = self.cp_model + + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() + + @thread_worker + def _do_train(): + # Swallow CancelledError here so it never reaches the worker's errored + # signal; real exceptions still propagate. + try: + with warnings.catch_warnings(): + warnings.simplefilter(action="ignore", category=FutureWarning) + cp_model.train(image_stack_norm, annot_data, memory_mode=mem_mode, + img_ids=img_name, in_channels=in_channels, skip_norm=False, + fe_use_device=fe_device, clf_use_device=clf_device, + cancel_token=cancel_token) + except CancelledError: + return None + + worker = _do_train() + worker.returned.connect(self._on_train_returned) + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + self._begin_worker('train', self.train_classifier_btn, cancel_token, worker, + desc='Training') + + def _on_train_returned(self, _result): + if self._op is not None and self._op.cancel_was_requested: + return + self._update_training_counts() + self.current_model_path = 'trained, unsaved' + self.trained = True + self._set_model_description() + self._pending_auto_seg = self.auto_seg + + def _handle_cancel_click(self, op_name): + """If a worker is active, request cancellation (only if op matches) and + tell the caller to return without starting new work.""" + if self._op is None: + return False + if self._op.name == op_name and not self._op.cancel_was_requested: + self._op.cancel_was_requested = True + self._op.cancel_token.cancel() + # Immediate feedback: the worker only stops at its next cancel + # checkpoint — an in-flight CatBoost fit or a single FE forward + # pass cannot be interrupted — so reflect that the cancel was + # registered and ignore further clicks until it takes effect. + self._op.button.setText('Cancelling…') + self._op.button.setEnabled(False) + if self._op.pbar is not None: + # No ellipsis here: napari appends ': ' after the description + self._op.pbar.set_description('Cancelling') + return True + + def _other_op_buttons(self, current_button): + """All existing operation-launching buttons except the given one. + Some of these buttons are only created in certain configurations + (advanced mode, Multifile tab), hence the getattr.""" + names = ['train_classifier_btn', 'segment_btn', 'segment_all_btn', + 'btn_add_features', 'btn_add_features_stack', + 'btn_train_on_selected', 'multifile_train_all_annot_btn', + 'multifile_preview_btn', 'multifile_segment_selected_btn'] + buttons = (getattr(self, n, None) for n in names) + return [b for b in buttons if b is not None and b is not current_button] + + def _begin_worker(self, name, button, cancel_token, worker, + desc='', total=0): + # Drain the delayed _on_select_layer QTimer before we start — otherwise + # it fires mid-op during layer-data assignment and resets the classifier. + # The old synchronous code got this flush for free from napari.utils.progress. + QApplication.processEvents() + # Create the progress bar on the main thread (QWidgets cannot be + # constructed from a worker thread on macOS — it raises NSInternalInconsistencyException). + pbar = progress(total=total, desc=desc) + worker.finished.connect(pbar.close) + if total: + worker.yielded.connect(pbar.increment_with_overflow) + self._op = _ActiveOp( + name=name, + cancel_token=cancel_token, + button=button, + button_orig_text=button.text(), + # Ops are mutually exclusive: all other op buttons are disabled + # while one is running (only the running op's button stays live, + # doubling as the Cancel button). + disabled_buttons=self._other_op_buttons(button), + pbar=pbar, + ) + button.setText('Cancel') + for b in self._op.disabled_buttons: + b.setEnabled(False) with warnings.catch_warnings(): warnings.simplefilter(action="ignore", category=FutureWarning) self.viewer.window._status_bar._toggle_activity_dock(True) + if self._sync_workers: + worker.run() + else: + worker.start() - with progress(total=0) as pbr: - pbr.set_description(f"Training") - img_name = self._get_selected_img().name - in_channels = self._parse_in_channels(self.input_channels) - # Train the model with the current image and annotations; skip normalization as it is done in the widget - _ = self.cp_model.train(image_stack_norm, annot, memory_mode=mem_mode, img_ids=img_name, - in_channels=in_channels, skip_norm=False, - fe_use_device=self.fe_device, clf_use_device=self.clf_device) - self._update_training_counts() - + def _on_worker_finished(self): + op = self._op + self._op = None + if op is None: + return + op.button.setText(op.button_orig_text) + op.button.setEnabled(True) # was disabled while 'Cancelling…' + # _reset_predict_buttons below re-decides segment/segment-all state based + # on self.trained; the train button has no such gating, so restoring it + # here unconditionally is what keeps it clickable after a predict run. + for b in op.disabled_buttons: + b.setEnabled(True) with warnings.catch_warnings(): warnings.simplefilter(action="ignore", category=FutureWarning) self.viewer.window._status_bar._toggle_activity_dock(False) - - # Set the current model path to 'trained, unsaved' and adjust the model description - self.current_model_path = 'trained, unsaved' - self.trained = True self._reset_predict_buttons() - self._set_model_description() + if op.cancel_was_requested: + show_info('Operation cancelled.') + if op.name in ('train', 'train_multiple') and self.current_model_path == 'in training': + self.current_model_path = 'not trained' if not self.trained else 'trained, unsaved' + self._set_model_description() + return + if op.name == 'train' and getattr(self, '_pending_auto_seg', False): + self._pending_auto_seg = False + if self.trained: + self._on_predict() - # Automatically segment the image if the option is activated - if self.auto_seg: - self._on_predict() + def _on_worker_errored(self, exc): + # CancelledError is swallowed inside each worker body, so only real + # failures reach here — napari's default error handler still displays + # the traceback; we just tidy up the 'in training' label. + if (self._op is not None and self._op.name in ('train', 'train_multiple') + and self.current_model_path == 'in training'): + self.current_model_path = 'not trained' if not self.trained else 'trained, unsaved' + self._set_model_description() # Predict def _on_predict(self, event=None): - """Predict the segmentation of the currently viewed frame based - on a classifier trained with annotations.""" + """Button slot: start single-frame prediction, or cancel the running one.""" + + if self._handle_cancel_click('predict'): + return if not (self.add_seg or self.add_probas): warnings.warn('Neither segmentation nor probabilities output selected to be added. Nothing to do.') return - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(True) - - with progress(total=0) as pbr: - pbr.set_description(f"Prediction") - - # Check dimensionality - img = self._get_selected_img() - data_dims = self._get_data_dims(img.data, img.ndim) if img is not None else None - if data_dims not in self.supported_data_dims: - warnings.warn(f'Non-supported image dimensions {data_dims}. Prediction not performed.') - return - - # Get the data - image_plane = self._get_current_plane_norm() - in_channels = self._parse_in_channels(self.input_channels) + img = self._get_selected_img() + data_dims = self._get_data_dims(img.data, img.ndim) if img is not None else None + if data_dims not in self.supported_data_dims: + warnings.warn(f'Non-supported image dimensions {data_dims}. Prediction not performed.') + return - # Predict image (use backend function which returns probabilities and segmentation); skip norm as it is done above - probas, segmentation = self.cp_model._predict(image_plane, add_seg=True, in_channels=in_channels, skip_norm=True, - use_dask=self.use_dask, fe_use_device=self.fe_device) + image_plane = self._get_current_plane_norm() + in_channels = self._parse_in_channels(self.input_channels) + use_dask = self.use_dask + fe_device = self.fe_device + cp_model = self.cp_model - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(False) + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() - # Get the current step in case of stacks + @thread_worker + def _do_predict(): + try: + return cp_model._predict(image_plane, add_seg=True, in_channels=in_channels, skip_norm=True, + use_dask=use_dask, fe_use_device=fe_device, + cancel_token=cancel_token) + except CancelledError: + return None + + # Capture the target slice now: the UI stays responsive during the + # prediction, so the user may scroll the stack before the result + # arrives — it must land on the slice that was predicted, not the one + # viewed at completion time. step = self.viewer.dims.current_step[-3] if data_dims in ['3D_single', '4D', '3D_RGB'] else None + self._pending_predict_ctx = (data_dims, step) + + worker = _do_predict() + worker.returned.connect(self._on_predict_returned) + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + # This slot serves both the Segment button and the Multifile preview + # button; show 'Cancel' on whichever one was actually clicked. + sender = self.sender() + button = sender if isinstance(sender, QPushButton) else self.segment_btn + self._begin_worker('predict', button, cancel_token, worker, + desc='Prediction') + + def _on_predict_returned(self, result): + cancelled = self._op is not None and self._op.cancel_was_requested + ctx = getattr(self, '_pending_predict_ctx', None) + self._pending_predict_ctx = None + if cancelled or result is None or ctx is None: + return + probas, segmentation = result + data_dims, step = ctx - # Add segmentation layer if enabled if self.add_seg: - # Check if we need to create a new segmentation layer self._check_create_segmentation_layer() - # Set the flag to False, so we don't create a new layer every time self.new_seg = False - - # Update segmentation layer if data_dims in ['2D', '2D_RGB', '3D_multi']: self.viewer.layers[self.seg_tag].data = segmentation elif data_dims in ['3D_single', '4D', '3D_RGB']: # seg has no channel dim -> z is first @@ -1824,55 +1987,66 @@ def _on_predict(self, event=None): # Case `data_dims is None` and other invalid cases are already caught above, so we don't need an else statement here self.viewer.layers[self.seg_tag].refresh() - # Add probabilities if enabled if self.add_probas: - # Check if we need to create a new probabilities layer num_classes = probas.shape[:1] self._check_create_probas_layer(num_classes) - # Set the flag to False, so we don't create a new layer every time self.new_proba = False - - # Update probabilities layer - if data_dims in ['2D', '2D_RGB', '3D_multi']: # No stack dim + if data_dims in ['2D', '2D_RGB', '3D_multi']: self.viewer.layers[self.proba_prefix].data = probas - elif data_dims in ['3D_single', '4D', '3D_RGB']: # (stack dim is second, probas first) + elif data_dims in ['3D_single', '4D', '3D_RGB']: self.viewer.layers[self.proba_prefix].data[:, step] = probas - # Case `data_dims is None` and other invalid cases are already caught above, so we don't need an else statement here self.viewer.layers[self.proba_prefix].refresh() def _on_get_feature_image(self, event=None): - """Get the feature image for the currently viewed frame based - on the current feature extractor and show it in a new layer.""" - - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(True) - - with progress(total=0) as pbr: - pbr.set_description(f"Feature extraction") + """Button slot: start feature extraction for the currently viewed frame, + or cancel the running one.""" + if self._handle_cancel_click('features'): + return - # Check dimensionality - img = self._get_selected_img() - data_dims = self._get_data_dims(img.data, img.ndim) if img is not None else None - if data_dims not in self.supported_data_dims: - warnings.warn(f'Non-supported image dimensions {data_dims}. Feature extraction not performed.') - return + img = self._get_selected_img() + data_dims = self._get_data_dims(img.data, img.ndim) if img is not None else None + if data_dims not in self.supported_data_dims: + warnings.warn(f'Non-supported image dimensions {data_dims}. Feature extraction not performed.') + return - # Get the data - image_plane = self._get_current_plane_norm() - in_channels = self._parse_in_channels(self.input_channels) + # Snapshot inputs on the main thread so the worker never touches UI state. + image_plane = self._get_current_plane_norm() + in_channels = self._parse_in_channels(self.input_channels) + pca, kmeans = self._check_parse_pca_kmeans() + fe_device = self.fe_device + cp_model = self.cp_model - # Check and parse PCA and Kmeans parameters - pca, kmeans = self._check_parse_pca_kmeans() + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() - # Get feature image; skip norm as it is done above - feature_image = self.cp_model.get_feature_image(image_plane, in_channels=in_channels, skip_norm=True, - pca_components=pca, kmeans_clusters=kmeans, - use_device=self.fe_device) + @thread_worker + def _do_features(): + try: + return cp_model.get_feature_image(image_plane, in_channels=in_channels, skip_norm=True, + pca_components=pca, kmeans_clusters=kmeans, + use_device=fe_device, cancel_token=cancel_token) + except CancelledError: + return None - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(False) + # Capture the target slice now (see _on_predict for why). + step = self.viewer.dims.current_step[-3] if data_dims in ['3D_single', '4D', '3D_RGB'] else None + self._pending_features_ctx = (data_dims, step, kmeans) + + worker = _do_features() + worker.returned.connect(self._on_features_returned) + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + self._begin_worker('features', self.btn_add_features, cancel_token, worker, + desc='Feature extraction') + + def _on_features_returned(self, result): + cancelled = self._op is not None and self._op.cancel_was_requested + ctx = getattr(self, '_pending_features_ctx', None) + self._pending_features_ctx = None + if cancelled or result is None or ctx is None: + return + data_dims, step, kmeans = ctx + feature_image = result # Check if we need to create a new features layer num_features = feature_image.shape[0] if not kmeans else 0 @@ -1884,16 +2058,15 @@ def _on_get_feature_image(self, event=None): if data_dims in ['2D', '2D_RGB', '3D_multi']: # No stack dim self.viewer.layers[self.features_prefix].data = feature_image elif data_dims in ['3D_single', '4D', '3D_RGB']: # stack dim is third last - step = self.viewer.dims.current_step[-3] self.viewer.layers[self.features_prefix].data[..., step, :, :] = feature_image # Case `data_dims is None` and other invalid cases are already caught above, so we don't need an else statement here self.viewer.layers[self.features_prefix].refresh() - def _on_predict_all(self): - """Predict the segmentation of all frames based - on a classifier model trained with annotations.""" - - # Get the data + def _on_predict_all(self, event=None): + """Button slot: start stack prediction, or cancel the running one.""" + if self._handle_cancel_click('predict_all'): + return + img = self._get_selected_img(check=True) # Check dimensionality @@ -1901,59 +2074,63 @@ def _on_predict_all(self): if data_dims not in ['3D_single', '3D_RGB', '4D']: warnings.warn(f'Image stack has wrong dimensionality ({data_dims}) for predicting stacks. Prediction not performed.') return - - # Create the segmentation layer if it is not already present - # (NOTE: probabilities layer is created in the prediction loop, as we need to know the number of classes) + + # Create seg layer up front so the worker can yield into it. The probas + # layer needs num_classes and is created when the first slice arrives. if self.add_seg: self._check_create_segmentation_layer() - # Set the flag to False, so we don't create a new layer every time self.new_seg = False - # Start prediction - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(True) - - # Get normalized stack data (entire stack, and stats prepared given the radio buttons) - image_stack_norm = self._get_data_channel_first_norm(img) # Normalize the entire stack - - # Step through the stack and predict each image + image_stack_norm = self._get_data_channel_first_norm(img) + in_channels = self._parse_in_channels(self.input_channels) + use_dask = self.use_dask + fe_device = self.fe_device + cp_model = self.cp_model num_steps = image_stack_norm.shape[-3] - for step in progress(range(num_steps)): - - # Take the slice of the 3rd last dimension (since images are C, Z, H, W or Z, H, W) - image = image_stack_norm[..., step, :, :] - - # Predict the current step; skip normalization as it is done above - in_channels = self._parse_in_channels(self.input_channels) - # Use the backend function which returns probabilities and segmentation - probas, seg = self.cp_model._predict(image, add_seg=True, in_channels=in_channels, skip_norm=True, - use_dask=self.use_dask, fe_use_device=self.fe_device) - - # In the first iteration, check if we need to create a new probas layer - # (we need the information about the number of classes) - if step == 0 and self.add_probas: - num_classes = probas.shape[0] - # Check if we need to create a new probabilities layer - self._check_create_probas_layer(num_classes) - # Set the flag to False, so we don't create a new layer every time - self.new_proba = False - - # Add the slices to the segmentation and probabilities layers - if self.add_seg: - self.viewer.layers[self.seg_tag].data[step] = seg - self.viewer.layers[self.seg_tag].refresh() - if self.add_probas: - self.viewer.layers[self.proba_prefix].data[..., step, :, :] = probas - self.viewer.layers[self.proba_prefix].refresh() - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(False) + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() + + @thread_worker + def _do_predict_all(): + try: + for step in range(num_steps): + cancel_token.raise_if_cancelled() + image = image_stack_norm[..., step, :, :] + probas, seg = cp_model._predict(image, add_seg=True, in_channels=in_channels, skip_norm=True, + use_dask=use_dask, fe_use_device=fe_device, + cancel_token=cancel_token) + yield step, probas, seg + except CancelledError: + # Any slices already yielded stay in the labels layer. + return + + worker = _do_predict_all() + worker.yielded.connect(self._on_predict_all_yielded) + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + self._begin_worker('predict_all', self.segment_all_btn, cancel_token, worker, + desc='Segmenting stack', total=num_steps) + + def _on_predict_all_yielded(self, value): + step, probas, seg = value + if step == 0 and self.add_probas: + num_classes = probas.shape[0] + self._check_create_probas_layer(num_classes) + self.new_proba = False + # Add the slices to the segmentation and probabilities layers + if self.add_seg: + self.viewer.layers[self.seg_tag].data[step] = seg + self.viewer.layers[self.seg_tag].refresh() + if self.add_probas: + self.viewer.layers[self.proba_prefix].data[..., step, :, :] = probas + self.viewer.layers[self.proba_prefix].refresh() - def _on_get_feature_image_all(self): - """Get the feature image for all frames based - on the current feature extractor and show it in a new layer.""" + def _on_get_feature_image_all(self, event=None): + """Button slot: start feature extraction for the whole stack, + or cancel the running one.""" + if self._handle_cancel_click('features_all'): + return # Get the data img = self._get_selected_img(check=True) @@ -1963,66 +2140,82 @@ def _on_get_feature_image_all(self): if data_dims not in ['3D_single', '3D_RGB', '4D']: warnings.warn(f'Image stack has wrong dimensionality ({data_dims}) for processing stacks. Feature extraction not performed.') return - - # Start feature extraction - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(True) - # Get normalized stack data (entire stack, and stats prepared given the radio buttons) + # Snapshot inputs on the main thread so the worker never touches UI state. image_stack_norm = self._get_data_channel_first_norm(img) # Normalize the entire stack pca, kmeans = self._check_parse_pca_kmeans() in_channels = self._parse_in_channels(self.input_channels) + fe_device = self.fe_device + cp_model = self.cp_model + num_steps = image_stack_norm.shape[-3] + + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() if kmeans: - # Get feature image for entire stack; skip norm as it is done above - feature_image = self.cp_model.get_feature_image(image_stack_norm, in_channels=in_channels, skip_norm=True, - pca_components=pca, kmeans_clusters=kmeans, - use_device=self.fe_device) - - # Check if we need to create a new features layer - # num_features = feature_image.shape[0] if not kmeans else 0 - # self._check_create_features_layer(num_features) - self._check_create_features_layer(0) - # Set the flag to False, so we don't create a new layer every time - self.new_features = False - # Update features layer - self.viewer.layers[self.features_prefix].data = feature_image + # Kmeans needs the entire stack at once, so there is no per-slice + # progress; the result arrives as a whole via `returned`. + @thread_worker + def _do_features_all(): + try: + return cp_model.get_feature_image(image_stack_norm, in_channels=in_channels, skip_norm=True, + pca_components=pca, kmeans_clusters=kmeans, + use_device=fe_device, cancel_token=cancel_token) + except CancelledError: + return None + + worker = _do_features_all() + worker.returned.connect(self._on_features_all_returned) + total = 0 else: # No kmeans, can do step-by-step to save memory (and show progress) - # Step through the stack and predict each image - num_steps = image_stack_norm.shape[-3] - for step in progress(range(num_steps)): - - # Take the slice of the 3rd last dimension (since images are C, Z, H, W or Z, H, W) - image = image_stack_norm[..., step, :, :] - - # Predict the current step; skip normalization as it is done above - # Get feature image; skip norm as it is done above - feature_image = self.cp_model.get_feature_image(image, in_channels=in_channels, skip_norm=True, - pca_components=pca, kmeans_clusters=kmeans, - use_device=self.fe_device) - - # In the first iteration, check if we need to create a new features layer - # (we need the information about the number of classes) - if step == 0: - # Check if we need to create a new features layer - num_features = feature_image.shape[0] if not kmeans else 0 - self._check_create_features_layer(num_features) - # Set the flag to False, so we don't create a new layer every time - self.new_features = False - - # Add the slices to the segmentation and probabilities layers - # if kmeans: - # self.viewer.layers[self.features_prefix].data[step] = feature_image - # self.viewer.layers[self.features_prefix].refresh() - # else: - self.viewer.layers[self.features_prefix].data[..., step, :, :] = feature_image - self.viewer.layers[self.features_prefix].refresh() - - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(False) + @thread_worker + def _do_features_all(): + try: + for step in range(num_steps): + cancel_token.raise_if_cancelled() + # Take the slice of the 3rd last dimension (since images are C, Z, H, W or Z, H, W) + image = image_stack_norm[..., step, :, :] + feature_image = cp_model.get_feature_image(image, in_channels=in_channels, skip_norm=True, + pca_components=pca, kmeans_clusters=kmeans, + use_device=fe_device, cancel_token=cancel_token) + yield step, feature_image + except CancelledError: + # Any slices already yielded stay in the features layer. + return + + worker = _do_features_all() + worker.yielded.connect(self._on_features_all_yielded) + total = num_steps + + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + self._begin_worker('features_all', self.btn_add_features_stack, cancel_token, worker, + desc='Extracting features', total=total) + + def _on_features_all_returned(self, result): + """Handles the kmeans (whole-stack) variant of _on_get_feature_image_all.""" + cancelled = self._op is not None and self._op.cancel_was_requested + if cancelled or result is None: + return + self._check_create_features_layer(0) + # Set the flag to False, so we don't create a new layer every time + self.new_features = False + # Update features layer + self.viewer.layers[self.features_prefix].data = result + self.viewer.layers[self.features_prefix].refresh() + + def _on_features_all_yielded(self, value): + step, feature_image = value + # In the first iteration, check if we need to create a new features layer + # (we need the information about the number of features) + if step == 0: + self._check_create_features_layer(feature_image.shape[0]) + # Set the flag to False, so we don't create a new layer every time + self.new_features = False + # Add the slice to the features layer + self.viewer.layers[self.features_prefix].data[..., step, :, :] = feature_image + self.viewer.layers[self.features_prefix].refresh() # Load/Save @@ -2300,6 +2493,7 @@ def _reset_attributes(self): self.cmap_flag = False # Flag to prevent infinite loops when changing colormaps self.labels_cmap = None # Colormap for the labels (annotations and segmentation) self._block_layer_select = True # Flag to block layer selection events temporarily + self._op: Optional[_ActiveOp] = None # Multifile attributes self._multifile_warned = False # Whether the user has already seen the "remove existing layers" warning (for Multifile) self._multifile_annotations_store = {} # Store for in-memory and saved annotations keyed by filename @@ -3578,12 +3772,15 @@ def _update_annotations_layers(self): # self.annotations_layer_selection_widget.choices = [(layer.name, layer) for layer in annot_layer_list] # return annot_layer_list - def _train_multiple(self, img_list, annot_list, id_list): - """Core training routine used by multiple callers. + def _train_multiple(self, img_list, annot_list, id_list, button=None): + """Core training routine used by multiple callers (train-on-selected, + multifile training). Runs the actual training on a cancellable worker; + callers must assemble the data lists on the main thread beforehand. id_list: list of str image ids/names img_list: list of numpy arrays (prepared via _get_data_channel_first) annot_list: list of numpy arrays (annotations masks) + button: the QPushButton that triggered the call (doubles as Cancel) """ if self.cp_model is None: warnings.warn('No model set. Cannot train.') @@ -3597,29 +3794,50 @@ def _train_multiple(self, img_list, annot_list, id_list): warnings.warn('Image and annotations lists must have identical lengths.') return - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(True) - - with progress(total=0) as pbr: - pbr.set_description("Training") - mem_mode = (self.cont_training == "Image" - or self.cont_training == "Global") - # Train; in this case, normalization is not skipped (but done in the ConvpaintModel) - in_channels = self._parse_in_channels(self.input_channels) - _ = self.cp_model.train(img_list, annot_list, memory_mode=mem_mode, img_ids=id_list, - in_channels=in_channels, skip_norm=False, - fe_use_device=self.fe_device, clf_use_device=self.clf_device) - self._update_training_counts() - - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(False) + self.current_model_path = 'in training' + self._set_model_description() - # Set the current model path to 'trained, unsaved' and adjust the model description + mem_mode = (self.cont_training == "Image" + or self.cont_training == "Global") + in_channels = self._parse_in_channels(self.input_channels) + # Copy annotations arrays: the UI stays responsive during training and + # some callers pass live layer data the user could keep painting into. + annot_list = [np.copy(a) for a in annot_list] + fe_device = self.fe_device + clf_device = self.clf_device + cp_model = self.cp_model + if button is None: + button = self.train_classifier_btn + + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() + + @thread_worker + def _do_train_multiple(): + try: + with warnings.catch_warnings(): + warnings.simplefilter(action="ignore", category=FutureWarning) + # Normalization is not skipped (but done in the ConvpaintModel) + cp_model.train(img_list, annot_list, memory_mode=mem_mode, img_ids=id_list, + in_channels=in_channels, skip_norm=False, + fe_use_device=fe_device, clf_use_device=clf_device, + cancel_token=cancel_token) + except CancelledError: + return None + + worker = _do_train_multiple() + worker.returned.connect(self._on_train_multiple_returned) + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + self._begin_worker('train_multiple', button, cancel_token, worker, + desc='Training') + + def _on_train_multiple_returned(self, _result): + if self._op is not None and self._op.cancel_was_requested: + return + self._update_training_counts() self.current_model_path = 'trained, unsaved' self.trained = True - self._reset_predict_buttons() self._set_model_description() ### ADVANCED TAB @@ -3664,6 +3882,8 @@ def _on_add_all_annot_layers(self): def _on_train_on_selected(self): """Train the model on the image and annotations layers currently selected in the layers widget (napari). Has been deprecated in favor of the "Multifile" Tab.""" + if self._handle_cancel_click('train_multiple'): + return # Get selected layers (arbitrary order) and sort them by their names layer_list = list(self.viewer.layers.selection) @@ -3690,7 +3910,7 @@ def _on_train_on_selected(self): id_list = [img.name for img in img_list] # Delegate core training to helper that can be reused - self._train_multiple(arr_list, annot_list, id_list) + self._train_multiple(arr_list, annot_list, id_list, button=self.btn_train_on_selected) def _update_training_counts(self): """Update the training counts (used with continuous_training/memory_mode) in the GUI.""" @@ -4255,6 +4475,9 @@ def _on_train_on_multifile(self): Ensures any current open multifile annotations is pushed to the in-memory store before assembling lists and delegating to `_train_multiple`. """ + if self._handle_cancel_click('train_multiple'): + return + # Ensure current open annotations (if any, and if it has annotations) is saved to the store fname = getattr(self, '_current_multifile_filename', None) if fname is not None and self.annot_tag in self.viewer.layers: @@ -4332,7 +4555,8 @@ def _on_train_on_multifile(self): # Train on all images # Delegate to core trainer - self._train_multiple(img_prepared, annots, filenames[:len(img_prepared)]) + self._train_multiple(img_prepared, annots, filenames[:len(img_prepared)], + button=self.multifile_train_all_annot_btn) def _on_segment_selected_multifile(self): """Segment selected files from the multifile list and save outputs to disk. @@ -4342,6 +4566,9 @@ def _on_segment_selected_multifile(self): - Loop over selected filenames, run backend predict and save segmentation TIFFs - Register saved paths in `_multifile_segmentation_store` and update table ticks. """ + if self._handle_cancel_click('segment_files'): + return + # Get selected rows try: sel = self.multifile_list.selectionModel().selectedRows() @@ -4408,62 +4635,97 @@ def _on_segment_selected_multifile(self): folder = Path(folder_text) if folder_text else None in_channels = self._parse_in_channels(self.input_channels) - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(True) + # Snapshot everything the worker needs; file reading, prediction and + # TIFF writing run on the worker thread, while the store/table-tick + # updates (Qt) happen on the main thread in the yielded handler. + cp_model = self.cp_model + use_dask = self.use_dask + fe_device = self.fe_device + is_rgb = cp_model.get_param('channel_mode') == 'rgb' + seg_tag = self.seg_tag - with progress(total=len(filenames)) as pbr: - pbr.set_description('Segmenting') + from .utils import CancelToken, CancelledError + cancel_token = CancelToken() + @thread_worker + def _do_segment_files(): segmented = 0 - for fname in filenames: - pbr.update(1) - try: - img_path = folder / fname - arr = imageio.imread(str(img_path)) - except Exception: - warnings.warn(f'Could not read image {fname}. Skipping.') - continue - - try: - is_rgb = self.cp_model.get_param('channel_mode') == 'rgb' - dims = arr.ndim if not is_rgb else arr.ndim - 1 # Account for channel dimension in RGB mode - prep = self._get_data_channel_first(arr, dims) if arr is not None else None - except Exception: - warnings.warn(f'Could not prepare image {fname} for prediction. Skipping.') - continue + try: + for fname in filenames: + cancel_token.raise_if_cancelled() + try: + arr = imageio.imread(str(folder / fname)) + except Exception: + warnings.warn(f'Could not read image {fname}. Skipping.') + yield fname, None + continue - try: - probas, seg = self.cp_model._predict(prep, add_seg=True, in_channels=in_channels, - skip_norm=False, use_dask=self.use_dask, - fe_use_device=self.fe_device) - except Exception: - warnings.warn(f'Prediction failed for {fname}. Skipping.') - continue + try: + dims = arr.ndim if not is_rgb else arr.ndim - 1 # Account for channel dimension in RGB mode + prep = self._get_data_channel_first(arr, dims) if arr is not None else None + except Exception: + warnings.warn(f'Could not prepare image {fname} for prediction. Skipping.') + yield fname, None + continue - try: - stem = Path(fname).stem - out_name = out_dir / f"{stem}_{self.seg_tag}.tif" - tifffile.imwrite(str(out_name), seg.astype(np.uint8)) - # Register saved segmentation and update table - self._multifile_segmentation_store[fname] = str(out_name) - self._update_multifile_seg_tick(fname) - segmented += 1 - except Exception: - warnings.warn(f'Could not write segmentation for {fname}.') + try: + probas, seg = cp_model._predict(prep, add_seg=True, in_channels=in_channels, + skip_norm=False, use_dask=use_dask, + fe_use_device=fe_device, + cancel_token=cancel_token) + except CancelledError: + raise + except Exception: + warnings.warn(f'Prediction failed for {fname}. Skipping.') + yield fname, None + continue - with warnings.catch_warnings(): - warnings.simplefilter(action="ignore", category=FutureWarning) - self.viewer.window._status_bar._toggle_activity_dock(False) + try: + stem = Path(fname).stem + out_name = out_dir / f"{stem}_{seg_tag}.tif" + tifffile.imwrite(str(out_name), seg.astype(np.uint8)) + segmented += 1 + yield fname, str(out_name) + except Exception: + warnings.warn(f'Could not write segmentation for {fname}.') + yield fname, None + except CancelledError: + # Files already segmented and saved stay on disk. + pass + return segmented + + self._pending_segment_files_ctx = (list(filenames), out_dir) + + worker = _do_segment_files() + worker.yielded.connect(self._on_segment_files_yielded) + worker.returned.connect(self._on_segment_files_returned) + worker.errored.connect(self._on_worker_errored) + worker.finished.connect(self._on_worker_finished) + self._begin_worker('segment_files', self.multifile_segment_selected_btn, cancel_token, worker, + desc='Segmenting', total=len(filenames)) + + def _on_segment_files_yielded(self, value): + fname, out_path = value + if out_path is not None: + # Register saved segmentation and update table + self._multifile_segmentation_store[fname] = str(out_path) + self._update_multifile_seg_tick(fname) + def _on_segment_files_returned(self, segmented): + ctx = getattr(self, '_pending_segment_files_ctx', None) + self._pending_segment_files_ctx = None + if ctx is None: + return + filenames, out_dir = ctx if segmented: show_info(f'Segmented {segmented} files and saved to {out_dir}') - else: + elif not (self._op is not None and self._op.cancel_was_requested): warnings.warn('No images were segmented.') # If segmented the opened image, open its segmentation current_open = getattr(self, '_current_multifile_filename', None) - if current_open in filenames and self.multifile_import_open_segmentations: + if (current_open in filenames and self.multifile_import_open_segmentations + and current_open in self._multifile_segmentation_store): self._multifile_open_segmentation(current_open) def _import_annot_and_seg(self): diff --git a/src/napari_convpaint/feature_extractor.py b/src/napari_convpaint/feature_extractor.py index 2fb2cbb..985c60b 100644 --- a/src/napari_convpaint/feature_extractor.py +++ b/src/napari_convpaint/feature_extractor.py @@ -2,7 +2,7 @@ import torch import warnings from .param import Param -from .utils import scale_img, rescale_features, reduce_to_patch_multiple, pad_to_shape, get_device_from_torch_model +from .utils import scale_img, rescale_features, reduce_to_patch_multiple, pad_to_shape, get_device_from_torch_model, check_cancel, cancel_scope class FeatureExtractor: def __init__(self, model_name="vgg16", model=None, **kwargs): @@ -300,7 +300,7 @@ def supported_devices(self): ### FEATURE EXTRACTION METHODS - def extract_features(self, data, param, device=torch.device("cpu")): + def extract_features(self, data, param, device=torch.device("cpu"), cancel_token=None): """ Extracts the features of an image (stack) with an arbitrary number of channels. This is the main method to call for feature extraction, which will handle scaling and rescaling of the features as needed. @@ -314,6 +314,13 @@ def extract_features(self, data, param, device=torch.device("cpu")): The parameters for the feature extraction. device : torch.device, optional The device on which to perform feature extraction. + cancel_token : CancelToken, optional + Cooperative cancellation token. If provided and cancelled (e.g. from + another thread), extraction aborts at the next checkpoint by raising + CancelledError. The token is installed as the ambient token for the + duration of the call, so subclass overrides don't need to accept or + forward it — the check_cancel() calls in the base-class loops pick + it up automatically. Returns: ---------- @@ -324,7 +331,8 @@ def extract_features(self, data, param, device=torch.device("cpu")): # self.move_model_to_device(device) # Extract features with scaling and rescaling as needed - features = self.extract_features_pyramid(data=data, param=param, patched=self.gives_patched_features(), device=device) + with cancel_scope(cancel_token): + features = self.extract_features_pyramid(data=data, param=param, patched=self.gives_patched_features(), device=device) return features @@ -363,8 +371,20 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic if not param.fe_scalings in self.get_proposed_scalings(): warnings.warn(f"The selected scalings {param.fe_scalings} are not in the proposed scalings {self.proposed_scalings}. Please check if this is intentional.") + # Post-processing of extracted features (rescaling, device transfer) works + # on one array per channel-series/layer; checking between elements keeps + # cancellation responsive even for huge feature images (e.g. many-channel + # inputs), where a single rescale/transfer can take seconds. + def _rescale_all(feats, shape): + out = [] + for f in feats: + check_cancel() + out.append(rescale_features(feature_img=f, target_shape=shape, order=param.fe_order)) + return out + # Iterate over the scales and extract features for each scale for s in param.fe_scalings: + check_cancel() # Downscale the image image_scaled = scale_img(data, s) @@ -396,11 +416,7 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic # NOTE: this should not be necessary if the inputs are already multiples of the patch size at all scales if patch_size > 1 and reduced_shape[2:] != pre_reduction_shape[2:] : # Step 1: rescale to the reduced (cropped to patch multiple) shape - features = [rescale_features( - feature_img=f, - target_shape=reduced_shape, - order=param.fe_order) - for f in features] + features = _rescale_all(features, reduced_shape) # Step 2: pad back to original pre_reduction_shape (but still downscaled) features = [pad_to_shape(f, pre_reduction_shape[2:] ) for f in features] @@ -408,18 +424,19 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic # Rescale to the full original shape target_shape = data.shape - features = [rescale_features( - feature_img=f, - target_shape=target_shape, - order=param.fe_order) - for f in features] + features = _rescale_all(features, target_shape) # If torch tensor is returned, convert to numpy array if isinstance(features[0], torch.Tensor): - # Detach, move to cpu, make np array - features = [feature.detach().cpu().numpy() for feature in features] - + # Detach, move to cpu, make np array (checking between transfers) + converted = [] + for feature in features: + check_cancel() + converted.append(feature.detach().cpu().numpy()) + features = converted + # Put together features for each input_channels procession (and layers if applicable) + check_cancel() # Last checkpoint before the (potentially large) concatenation features = np.concatenate(features, axis=0) # If use_min_features is True, shorten features @@ -486,12 +503,12 @@ def extract_features_from_multichannel_stack(self, image, rgb_data=False, device # For each channel, create a replicate with the needed number of input channels fe_input_channels = min(fe_input_channels) channel_series = [np.tile(ch, (fe_input_channels, 1, 1, 1)) for ch in image] - - # Get outputs for each channel_series + + # Get outputs for each channel_series. The per-Z-plane check inside + # extract_features_from_stack fires almost immediately, so no extra + # check is needed at this loop boundary. all_outputs = [] for channel in channel_series: - # Output is either a single array or a list of features, - # possibly from different layers (and thus with different sizes) output = self.extract_features_from_stack(channel, device=device) # Make one list of all outputs (aligning different channel_series and layers) if isinstance(output, list): @@ -526,6 +543,7 @@ def extract_features_from_stack(self, image, device=torch.device("cpu")): all_features = [] # Go through the stack, and get features for each plane for z in range(image.shape[1]): + check_cancel() features = self.extract_features_from_plane(image[:,z], device=device) all_features.append(features) diff --git a/src/napari_convpaint/feature_extractors/nnlayers.py b/src/napari_convpaint/feature_extractors/nnlayers.py index 49a1880..c56596f 100644 --- a/src/napari_convpaint/feature_extractors/nnlayers.py +++ b/src/napari_convpaint/feature_extractors/nnlayers.py @@ -1,7 +1,8 @@ import numpy as np import torch from torch import nn -from ..utils import get_device_from_torch_model, guided_model_download +from ..utils import get_device_from_torch_model, guided_model_download, check_cancel + def import_models(): try: @@ -254,7 +255,7 @@ def extract_features_from_stack(self, image, device=torch.device("cpu")): pass # Stop at hook except Exception as ex: raise ex - + # Move the z dimension back to the second position (and features to first) outputs = [o.permute(1, 0, 2, 3) for o in self.outputs] @@ -265,12 +266,15 @@ def __call__(self, tensor_image): return self.model(tensor_image_dev) def hook_normal(self, module, input, output): - # print("extracting with normal layer") self.outputs.append(output) + # Checking the ambient token between hooked layers lets a cancel take + # effect mid-forward-pass; without it, heavy VGG16 configs only cancel + # after the whole forward completes. + check_cancel() def hook_last(self, module, input, output): - # print("extracting with last layer") self.outputs.append(output) + check_cancel() assert False def register_hooks(self, selected_layers): # , selected_layer_pos): diff --git a/src/napari_convpaint/utils.py b/src/napari_convpaint/utils.py index 84a6366..0e1b874 100644 --- a/src/napari_convpaint/utils.py +++ b/src/napari_convpaint/utils.py @@ -1,4 +1,6 @@ +import contextvars import warnings +from contextlib import contextmanager import torch import numpy as np from scipy.ndimage import gaussian_filter @@ -18,6 +20,59 @@ # napari.utils (see details below) +### Cooperative cancellation + +class CancelledError(Exception): + pass + + +class CancelToken: + def __init__(self): + self._cancelled = False + + def cancel(self): + self._cancelled = True + + @property + def cancelled(self): + return self._cancelled + + def raise_if_cancelled(self): + if self._cancelled: + raise CancelledError() + + +# The active token is carried in a ContextVar rather than threaded through every +# method signature. Entry points (ConvpaintModel.train/segment/... ) install the +# token with cancel_scope(); any code below them — including custom +# FeatureExtractor subclasses with pre-cancellation signatures — is covered by +# the plain check_cancel() calls in the base-class loops without needing a +# cancel_token parameter of its own. ContextVars are per-thread, so the token +# installed by a worker thread is invisible to other threads; the shared +# CancelToken object is what crosses threads (cancel() from the GUI thread, +# checks in the worker). +_current_cancel_token = contextvars.ContextVar("convpaint_cancel_token", default=None) + + +@contextmanager +def cancel_scope(cancel_token): + """Install `cancel_token` (may be None) as the ambient token for the duration.""" + reset_token = _current_cancel_token.set(cancel_token) + try: + yield + finally: + _current_cancel_token.reset(reset_token) + + +def check_cancel(cancel_token=None): + """Raise CancelledError if the given token — or, when None, the ambient + token installed by the innermost cancel_scope() — has been cancelled.""" + if cancel_token is None: + cancel_token = _current_cancel_token.get() + if cancel_token is not None: + cancel_token.raise_if_cancelled() + + ### PCA and Kmeans on feature images def apply_pca_to_f_image(feature_img, n_components, norm=True):