diff --git a/src/napari_convpaint/_tests/test_feature_cache.py b/src/napari_convpaint/_tests/test_feature_cache.py
new file mode 100644
index 0000000..5ba78d8
--- /dev/null
+++ b/src/napari_convpaint/_tests/test_feature_cache.py
@@ -0,0 +1,347 @@
+"""Tests for the bounded feature cache (storage/eviction/budget logic)."""
+import numpy as np
+
+from napari_convpaint.feature_cache import FeatureCache
+
+
+def _arr(mb):
+ """A float32 array of approximately `mb` megabytes."""
+ n = int(mb * 1e6 / 4)
+ return np.zeros(n, dtype=np.float32)
+
+
+def test_hit_and_miss():
+ c = FeatureCache(max_bytes=100 * 10**6, headroom_frac=0.0)
+ assert c.get(("img", 0, "sig")) is None
+ payload = _arr(1)
+ c.put(("img", 0, "sig"), payload)
+ got = c.get(("img", 0, "sig"))
+ assert got is payload
+ assert c.stats()["hits"] == 1
+ assert c.stats()["misses"] == 1
+
+
+def test_lru_eviction_by_cap():
+ # Cap ~2.5 MB; each entry ~1 MB -> at most 2 fit, oldest evicted.
+ c = FeatureCache(max_bytes=int(2.5 * 10**6), headroom_frac=0.0)
+ c.put(("a",), _arr(1))
+ c.put(("b",), _arr(1))
+ assert len(c) == 2
+ c.put(("c",), _arr(1)) # evicts "a" (LRU)
+ assert len(c) == 2
+ assert c.get(("a",)) is None
+ assert c.get(("b",)) is not None
+ assert c.get(("c",)) is not None
+
+
+def test_lru_touch_on_get_protects_entry():
+ c = FeatureCache(max_bytes=int(2.5 * 10**6), headroom_frac=0.0)
+ c.put(("a",), _arr(1))
+ c.put(("b",), _arr(1))
+ assert c.get(("a",)) is not None # touch "a" -> now "b" is LRU
+ c.put(("c",), _arr(1)) # should evict "b", not "a"
+ assert c.get(("a",)) is not None
+ assert c.get(("b",)) is None
+
+
+def test_single_oversize_payload_is_not_cached():
+ c = FeatureCache(max_bytes=1 * 10**6, headroom_frac=0.0)
+ c.put(("big",), _arr(5)) # 5 MB into a 1 MB cap -> skipped, not cached
+ assert len(c) == 0
+ assert c.get(("big",)) is None
+
+
+def test_overwrite_updates_size():
+ c = FeatureCache(max_bytes=100 * 10**6, headroom_frac=0.0)
+ c.put(("k",), _arr(1))
+ b0 = c.nbytes
+ c.put(("k",), _arr(3)) # replace with a bigger payload
+ assert c.nbytes > b0
+ assert len(c) == 1
+
+
+def test_clear():
+ c = FeatureCache(max_bytes=100 * 10**6, headroom_frac=0.0)
+ c.put(("a",), _arr(1))
+ c.put(("b",), _arr(1))
+ c.clear()
+ assert len(c) == 0
+ assert c.nbytes == 0
+
+
+def test_disabled_cache_is_noop():
+ c = FeatureCache(max_bytes=100 * 10**6, headroom_frac=0.0, enabled=False)
+ c.put(("a",), _arr(1))
+ assert c.get(("a",)) is None
+ assert len(c) == 0
+
+
+def test_list_payload_size_accounted():
+ c = FeatureCache(max_bytes=int(2.5 * 10**6), headroom_frac=0.0)
+ c.put(("a",), [_arr(1), _arr(1)]) # ~2 MB as a list of arrays
+ assert len(c) == 1
+ c.put(("b",), _arr(1)) # pushes over 2.5 MB -> evicts "a"
+ assert c.get(("a",)) is None
+
+
+def test_model_feature_cache_identical_and_reuses():
+ """With the cache on, a re-extraction of the same image reuses features and
+ produces bit-identical output vs the cache off."""
+ import warnings
+ from napari_convpaint.convpaint_model import ConvpaintModel
+
+ rng = np.random.default_rng(0)
+ img = rng.random((64, 64), dtype=np.float32)
+ annot = np.zeros((64, 64), dtype=np.uint8)
+ annot[10:20, 10:20] = 1
+ annot[40:50, 40:50] = 2
+
+ def run(enable):
+ m = ConvpaintModel(fe_name="gaussian_features")
+ m.set_params(channel_mode="single")
+ if enable:
+ m.enable_feature_cache(True)
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ m.train(img, annot)
+ seg1 = np.asarray(m.segment(img))
+ seg2 = np.asarray(m.segment(img))
+ return m, seg1, seg2
+
+ m_off, off1, off2 = run(False)
+ m_on, on1, on2 = run(True)
+ assert np.array_equal(off1, on1)
+ assert np.array_equal(off2, on2)
+ # cache actually stored and served something
+ assert m_on._feature_cache.stats()["hits"] >= 1
+
+
+def test_disk_spillover_serves_ram_evicted_entries():
+ """RAM-evicted entries spill to disk and are served from there (bit-identical)."""
+ c = FeatureCache(max_bytes=int(2.5 * 10**6), headroom_frac=0.0,
+ disk_max_bytes=100 * 10**6)
+ a = _arr(1); b = _arr(1); d = _arr(1)
+ c.put(("a",), a); c.put(("b",), b) # RAM full (2 entries)
+ c.put(("c",), d) # evicts "a" from RAM -> spills to disk
+ assert len(c) == 2 and c.stats()["disk_entries"] == 1
+ got = c.get(("a",)) # RAM miss -> disk hit
+ assert got is not None and np.array_equal(got, a) # round-trips bit-identical
+ assert c.stats()["disk_hits"] == 1
+
+
+def test_disk_lru_eviction_and_total_miss():
+ c = FeatureCache(max_bytes=int(1.5 * 10**6), headroom_frac=0.0,
+ disk_max_bytes=int(1.5 * 10**6)) # RAM holds 1, disk holds 1
+ c.put(("a",), _arr(1)); c.put(("b",), _arr(1)) # a -> disk, b in RAM
+ c.put(("c",), _arr(1)) # b -> disk (evicts a from disk), c in RAM
+ assert c.get(("a",)) is None # a fell off disk entirely -> recompute
+ assert c.get(("b",)) is not None # b on disk
+ assert c.get(("c",)) is not None # c in RAM
+
+
+def test_disk_disabled_by_default():
+ c = FeatureCache(max_bytes=int(1.5 * 10**6), headroom_frac=0.0) # no disk
+ c.put(("a",), _arr(1)); c.put(("b",), _arr(1)) # a evicted, dropped (no disk)
+ assert c.get(("a",)) is None and c.stats()["disk_entries"] == 0
+
+
+def test_clear_removes_disk_tier_and_tempdir():
+ import os
+ c = FeatureCache(max_bytes=int(1.5 * 10**6), headroom_frac=0.0,
+ disk_max_bytes=100 * 10**6)
+ c.put(("a",), _arr(1)); c.put(("b",), _arr(1)) # a on disk
+ disk_dir = c._disk_dir
+ assert disk_dir is not None and os.path.isdir(disk_dir)
+ c.clear()
+ assert c.stats()["disk_entries"] == 0 and c.disk_nbytes == 0
+ c.close()
+ assert not os.path.isdir(disk_dir) # temp dir removed
+
+
+def test_disk_bytes_never_exceeds_cap():
+ """Stress: many puts must never push the disk tier over its byte cap."""
+ cap = int(3.5 * 10**6) # ~3 entries of 1 MB
+ c = FeatureCache(max_bytes=int(1.5 * 10**6), headroom_frac=0.0, disk_max_bytes=cap)
+ for i in range(20):
+ c.put((i,), _arr(1))
+ assert c.disk_nbytes <= cap # invariant holds after every put
+ c.close()
+
+
+# --- integration with the model-level cache protocol -----------------------
+
+def test_oversized_payload_goes_to_disk_tier():
+ from napari_convpaint.feature_cache import FeatureCache
+ c = FeatureCache(max_bytes=1024 * 1024, headroom_frac=0.0,
+ disk_max_bytes=64 * 1024 * 1024)
+ try:
+ payload = np.zeros(2 * 1024 * 1024, dtype=np.uint8) # 2 MB > 1 MB RAM cap
+ c.put(('big',), payload)
+ assert len(c) == 0
+ assert c.stats()['disk_entries'] == 1
+ got = c.get(('big',))
+ assert got is not None and got.nbytes == payload.nbytes
+ finally:
+ c.close()
+
+
+def test_spill_ok_false_never_touches_disk():
+ from napari_convpaint.feature_cache import FeatureCache
+ c = FeatureCache(max_bytes=1024 * 1024, headroom_frac=0.0,
+ disk_max_bytes=64 * 1024 * 1024)
+ try:
+ c.put(('a',), np.zeros(600 * 1024, dtype=np.uint8), spill_ok=False)
+ c.put(('b',), np.zeros(600 * 1024, dtype=np.uint8)) # evicts 'a' -> dropped
+ assert c.stats()['disk_entries'] == 0
+ assert c.get(('a',)) is None
+ c.put(('huge',), np.zeros(2 * 1024 * 1024, dtype=np.uint8), spill_ok=False)
+ assert c.get(('huge',)) is None
+ assert c.stats()['disk_entries'] == 0
+ finally:
+ c.close()
+
+
+def test_hookmodel_opts_out_of_disk_spill():
+ from napari_convpaint.feature_extractor import FeatureExtractor
+ assert FeatureExtractor.cache_spill_to_disk(object()) is True
+ from napari_convpaint.feature_extractors.nnlayers import Hookmodel
+ assert Hookmodel.cache_spill_to_disk(object()) is False
+
+
+def test_cache_key_includes_fe_instance_state():
+ from napari_convpaint.convpaint_model import ConvpaintModel
+ cp = ConvpaintModel('gaussian')
+ sig_before = cp._fe_cache_signature(cp._param)
+ cp.fe_model.sigma = cp.fe_model.sigma + 1
+ assert cp._fe_cache_signature(cp._param) != sig_before
+ # generic hook: any change in reported extra state must change the key
+ orig = cp.fe_model.cache_extra_state
+ cp.fe_model.cache_extra_state = lambda: ('jafar_scalings', (1, 8))
+ try:
+ assert cp._fe_cache_signature(cp._param) != sig_before
+ finally:
+ cp.fe_model.cache_extra_state = orig
+
+
+def test_cached_prediction_bit_identical_and_hits():
+ import warnings as _w
+ from napari_convpaint.convpaint_model import ConvpaintModel
+ rng = np.random.RandomState(0)
+ img = rng.rand(1, 96, 96).astype(np.float32)
+ annot = np.zeros((1, 96, 96), dtype=np.uint8)
+ annot[0, :12, :12] = 1
+ annot[0, -12:, -12:] = 2
+ with _w.catch_warnings():
+ _w.simplefilter('ignore')
+ cp = ConvpaintModel('gaussian')
+ fc = cp.enable_feature_cache(max_bytes=64 * 1024 * 1024)
+ cp.train(img, annot)
+ seg_first = cp.segment(img)
+ hits_before = fc.stats()['hits']
+ seg_second = cp.segment(img)
+ assert fc.stats()['hits'] > hits_before # second pass hits
+ assert np.array_equal(seg_first, seg_second)
+ # peek semantics
+ assert cp._predict(rng.rand(1, 96, 96).astype(np.float32), cache_only=True) is None
+ assert cp._predict(img, cache_only=True) is not None
+ # uncached model produces the identical segmentation
+ cp2 = ConvpaintModel('gaussian')
+ cp2.train(img, annot)
+ assert np.array_equal(seg_second, cp2.segment(img))
+
+
+def test_thread_safety_under_concurrent_use():
+ """Hammer the cache from worker threads while the "GUI" thread clears it and
+ changes limits (exactly what the napari widget does during a threaded op).
+ Correctness bar: no exceptions and consistent bookkeeping afterwards."""
+ import threading
+
+ c = FeatureCache(max_bytes=int(3 * 10**6), headroom_frac=0.0,
+ disk_max_bytes=int(5 * 10**6))
+ errors = []
+ start = threading.Barrier(5)
+
+ def worker(tid):
+ try:
+ start.wait()
+ for i in range(200):
+ key = ("img", tid, i % 7)
+ if c.get(key) is None:
+ c.put(key, _arr(0.1))
+ len(c), c.stats()
+ except Exception as e: # pragma: no cover - only on regression
+ errors.append(e)
+
+ def gui():
+ try:
+ start.wait()
+ for i in range(100):
+ c.set_max_bytes(int((2 + i % 3) * 10**6))
+ c.set_disk_max_bytes(int((i % 2) * 5 * 10**6))
+ c.stats()
+ if i % 10 == 0:
+ c.clear()
+ except Exception as e: # pragma: no cover - only on regression
+ errors.append(e)
+
+ threads = [threading.Thread(target=worker, args=(t,)) for t in range(4)]
+ threads.append(threading.Thread(target=gui))
+ for t in threads:
+ t.start()
+ for t in threads:
+ t.join()
+
+ assert errors == []
+ # Bookkeeping must be consistent: recompute sizes from the stores.
+ assert c.nbytes == sum(item[1] for item in c._store.values())
+ assert c.disk_nbytes == sum(item[1] for item in c._disk_store.values())
+ assert c.nbytes <= c.stats()["max_bytes"]
+ c.close()
+
+
+def test_nn_fe_cache_hit_matches_fresh_and_uses_torch_payload():
+ """NN FEs keep their native features on-device (torch); the cache payload
+ is cast to numpy for storage but remembers it was torch, so hits are
+ lifted back and reconstructed with the SAME torch backend as fresh
+ extractions. Guards against a hit/miss backend split (skimage vs torch)
+ which would make cache-enabled extraction both slow (CPU rescale) and
+ potentially non-identical to fresh results."""
+ import warnings as _w
+ from napari_convpaint.convpaint_model import ConvpaintModel
+ rng = np.random.RandomState(0)
+ img = rng.rand(1, 64, 64).astype(np.float32)
+ with _w.catch_warnings():
+ _w.simplefilter('ignore')
+ cp = ConvpaintModel(fe_name='vgg16')
+ cp.set_params(fe_scalings=[1, 2])
+ feat_off = cp.get_feature_image(img) # cache disabled: fresh
+ fc = cp.enable_feature_cache(max_bytes=512 * 10**6)
+ feat_miss = cp.get_feature_image(img) # miss: fills cache
+ feat_hit = cp.get_feature_image(img) # hit: from payload
+ assert fc.stats()['hits'] >= 1
+ assert np.array_equal(feat_off, feat_miss), "cache-on (miss) differs from cache-off"
+ assert np.array_equal(feat_miss, feat_hit), "cache hit differs from miss"
+ payload = next(iter(fc._store.values()))[0]
+ assert payload['was_torch'] is True
+ for features, _, _ in payload['scales']: # stored form is numpy
+ assert all(isinstance(f, np.ndarray) for f in features)
+
+
+def test_numpy_fe_payload_stays_numpy_and_identical():
+ """Numpy-native FEs (e.g. gaussian) must NOT be lifted to torch on a hit —
+ their fresh path is skimage, and hit/miss must keep sharing it."""
+ import warnings as _w
+ from napari_convpaint.convpaint_model import ConvpaintModel
+ rng = np.random.RandomState(0)
+ img = rng.rand(1, 96, 96).astype(np.float32)
+ with _w.catch_warnings():
+ _w.simplefilter('ignore')
+ cp = ConvpaintModel(fe_name='gaussian_features')
+ feat_off = cp.get_feature_image(img)
+ fc = cp.enable_feature_cache(max_bytes=256 * 10**6)
+ feat_miss = cp.get_feature_image(img)
+ feat_hit = cp.get_feature_image(img)
+ assert np.array_equal(feat_off, feat_miss)
+ assert np.array_equal(feat_miss, feat_hit)
+ payload = next(iter(fc._store.values()))[0]
+ assert payload['was_torch'] is False
diff --git a/src/napari_convpaint/_tests/test_widget_layout.py b/src/napari_convpaint/_tests/test_widget_layout.py
new file mode 100644
index 0000000..f2ea93f
--- /dev/null
+++ b/src/napari_convpaint/_tests/test_widget_layout.py
@@ -0,0 +1,230 @@
+"""Regression tests for the widget's tab/scroll layout."""
+from qtpy.QtWidgets import QScrollArea
+from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+
+def test_widget_opens_on_home_tab_and_tabs_scroll(make_napari_viewer, qtbot):
+ viewer = make_napari_viewer()
+ widget = ConvpaintWidget(viewer)
+ if hasattr(widget, 'ensure_init'):
+ widget.ensure_init()
+ widget.show()
+ qtbot.waitExposed(widget)
+ # A fresh widget must open on the first (Home) tab — the scroll-wrap
+ # remove/insert dance moves the current index around during construction.
+ assert widget.tabs.currentIndex() == 0
+ # Every tab is wrapped in a scroll area so content can't be cut off.
+ for i in range(widget.tabs.count()):
+ assert isinstance(widget.tabs.widget(i), QScrollArea), widget.tabs.tab_names[i]
+
+
+def test_output_layers_do_not_steal_selection(make_napari_viewer):
+ """Creating segmentation/probabilities/features layers must leave the
+ active layer untouched — napari selects new layers by default, which
+ would steer the user's brush into the output layer instead of the
+ annotations layer they were painting on."""
+ import numpy as np
+ from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+ viewer = make_napari_viewer()
+ w = ConvpaintWidget(viewer)
+ w.ensure_init()
+ viewer.add_image(np.random.random((64, 64)), name='img')
+ w._on_add_annot_layer()
+ annot = viewer.layers['annotations']
+ viewer.layers.selection.active = annot
+
+ w._check_create_segmentation_layer()
+ assert viewer.layers.selection.active is annot
+ w._check_create_probas_layer(2)
+ assert viewer.layers.selection.active is annot
+ w._check_create_features_layer(4)
+ assert viewer.layers.selection.active is annot
+ assert {'segmentation'} <= {l.name for l in viewer.layers}
+
+
+def test_classes_tab_value_model(make_napari_viewer):
+ """Value-based class rows: placeholder icons before any layer; 'Remove
+ class' targets the SELECTED class (grayed without a selection or at the
+ two-class floor) and never renumbers others; sync adds exactly the painted
+ values (sparse — no gap filling)."""
+ import numpy as np
+ from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+ viewer = make_napari_viewer()
+ w = ConvpaintWidget(viewer)
+ w.ensure_init()
+
+ # Startup: NO classes — two empty placeholder slots, Remove grayed
+ assert w.class_rows == []
+ assert len(w._placeholder_rows) == 2
+ assert all(not r['icon'].pixmap().isNull() for r in w._placeholder_rows)
+ assert not w.remove_class_btn.isEnabled()
+ # Create the two classes the rest of this test works with
+ w._on_add_class(text='Background', value=1)
+ w._on_add_class(text='Foreground', value=2)
+
+ viewer.add_image(np.random.random((64, 64)), name='img')
+ w._on_add_annot_layer()
+ annot = viewer.layers['annotations']
+
+ # A selected real class is removable even at two classes (no floor)
+ annot.selected_label = 1
+ assert w.remove_class_btn.isEnabled()
+
+ # Sparse sync: painting 7 adds EXACTLY value 7 (no gap rows 3..6)
+ annot.data[5:10, 5:10] = 7
+ w._on_sync_classes_from_annotations()
+ assert [r['value'] for r in w.class_rows] == [1, 2, 7]
+
+ # Remove the selected middle-by-value class: value 2 goes, 7 keeps its value
+ annot.selected_label = 2
+ assert w.remove_class_btn.isEnabled()
+ w._on_remove_class()
+ assert [r['value'] for r in w.class_rows] == [1, 7]
+ # ...and its annotations were erased, others kept
+ assert not (annot.data == 2).any()
+ assert (annot.data == 7).any()
+
+ # Selection now points at a value with no row -> Remove grayed
+ assert annot.selected_label == 2
+ assert not w.remove_class_btn.isEnabled()
+
+ # Add class assigns max+1
+ w._on_add_class()
+ assert [r['value'] for r in w.class_rows] == [1, 7, 8]
+
+
+def test_classes_tab_no_floor_and_placeholders(make_napari_viewer):
+ """No class-count floor; placeholders are pure UI (not clickable classes)
+ that pad the display to two slots; typing into a placeholder's name field
+ creates the class in place."""
+ import numpy as np
+ from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+ viewer = make_napari_viewer()
+ w = ConvpaintWidget(viewer)
+ w.ensure_init()
+ viewer.add_image(np.random.random((64, 64)), name='img')
+ w._on_add_annot_layer()
+ annot = viewer.layers['annotations']
+ w._on_add_class(text='Background', value=1)
+ w._on_add_class(text='Foreground', value=2)
+
+ # Remove both classes, one by one — no floor
+ annot.selected_label = 2
+ w._on_remove_class()
+ assert [r['value'] for r in w.class_rows] == [1]
+ assert len(w._placeholder_rows) == 1
+ annot.selected_label = 1
+ w._on_remove_class()
+ assert w.class_rows == []
+ assert len(w._placeholder_rows) == 2
+
+ # Placeholders: editable italic 'add class' fields, no class behavior
+ for ph in w._placeholder_rows:
+ assert ph['name'].isEnabled()
+ assert ph['name'].placeholderText() == 'add class'
+ assert 'italic' in ph['name'].styleSheet()
+ # Selecting a placeholder's value (e.g. via the label spinbox) grays Remove
+ annot.selected_label = w._placeholder_rows[0]['value']
+ assert not w.remove_class_btn.isEnabled()
+
+ # Typing into a placeholder creates the class with that text, in place
+ ph = w._placeholder_rows[0]
+ v = ph['value']
+ w._on_placeholder_name_edited(ph, 'Nu')
+ assert [r['value'] for r in w.class_rows] == [v]
+ assert w._row_for_value(v)['name'].text() == 'Nu'
+
+
+def test_classes_selectable_without_annotations_layer(make_napari_viewer):
+ """Deleting the annotations layer must not strand the class list: the
+ widget keeps its own selection memory, so swatch clicks still select and
+ Remove still works with no layer present."""
+ import numpy as np
+ from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+ viewer = make_napari_viewer()
+ w = ConvpaintWidget(viewer)
+ w.ensure_init()
+ viewer.add_image(np.random.random((64, 64)), name='img')
+ w._on_add_annot_layer()
+ annot = viewer.layers['annotations']
+ w._on_add_class(text='Background', value=1)
+ w._on_add_class(text='Foreground', value=2)
+ annot.selected_label = 2 # mirrored into widget memory
+ viewer.layers.remove('annotations') # layer gone
+
+ assert w._selected_class_value() == 2 # selection survives
+ w._update_selected_class_highlight()
+ assert w.remove_class_btn.isEnabled()
+ w._on_remove_class() # removable without a layer
+ assert [r['value'] for r in w.class_rows] == [1]
+
+ # Swatch click selects widget-side with no layer at all
+ w._on_class_swatch_clicked(1)
+ assert w._selected_class_value() == 1
+ assert w.remove_class_btn.isEnabled()
+
+
+def test_selected_class_highlight_follows_selected_label(make_napari_viewer):
+ """The Classes tab outlines the row matching the annotations layer's
+ selected label (by VALUE), in both directions."""
+ import numpy as np
+ from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+ viewer = make_napari_viewer()
+ w = ConvpaintWidget(viewer)
+ w.ensure_init()
+ viewer.add_image(np.random.random((64, 64)), name='img')
+ w._on_add_annot_layer()
+ annot = viewer.layers['annotations']
+ w._on_add_class(value=1)
+ w._on_add_class(value=2)
+
+ def outlined():
+ return [r['value'] for r in w.class_rows
+ if 'transparent' not in r['icon'].styleSheet()]
+
+ annot.selected_label = 2
+ assert outlined() == [2]
+ annot.selected_label = 1
+ assert outlined() == [1]
+
+
+def test_class_value_limit(make_napari_viewer, tmp_path):
+ """Class values are capped at 255 (annotation/segmentation data are
+ uint8): adding beyond the limit is refused, the Add button grays out at
+ the limit, and a CSV with too-high values fails without side effects."""
+ import pytest
+ from napari_convpaint.convpaint_widget import ConvpaintWidget
+
+ viewer = make_napari_viewer()
+ w = ConvpaintWidget(viewer)
+ w.ensure_init()
+
+ # Adding a value above the limit is a no-op
+ w._on_add_class(value=300)
+ assert w._class_values() == []
+
+ # At the limit the Add button grays out; below it stays enabled
+ w._on_add_class(value=254)
+ assert w.add_class_btn.isEnabled()
+ w._on_add_class(value=255)
+ assert w._class_values() == [254, 255]
+ assert not w.add_class_btn.isEnabled()
+
+ # A CSV holding a too-high value raises and leaves the classes untouched
+ bad = tmp_path / 'bad.csv'
+ bad.write_text('index,name\n1,ok\n300,too high\n')
+ with pytest.raises(ValueError, match='255'):
+ w.import_class_names_csv(str(bad))
+ assert w._class_values() == [254, 255]
+
+ # A valid sparse CSV still round-trips
+ good = tmp_path / 'good.csv'
+ good.write_text('index,name\n1,bg\n7,rare\n')
+ w.import_class_names_csv(str(good))
+ assert w._class_values() == [1, 7]
+ assert w.add_class_btn.isEnabled()
diff --git a/src/napari_convpaint/convpaint_model.py b/src/napari_convpaint/convpaint_model.py
index 6c13f94..57e1653 100644
--- a/src/napari_convpaint/convpaint_model.py
+++ b/src/napari_convpaint/convpaint_model.py
@@ -110,6 +110,7 @@ def __init__(self, alias=None, model_path=None, param=None, fe_name=None, **kwar
self.num_features = 0
self._fe_locked_device = None
self._clf_locked_device = None
+ self._feature_cache = None # created by enable_feature_cache
self._params_to_reset_training = ['channel_mode',
'normalize',
# 'image_downsample',
@@ -718,7 +719,7 @@ def get_fe_defaults(self):
new_param : Param
Convpaint model defaults adjusted to the feature extractor defaults
"""
- cpm_defaults = ConvpaintModel.get_default_params() # Get ConvPaint defaults
+ cpm_defaults = ConvpaintModel.get_default_params() # Get Convpaint defaults
new_param = self.fe_model.get_default_params(cpm_defaults) # Overwrite defaults defined in the FE model
return new_param
@@ -913,13 +914,93 @@ def get_feature_image(self, data,
)
return features
-
+
+### FEATURE CACHE (opt-in; see feature_cache.py)
+
+ def enable_feature_cache(self, enabled=True, max_bytes=None, disk_max_bytes=0):
+ """Turn on whole-image feature caching. When on, the (resolution-
+ independent) native features of an extracted image are cached and reused
+ the next time the *same* image is processed with the same FE settings —
+ e.g. re-segmenting while refining scribbles, or the train->predict of one
+ image — instead of recomputing them. Cache entries are content-addressed
+ (a hash of the prepared image), so it is self-invalidating: a changed
+ image simply misses. Bounded by a RAM budget (``max_bytes``); RAM-evicted
+ entries spill to disk up to ``disk_max_bytes`` (0 = off), which lets a
+ stack too large for RAM still benefit on the next iteration (loading a
+ cached slice from disk is much faster than recomputing it). Off by
+ default (opt-in)."""
+ from .feature_cache import FeatureCache
+ self._feature_cache = FeatureCache(max_bytes=max_bytes, enabled=enabled,
+ disk_max_bytes=disk_max_bytes)
+ return self._feature_cache
+
+ def _fe_cache_signature(self, param):
+ """The FE-relevant part of the cache key: parameters whose change
+ invalidates extracted features (reusing the model's own train-reset set),
+ plus image_downsample and the FE's patch size."""
+ def _hashable(v):
+ # fe_scalings / fe_layers are lists (and FE extra state may nest
+ # lists in tuples) -> make them hashable for the key.
+ if isinstance(v, (list, tuple)):
+ return tuple(_hashable(x) for x in v)
+ return v
+ keys = self._params_to_reset_training
+ sig = tuple((k, _hashable(getattr(param, k, None))) for k in keys)
+ return sig + (("image_downsample", getattr(param, "image_downsample", 1)),
+ ("patch_size", self.fe_model.get_patch_size()),
+ # FE instance state outside the Param (e.g. jafar_scalings,
+ # gaussian sigma) — without it, changing that state would
+ # serve stale cached features.
+ ("fe_extra", _hashable(self.fe_model.cache_extra_state())))
+
+ @staticmethod
+ def _data_hash(d):
+ """Content hash of a prepared image tile, so train/predict of the same
+ pixels share a cache entry without threading an id through the pipeline."""
+ import hashlib
+ arr = np.ascontiguousarray(d)
+ h = hashlib.blake2b(arr.view(np.uint8), digest_size=16)
+ h.update(str(arr.shape).encode())
+ h.update(str(arr.dtype).encode())
+ return h.hexdigest()
+
+ def _extract_pyramid_cached(self, d, param, keep_patched, device, cache_only=False):
+ """Extract the feature pyramid for one image, consulting the feature
+ cache. Behaviour with the cache disabled (the default) is exactly
+ extract_features_pyramid; enabled, it caches/reuses the native features
+ (bit-identical output, since the pyramid split is exact).
+
+ ``cache_only=True`` is a peek: return the reconstructed features only if
+ they are already cached, else return None WITHOUT running the (expensive)
+ feature extractor. This lets a stack prediction serve already-cached
+ slices first, before a sequential scan evicts them (see the widget)."""
+ cache = self._feature_cache
+ fe = self.fe_model
+ if cache is None or not cache.enabled or not fe.supports_feature_cache(param):
+ return None if cache_only else fe.extract_features_pyramid(d, param, patched=keep_patched, device=device)
+ key = (self._data_hash(d), self._fe_cache_signature(param))
+ payload = cache.get(key)
+ if payload is not None:
+ # Hit: reconstruct on `device` (the payload is lifted back to torch
+ # if that is the FE's native form) — same backend as a fresh
+ # extraction, so hits are as fast as (and identical to) misses.
+ return fe.features_from_cacheable(payload, d.shape, param,
+ patched=keep_patched, device=device)
+ if cache_only:
+ return None
+ # Miss: one extraction pass yields both the features (reconstructed
+ # from the on-device native form) and the numpy payload to store.
+ features, payload = fe.cacheable_repr_and_features(d, param, device,
+ patched=keep_patched)
+ cache.put(key, payload, spill_ok=fe.cache_spill_to_disk())
+ return features
+
### BACKEND METHOD FOR FEATURE EXTRACTION
def _get_features(self, data, annotations=None, restore_input_form=True,
memory_mode=False, img_ids=None,
in_channels=None, skip_norm=False, use_device=None,
- pca_components=0, kmeans_clusters=0):
+ pca_components=0, kmeans_clusters=0, cache_only=False):
"""
Returns the features of images extracted by the feature extractor model.
@@ -1123,13 +1204,16 @@ 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,
- device=fe_runtime_device)
+ features = [self._extract_pyramid_cached(
+ d, params_for_extract, keep_patched, fe_runtime_device,
+ cache_only=cache_only)
for d in data]
-
+
+ # cache_only peek: if any image's features are not already cached, signal
+ # a miss so the caller can defer this image to the compute pass.
+ if cache_only and any(f is None for f in features):
+ return None
+
if pca_components:
features = [utils.apply_pca_to_f_image(f, n_components=pca_components)
for f in features]
@@ -1504,7 +1588,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, cache_only=False):
"""
Backend method to predict images as a whole or tiling and parallelizing the prediction.
@@ -1537,10 +1621,17 @@ def _predict(self, data, add_seg=False, in_channels=None, skip_norm=False, use_d
# Get class probabilities, using tiling if enabled
if self._param.tile_image:
+ if cache_only:
+ # The tiled path doesn't support the peek, so it counts as a miss.
+ return None
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
+ # cache_only is a peek: only serve images whose features are already
+ # cached; None on any miss defers them to the caller's compute pass.
+ probas = self._predict_image(data, return_proba=True, fe_use_device=fe_use_device, cache_only=cache_only) # Can handle lists directly
+ if probas is None:
+ return None
# Restore input dimensionality (especially see if we want to remove z dimension)
probas = [self._restore_dims(probas[i], input_shapes[i])
@@ -1559,7 +1650,7 @@ def _predict(self, data, add_seg=False, in_channels=None, skip_norm=False, use_d
else:
return probas
- def _predict_image(self, image, return_proba=True, feature_img=None, fe_use_device=None):
+ def _predict_image(self, image, return_proba=True, feature_img=None, fe_use_device=None, cache_only=False):
"""
Backend method to predict images without tiling and parallelization.
Returns the class probabilities and optionally the segmentation of the images.
@@ -1580,7 +1671,10 @@ def _predict_image(self, image, return_proba=True, feature_img=None, fe_use_devi
restore_input_form=False,
in_channels=None, # already extracted outside
skip_norm=True, # already normalized outside
- use_device=fe_use_device)
+ use_device=fe_use_device,
+ cache_only=cache_only)
+ if feature_img is None: # cache_only peek: features not cached
+ return None
num_f = feature_img[0].shape[0] if isinstance(feature_img, list) else feature_img.shape[0]
num_f_clf = self.num_features
diff --git a/src/napari_convpaint/convpaint_widget.py b/src/napari_convpaint/convpaint_widget.py
index 649f069..84d1581 100644
--- a/src/napari_convpaint/convpaint_widget.py
+++ b/src/napari_convpaint/convpaint_widget.py
@@ -41,6 +41,11 @@ class ConvpaintWidget(QWidget):
by default.
"""
+ # Highest usable class label value: annotation layers are uint8 and the
+ # predicted segmentation is cast to uint8 (required by the smoothening
+ # filter), so larger values would wrap around.
+ MAX_CLASS_VALUE = 255
+
### Define the basic structure of the widget
def __init__(self, napari_viewer, parent=None, third_party=False):
@@ -75,29 +80,39 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.tab_names += ['Multifile']
tab_layouts = [None if name not in ['Models', 'Multifile'] else QGridLayout() for name in self.tab_names]
self.tabs = TabSet(self.tab_names, tab_layouts=tab_layouts) # [None, None, QGridLayout()])
+ # Left-aligned tabs; scroll buttons let the bar collapse gracefully
+ # when the dock is narrow.
tab_bar = self.tabs.tabBar()
- tab_bar.setSizePolicy(tab_bar.sizePolicy().horizontalPolicy(), tab_bar.sizePolicy().verticalPolicy())
-
- # Create docs button
- docs_button = QtWidgets.QToolButton()
- docs_button.setText("Documentation")
- docs_button.setStyleSheet("QToolButton {color: #999; text-decoration: underline; margin-left: 4px; margin-right: 8px}")
- docs_button.clicked.connect(lambda: QtGui.QDesktopServices.openUrl(QUrl("https://guiwitz.github.io/napari-convpaint/book/Landing.html")))
- docs_button.setToolTip("Open the documentation in your default browser.")
-
- # Create a widget to hold tab bar and button side by side
- tab_header_widget = QWidget()
- tab_header_layout = QtWidgets.QHBoxLayout(tab_header_widget)
- tab_header_layout.setContentsMargins(0, 0, 0, 0)
- tab_header_layout.setSpacing(0)
+ tab_bar.setUsesScrollButtons(True)
- tab_header_layout.addWidget(tab_bar)
- tab_header_layout.addWidget(docs_button)
+ # (Do NOT reparent the tab bar into a custom header row: QTabWidget
+ # keeps managing its bar's geometry on every resize and re-centers it,
+ # fighting any outside layout. The docs link lives on the Home tab.)
# Add to your main layout
- self.main_layout.addWidget(tab_header_widget)
self.main_layout.addWidget(self.tabs)
-
+
+ # Remove the dead space around the tab content: no pane frame, tabs
+ # left-aligned on the bar row, and a tight top margin on each tab page.
+ self.main_layout.setSpacing(0)
+ # Tight outer margins so the widget sits in its dock like napari's own
+ # panels (the default ~20px on every side reads as extra indentation
+ # compared to e.g. the layer controls); top matches the 4px gap
+ # between the tab bar and the first item.
+ self.main_layout.setContentsMargins(6, 4, 6, 6)
+ self._style_tabs()
+ self.viewer.events.theme.connect(self._style_tabs)
+ self.viewer.events.theme.connect(self._update_selected_class_highlight)
+ for i in range(self.tabs.count()):
+ page_layout = self.tabs.widget(i).layout()
+ if page_layout is not None:
+ # Zero left margin: the first tab starts exactly at the bar's
+ # left edge (x=0, measured), and a group box draws its frame at
+ # its widget edge — so any left page margin shows up as
+ # misalignment between tab headers and content. The small right
+ # margin keeps a gap between items and the vertical scrollbar.
+ page_layout.setContentsMargins(0, 4, 6, 8)
+
# Align rows in some tabs on top
for tab_name in ['Home', 'Models', 'Advanced']:
if tab_name in self.tabs.tab_names:
@@ -108,21 +123,33 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Create groups and separate labels
self.model_group = VHGroup('Model', orientation='G')
self.layer_selection_group = VHGroup('Layer selection', orientation='G')
- self.image_processing_group = VHGroup('Image type and normalization', orientation='G')
- self.train_group = VHGroup('Train/Segment', orientation='G')
+ self.image_processing_group = VHGroup('Image type && Normalization', orientation='G')
+ self.train_group = VHGroup('Train / Segment', orientation='G')
# self.segment_group = VHGroup('Segment', orientation='G')
# self.load_save_group = VHGroup('Load/Save', orientation='G')
- self.acceleration_group = VHGroup('Acceleration and post-processing', orientation='G')
+ self.acceleration_group = VHGroup('Acceleration && Post-processing', orientation='G')
# Create the shortcuts info
- shortcuts_text1 = 'Shift+a: Toggle annotations\nShift+s: Train\nShift+d: Predict\nShift+f: Toggle prediction'
- shortcuts_text2 = 'Shift+q: Set annotations label 1\nShift+w: Set annotations label 2\nShift+e: Set annotations label 3\nShift+r: Set annotations label 4'
+ shortcuts_text1 = 'Shift+a: Toggle annot.\nShift+s: Train\nShift+d: Predict\nShift+f: Toggle prediction'
+ shortcuts_text2 = 'Shift+q: Select 1st class\nShift+w: Select 2nd class\nShift+e: Select 3rd class\nShift+r: Select 4th class'
shortcuts_label1 = QLabel(shortcuts_text1)
shortcuts_label2 = QLabel(shortcuts_text2)
shortcuts_label1.setStyleSheet(style_for_shortcut_info)
shortcuts_label2.setStyleSheet(style_for_shortcut_info)
+ # Docs link above the shortcut hints
+ docs_link = QLabel(
+ 'Information and tutorials in '
+ 'documentation') # same grey as the hint text
+ docs_link.setOpenExternalLinks(True)
+ docs_link.setToolTip("Open the documentation in your default browser.")
+ docs_link.setStyleSheet(style_for_shortcut_info)
shortcuts_grid = QGridLayout()
- shortcuts_grid.addWidget(shortcuts_label1, 0, 0)
- shortcuts_grid.addWidget(shortcuts_label2, 0, 1)
+ # No grid margins: the hints are plain info text, so the usual widget
+ # margins just read as a large gap below the last group box.
+ shortcuts_grid.setContentsMargins(4, 0, 4, 0)
+ shortcuts_grid.addWidget(docs_link, 0, 0, 1, 2)
+ shortcuts_grid.addWidget(shortcuts_label1, 1, 0)
+ shortcuts_grid.addWidget(shortcuts_label2, 1, 1)
shortcuts_widget = QWidget()
shortcuts_widget.setLayout(shortcuts_grid)
@@ -144,6 +171,8 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Add buttons for "Model" group
# Current model description label
self.model_description1 = QLabel('None')
+ # Word wrap so the summary never dictates the dock's minimum width
+ self.model_description1.setWordWrap(True)
self.model_group.glayout.addWidget(self.model_description1, 0,0,1,2)
# Save and load model buttons
self.save_model_btn = QPushButton('Save model')
@@ -166,7 +195,7 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.image_layer_label = QLabel('Image layer')
self.layer_selection_group.glayout.addWidget(self.image_layer_label, 0,0,1,1)
self.layer_selection_group.glayout.addWidget(self.image_layer_selection_widget.native, 0,1,1,1)
- self.annotations_layer_label = QLabel('annotations layer')
+ self.annotations_layer_label = QLabel('Annotations layer')
self.layer_selection_group.glayout.addWidget(self.annotations_layer_label, 1,0,1,1)
self.layer_selection_group.glayout.addWidget(self.annotations_layer_selection_widget.native, 1,1,1,1)
@@ -178,9 +207,9 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Add buttons for "Image Processing" group
# Radio buttons for "Data Dimensions"
self.button_group_channels = QButtonGroup()
- self.radio_single_channel = QRadioButton('Single channel image')
- self.radio_multi_channel = QRadioButton('Multichannel image')
- self.radio_rgb = QRadioButton('RGB image')
+ self.radio_single_channel = QRadioButton('Single channel img')
+ self.radio_multi_channel = QRadioButton('Multichannel img')
+ self.radio_rgb = QRadioButton('RGB img')
self.radio_single_channel.setChecked(True)
self.channel_buttons = [self.radio_single_channel, self.radio_multi_channel, self.radio_rgb]
for x in self.channel_buttons: x.setEnabled(False)
@@ -198,16 +227,23 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# "Normalize" radio buttons
self.button_group_normalize = QButtonGroup()
self.radio_no_normalize = QRadioButton('No normalization')
- self.radio_normalize_over_stack = QRadioButton('Normalize over stack')
- self.radio_normalize_by_image = QRadioButton('Normalized by plane')
+ self.radio_normalize_over_stack = QRadioButton('Norm. over stack')
+ self.radio_normalize_by_image = QRadioButton('Norm. by plane')
self.radio_normalize_over_stack.setChecked(True)
self.norm_buttons = [self.radio_no_normalize, self.radio_normalize_over_stack, self.radio_normalize_by_image]
self.button_group_normalize.addButton(self.radio_no_normalize, id=1)
self.button_group_normalize.addButton(self.radio_normalize_over_stack, id=2)
self.button_group_normalize.addButton(self.radio_normalize_by_image, id=3)
- self.image_processing_group.glayout.addWidget(self.radio_no_normalize, 0,2,1,1)
- self.image_processing_group.glayout.addWidget(self.radio_normalize_over_stack, 1,2,1,1)
- self.image_processing_group.glayout.addWidget(self.radio_normalize_by_image, 2,2,1,1)
+ # Left-align the right radio column within its cells
+ self.image_processing_group.glayout.addWidget(self.radio_no_normalize, 0,2,1,1, Qt.AlignLeft)
+ self.image_processing_group.glayout.addWidget(self.radio_normalize_over_stack, 1,2,1,1, Qt.AlignLeft)
+ self.image_processing_group.glayout.addWidget(self.radio_normalize_by_image, 2,2,1,1, Qt.AlignLeft)
+ # Extra width goes to the two radio columns, not the divider column —
+ # otherwise the divider's cell grows and pushes the right column away
+ # from it (looks centered instead of left-aligned).
+ self.image_processing_group.glayout.setColumnStretch(0, 1)
+ self.image_processing_group.glayout.setColumnStretch(1, 0)
+ self.image_processing_group.glayout.setColumnStretch(2, 1)
# Add buttons for "Train/Segment" group
self.train_classifier_btn = QPushButton('Train')
@@ -226,35 +262,37 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# "Tile annotations" checkbox
self.check_tile_annotations = QCheckBox('Tile annotations for training')
self.check_tile_annotations.setChecked(False)
- self.acceleration_group.glayout.addWidget(self.check_tile_annotations, 0,0,1,1)
+ # Stacked vertically: side by side these two are the widest row of the
+ # Home tab and would dictate the dock's minimum width.
+ self.acceleration_group.glayout.addWidget(self.check_tile_annotations, 0,0,1,2)
# "Tile image" checkbox
self.check_tile_image = QCheckBox('Tile image for segmentation')
self.check_tile_image.setChecked(False)
- self.acceleration_group.glayout.addWidget(self.check_tile_image, 0,1,1,1)
+ self.acceleration_group.glayout.addWidget(self.check_tile_image, 1,0,1,2)
# Use Device/GPU dropdown
self.device_options_default = ['auto', 'gpu', 'cpu']
self.device_options_gpu_only_clf = ['auto', 'gpu (only classifier)', 'cpu']
self.device_dropdown = QComboBox()
self.device_dropdown.addItems(self.device_options_default)
self.device_label = QLabel('Device (GPU/CPU)')
- self.acceleration_group.glayout.addWidget(self.device_label, 1,0,1,1)
- self.acceleration_group.glayout.addWidget(self.device_dropdown, 1,1,1,1)
+ self.acceleration_group.glayout.addWidget(self.device_label, 2,0,1,1)
+ self.acceleration_group.glayout.addWidget(self.device_dropdown, 2,1,1,1)
# "Downsample" spinbox
self.spin_downsample = QSpinBox()
self.spin_downsample.setMinimum(-20)
self.spin_downsample.setMaximum(20)
self.spin_downsample.setValue(1)
self.downsample_label = QLabel('Downsample input')
- self.acceleration_group.glayout.addWidget(self.downsample_label, 2,0,1,1)
- self.acceleration_group.glayout.addWidget(self.spin_downsample, 2,1,1,1)
+ self.acceleration_group.glayout.addWidget(self.downsample_label, 3,0,1,1)
+ self.acceleration_group.glayout.addWidget(self.spin_downsample, 3,1,1,1)
# "Smoothen output" spinbox
self.spin_smoothen = QSpinBox()
self.spin_smoothen.setMinimum(1)
self.spin_smoothen.setMaximum(20)
self.spin_smoothen.setValue(1)
self.smoothen_label = QLabel('Smoothen output')
- self.acceleration_group.glayout.addWidget(self.smoothen_label, 3,0,1,1)
- self.acceleration_group.glayout.addWidget(self.spin_smoothen, 3,1,1,1)
+ self.acceleration_group.glayout.addWidget(self.smoothen_label, 4,0,1,1)
+ self.acceleration_group.glayout.addWidget(self.spin_smoothen, 4,1,1,1)
# === MODEL TAB ===
@@ -271,6 +309,8 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Current model
self.model_description2 = QLabel('None')
+ # Word wrap so the summary never dictates the dock's minimum width
+ self.model_description2.setWordWrap(True)
self.current_model_group.glayout.addWidget(self.model_description2, 0, 0, 1, 2)
# Add "FE architecture" combo box to FE group
@@ -280,6 +320,8 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Add "FE description" label to FE group
self.FE_description = QLabel('None')
self.FE_description.setWordWrap(True)
+ # Info-text styling (italic, dimmed), like the notes on other tabs
+ self.FE_description.setStyleSheet(style_for_infos)
self.fe_group.glayout.addWidget(self.FE_description, 2, 0, 1, 2)
# Add "FE layers" list to FE group
@@ -370,28 +412,22 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# class_names_text.setStyleSheet("font-size: 11px; color: rgba(120, 120, 120, 70%)")#; font-style: italic")
self.classes_layout.addWidget(class_names_text, 0, 0, 1, 10)
- # Add buttons ("add class", "remove class" and reset)
+ # Add buttons ("add class", "remove class", import/export and reset)
self.add_class_btn = QPushButton('Add class')
- self.classes_layout.addWidget(self.add_class_btn, len(self.initial_names)+1, 0, 1, 5)
self.remove_class_btn = QPushButton('Remove class')
- self.classes_layout.addWidget(self.remove_class_btn, len(self.initial_names)+1, 5, 1, 5)
- # Minimal import/export buttons (CSV)
self.export_class_names_btn = QPushButton('Export class names (csv)')
- self.classes_layout.addWidget(self.export_class_names_btn, len(self.initial_names)+2, 0, 1, 5)
self.import_class_names_btn = QPushButton('Import class names (csv/txt)')
- self.classes_layout.addWidget(self.import_class_names_btn, len(self.initial_names)+2, 5, 1, 5)
- # Reset to initial state
self.reset_class_names_btn = QPushButton('Reset to default')
- self.classes_layout.addWidget(self.reset_class_names_btn, len(self.initial_names)+3, 0, 1, 10)
+ self.sync_classes_btn = QPushButton('Add classes from annotations')
self.btn_class_distribution_annot = QPushButton('Show class distribution (in annotation)')
- self.classes_layout.addWidget(self.btn_class_distribution_annot, len(self.initial_names)+4, 0, 1, 10)
+ self._place_class_buttons(2) # two (placeholder) slots at startup
# Create the class names
self._create_default_class_names()
# Add the widget to the tab
- self.classes_layout.setColumnStretch(1, 1)
- self.classes_layout.setColumnStretch(5, 1)
+ self.classes_layout.setColumnStretch(2, 1)
+ self.classes_layout.setColumnStretch(6, 1)
self.tabs.add_named_tab('Classes', self.classes_widget)
# === ADVANCED TAB ===
@@ -399,22 +435,18 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
if 'Advanced' in self.tab_names:
# Create group boxes
self.advanced_note_group = VHGroup('Important note', orientation='G')
- self.advanced_appearance_group = VHGroup('Appearance', orientation='G')
- self.advanced_labels_group = VHGroup('Layers handling', orientation='G')
+ self.advanced_labels_group = VHGroup('Layers handling && Appearance', orientation='G')
self.advanced_training_group = VHGroup('Training', orientation='G')
# self.advanced_multifile_group = VHGroup('Multifile Training', orientation='G')
- self.advanced_prediction_group = VHGroup('Prediction', orientation='G')
self.advanced_input_group = VHGroup('Input', orientation='G')
self.advanced_output_group = VHGroup('Output', orientation='G')
self.advanced_unsupervised_group = VHGroup('Unsupervised extraction (without annotations)', orientation='G')
- # Add groups to the tab
+ # Add groups to the tab (the 'Performance' group is added below)
self.tabs.add_named_tab('Advanced', self.advanced_note_group.gbox)
- self.tabs.add_named_tab('Advanced', self.advanced_appearance_group.gbox)
self.tabs.add_named_tab('Advanced', self.advanced_labels_group.gbox)
self.tabs.add_named_tab('Advanced', self.advanced_training_group.gbox)
# self.tabs.add_named_tab('Advanced', self.advanced_multifile_group.gbox)$
- self.tabs.add_named_tab('Advanced', self.advanced_prediction_group.gbox)
self.tabs.add_named_tab('Advanced', self.advanced_input_group.gbox)
self.tabs.add_named_tab('Advanced', self.advanced_output_group.gbox)
self.tabs.add_named_tab('Advanced', self.advanced_unsupervised_group.gbox)
@@ -427,10 +459,10 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.advanced_note.setWordWrap(True)
self.advanced_note_group.glayout.addWidget(self.advanced_note, 0, 0, 1, 2)
- # Appearance: show/hide tooltips
+ # Show/hide tooltips
self.check_show_tooltips = QCheckBox('Show tooltips')
self.check_show_tooltips.setChecked(True)
- self.advanced_appearance_group.glayout.addWidget(self.check_show_tooltips, 0, 0, 1, 1)
+ self.advanced_labels_group.glayout.addWidget(self.check_show_tooltips, 5, 0, 1, 1)
# Wire the checkbox to toggle the promoted widgets' tooltips
self.check_show_tooltips.toggled.connect(lambda checked: self._setup_init_tooltips() if checked else self._remove_init_tooltips())
@@ -444,14 +476,17 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.check_keep_layers.setChecked(self.keep_layers)
self.advanced_labels_group.glayout.addWidget(self.check_keep_layers, 1, 1, 1, 1)
- # Button for adding annotations layers for selected images
- self.btn_add_all_annot_layers = QPushButton('Add for all selected')
- self.advanced_labels_group.glayout.addWidget(self.btn_add_all_annot_layers, 2, 0, 1, 1)
-
# Checkbox for auto-selecting annotations layers
- self.check_auto_select_annot = QCheckBox('Auto-select annotations layer')
+ self.check_auto_select_annot = QCheckBox('Auto-select annot. layer')
self.check_auto_select_annot.setChecked(self.auto_select_annot)
- self.advanced_labels_group.glayout.addWidget(self.check_auto_select_annot, 2, 1, 1, 1)
+ self.advanced_labels_group.glayout.addWidget(self.check_auto_select_annot, 2, 0, 1, 2)
+
+ # Button for adding annotations layers for selected images
+ self.btn_add_all_annot_layers = QPushButton('Add annot. layers for all selected images')
+ self.advanced_labels_group.glayout.addWidget(self.btn_add_all_annot_layers, 3, 0, 1, 2)
+
+ # --- dashed divider between the layer settings and appearance ---
+ self.advanced_labels_group.glayout.addWidget(self._dashed_divider(), 4, 0, 1, 2)
# Textbox to define the prefix for the annotations layers; NOTE: DISABLED FOR NOW
# self.text_annot_prefix = QtWidgets.QLineEdit()
@@ -486,20 +521,16 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Label for number of trainings performed
self.label_training_count = QLabel('')
- self.advanced_training_group.glayout.addWidget(self.label_training_count, 3, 0, 1, 2)
+ self.advanced_training_group.glayout.addWidget(self.label_training_count, 3, 0, 1, 4)
# Button to display a diagram of class distribution
self.btn_class_distribution_trained = QPushButton('Show class distr. (trained)')
- self.advanced_training_group.glayout.addWidget(self.btn_class_distribution_trained, 3, 2, 1, 2)
+ # Below the counts label (side by side they would be the widest row)
+ self.advanced_training_group.glayout.addWidget(self.btn_class_distribution_trained, 4, 0, 1, 4)
# Reset training button
self.btn_reset_training = QPushButton('Reset continuous training')
- self.advanced_training_group.glayout.addWidget(self.btn_reset_training, 4, 0, 1, 4)
-
- # Dask option
- self.check_use_dask = QCheckBox('Use Dask when tiling image for segmentation')
- self.check_use_dask.setChecked(self.use_dask)
- self.advanced_prediction_group.glayout.addWidget(self.check_use_dask, 0, 0, 1, 1)
+ self.advanced_training_group.glayout.addWidget(self.btn_reset_training, 5, 0, 1, 4)
# Input channels option
self.text_input_channels = QtWidgets.QLineEdit()
@@ -511,7 +542,7 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
# Button to switch first to axes
self.btn_switch_axes = QPushButton('Switch channels axis')
- self.advanced_input_group.glayout.addWidget(self.btn_switch_axes, 1, 0, 1, 2)
+ self.advanced_input_group.glayout.addWidget(self.btn_switch_axes, 1, 0, 1, 4)
# Checkbox for adding segmentation
self.check_add_seg = QCheckBox('Segmentation')
@@ -547,14 +578,72 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.advanced_unsupervised_group.glayout.addWidget(self.kmeans_label, 1, 0, 1, 2)
self.advanced_unsupervised_group.glayout.addWidget(self.text_features_kmeans, 1, 2, 1, 2)
+ # === ADVANCED TAB: PERFORMANCE SECTION (feature caching + Dask) ===
+
+ if 'Advanced' in self.tab_names:
+ self.advanced_cache_group = VHGroup('Performance', orientation='G')
+ self.tabs.add_named_tab('Advanced', self.advanced_cache_group.gbox)
+
+ # Explanatory note
+ cache_note = QLabel(
+ "Reuse extracted features when segmenting or training the same image "
+ "repeatedly (e.g. while refining annotations), instead of recomputing "
+ "them. Bounded by the memory limit below; on stacks/movies the oldest "
+ "cached slices are dropped first.")
+ cache_note.setStyleSheet(style_for_infos)
+ cache_note.setWordWrap(True)
+ self.advanced_cache_group.glayout.addWidget(cache_note, 0, 0, 1, 3)
+
+ # Enable/disable checkbox
+ self.check_use_cache = QCheckBox('Enable feature caching')
+ self.check_use_cache.setChecked(self.cache_enabled)
+ self.advanced_cache_group.glayout.addWidget(self.check_use_cache, 1, 0, 1, 3)
+
+ # Max RAM spinbox (MB)
+ self.cache_max_ram_label = QLabel('Max cache RAM (MB)')
+ self.advanced_cache_group.glayout.addWidget(self.cache_max_ram_label, 2, 0, 1, 2)
+ self.cache_max_ram_spinbox = QSpinBox()
+ self.cache_max_ram_spinbox.setRange(64, 1024 * 1024) # 64 MB .. 1 TB
+ self.cache_max_ram_spinbox.setSingleStep(256)
+ self.cache_max_ram_spinbox.setValue(self.cache_max_mb)
+ self.advanced_cache_group.glayout.addWidget(self.cache_max_ram_spinbox, 2, 2, 1, 1)
+
+ # Max disk spinbox (MB) — features evicted from RAM spill here instead
+ # of being recomputed. 0 disables disk spillover.
+ self.cache_max_disk_label = QLabel('Max cache disk (MB)')
+ self.advanced_cache_group.glayout.addWidget(self.cache_max_disk_label, 3, 0, 1, 2)
+ self.cache_max_disk_spinbox = QSpinBox()
+ self.cache_max_disk_spinbox.setRange(0, 8 * 1024 * 1024) # 0 .. 8 TB
+ self.cache_max_disk_spinbox.setSingleStep(1024)
+ self.cache_max_disk_spinbox.setValue(self.cache_disk_max_mb)
+ self.advanced_cache_group.glayout.addWidget(self.cache_max_disk_spinbox, 3, 2, 1, 1)
+
+ # Current cache size label (RAM + disk)
+ self.cache_size_label = QLabel('Current cache size: 0 MB')
+ self.advanced_cache_group.glayout.addWidget(self.cache_size_label, 4, 0, 1, 3)
+
+ # --- dashed divider between the caching and Dask parts ---
+ self.advanced_cache_group.glayout.addWidget(self._dashed_divider(), 5, 0, 1, 3)
+
+ # Dask option (applies to tiled segmentation)
+ dask_note = QLabel(
+ "Distribute the tiles of a tiled segmentation to parallel Dask "
+ "workers (only applies when 'Tile for segmentation' is enabled).")
+ dask_note.setStyleSheet(style_for_infos)
+ dask_note.setWordWrap(True)
+ self.advanced_cache_group.glayout.addWidget(dask_note, 6, 0, 1, 3)
+ self.check_use_dask = QCheckBox('Use Dask')
+ self.check_use_dask.setChecked(self.use_dask)
+ self.advanced_cache_group.glayout.addWidget(self.check_use_dask, 7, 0, 1, 3)
+
# === MULTIFILE TAB ===
if 'Multifile' in self.tab_names:
# Create three groups for the Multifile tab to match other tabs' style
self.multifile_files_group = VHGroup('Files', orientation='G')
- self.multifile_train_group = VHGroup('Train/Segment', orientation='G')
- self.multifile_reset_group = VHGroup('Clear/Close', orientation='G')
- self.multifile_export_import_group = VHGroup('Export/Import', orientation='G')
+ self.multifile_train_group = VHGroup('Train / Segment', orientation='G')
+ self.multifile_reset_group = VHGroup('Clear / Close', orientation='G')
+ self.multifile_export_import_group = VHGroup('Export / Import', orientation='G')
self.multifile_settings_group = VHGroup('Preferences', orientation='G')
# Add groups to the Multifile tab
@@ -580,6 +669,14 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.multifile_list = QTableWidget()
self.multifile_list.setColumnCount(3)
self.multifile_list.setHorizontalHeaderLabels(['Annot.', 'Image Filename', 'Segm.'])
+ # Flat header and list in the ACTIVE napari theme's colors; re-apply
+ # whenever the theme changes (dark <-> light).
+ self._style_multifile_list()
+ self.viewer.events.theme.connect(self._style_multifile_list)
+ # Same font as the rest of the plugin (tables default to the
+ # platform's smaller 'small-widget' font on some systems).
+ self.multifile_list.setFont(self.font())
+ self.multifile_list.horizontalHeader().setFont(self.font())
# Align the 'Image Filename' header label to the left for readability
try:
header_item = self.multifile_list.horizontalHeaderItem(1)
@@ -604,21 +701,25 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.multifile_list.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeToContents)
self.multifile_files_group.glayout.addWidget(self.multifile_list, 1, 0, 1, 3)
+ # 2+1 rows: three buttons side by side would be the widest row of
+ # the tab and dictate the dock's minimum width.
self.multifile_clear_annotations_btn = QPushButton('Clear selected annot.')
self.multifile_reset_group.glayout.addWidget(self.multifile_clear_annotations_btn, 1, 0, 1, 1)
- self.multifile_reset_folder_btn = QPushButton('Close folder')
- self.multifile_reset_group.glayout.addWidget(self.multifile_reset_folder_btn, 1, 1, 1, 1)
self.multifile_clear_segmentations_btn = QPushButton('Clear selected segm.')
- self.multifile_reset_group.glayout.addWidget(self.multifile_clear_segmentations_btn, 1, 2, 1, 1)
+ self.multifile_reset_group.glayout.addWidget(self.multifile_clear_segmentations_btn, 1, 1, 1, 1)
+ self.multifile_reset_folder_btn = QPushButton('Close folder')
+ self.multifile_reset_group.glayout.addWidget(self.multifile_reset_folder_btn, 2, 0, 1, 2)
# --- Train/Segment group: action buttons (placeholders for now)
- self.multifile_train_all_annot_btn = QPushButton('Train on annotated')
+ self.multifile_train_all_annot_btn = QPushButton('Train on annot.')
self.multifile_preview_btn = QPushButton('Preview segmentation')
self.multifile_segment_selected_btn = QPushButton('Segment selected')
+ # 2+1 rows: three buttons side by side would be the widest row of
+ # the tab and dictate the dock's minimum width.
self.multifile_train_group.glayout.addWidget(self.multifile_train_all_annot_btn, 0, 0, 1, 1)
self.multifile_train_group.glayout.addWidget(self.multifile_preview_btn, 0, 1, 1, 1)
- self.multifile_train_group.glayout.addWidget(self.multifile_segment_selected_btn, 0, 2, 1, 1)
+ self.multifile_train_group.glayout.addWidget(self.multifile_segment_selected_btn, 1, 0, 1, 2)
# --- Import/Export group: action buttons (placeholders for now)
self.multifile_export_annot_btn = QPushButton('Export annotations')
@@ -647,13 +748,36 @@ def __init__(self, napari_viewer, parent=None, third_party=False):
self.multifile_segmentation_suffix_txt.setText('segmentation')
self.multifile_settings_group.glayout.addWidget(self.multifile_annotations_suffix_txt, 1, 1, 1, 1)
self.multifile_settings_group.glayout.addWidget(self.multifile_segmentation_suffix_txt, 1, 2, 1, 1)
-
+
+ # === Match napari's control density ===
+ # Qt's default grid spacing is looser than napari's own panels; tighten
+ # the vertical spacing between rows inside all group boxes (and the
+ # Classes grid, which lives in a plain widget).
+ for gbox in self.findChildren(QtWidgets.QGroupBox):
+ gbox_layout = gbox.layout()
+ if isinstance(gbox_layout, QGridLayout):
+ gbox_layout.setVerticalSpacing(4)
+ gbox_layout.setHorizontalSpacing(4)
+ if hasattr(self, 'classes_layout'):
+ self.classes_layout.setVerticalSpacing(4)
+ self.classes_layout.setHorizontalSpacing(4)
+
+ # === Make all tabs scrollable ===
+ # All tab content is added by now — wrap every tab in a scroll area so
+ # nothing can be cut off on small screens.
+ for tab_name in self.tab_names:
+ self._make_tab_scrollable(tab_name)
+ # The remove/insert dance above moves the current-tab index around;
+ # make sure a fresh widget always opens on the first (Home) tab.
+ self.tabs.setCurrentIndex(0)
+
# === Show tooltips by default ===
self._setup_init_tooltips()
# Set device dropdown tooltip separately, as we want to show these dynamically and permanently, even when the "Show tooltips" checkbox is unchecked
self.device_dropdown.setToolTip('Select device policy for feature extraction and classifier.')
+
def _setup_init_tooltips(self):
# Set tooltip for the tabs
@@ -726,10 +850,11 @@ def _setup_init_tooltips(self):
# Classes tab
if 'Classes' in self.tab_names:
self.add_class_btn.setToolTip('Add a class name to the list.')
- self.remove_class_btn.setToolTip('Remove a class from the list. Note that this will also delete the corresponding annotations from the annotations layer, if they exist.')
+ self.remove_class_btn.setToolTip('Remove the SELECTED class (the one outlined / active in the annotations layer). Also deletes its annotations, if any. Other classes keep their label values. Grayed out when the selection is not a class.')
self.export_class_names_btn.setToolTip('Export class names as a csv file.')
self.import_class_names_btn.setToolTip('Import class names from a csv or txt file.')
- self.reset_class_names_btn.setToolTip('Reset the list of class names to "Background" and "Foreground".')
+ self.reset_class_names_btn.setToolTip('Clear all classes and start over with two empty slots.')
+ self.sync_classes_btn.setToolTip('Add class rows for label values already painted in the selected annotations layer (never removes rows).')
self.btn_class_distribution_annot.setToolTip('Show a diagram of the class distribution in the annotations layer.')
# Advanced tab
@@ -811,7 +936,7 @@ def _remove_init_tooltips(self):
# Classes tab
if 'Classes' in self.tab_names:
- for w in [self.add_class_btn, self.remove_class_btn,
+ for w in [self.add_class_btn, self.remove_class_btn, self.sync_classes_btn,
self.export_class_names_btn, self.import_class_names_btn,
self.reset_class_names_btn, self.btn_class_distribution_annot]:
w.setToolTip('')
@@ -864,6 +989,70 @@ def _import_convpaint_model_class(self):
from .convpaint_model import ConvpaintModel
self._cpm_class = ConvpaintModel
+ def _make_tab_scrollable(self, tab_name):
+ """Wrap a tab's content in a scroll area so it cannot be cut off on
+ small screens. Must be called AFTER everything has been added to the
+ tab (add_named_tab resolves the tab's widget by index, which becomes
+ the scroll area after wrapping)."""
+ idx = self.tabs.tab_names.index(tab_name)
+ content = self.tabs.widget(idx)
+ # Detach the page BEFORE handing it to the scroll area: setWidget()
+ # reparents it, which would already remove it from the tab widget and
+ # shift the indices under removeTab().
+ self.tabs.removeTab(idx)
+ scroll = QtWidgets.QScrollArea()
+ scroll.setWidgetResizable(True)
+ scroll.setFrameShape(QtWidgets.QFrame.NoFrame)
+ # No width floor: the dock may be made thinner than any tab's content
+ # (the widest tab must not dictate the plugin's minimum width). A tab
+ # whose content doesn't fit gets a horizontal scrollbar on demand
+ # instead of clipping.
+ scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
+ scroll.setWidget(content)
+ self.tabs.insertTab(idx, scroll, tab_name)
+
+ def _apply_feature_cache(self, *args, recreate=False):
+ """Apply the caching settings from the GUI controls to the active model
+ (connected directly to the controls' change signals). The cache_*
+ attributes hold the pre-GUI defaults and simply mirror the controls
+ afterwards. Pass recreate=True right after the model is (re)created;
+ otherwise the existing cache is updated in place so its entries survive
+ a settings change (disabling clears it, freeing RAM+disk)."""
+ model = getattr(self, "cp_model", None)
+ if model is None:
+ return
+ if hasattr(self, "check_use_cache"):
+ self.cache_enabled = self.check_use_cache.isChecked()
+ self.cache_max_mb = self.cache_max_ram_spinbox.value()
+ self.cache_disk_max_mb = self.cache_max_disk_spinbox.value()
+ # Use decimal MB (1e6) here to match the size shown in the label (also
+ # /1e6), so the number the user types is exactly the max size displayed.
+ max_bytes = int(self.cache_max_mb) * 1_000_000
+ disk_max_bytes = int(self.cache_disk_max_mb) * 1_000_000
+ fc = model._feature_cache
+ if fc is None or recreate:
+ model.enable_feature_cache(enabled=self.cache_enabled, max_bytes=max_bytes,
+ disk_max_bytes=disk_max_bytes)
+ else:
+ fc.set_max_bytes(max_bytes)
+ fc.set_disk_max_bytes(disk_max_bytes)
+ fc.set_enabled(self.cache_enabled)
+ self._refresh_cache_size_label()
+
+ def _refresh_cache_size_label(self):
+ if not hasattr(self, "cache_size_label"):
+ return
+ model = getattr(self, "cp_model", None)
+ fc = model._feature_cache if model is not None else None
+ if fc is None:
+ self.cache_size_label.setText('Current cache size: 0 MB')
+ return
+ s = fc.stats() # one lock acquisition for all fields
+ text = (f'Current cache size: RAM {s["bytes"] / 1e6:.0f} MB ({s["entries"]}), '
+ f'disk {s["disk_bytes"] / 1e6:.0f} MB ({s["disk_entries"]})')
+ if text != self.cache_size_label.text():
+ self.cache_size_label.setText(text)
+
def _late_init(self):
"""Populate UI widgets with defaults from ConvpaintModel, set up connections, and reset model.
This is called after the GUI is shown to ensure that all components are properly initialized."""
@@ -871,6 +1060,7 @@ def _late_init(self):
# === MODEL DEFAULTS & WIDGET POPULATION ===
self._import_convpaint_model_class()
self.cp_model = self._cpm_class()
+ self._apply_feature_cache(recreate=True)
# Get default parameters to set in widget
self.default_cp_param = self._cpm_class.get_default_params()
# Use variables of main model as temp variables for the Models tab, as it is the one model used at that time
@@ -909,10 +1099,12 @@ def _late_init(self):
self.viewer.bind_key('Shift+s', self._on_train, overwrite=True)
self.viewer.bind_key('Shift+d', self._on_predict, overwrite=True)
self.viewer.bind_key('Shift+f', self.toggle_prediction, overwrite=True)
- self.viewer.bind_key('Shift+q', lambda event=None: self.set_annot_label_class(1, event), overwrite=True)
- self.viewer.bind_key('Shift+w', lambda event=None: self.set_annot_label_class(2, event), overwrite=True)
- self.viewer.bind_key('Shift+e', lambda event=None: self.set_annot_label_class(3, event), overwrite=True)
- self.viewer.bind_key('Shift+r', lambda event=None: self.set_annot_label_class(4, event), overwrite=True)
+ # Shortcuts select the Nth class ROW (its actual label value may differ
+ # from N once classes are sparse).
+ self.viewer.bind_key('Shift+q', lambda event=None: self._set_annot_label_by_row(0, event), overwrite=True)
+ self.viewer.bind_key('Shift+w', lambda event=None: self._set_annot_label_by_row(1, event), overwrite=True)
+ self.viewer.bind_key('Shift+e', lambda event=None: self._set_annot_label_by_row(2, event), overwrite=True)
+ self.viewer.bind_key('Shift+r', lambda event=None: self._set_annot_label_by_row(3, event), overwrite=True)
### Define the connections between the widget elements
@@ -945,6 +1137,7 @@ def _add_connections(self):
self.image_layer_selection_widget.changed.connect(self._delayed_on_select_layer)
self.annotations_layer_selection_widget.native.activated.connect(self._on_select_annot)
self.annotations_layer_selection_widget.changed.connect(self._on_select_annot)
+ self.annotations_layer_selection_widget.changed.connect(self._update_selected_class_highlight)
self.add_layers_btn.clicked.connect(self._on_add_annot_layer)
# Image Processing; only trigger from buttons that are activated (checked)
@@ -1018,9 +1211,10 @@ def _add_connections(self):
self.export_class_names_btn.clicked.connect(lambda: self._export_class_names_dialog())
self.import_class_names_btn.clicked.connect(lambda: self._import_class_names_dialog())
self.reset_class_names_btn.clicked.connect(self._on_reset_class_names)
+ self.sync_classes_btn.clicked.connect(self._on_sync_classes_from_annotations)
- for class_name in self.class_names:
- class_name.textChanged.connect(self._update_class_names)
+ for r in self.class_rows:
+ r['name'].textChanged.connect(self._update_class_names)
if self.annotations_layer_selection_widget.value is not None:
labels_layer = self.annotations_layer_selection_widget.value
labels_layer.events.colormap.connect(self._on_change_annot_cmap)
@@ -1055,6 +1249,17 @@ def _add_connections(self):
self.check_use_dask.stateChanged.connect(lambda: setattr(
self, 'use_dask', self.check_use_dask.isChecked()))
+ if hasattr(self, 'check_use_cache'):
+ # All three controls apply the full settings set in one go.
+ self.check_use_cache.stateChanged.connect(self._apply_feature_cache)
+ self.cache_max_ram_spinbox.valueChanged.connect(self._apply_feature_cache)
+ self.cache_max_disk_spinbox.valueChanged.connect(self._apply_feature_cache)
+ # Keep the "current cache size" label live.
+ self._cache_size_timer = QTimer(self)
+ self._cache_size_timer.setInterval(1000)
+ self._cache_size_timer.timeout.connect(self._refresh_cache_size_label)
+ self._cache_size_timer.start()
+
self.text_input_channels.textChanged.connect(lambda: setattr(
self, 'input_channels', self.text_input_channels.text()))
@@ -1135,6 +1340,11 @@ def toggle_prediction(self, event=None):
else:
self.viewer.layers[self.seg_tag].visible = False
+ def _set_annot_label_by_row(self, row_idx, event=None):
+ """Shortcut target: select the label value of the row_idx-th class row."""
+ if row_idx < len(self.class_rows):
+ self.set_annot_label_class(self.class_rows[row_idx]['value'], event)
+
def set_annot_label_class(self, x, event=None):
"""Set the label class of the annotations layer."""
annot_layer = self.annotations_layer_selection_widget.value
@@ -1155,11 +1365,20 @@ def set_annot_label_class(self, x, event=None):
# Classes
+ def _clear_class_rows(self):
+ """Remove and delete all class-row widgets."""
+ for r in self.class_rows:
+ for w in (r['icon'], r['value_lbl'], r['name']):
+ self.classes_layout.removeWidget(w)
+ w.deleteLater()
+ self.class_rows.clear()
+
def _create_default_class_names(self):
- """Create the default class names and icons in the layout."""
- # Start with default class names
- for name in self.initial_names:
- self._on_add_class(text=name)
+ """Initialize the class list to its default state: NO classes, two
+ empty placeholder slots (type a name into one to create a class)."""
+ self._rebuild_class_rows_layout()
+ self._update_class_names()
+ self._update_selected_class_highlight()
# Add default annot and seg layers if they exist
if self.annotations_layer_selection_widget.value is not None:
self.annot_layers.add(self.annotations_layer_selection_widget.value)
@@ -1172,19 +1391,12 @@ def _on_reset_class_names(self):
"""Reset the class names to the default ones and update all annotations and segmentation layers."""
# Remove and delete all class name widgets and icons
- for name in self.class_names:
- self.classes_layout.removeWidget(name)
- name.deleteLater()
- for icon in self.class_icons:
- self.classes_layout.removeWidget(icon)
- icon.deleteLater()
-
- self.class_names.clear()
- self.class_icons.clear()
+ self._clear_class_rows()
# Remove the buttons from the layout
self.classes_layout.removeWidget(self.add_class_btn)
self.classes_layout.removeWidget(self.remove_class_btn)
+ self.classes_layout.removeWidget(self.sync_classes_btn)
self.classes_layout.removeWidget(self.reset_class_names_btn)
self.classes_layout.removeWidget(self.btn_class_distribution_annot)
@@ -1192,134 +1404,314 @@ def _on_reset_class_names(self):
self._create_default_class_names()
# Re-add the buttons below the class names
- self.classes_layout.addWidget(self.add_class_btn, len(self.class_names)+1, 0, 1, 5)
- self.classes_layout.addWidget(self.remove_class_btn, len(self.class_names)+1, 5, 1, 5)
- self.classes_layout.addWidget(self.export_class_names_btn, len(self.class_names)+2, 0, 1, 5)
- self.classes_layout.addWidget(self.import_class_names_btn, len(self.class_names)+2, 5, 1, 5)
- self.classes_layout.addWidget(self.reset_class_names_btn, len(self.class_names)+3, 0, 1, 10)
- self.classes_layout.addWidget(self.btn_class_distribution_annot, len(self.class_names)+4, 0, 1, 10)
-
- def _on_add_class(self, text=None):
- """Add a new class name and icon to the layout and update all annotations and segmentation layers."""
-
- # Create a new class name
- new_name = QtWidgets.QLineEdit()
- new_name.setStyleSheet("font-size: 12px;")
- self.class_names.append(new_name)
- class_num = len(self.class_names) # Class number is the length of the list
- # Add the new name to the layout
- self.classes_layout.addWidget(new_name, class_num, 1, 1, 9)
- # Set the text of the new name
- text_str = text if text is not None else f'Class {class_num}'
- new_name.setText(text_str)
-
- # Change "clear" button to the last name and connect it to deleting the entire entry (instead of only text)
- # new_name.setClearButtonEnabled(True)
- # new_name.textChanged.connect(self.remove_class_name)
- # self.class_names[-2].setClearButtonEnabled(False)
-
- # Connect the new name to the update function
- new_name.textChanged.connect(self._update_class_names)
-
- # Add a new icon
- new_icon = QtWidgets.QLabel()
- self.class_icons.append(new_icon)
- self.classes_layout.addWidget(new_icon, class_num, 0)
- new_icon.mousePressEvent = lambda event: self._set_all_labels_classes(class_num, event)
-
- # Update the icon with the color of the last label and all class names
- self._update_class_icons(class_num)
- self._update_class_names()
+ self._place_class_buttons(len(self.class_rows))
+
+ def _style_tabs(self, event=None):
+ """(Re-)apply the tab-bar style with the active theme's colors as SOLID
+ fills — napari's own tab rule paints a vertical gradient, which clashes
+ with the joined segmented-control look. Connected to viewer.events.theme
+ so dark <-> light switches restyle."""
+ def _hex(color):
+ as_hex = getattr(color, 'as_hex', None)
+ return as_hex() if callable(as_hex) else str(color)
+ try:
+ from napari.utils.theme import get_theme
+ theme = get_theme(self.viewer.theme)
+ if isinstance(theme, dict):
+ fg, cur = _hex(theme['foreground']), _hex(theme['current'])
+ else:
+ fg, cur = _hex(theme.foreground), _hex(theme.current)
+ except Exception:
+ fg, cur = '#414851', '#0f6285' # napari dark
+ self.tabs.setStyleSheet(
+ "QTabWidget::pane { border: 0; margin: 0; padding: 0; } "
+ "QTabWidget::tab-bar { alignment: left; } "
+ # Joined segmented-control look: adjacent tabs share square inner
+ # corners (rounded inner corners leave notches that expose
+ # tab-colored nubs of the neighbor when the bar is squeezed);
+ # only the outer corners of the first/last tab stay rounded.
+ "QTabBar { background: transparent; } "
+ f"QTabBar::tab {{ margin-right: 0px; border-radius: 0px; background: {fg}; }} "
+ f"QTabBar::tab:selected {{ background: {cur}; }} "
+ "QTabBar::tab:first { border-top-left-radius: 4px; border-bottom-left-radius: 4px; } "
+ "QTabBar::tab:last { border-top-right-radius: 4px; border-bottom-right-radius: 4px; } "
+ "QTabBar::tab:only-one { border-radius: 4px; }")
+
+ def _style_multifile_list(self, event=None):
+ """(Re-)apply the active napari theme's colors to the multifile list.
+ Palette roles can't be used: napari themes via stylesheet only, so the
+ Qt palette keeps the platform's light look (a Win95-style bevel).
+ Connected to viewer.events.theme so dark <-> light switches restyle."""
+ def _hex(color):
+ as_hex = getattr(color, 'as_hex', None)
+ return as_hex() if callable(as_hex) else str(color)
+ try:
+ from napari.utils.theme import get_theme
+ theme = get_theme(self.viewer.theme)
+ if isinstance(theme, dict):
+ bg, fg, txt = (_hex(theme['background']), _hex(theme['foreground']),
+ _hex(theme['text']))
+ else:
+ bg, fg, txt = _hex(theme.background), _hex(theme.foreground), _hex(theme.text)
+ except Exception:
+ bg, fg, txt = '#262930', '#414851', '#f0f1f2' # napari dark
+ self.multifile_list.setStyleSheet(
+ f"QHeaderView::section {{ background-color: {fg}; color: {txt};"
+ f" border: none; border-right: 1px solid {bg}; padding: 3px 6px; }}"
+ f"QHeaderView::section:first {{ border-top-left-radius: 4px; }}"
+ f"QHeaderView::section:last {{ border-top-right-radius: 4px; border-right: none; }}"
+ f"QTableWidget {{ border: 1px solid {fg}; border-radius: 4px;"
+ f" gridline-color: {fg}; background-color: {bg}; }}"
+ f"QTableCornerButton::section {{ background-color: {fg}; border: none; }}")
- # Move the add, remove and reset buttons one down
- self.classes_layout.removeWidget(self.add_class_btn)
- self.classes_layout.removeWidget(self.remove_class_btn)
- self.classes_layout.removeWidget(self.reset_class_names_btn)
- self.classes_layout.removeWidget(self.btn_class_distribution_annot)
- self.classes_layout.addWidget(self.add_class_btn, class_num+1, 0, 1, 5)
- self.classes_layout.addWidget(self.remove_class_btn, class_num+1, 5, 1, 5)
- self.classes_layout.addWidget(self.export_class_names_btn, class_num+2, 0, 1, 5)
- self.classes_layout.addWidget(self.import_class_names_btn, class_num+2, 5, 1, 5)
- self.classes_layout.addWidget(self.reset_class_names_btn, class_num+3, 0, 1, 10)
- self.classes_layout.addWidget(self.btn_class_distribution_annot, class_num+4, 0, 1, 10)
+ @staticmethod
+ def _dashed_divider():
+ """A thin dashed horizontal line used to separate subsections,
+ with breathing room above and below."""
+ divider = QtWidgets.QFrame()
+ divider.setFixedHeight(13) # 6px margin + 1px line + 6px margin
+ divider.setStyleSheet(
+ "border: none; border-top: 1px dashed rgba(120, 120, 120, 50%); margin: 6px 0;")
+ return divider
+
+ def _place_class_buttons(self, n_classes):
+ """(Re-)place the static buttons below the class-name rows.
+ Export/import are stacked vertically: side by side they would be the
+ widest row of the Classes tab and dictate the dock's minimum width."""
+ # Dividers are created once and re-placed on layout rebuilds.
+ if not hasattr(self, '_classes_divider1'):
+ self._classes_divider1 = self._dashed_divider()
+ self._classes_divider2 = self._dashed_divider()
+ self.classes_layout.addWidget(self.add_class_btn, n_classes+1, 0, 1, 5)
+ self.classes_layout.addWidget(self.remove_class_btn, n_classes+1, 5, 1, 5)
+ self.classes_layout.addWidget(self.sync_classes_btn, n_classes+2, 0, 1, 10)
+ self.classes_layout.addWidget(self._classes_divider1, n_classes+3, 0, 1, 10)
+ self.classes_layout.addWidget(self.export_class_names_btn, n_classes+4, 0, 1, 10)
+ self.classes_layout.addWidget(self.import_class_names_btn, n_classes+5, 0, 1, 10)
+ self.classes_layout.addWidget(self._classes_divider2, n_classes+6, 0, 1, 10)
+ self.classes_layout.addWidget(self.reset_class_names_btn, n_classes+7, 0, 1, 10)
+ self.classes_layout.addWidget(self.btn_class_distribution_annot, n_classes+8, 0, 1, 10)
+
+ def _class_values(self):
+ return [r['value'] for r in self.class_rows]
+
+ def _row_for_value(self, value):
+ for r in self.class_rows:
+ if r['value'] == value:
+ return r
+ return None
+
+ def _class_name_for_value(self, value):
+ """Display name for a label value; a plain 'Class N' for values that
+ have no row (they can exist internally, e.g. in dense layer props)."""
+ row = self._row_for_value(value)
+ return row['name'].text() if row is not None else f'Class {value}'
+
+ def _selected_class_value(self):
+ """The selected class value: the annotations layer's selected label
+ while a layer exists, else the widget's own last selection."""
+ annot = self.annotations_layer_selection_widget.value
+ if annot is not None:
+ return getattr(annot, 'selected_label', None)
+ return self._ui_selected_class
+
+ def _create_class_row(self, value, text=None):
+ """Build the widgets for one class row and register it (sorted by value).
+ Layout placement happens in _rebuild_class_rows_layout."""
+ icon = QtWidgets.QLabel()
+ icon.mousePressEvent = lambda event, v=value: self._on_class_swatch_clicked(v)
+ value_lbl = QtWidgets.QLabel(str(value))
+ value_lbl.setStyleSheet("font-size: 12px; color: rgba(120, 120, 120, 80%);")
+ value_lbl.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
+ value_lbl.setToolTip('The label value this class paints and trains as.')
+ name = QtWidgets.QLineEdit()
+ name.setStyleSheet("font-size: 12px;")
+ name.setText(text if text is not None else f'Class {value}')
+ name.textChanged.connect(self._update_class_names)
+ row = {'value': value, 'icon': icon, 'value_lbl': value_lbl, 'name': name}
+ self.class_rows.append(row)
+ self.class_rows.sort(key=lambda r: r['value'])
+ return row
+
+ def _rebuild_class_rows_layout(self):
+ """(Re-)place all class rows (sorted by value) and the buttons below
+ them, padding the display to two slots with PLACEHOLDER rows (ghost
+ value, striped swatch, disabled name) — training needs two classes, so
+ the tab keeps that shape, but the slots are not real classes and the
+ list may hold fewer. Rebuilding wholesale keeps insertion/removal
+ simple; the list is small."""
+ for ph in getattr(self, '_placeholder_rows', []):
+ for w in (ph['icon'], ph['value_lbl'], ph['name']):
+ self.classes_layout.removeWidget(w)
+ w.deleteLater()
+ self._placeholder_rows = []
+ next_free = max(self._class_values(), default=0) + 1
+ n_placeholders = max(0, 2 - len(self.class_rows))
+ for i in range(n_placeholders):
+ v = next_free + i
+ # Pure UI placeholders: the swatch is NOT clickable and nothing
+ # here touches the layers — they are not classes, just the shape
+ # of the two slots training will need. Typing a name creates the
+ # real class in place.
+ icon = QtWidgets.QLabel()
+ icon.setPixmap(self._placeholder_pixmap())
+ # Same transparent border the highlight system puts on real
+ # swatches, so placeholder and class swatches render equal-sized.
+ icon.setStyleSheet('border: 2px solid transparent;')
+ icon.setToolTip('Not a class yet — type a name to add it.')
+ value_lbl = QtWidgets.QLabel(str(v))
+ value_lbl.setStyleSheet("font-size: 12px; color: rgba(120, 120, 120, 50%);")
+ value_lbl.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
+ name = QtWidgets.QLineEdit()
+ name.setStyleSheet("font-size: 12px; font-style: italic;")
+ name.setPlaceholderText('add class')
+ ph = {'value': v, 'icon': icon, 'value_lbl': value_lbl, 'name': name}
+ name.textEdited.connect(lambda text, ph=ph: self._on_placeholder_name_edited(ph, text))
+ self._placeholder_rows.append(ph)
+ display_rows = self.class_rows + self._placeholder_rows
+ for grid_row, r in enumerate(display_rows, start=1):
+ self.classes_layout.addWidget(r['value_lbl'], grid_row, 0)
+ self.classes_layout.addWidget(r['icon'], grid_row, 1)
+ self.classes_layout.addWidget(r['name'], grid_row, 2, 1, 8)
+ for btn in (self.add_class_btn, self.remove_class_btn, self.sync_classes_btn,
+ self.reset_class_names_btn, self.btn_class_distribution_annot):
+ self.classes_layout.removeWidget(btn)
+ self._place_class_buttons(len(display_rows))
+
+ def _on_placeholder_name_edited(self, ph, text):
+ """First keystroke in a placeholder's name field turns it into a real
+ class with that text; focus moves to the created row's name field so
+ typing continues seamlessly."""
+ if ph not in self._placeholder_rows:
+ return
+ value = ph['value']
+ self._on_add_class(text=text, value=value)
+ row = self._row_for_value(value)
+ if row is not None:
+ field = row['name']
+ def _refocus(field=field):
+ field.setFocus()
+ field.setCursorPosition(len(field.text()))
+ # Deferred: the placeholder field we are typing in is deleted via
+ # deleteLater once this handler returns, and destroying a focused
+ # widget steals focus back — refocus after that has settled.
+ QTimer.singleShot(0, _refocus)
+
+ def _on_add_class(self, text=None, value=None):
+ """Add a class row. Without an explicit value, the next value after the
+ current maximum is used (values are never reused automatically).
+ Values above MAX_CLASS_VALUE are rejected."""
+ if value is None:
+ value = max(self._class_values(), default=0) + 1
+ if value > self.MAX_CLASS_VALUE:
+ show_info(f'Class values above {self.MAX_CLASS_VALUE} are not supported.')
+ return
+ if self._row_for_value(value) is not None:
+ return
+ self._create_class_row(value, text)
+ self._rebuild_class_rows_layout()
+ self._update_class_icons()
+ self._update_class_names()
+ self._update_selected_class_highlight()
def _on_remove_class(self, del_annots=True, event=None):
- """Remove the last class name and icon from the layout and update all annotations and segmentation layers."""
- last_name_idx = len(self.class_names)
- if last_name_idx > 2:
- # Remove the annotations from all annotations layers (do NOT do it in segmentation layers, as this would leave holes)
- if del_annots:
- for layer in self.annot_layers:
- if layer is not None and layer.name in self.viewer.layers:
- # Get the annotations image and remove the last label from it
- label_img = layer.data
- label_img[label_img == last_name_idx] = 0
- # Update the layer to show changes immediately
- layer.refresh()
- # Remove the last label and icon from the layout
- self.class_names[-1].deleteLater()
- self.class_icons[-1].deleteLater()
- self.class_names.pop()
- self.class_icons.pop()
- # Move the buttons one up
- self.classes_layout.removeWidget(self.add_class_btn)
- self.classes_layout.removeWidget(self.remove_class_btn)
- self.classes_layout.removeWidget(self.reset_class_names_btn)
- self.classes_layout.removeWidget(self.btn_class_distribution_annot)
- self.classes_layout.addWidget(self.add_class_btn, len(self.class_names)+1, 0, 1, 5)
- self.classes_layout.addWidget(self.remove_class_btn, len(self.class_names)+1, 5, 1, 5)
- self.classes_layout.addWidget(self.export_class_names_btn, len(self.class_names)+2, 0, 1, 5)
- self.classes_layout.addWidget(self.import_class_names_btn, len(self.class_names)+2, 5, 1, 5)
- self.classes_layout.addWidget(self.reset_class_names_btn, len(self.class_names)+3, 0, 1, 10)
- self.classes_layout.addWidget(self.btn_class_distribution_annot, len(self.class_names)+4, 0, 1, 10)
- # Update the icons and class names
- self._update_class_names()
- else:
- show_info('You need at least two classes.')
-
- def _update_class_icons(self, class_num=None, event=None):
- """Update the class icons with the colors of the class names.
- If class_num is given, only update the icon of that class."""
+ """Remove the SELECTED class (the row matching the annotations layer's
+ selected label). Other classes keep their values — no renumbering.
+ The button is disabled when this cannot apply; the guards below are
+ backstops for programmatic calls."""
+ value = self._selected_class_value()
+ row = self._row_for_value(value) if value is not None else None
+ if row is None:
+ show_info('No class selected. Click a class color (or use the label controls) first.')
+ return
+ if del_annots:
+ # Remove this class' annotations (NOT from segmentation layers, as
+ # that would leave holes)
+ for layer in self.annot_layers:
+ if layer is not None and layer.name in self.viewer.layers:
+ label_img = layer.data
+ label_img[label_img == value] = 0
+ layer.refresh()
+ for w in (row['icon'], row['value_lbl'], row['name']):
+ self.classes_layout.removeWidget(w)
+ w.deleteLater()
+ self.class_rows.remove(row)
+ self._rebuild_class_rows_layout()
+ self._update_class_icons()
+ self._update_class_names()
+ self._update_selected_class_highlight()
+ def _update_class_icons(self, value=None, event=None):
+ """Update the class icons with their label values' colors.
+ Without an annotations layer there is no colormap yet — show a striped
+ gray placeholder instead of an empty box."""
+ rows = self.class_rows if value is None else filter(None, [self._row_for_value(value)])
if self.labels_cmap is None:
+ for r in rows:
+ r['icon'].setPixmap(self._placeholder_pixmap())
+ r['icon'].setToolTip('Class colors appear once an annotations layer exists.')
return
-
cmap = self.labels_cmap.copy()
+ for r in rows:
+ col = cmap.map(r['value'])
+ r['icon'].setPixmap(self.get_pixmap(col))
+ r['icon'].setToolTip('')
- if class_num is not None:
- col = cmap.map(class_num)
- pixmap = self.get_pixmap(col)
- self.class_icons[class_num-1].setPixmap(pixmap)
- self.class_icons[class_num-1].mousePressEvent = lambda event: self._set_all_labels_classes(class_num, event)
-
- # Update all icons with the colors of the class names
- else:
- for i, _ in enumerate(self.class_names):
- cl = i+1
- col = cmap.map(cl)
- pixmap = self.get_pixmap(col)
- self.class_icons[i].setPixmap(pixmap)
- # Bind clicking on the icon to selecting the label
- # self.class_icons[i].mousePressEvent = lambda event, idx=i: self._set_all_labels_classes(idx+1, event)
+ def _on_class_swatch_clicked(self, value):
+ """Select a class from its swatch: remember it widget-side and push it
+ to the layers' selected label when layers exist."""
+ self._ui_selected_class = value
+ self._set_all_labels_classes(value)
+ self._update_selected_class_highlight()
def _set_all_labels_classes(self, x, event=None):
"""Set the selected label of all annotations and segmentation layers to x."""
- # For all annotations and segmentation layers added previously, set the selected label to x
labels_layers = self.annot_layers.union(self.seg_layers)
for l in labels_layers:
if l is not None and l.name in self.viewer.layers:
self.viewer.layers[l.name].selected_label = x
def _update_class_names(self, event=None):
- """Update the class names for all annotations and segmentation layers."""
- # For all annotations and segmentation layers, set the class names (= layer property) to the ones defined in the widget
- class_names = ["No label"] + [label.text() for label in self.class_names]
+ """Update the class names for all annotations and segmentation layers.
+ The layer property list is DENSE (index = label value, up to the highest
+ class), so values without a row get a plain internal placeholder name."""
+ max_v = max(self._class_values(), default=0)
+ class_names = ["No label"] + [self._class_name_for_value(v) for v in range(1, max_v + 1)]
props = {"Class": class_names}
labels_layers = self.annot_layers.union(self.seg_layers)
for l in labels_layers:
if l is not None and l.name in self.viewer.layers:
self.viewer.layers[l.name].properties = props
+ @staticmethod
+ def _placeholder_pixmap():
+ """Striped gray placeholder for class-color icons while no annotations
+ layer (hence no labels colormap) exists yet."""
+ pixmap = QtGui.QPixmap(20, 20)
+ pixmap.fill(QtGui.QColor(150, 150, 150))
+ painter = QtGui.QPainter(pixmap)
+ painter.setPen(QtGui.QPen(QtGui.QColor(115, 115, 115), 3))
+ for off in range(-20, 40, 8):
+ painter.drawLine(off, 20, off + 20, 0)
+ painter.end()
+ return pixmap
+
+ def _on_sync_classes_from_annotations(self, event=None):
+ """Add class rows for label values already painted in the selected
+ annotations layer. Explicit bottom-up sync, add-only, and exact: values
+ are per-row, so painting 1, 2 and 7 yields exactly those three rows."""
+ annot = self.annotations_layer_selection_widget.value
+ if annot is None:
+ show_info('No annotations layer selected.')
+ return
+ painted = set(int(v) for v in np.unique(annot.data)) - {0}
+ missing = sorted(painted - set(self._class_values()))
+ if not missing:
+ show_info('The annotations contain no classes beyond the current list.')
+ return
+ for v in missing:
+ self._on_add_class(value=v)
+
@staticmethod
def get_pixmap(color):
"""Convert a color (array or list) to a QPixmap for displaying as icon."""
@@ -1377,8 +1769,57 @@ def _update_cmaps(self, source_layer=None):
def _connect_all_cmaps(self):
"""Connect colormap changes for all annotations layers."""
labels_layers = self.annot_layers.union(self.seg_layers)
+ if not hasattr(self, '_selected_label_connected'):
+ self._selected_label_connected = set()
for l in labels_layers:
l.events.colormap.connect(lambda event: self._update_cmaps(source_layer=event.source))
+ # Follow the layer's active label so the Classes tab can highlight
+ # the corresponding row (guard: this method re-runs on every layer
+ # addition, and event connections are not deduplicated).
+ if id(l) not in self._selected_label_connected:
+ l.events.selected_label.connect(self._on_selected_label_changed)
+ self._selected_label_connected.add(id(l))
+
+ def _on_selected_label_changed(self, event=None):
+ """Reflect a layer's selected-label change in the Classes tab (only for
+ the annotations layer the user is painting into)."""
+ if event is not None and event.source != self.annotations_layer_selection_widget.value:
+ return
+ value = self._selected_class_value()
+ if value is not None:
+ self._ui_selected_class = value # survives layer deletion
+ self._update_selected_class_highlight()
+
+ def _update_selected_class_highlight(self, event=None):
+ """Outline the class row matching the annotations layer's selected
+ label — the tab-side view of napari's active-label notion (class icons
+ and the Shift+Q/W/E/R shortcuts set it, the layer controls change it).
+ All icons carry an equal-width transparent border so highlighting never
+ shifts the layout."""
+ selected = self._selected_class_value()
+ def _hex(color):
+ as_hex = getattr(color, 'as_hex', None)
+ return as_hex() if callable(as_hex) else str(color)
+ try:
+ from napari.utils.theme import get_theme
+ theme = get_theme(self.viewer.theme)
+ txt = _hex(theme['text']) if isinstance(theme, dict) else _hex(theme.text)
+ except Exception:
+ txt = '#f0f1f2' # napari dark
+ for r in self.class_rows:
+ color = txt if r['value'] == selected else 'transparent'
+ r['icon'].setStyleSheet(f'border: 2px solid {color};')
+ # 'Remove class' removes the selected class; gray it out when the
+ # selection is no real class (e.g. a placeholder slot, or nothing).
+ self.remove_class_btn.setEnabled(
+ selected is not None and self._row_for_value(selected) is not None)
+ # 'Add class' uses the next value after the current maximum; gray it
+ # out when that value would exceed the label-value limit.
+ at_limit = max(self._class_values(), default=0) >= self.MAX_CLASS_VALUE
+ self.add_class_btn.setEnabled(not at_limit)
+ self.add_class_btn.setToolTip(
+ f'No class values left (limit: {self.MAX_CLASS_VALUE}).' if at_limit
+ else 'Add a class name to the list.')
def _on_change_annot_cmap(self, event=None):
"""Update class icons and segmentation colormap when annotations colormap changes."""
@@ -1430,23 +1871,24 @@ def export_class_names_csv(self, file_path):
import csv
# Build list of label names: include "No label" at index 0
- class_names = ["No label"] + [w.text() for w in self.class_names]
-
with open(file_path, "w", newline="", encoding="utf-8") as fh:
writer = csv.writer(fh)
writer.writerow(["index", "name"])
- for idx, name in enumerate(class_names):
- writer.writerow([idx, name])
+ writer.writerow([0, "No label"])
+ for r in self.class_rows: # index = the row's actual label value
+ writer.writerow([r['value'], r['name'].text()])
def import_class_names_csv(self, file_path):
"""
Load class names from a CSV produced by `export_class_names_csv`.
Behavior:
- - Expects rows with columns `index,label` (header optional).
- - Ignores the index==0 row ("No label").
- - Resets widget to defaults (2 labels), then adds extra labels if CSV contains more than 2 labels.
- - If CSV provides fewer than 2 labels, remaining labels are set to "Class N" (e.g. "Class 2").
+ - Rows with columns `index,name` (header optional): the index becomes
+ the class' actual label value — sparse values are preserved.
+ - The index==0 row ("No label") is ignored.
+ - Name-only rows get sequential values after the highest explicit one.
+ - Values above MAX_CLASS_VALUE raise a ValueError (annotation and
+ segmentation data are uint8), leaving the current classes untouched.
"""
if file_path is None:
raise ValueError("file_path must be provided")
@@ -1477,54 +1919,50 @@ def import_class_names_csv(self, file_path):
elif len(first) >= 1 and first[0].strip().lower() == 'name' and len(first) == 1:
rows = rows[1:]
- num_appended = 0
+ # Parse into (value, name) pairs. Explicit indices become the actual
+ # label values (rows may be sparse); name-only rows get sequential
+ # values after the highest parsed one.
+ pairs = []
for row in rows:
if not row:
continue
- # If two columns, treat as index,name
if len(row) >= 2:
- # try parse index; if non-numeric, treat first column as name
try:
idx = int(row[0])
except Exception:
- parsed.append(row[0].strip())
- num_appended += 1
+ pairs.append((None, row[0].strip()))
continue
- name = row[1].strip()
if idx == 0:
- # skip "No label"
- continue
- if idx != num_appended + 1: # + 1 because index 0 is "No label"
- warnings.warn(f"Row {num_appended + 1} of named classes has index {idx}, meaning it is not increasing sequentially. Using the row number as index instead.")
- parsed.append(name)
- num_appended += 1
- # Single column: treat as names
+ continue # skip "No label"
+ pairs.append((idx, row[1].strip()))
else:
val = row[0].strip()
- # ignore empty rows
if val != "":
- parsed.append(val)
- num_appended += 1
-
- # Reset to defaults (this will create 2 labels)
- self._on_reset_class_names()
-
- # Add extra labels if CSV contains more than 2
- n_parsed = len(parsed)
- if n_parsed > 2:
- for i in range(n_parsed - 2):
- # add extra empty labels; we'll set texts in the unified loop below
- self._on_add_class()
- # If there are fewer than 2 in the import, rename the 2 base labels to "pad" the imported ones
- elif n_parsed < 2:
- # Rename the 2 base labels to "pad" the imported ones
- self.class_names[0].setText(f'Class 1')
- self.class_names[1].setText(f'Class 2')
-
- # Overwrite as many label texts as available; pad missing up to 2
- for i in range(0, n_parsed):
- self.class_names[i].setText(parsed[i])
-
+ pairs.append((None, val))
+ for name in parsed: # single-row "A,B,C" files: names only
+ pairs.append((None, name))
+ # Assign sequential values to name-only entries
+ next_v = max((v for v, _ in pairs if v is not None), default=0)
+ resolved = []
+ for v, name in pairs:
+ if v is None:
+ next_v += 1
+ v = next_v
+ resolved.append((v, name))
+
+ # Reject values beyond the label-value limit before touching the
+ # current classes, so a bad file leaves the widget unchanged.
+ too_high = sorted(v for v, _ in resolved if v > self.MAX_CLASS_VALUE)
+ if too_high:
+ raise ValueError(
+ f'Class values above {self.MAX_CLASS_VALUE} are not supported '
+ f'(annotation data is uint8); got {too_high[0]}'
+ + (f' and {len(too_high) - 1} more' if len(too_high) > 1 else '') + '.')
+
+ # Rebuild the class rows from the imported values
+ self._clear_class_rows()
+ for v, name in sorted(dict(resolved).items()): # last name wins per value
+ self._on_add_class(text=name, value=v)
# Sync labels to layers
self._update_class_names()
# Keep icons/cmaps in sync (no color data is read or written)
@@ -1572,6 +2010,31 @@ def _on_layer_removed(self, event=None):
self.annot_layers = {l for l in self.annot_layers if l is None or l.name in self.viewer.layers}
self.seg_layers = {l for l in self.seg_layers if l is None or l.name in self.viewer.layers}
+ # Clear the feature cache only when the LAST user image layer is removed.
+ # The cache is content-addressed (a removed image's entries simply stop
+ # hitting and age out via LRU), so clearing on every removal would throw
+ # away valid entries for the images still open — including when the
+ # plugin itself removes/recreates its own probabilities/features layers
+ # (e.g. after a class-count change), which must never wipe the cache.
+ removed = getattr(event, 'value', None) if event is not None else None
+
+ def _is_plugin_image(name):
+ # Live plugin layers are named exactly proba_prefix/features_prefix;
+ # backups renamed on image switch get a '_' suffix.
+ return any(name == p or name.startswith(p + '_')
+ for p in (self.proba_prefix, self.features_prefix))
+
+ if (isinstance(removed, napari.layers.Image)
+ and not _is_plugin_image(removed.name)):
+ user_images_left = any(
+ isinstance(l, napari.layers.Image) and not _is_plugin_image(l.name)
+ for l in self.viewer.layers)
+ if not user_images_left:
+ fc = getattr(getattr(self, 'cp_model', None), '_feature_cache', None)
+ if fc is not None:
+ fc.clear()
+ self._refresh_cache_size_label()
+
# Layer selection
def _on_select_layer(self, newtext=None):
@@ -1750,9 +2213,15 @@ def _on_train(self, event=None):
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
+ # skip_norm: the widget already normalized the stack
+ # (image_stack_norm), and prediction passes skip_norm=True on the
+ # same pre-normalized data. Matching it here keeps normalization
+ # single-pass (data-dependent modes like percentile must not be
+ # applied twice) and keeps train/predict features identical — so
+ # they can share feature-cache entries (keys are content hashes of
+ # the prepared image).
_ = self.cp_model.train(image_stack_norm, annot, memory_mode=mem_mode, img_ids=img_name,
- in_channels=in_channels, skip_norm=False,
+ in_channels=in_channels, skip_norm=True,
fe_use_device=self.fe_device, clf_use_device=self.clf_device)
self._update_training_counts()
@@ -1917,35 +2386,58 @@ def _on_predict_all(self):
# 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
+ # Step through the stack and predict each image.
num_steps = image_stack_norm.shape[-3]
- for step in progress(range(num_steps)):
+ in_channels = self._parse_in_channels(self.input_channels)
- # Take the slice of the 3rd last dimension (since images are C, Z, H, W or Z, H, W)
+ def _predict_and_write(step, cache_only):
+ """Predict one slice and write it to the layers. With cache_only=True,
+ only slices whose features are already cached are predicted (returns
+ False on a miss, without running the extractor). Returns True if
+ written."""
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
+ out = 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, cache_only=cache_only)
+ if out is None: # cache_only peek: this slice is not cached yet
+ return False
+ probas, seg = out
+ if self.add_probas:
+ # Creates the probabilities layer on the first actual prediction
+ # (we need the class count); with cache-first ordering this may
+ # not be step 0. A no-op once new_proba is cleared.
+ self._check_create_probas_layer(probas.shape[0])
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()
+ return True
+
+ fc = self.cp_model._feature_cache
+ cache_primed = (fc is not None and fc.enabled
+ and (len(fc) + fc.stats()["disk_entries"]) > 0)
+ done = [False] * num_steps
+ with progress(total=num_steps) as pbr:
+ pbr.set_description("Predicting")
+ if cache_primed:
+ # Cache-first ordering: serve slices already in the cache before
+ # computing the rest. A plain sequential scan over a stack larger
+ # than the cache evicts the very slices the next pass needs first
+ # (classic LRU thrash) — so cached slices would be recomputed for
+ # no benefit. Predicting cached slices first guarantees they are
+ # used before the compute pass evicts them. (Skipped when the
+ # cache is empty — nothing to serve first.)
+ for step in range(num_steps): # phase 1: already-cached slices
+ if _predict_and_write(step, cache_only=True):
+ done[step] = True
+ pbr.update(1)
+ for step in range(num_steps): # phase 2: compute the rest
+ if not done[step]:
+ _predict_and_write(step, cache_only=False)
+ pbr.update(1)
with warnings.catch_warnings():
warnings.simplefilter(action="ignore", category=FutureWarning)
@@ -2109,6 +2601,7 @@ def _on_load_model(self, event=None, save_file=None):
# Load the model (Note: done after updating GUI, since GUI updates might reset clf or change model)
self.cp_model = new_model
+ self._apply_feature_cache(recreate=True)
self.cp_model._param = new_param
temp_fe_model = self._cpm_class.create_fe(new_param.fe_name)
self.temp_fe_description = temp_fe_model.get_description()
@@ -2176,12 +2669,7 @@ def _on_reset_convpaint(self, event=None):
# Remove class names (note, resetting of class names needs to be split because of the handling of the attributes)
if 'Classes' in self.tab_names:
- for name in self.class_names:
- self.classes_layout.removeWidget(name)
- name.deleteLater()
- for icon in self.class_icons:
- self.classes_layout.removeWidget(icon)
- icon.deleteLater()
+ self._clear_class_rows()
# Reset the model to default
self._reset_model()
@@ -2216,12 +2704,7 @@ def _on_reset_convpaint(self, event=None):
self._create_default_class_names()
# Re-add the buttons below the class names
- self.classes_layout.addWidget(self.add_class_btn, len(self.class_names)+1, 0, 1, 5)
- self.classes_layout.addWidget(self.remove_class_btn, len(self.class_names)+1, 5, 1, 5)
- self.classes_layout.addWidget(self.export_class_names_btn, len(self.class_names)+2, 0, 1, 5)
- self.classes_layout.addWidget(self.import_class_names_btn, len(self.class_names)+2, 5, 1, 5)
- self.classes_layout.addWidget(self.reset_class_names_btn, len(self.class_names)+3, 0, 1, 10)
- self.classes_layout.addWidget(self.btn_class_distribution_annot, len(self.class_names)+4, 0, 1, 10)
+ self._place_class_buttons(len(self.class_rows))
if 'Multifile' in self.tab_names:
self._reset_multifile_folder()
@@ -2282,6 +2765,9 @@ def _reset_attributes(self):
self.features_prefix = 'features' # Prefix for the feature image layer name
self.cont_training = "Image" # Update features for subsequent training ("Image" or "Off" or "Global")
self.use_dask = False # Use Dask for parallel processing
+ self.cache_enabled = True # Reuse extracted features when re-segmenting / re-training the same image
+ self.cache_max_mb = 2048 # Max RAM (MB) the feature cache may use (moderate default)
+ self.cache_disk_max_mb = 8192 # Max disk (MB) for spilled features (0 = disk spillover off)
self.fe_device = 'auto' # Device to use for the FE (if applicable); 'auto' will use GPU if available, otherwise CPU
self.clf_device = 'auto' # Device to use for the classifier (if applicable); 'auto' will use GPU if available, otherwise CPU
self.input_channels = "" # Input channels for the model (as txt, will be parsed)
@@ -2292,11 +2778,18 @@ def _reset_attributes(self):
self.new_features = True
self.features_pca_components = "0" # Number of PCA components for feature image (0 = no PCA)
self.features_kmeans_clusters = "0" # Number of k-means clusters for feature image (0 = no k-means)
- self.initial_names = ['Background', 'Foreground']
self.annot_layers = set() # List of annotations layers
self.seg_layers = set() # List of segmentation layers
- self.class_names = [] # List of class names
- self.class_icons = [] # List of class icons
+ # Class rows: each is {'value': int, 'icon': QLabel, 'value_lbl': QLabel,
+ # 'name': QLineEdit}, kept sorted by value. The label VALUE is explicit
+ # per row — rows may be sparse (e.g. classes 1, 2, 7) and removing a
+ # middle class never renumbers the others.
+ self.class_rows = []
+ # Widget-side memory of the selected class value: mirrors the
+ # annotations layer's selected_label while one exists, and keeps the
+ # last selection alive when the layer is deleted — classes stay
+ # selectable/removable without any layer.
+ self._ui_selected_class = None
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
@@ -2475,6 +2968,7 @@ def _on_set_fe_model(self, event=None):
# Create a new model with the new FE
self.cp_model = self._cpm_class(param=new_param)
+ self._apply_feature_cache(recreate=True)
self._reset_device_options()
self._reset_clf() # Call to take all actions needed after resetting the clf
# Reset the features for continuous training
@@ -2630,13 +3124,22 @@ def _add_empty_annot(self, event=None, force_add=True, from_multifile=False):
# Track annotations data changes to keep in-memory store in sync (for Multifile)
self.store_annot = from_multifile # Only store if the annot was added from Multifile, to avoid storing unnecessarily when not using Multifile
+ def _restore_active_layer(self, prev_active):
+ """Re-activate `prev_active` after adding an output layer. napari
+ selects newly added layers, which would silently steer the user's
+ brush into the segmentation/probabilities/features layer instead of
+ the annotations layer they were painting on."""
+ if prev_active is not None and prev_active in self.viewer.layers:
+ self.viewer.layers.selection.active = prev_active
+
def _check_create_segmentation_layer(self):
"""Check if segmentation layer exists and create it if not."""
-
+
img = self._get_selected_img(check=True)
if img is None:
warnings.warn('No image selected. No layers added.')
return
+ prev_active = self.viewer.layers.selection.active
layer_shape = self._get_annot_shape(img)
num_spatial = len(layer_shape)
transform_kwargs = self._get_layer_transform_kwargs(img, num_spatial_dims=num_spatial, num_leading_dims=0)
@@ -2666,6 +3169,7 @@ def _check_create_segmentation_layer(self):
# Add it to the list of layers where class names shall be updated
self.seg_layers.add(self.viewer.layers[self.seg_tag])
self.update_all_class_names_and_cmaps()
+ self._restore_active_layer(prev_active)
def _check_create_probas_layer(self, num_classes):
"""Check if class probabilities layer exists and create it if not."""
@@ -2675,6 +3179,7 @@ def _check_create_probas_layer(self, num_classes):
warnings.warn('No image selected. No layers added.')
return
+ prev_active = self.viewer.layers.selection.active
spatial_dims = self._get_annot_shape(img)
if isinstance(num_classes, int):
num_classes = (num_classes,)
@@ -2708,6 +3213,7 @@ def _check_create_probas_layer(self, num_classes):
self.viewer.layers[self.proba_prefix].colormap = "turbo"
# Save information about the probabilities layer to be able to rename it later
self._set_old_proba_tag()
+ self._restore_active_layer(prev_active)
def _check_create_features_layer(self, num_features):
"""Check if feature image layer exists and create it if not."""
@@ -2717,6 +3223,7 @@ def _check_create_features_layer(self, num_features):
warnings.warn('No image selected. No layers added.')
return
+ prev_active = self.viewer.layers.selection.active
spatial_dims = self._get_annot_shape(img)
# Create a new features layer if it doesn't exist yet or we need a new one
@@ -2760,6 +3267,7 @@ def _check_create_features_layer(self, num_features):
)
# Save information about the features layer to be able to rename it later
self._set_old_features_tag()
+ self._restore_active_layer(prev_active)
def _rename_annot_for_backup(self):
"""Name the annotations with a unique name according to its image,
@@ -3699,7 +4207,7 @@ def _update_training_counts(self):
pix = len(self.cp_model.table)
imgs = len(np.unique(self.cp_model.table['img_id']))
lbls = len(np.unique(self.cp_model.table['label']))
- self.label_training_count.setText(f'{pix} pixels, {imgs} image{"s"*(imgs>1)}, {lbls} labels')
+ self.label_training_count.setText(f'{pix} px / {imgs} img{"s"*(imgs>1)} / {lbls} labels')
def _on_show_class_distribution(self, trained_data=False):
"""Show the class distribution of the data used with continuous_training/memory_mode (saved in self.cp_model.table)
@@ -3737,8 +4245,8 @@ def _on_show_class_distribution(self, trained_data=False):
percs = counts / np.sum(counts) * 100
# Get class display names from a list, assuming class numbers start at 1
- if self.class_names is not None and self.class_names:
- class_names = [self.class_names[c - 1].text() if 1 <= c <= len(self.class_names) else str(c) for c in classes]
+ if self.class_rows:
+ class_names = [self._class_name_for_value(c) for c in classes]
# Create label strings for the pie chart
pie_labels = [f'{count} ({perc:.1f}%)' for count, perc in zip(counts, percs)]
diff --git a/src/napari_convpaint/feature_cache.py b/src/napari_convpaint/feature_cache.py
new file mode 100644
index 0000000..77b165b
--- /dev/null
+++ b/src/napari_convpaint/feature_cache.py
@@ -0,0 +1,362 @@
+"""Bounded, FE-pluggable feature cache for the interactive annotate→predict loop.
+
+The expensive step in interactive segmentation is feature extraction; when the
+same image (or z-slice / movie frame) is segmented repeatedly while refining
+scribbles, its features can be reused instead of recomputed. This module holds a
+generic, feature-extractor-agnostic cache: it stores an *opaque payload* defined
+by each FE (e.g. DINO patch tokens — tiny and lossless to upsample), keyed by
+`(img_id, slice, FE-signature)`, and owns everything storage-related — LRU
+eviction and, crucially, a memory budget so caching a 100-slice stack or a
+300-frame movie can never grow unbounded and crash the kernel.
+
+Triage principle (never OOM): before storing an entry, its size is checked
+against a live budget = `min(configured cap, available_RAM − headroom)`. If it
+does not fit, the least-recently-used entries are evicted; if it still does not
+fit (a single payload larger than the budget), it is simply not cached and the
+caller recomputes. The cache never exceeds the budget, so it degrades to
+recomputation rather than pushing the system into swap.
+
+FE-specific behaviour lives in the FeatureExtractor (see the
+`cacheable_repr_and_features` / `features_from_cacheable` /
+`supports_feature_cache` protocol on the base class): the cache never needs to know what is inside a
+payload.
+"""
+from __future__ import annotations
+
+import functools
+import os
+import pickle
+import shutil
+import tempfile
+import threading
+from collections import OrderedDict
+
+try:
+ import psutil
+ _HAVE_PSUTIL = True
+except Exception: # pragma: no cover - psutil is optional
+ _HAVE_PSUTIL = False
+
+
+# Fallback absolute cap when available-RAM cannot be queried (no psutil).
+_DEFAULT_MAX_BYTES = 2 * 1024 ** 3 # 2 GiB
+# Keep at least this fraction of currently-available RAM free (never consume it
+# all with cache), as a safety headroom against OOM.
+_DEFAULT_HEADROOM_FRAC = 0.25
+# Never write disk-cache data that would leave less than this much free on the
+# target filesystem — so the disk cache can't fill the user's disk.
+_DISK_HEADROOM_BYTES = 2 * 1024 ** 3 # 2 GiB
+
+
+def _payload_nbytes(payload) -> int:
+ """Best-effort byte size of an opaque payload (array, list/tuple of arrays,
+ or anything exposing .nbytes, recursing into lists/tuples/dicts).
+ Unknown → 0 (treated as free)."""
+ if payload is None:
+ return 0
+ if hasattr(payload, "nbytes"):
+ return int(payload.nbytes)
+ if isinstance(payload, (list, tuple)):
+ return sum(_payload_nbytes(p) for p in payload)
+ if isinstance(payload, dict):
+ return sum(_payload_nbytes(v) for v in payload.values())
+ return 0
+
+
+def _locked(method):
+ """Run `method` under the cache's re-entrant lock. The cache is shared
+ across threads (e.g. the napari GUI thread changing limits or clearing
+ while a worker thread is inside get/put), so every public entry point must
+ hold the lock; private helpers are only called from within one."""
+ @functools.wraps(method)
+ def wrapper(self, *args, **kwargs):
+ with self._lock:
+ return method(self, *args, **kwargs)
+ return wrapper
+
+
+class FeatureCache:
+ """LRU feature cache bounded by a memory budget.
+
+ Thread-safe: all public methods take a re-entrant lock, so a worker thread
+ can extract/cache while the GUI thread clears the cache or changes limits.
+
+ Parameters
+ ----------
+ max_bytes : int or None
+ Hard cap on the cache's own size. None → an automatic cap derived from
+ system RAM (a quarter of total, or `_DEFAULT_MAX_BYTES` without psutil).
+ headroom_frac : float
+ Fraction of *currently available* RAM to always keep free. The live
+ budget is `available_RAM * (1 - headroom_frac)`; entries are never added
+ (and are evicted) to respect it, so the cache cannot trigger OOM.
+ enabled : bool
+ Master switch; when False, get() always misses and put() is a no-op.
+ """
+
+ def __init__(self, max_bytes: int | None = None,
+ headroom_frac: float = _DEFAULT_HEADROOM_FRAC,
+ enabled: bool = True,
+ disk_max_bytes: int = 0):
+ self._lock = threading.RLock()
+ self._store: "OrderedDict[tuple, tuple]" = OrderedDict() # key -> (payload, nbytes, spill_ok)
+ self._total_bytes = 0
+ self._headroom_frac = float(headroom_frac)
+ self.enabled = bool(enabled)
+ if max_bytes is None:
+ if _HAVE_PSUTIL:
+ max_bytes = int(psutil.virtual_memory().total * 0.25)
+ else:
+ max_bytes = _DEFAULT_MAX_BYTES
+ self._max_bytes = int(max_bytes)
+ self.hits = 0
+ self.misses = 0
+ # --- disk spillover: a second, larger LRU tier on disk ---
+ # RAM-evicted entries spill here (pickled) instead of being dropped, up to
+ # `disk_max_bytes` (0 = off). get() checks RAM then disk. This mainly helps
+ # FEs whose cacheable payload can't be compressed to a small form (e.g. an
+ # upsampling FE that emits full-resolution features), where per-slice
+ # payloads are large and few fit in RAM: reading one back from disk is far
+ # cheaper than recomputing it, so a stack too large for the RAM tier still
+ # benefits on the next iteration.
+ self._disk_max_bytes = int(disk_max_bytes)
+ self._disk_dir = None # temp dir, made lazily on first spill
+ self._disk_store: "OrderedDict[tuple, tuple]" = OrderedDict() # key -> (path, nbytes)
+ self._disk_bytes = 0
+ self._disk_seq = 0
+ self.disk_hits = 0
+
+ # -- budget helpers ----------------------------------------------------
+
+ def _available_bytes(self) -> int:
+ if _HAVE_PSUTIL:
+ return int(psutil.virtual_memory().available)
+ # No psutil: rely solely on the configured cap (assume plenty free).
+ return self._max_bytes
+
+ def _fits(self, nbytes: int, available: int | None = None) -> bool:
+ """Whether adding `nbytes` keeps the cache under its cap AND leaves the
+ configured headroom of currently-available RAM free. Pass `available`
+ to reuse one RAM snapshot across repeated checks (e.g. put's eviction
+ loop) instead of re-querying psutil per call."""
+ if self._total_bytes + nbytes > self._max_bytes:
+ return False
+ if available is None:
+ available = self._available_bytes()
+ # available RAM already accounts for the cache's current allocation, so
+ # only the *new* bytes reduce it further.
+ return available - nbytes >= self._headroom_frac * available
+
+ # -- public API --------------------------------------------------------
+
+ @_locked
+ def get(self, key):
+ """Return the cached payload for `key`, or None. Checks RAM, then the disk
+ tier. A disk hit returns the loaded payload but leaves it on disk (no
+ promote-and-thrash): the RAM tier holds the most recent entries, disk the
+ older ones."""
+ if not self.enabled:
+ return None
+ item = self._store.get(key)
+ if item is not None:
+ self._store.move_to_end(key) # most-recently-used
+ self.hits += 1
+ return item[0]
+ payload = self._load_from_disk(key) # RAM miss -> try disk
+ if payload is not None:
+ self.hits += 1
+ self.disk_hits += 1
+ return payload
+ self.misses += 1
+ return None
+
+ @_locked
+ def put(self, key, payload, nbytes: int | None = None, spill_ok: bool = True):
+ """Store `payload` under `key` if it fits the budget; else evict LRU and
+ retry. A payload that can never fit the RAM tier goes straight to the
+ disk tier (if `spill_ok` and it fits the disk budget); otherwise caching
+ is skipped and the caller recomputes. `spill_ok=False` keeps a payload
+ out of the disk tier entirely (RAM only, evict = drop) — for FEs whose
+ payloads are huge relative to their recompute cost."""
+ if not self.enabled or payload is None:
+ return
+ if nbytes is None:
+ nbytes = _payload_nbytes(payload)
+ # Overwrite of an existing key: drop the old size first.
+ if key in self._store:
+ self._total_bytes -= self._store.pop(key)[1]
+ # A single payload larger than the whole RAM cap can never be held in
+ # RAM — route it directly to the disk tier instead of dropping it.
+ if nbytes > self._max_bytes:
+ if spill_ok:
+ self._spill_to_disk(key, payload, nbytes)
+ return
+ # Evict least-recently-used until the new entry fits. One RAM snapshot
+ # serves the whole loop: eviction only increases availability, so the
+ # snapshot errs conservative.
+ available = self._available_bytes()
+ while self._store and not self._fits(nbytes, available):
+ self._evict_one()
+ if not self._fits(nbytes, available):
+ # The live headroom refuses it even with the RAM tier empty; the
+ # disk tier can still hold it.
+ if spill_ok:
+ self._spill_to_disk(key, payload, nbytes)
+ return
+ # The key is absent at this point (popped above if present), so
+ # assignment appends at the MRU end.
+ self._store[key] = (payload, nbytes, spill_ok)
+ self._total_bytes += nbytes
+
+ def _evict_one(self):
+ key, (payload, nbytes, spill_ok) = self._store.popitem(last=False) # LRU = oldest
+ self._total_bytes -= nbytes
+ # Spill to the disk tier instead of dropping (if disk spillover is on).
+ if spill_ok:
+ self._spill_to_disk(key, payload, nbytes)
+
+ # -- disk tier ---------------------------------------------------------
+
+ def _ensure_disk_dir(self) -> str:
+ if self._disk_dir is None:
+ self._disk_dir = tempfile.mkdtemp(prefix="convpaint_fcache_")
+ os.makedirs(self._disk_dir, exist_ok=True)
+ return self._disk_dir
+
+ def _safe_remove(self, path):
+ try:
+ os.remove(path)
+ except OSError:
+ pass
+
+ def _remove_disk_entry(self, key):
+ item = self._disk_store.pop(key, None)
+ if item is not None:
+ path, nbytes = item
+ self._disk_bytes -= nbytes
+ self._safe_remove(path)
+
+ def _evict_disk_one(self):
+ key = next(iter(self._disk_store)) # LRU = oldest
+ self._remove_disk_entry(key)
+
+ def _spill_to_disk(self, key, payload, nbytes):
+ """Write an RAM-evicted payload to the disk tier (LRU-bounded). No-op if
+ disk spillover is off or the single payload exceeds the disk cap."""
+ if self._disk_max_bytes <= 0 or nbytes > self._disk_max_bytes:
+ return
+ self._remove_disk_entry(key) # replace any stale copy
+ while self._disk_store and self._disk_bytes + nbytes > self._disk_max_bytes:
+ self._evict_disk_one()
+ if self._disk_bytes + nbytes > self._disk_max_bytes:
+ return
+ # Also never fill the actual filesystem below the free-space headroom.
+ try:
+ free = shutil.disk_usage(self._ensure_disk_dir()).free
+ if free - nbytes < _DISK_HEADROOM_BYTES:
+ return
+ except OSError:
+ pass
+ self._disk_seq += 1
+ path = os.path.join(self._ensure_disk_dir(), f"{self._disk_seq}.pkl")
+ try:
+ with open(path, "wb") as f:
+ pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
+ except Exception:
+ self._safe_remove(path)
+ return
+ self._disk_store[key] = (path, nbytes)
+ self._disk_bytes += nbytes
+
+ def _load_from_disk(self, key):
+ item = self._disk_store.get(key)
+ if item is None:
+ return None
+ path, _ = item
+ try:
+ with open(path, "rb") as f:
+ payload = pickle.load(f)
+ except Exception:
+ self._remove_disk_entry(key)
+ return None
+ self._disk_store.move_to_end(key) # mark recently used
+ return payload
+
+ def _clear_disk(self):
+ for path, _ in list(self._disk_store.values()):
+ self._safe_remove(path)
+ self._disk_store.clear()
+ self._disk_bytes = 0
+
+ # -- invalidation / limits --------------------------------------------
+
+ @_locked
+ def clear(self):
+ """Drop all entries (both RAM and disk tiers). Entries are
+ content-addressed and never go stale; clearing only frees memory."""
+ self._store.clear()
+ self._total_bytes = 0
+ self._clear_disk()
+
+ @_locked
+ def set_max_bytes(self, max_bytes: int):
+ """Change the RAM cap in place, evicting (spilling) LRU entries if over."""
+ self._max_bytes = int(max_bytes)
+ while self._store and self._total_bytes > self._max_bytes:
+ self._evict_one()
+
+ @_locked
+ def set_disk_max_bytes(self, disk_max_bytes: int):
+ """Change the disk cap in place, evicting disk LRU if over (0 = off)."""
+ self._disk_max_bytes = int(disk_max_bytes)
+ if self._disk_max_bytes <= 0:
+ self._clear_disk()
+ else:
+ while self._disk_store and self._disk_bytes > self._disk_max_bytes:
+ self._evict_disk_one()
+
+ @_locked
+ def set_enabled(self, enabled: bool):
+ """Enable/disable in place; disabling clears both tiers to free space."""
+ self.enabled = bool(enabled)
+ if not self.enabled:
+ self.clear()
+
+ @_locked
+ def close(self):
+ """Free the disk tier and remove the temp directory this cache created."""
+ self._clear_disk()
+ if self._disk_dir:
+ shutil.rmtree(self._disk_dir, ignore_errors=True)
+ self._disk_dir = None
+
+ def __del__(self): # pragma: no cover - best-effort cleanup
+ try:
+ self.close()
+ except Exception:
+ pass
+
+ @property
+ def nbytes(self) -> int:
+ return self._total_bytes
+
+ @property
+ def disk_nbytes(self) -> int:
+ return self._disk_bytes
+
+ @_locked
+ def __len__(self):
+ return len(self._store)
+
+ @_locked
+ def stats(self) -> dict:
+ return {
+ "entries": len(self._store),
+ "bytes": self._total_bytes,
+ "disk_entries": len(self._disk_store),
+ "disk_bytes": self._disk_bytes,
+ "disk_hits": self.disk_hits,
+ "hits": self.hits,
+ "misses": self.misses,
+ "max_bytes": self._max_bytes,
+ }
diff --git a/src/napari_convpaint/feature_extractor.py b/src/napari_convpaint/feature_extractor.py
index 2fb2cbb..364e2c4 100644
--- a/src/napari_convpaint/feature_extractor.py
+++ b/src/napari_convpaint/feature_extractor.py
@@ -4,6 +4,13 @@
from .param import Param
from .utils import scale_img, rescale_features, reduce_to_patch_multiple, pad_to_shape, get_device_from_torch_model
+def _concat(arrays, axis=0):
+ """Concatenate along `axis`, staying on-device for torch tensors."""
+ if len(arrays) and isinstance(arrays[0], torch.Tensor):
+ return torch.cat(arrays, dim=axis)
+ return np.concatenate(arrays, axis=axis)
+
+
class FeatureExtractor:
def __init__(self, model_name="vgg16", model=None, **kwargs):
"""
@@ -30,6 +37,11 @@ def __init__(self, model_name="vgg16", model=None, **kwargs):
# For such FEs, tile_annotations / tile_image cannot match whole-image features at any finite padding
self.tile_block_size = None # If not None, this block size is used for tiling the image at segmentation
self.num_input_channels = [1]
+ # Number of features each hooked layer contributes (only Hookmodel sets a
+ # real value); default None so `fe_use_min_features` degrades gracefully
+ # (warns and uses all features) for FEs that don't track it, instead of
+ # raising AttributeError.
+ self.features_per_layer = None
self.norm_mode = "default" # or "imagenet" or "percentile"
self.rgb_input = False # Whether the model takes RGB input or not
self.proposed_scalings = [[1],
@@ -292,7 +304,7 @@ def supported_devices(self):
The list of devices that the feature extractor supports.
"""
if self.model is not None and hasattr(self.model, "to"):
- return [torch.device("cuda"), torch.device("mps"), [torch.device("cpu")]]
+ return [torch.device("cuda"), torch.device("mps"), torch.device("cpu")]
else:
return [torch.device("cpu")]
@@ -357,21 +369,37 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic
The extracted features of the image as a single array with [nb_features, Z, H, W]
"""
- features_all_scales = []
+ native = self._pyramid_native(data, param, device)
+ return self._pyramid_reconstruct(native, data.shape, param, patched)
+
+ def _pyramid_native(self, data, param, device=torch.device("cpu")):
+ """Expensive half of the pyramid: extract the per-scale *native*
+ (pre-rescale) features. Returns one entry per scale of
+ ``(features_list, pre_reduction_shape, reduced_shape)`` that
+ ``_pyramid_reconstruct`` turns into the final feature stack.
+ The features stay in whatever form the extractor produced them — for
+ NN FEs, torch tensors on the extraction device — so
+ ``_pyramid_reconstruct`` rescales on-device. Only the feature-cache
+ payload is cast to device-independent CPU numpy
+ (``_native_to_payload``).
+
+ This is the (pre-cast) cacheable representation: it is independent of
+ the requested output resolution (``patched``), and for patched FEs
+ (ViTs) it is the small patch-grid features rather than the
+ full-resolution stack. See the feature-cache protocol below."""
# Check if the given selection of scalings is in the proposed scalings, and if not, give a warning
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.")
- # Iterate over the scales and extract features for each scale
+ native = []
for s in param.fe_scalings:
-
# Downscale the image
image_scaled = scale_img(data, s)
# Make sure the downscaled part is a multiple of the patch size
patch_size = self.get_patch_size()
- pre_reduction_shape = image_scaled.shape
+ pre_reduction_shape = image_scaled.shape
# NOTE: reduce_to_patch_multiple should not do anything if the inputs are already multiples of the patch size at all scales
image_scaled = reduce_to_patch_multiple(image_scaled, patch_size)
reduced_shape = image_scaled.shape
@@ -383,14 +411,24 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic
# In case the features are not a list, but a single array, make it a list
if not isinstance(features, list):
features = [features]
-
+ native.append((features, pre_reduction_shape, reduced_shape))
+ return native
+
+ def _pyramid_reconstruct(self, native, data_shape, param, patched=True):
+ """Cheap half of the pyramid: rescale the native per-scale features to the
+ requested resolution and concatenate across channel-series and scales.
+ ``native`` is the payload from ``_pyramid_native``; ``data_shape`` is the
+ original (pre-scaling) input shape [C, Z, H, W]."""
+ patch_size = self.get_patch_size()
+ features_all_scales = []
+ for features, pre_reduction_shape, reduced_shape in native:
# Resize the features from this downscaling to the size of the (possibly patched) input
# NOTE: this shouldn't do anything for scaling of 1, unless we want to "unpatch"
if patched:
- target_shape = (data.shape[0],
- data.shape[1],
- data.shape[2]//self.get_patch_size(),
- data.shape[3]//self.get_patch_size())
+ target_shape = (data_shape[0],
+ data_shape[1],
+ data_shape[2]//patch_size,
+ data_shape[3]//patch_size)
else:
# When not patched, but patch_size > 1, we handle possible cropping due to reduce_to_patch_multiple
# NOTE: this should not be necessary if the inputs are already multiples of the patch size at all scales
@@ -406,7 +444,7 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic
features = [pad_to_shape(f, pre_reduction_shape[2:] ) for f in features]
# Rescale to the full original shape
- target_shape = data.shape
+ target_shape = data_shape
features = [rescale_features(
feature_img=f,
@@ -414,13 +452,10 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic
order=param.fe_order)
for f in features]
- # 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]
-
- # Put together features for each input_channels procession (and layers if applicable)
- features = np.concatenate(features, axis=0)
+ # Put together features for each input_channels procession (and
+ # layers if applicable), staying on the extraction device for torch
+ # tensors; the single host transfer happens at the very end.
+ features = _concat(features)
# If use_min_features is True, shorten features
if param.fe_use_min_features:
@@ -432,12 +467,89 @@ def extract_features_pyramid(self, data, param, patched=True, device=torch.devic
# Add the features to the list of features
features_all_scales.append(features)
-
- # Concatenate the features from all scales along the first axis
- features_all_scales = np.concatenate(features_all_scales, axis=0)
+
+ # Concatenate all scales along the first axis, then move to CPU numpy
+ # in a single host transfer.
+ features_all_scales = _concat(features_all_scales)
+ if isinstance(features_all_scales, torch.Tensor):
+ features_all_scales = features_all_scales.detach().cpu().numpy()
return features_all_scales
+### FEATURE CACHING PROTOCOL (optional per-FE optimization; see feature_cache.py)
+
+ def supports_feature_cache(self, param):
+ """Whether whole-image feature caching is worthwhile for this FE. FEs that
+ are cheap to recompute or whose extraction doesn't fit the pyramid split
+ can return False to opt out."""
+ return True
+
+ def cache_spill_to_disk(self):
+ """Whether this FE's payloads may be spilled to the disk cache tier on
+ RAM eviction. FEs whose payloads are huge relative to their recompute
+ cost (e.g. CNN feature maps) should return False: they stay cacheable in
+ RAM but are dropped instead of pickled to disk."""
+ return True
+
+ def cache_extra_state(self):
+ """Extraction-relevant state that lives on the FE instance (rather than
+ in the Param), to be mixed into the feature-cache key. Any FE whose output
+ depends on constructor/instance state (e.g. sigmas, scalings moved out of
+ the Param by `get_enforced_params`) must return it here, or stale cached
+ features will be served after that state changes. Return value must be
+ hashable (or None)."""
+ return None
+
+ @staticmethod
+ def _native_to_payload(native):
+ """Cast a native pyramid (whose feature arrays may be on-device torch
+ tensors) to a device-independent cache payload: CPU numpy arrays, safe
+ to hold in RAM and pickle to the disk tier. ``was_torch`` records the
+ native form, so reconstruction from the cache can lift the payload back
+ to tensors and use the same rescale backend as a fresh extraction —
+ cached results must be identical to fresh ones."""
+ was_torch = any(isinstance(f, torch.Tensor)
+ for features, _, _ in native for f in features)
+ scales = [([f.detach().cpu().numpy() if isinstance(f, torch.Tensor) else f
+ for f in features], pre_shape, red_shape)
+ for features, pre_shape, red_shape in native]
+ return {"scales": scales, "was_torch": was_torch}
+
+ def cacheable_repr_and_features(self, data, param, device=torch.device("cpu"), patched=True):
+ """Compute the features AND the cache payload in one extraction pass.
+ This is THE extension point for FEs with a custom payload (override
+ together with `features_from_cacheable`); the cache's miss path calls
+ only this method.
+
+ The reconstruction runs from the on-device native form (fast torch
+ rescale for NN FEs); only the stored payload is cast to CPU numpy."""
+ native = self._pyramid_native(data, param, device)
+ payload = self._native_to_payload(native)
+ features = self._pyramid_reconstruct(native, data.shape, param, patched)
+ return features, payload
+
+ def features_from_cacheable(self, payload, data_shape, param, patched=True, device=None):
+ """Reconstruct the features the pipeline needs from a cached payload.
+ If the payload originated from torch tensors, lift it back onto
+ `device` first, so cache hits use the same (fast, on-device) rescale
+ backend as fresh extractions — both for speed and so hit and miss
+ results are identical.
+
+ Contract: pass the SAME resolved device the fresh extraction would use
+ (as `_extract_pyramid_cached` does). `device=None` reconstructs on the
+ CPU, which matches a CPU extraction but not bit-exactly a GPU one
+ (torch's CPU and GPU interpolation kernels may differ at order>0)."""
+ native = payload["scales"]
+ if payload.get("was_torch"):
+ # NOTE: device=None falls back to CPU (see docstring contract) — the
+ # caller (_extract_pyramid_cached) always passes the extraction device,
+ # so a hit reconstructs on the same backend as the miss that stored it.
+ lift_device = device if device is not None else "cpu"
+ native = [([torch.from_numpy(f).to(lift_device)
+ for f in features], pre_shape, red_shape)
+ for features, pre_shape, red_shape in native]
+ return self._pyramid_reconstruct(native, data_shape, param, patched)
+
def extract_features_from_multichannel_stack(self, image, rgb_data=False, device=torch.device("cpu")):
"""
Extracts the features of an image (stack) with an arbitrary number of channels.
diff --git a/src/napari_convpaint/feature_extractors/combo_fe.py b/src/napari_convpaint/feature_extractors/combo_fe.py
index 5f89fa6..8c528f6 100644
--- a/src/napari_convpaint/feature_extractors/combo_fe.py
+++ b/src/napari_convpaint/feature_extractors/combo_fe.py
@@ -109,6 +109,12 @@ def gives_patched_features(self):
# So, the combo FE itself is not patched, even if it works with a patch_size to comply with the models
return False
+ def supports_feature_cache(self, param):
+ # ComboFeatures overrides extract_features_pyramid to combine two sub-FEs,
+ # so it does not go through the base _pyramid_native/_pyramid_reconstruct
+ # split the cache relies on. Opt out of caching for now (v1).
+ return False
+
def extract_features_pyramid(self, image, param, patched=False, device=None):
def1 = self.model1.get_default_params(param)
features1 = self.model1.extract_features_pyramid(image, def1, patched=False, device=device)
diff --git a/src/napari_convpaint/feature_extractors/dino_jafar.py b/src/napari_convpaint/feature_extractors/dino_jafar.py
index 51af56b..6483337 100644
--- a/src/napari_convpaint/feature_extractors/dino_jafar.py
+++ b/src/napari_convpaint/feature_extractors/dino_jafar.py
@@ -56,8 +56,8 @@ def import_vitwrapper_jafar():
class DinoJafarFeatures(FeatureExtractor):
"""
- DINO + JAFAR upsampler feature extractor integrated with ConvPaint.
- Expects that ConvPaint already padded/cropped images so H,W are multiples
+ DINO + JAFAR upsampler feature extractor integrated with Convpaint.
+ Expects that Convpaint already padded/cropped images so H,W are multiples
of self.patch_size. Provides dynamic patch size: large images use sliding
patches with overlap; smaller images shrink patch size to the largest
multiple of the backbone patch size that fits within min(H,W).
@@ -84,6 +84,11 @@ def __init__(self, model_name="dinov2_small-reg_jafar", **kwargs):
[1, 8],
[1, 8, self.patch_size],
]
+ # Internal JAFAR upsampling scales; normally (re)set from fe_scalings in
+ # get_enforced_params before extraction, but default it here so direct FE
+ # use (extract_features_from_plane without going through ConvpaintModel)
+ # doesn't hit an AttributeError.
+ self.jafar_scalings = [1]
# Parent .create_model() saves tuple (hr_head, backbone) in self.model
self.model, self.backbone = self.model
@@ -170,6 +175,13 @@ def get_enforced_params(self, param=None):
#param.fe_scalings = [4]
return param
+ def cache_extra_state(self):
+ # The user's fe_scalings are moved out of the Param (forced to [1]) into
+ # self.jafar_scalings by get_enforced_params, and the cached payload
+ # bakes them in — so they must be part of the cache key or changing the
+ # scalings would silently serve stale features.
+ return ("jafar_scalings", tuple(self.jafar_scalings))
+
# ------------------------------------------------------------------ #
# Public extraction entry points
# ------------------------------------------------------------------ #
diff --git a/src/napari_convpaint/feature_extractors/gaussian.py b/src/napari_convpaint/feature_extractors/gaussian.py
index 9ea4ae1..da0f908 100644
--- a/src/napari_convpaint/feature_extractors/gaussian.py
+++ b/src/napari_convpaint/feature_extractors/gaussian.py
@@ -27,6 +27,11 @@ def get_default_params(self, param=None):
param.fe_layers = None
return param
+ def cache_extra_state(self):
+ # sigma lives on the instance, not the Param — it must enter the cache
+ # key or a model rebuilt with a different sigma could hit stale entries.
+ return ("sigma", self.sigma)
+
def extract_features_from_plane(self, image, device=None):
# Given that we get single-channel images as input:
diff --git a/src/napari_convpaint/feature_extractors/nnlayers.py b/src/napari_convpaint/feature_extractors/nnlayers.py
index 49a1880..cef8e74 100644
--- a/src/napari_convpaint/feature_extractors/nnlayers.py
+++ b/src/napari_convpaint/feature_extractors/nnlayers.py
@@ -185,6 +185,12 @@ def gives_patched_features(self):
# So we return False here, even if the patch size is >1.
return False
+ def cache_spill_to_disk(self):
+ # CNN native payloads are per-pixel multi-scale feature maps — often
+ # hundreds of MB per slice, far more than the recompute cost justifies
+ # pickling to disk. Keep them RAM-cacheable only.
+ return False
+
def _compute_nn_properties(self):
"""Walk the network in execution order, accumulating receptive field,
total stride, and whether any global-context op was encountered, up to
diff --git a/src/napari_convpaint/utils.py b/src/napari_convpaint/utils.py
index 84a6366..da3c877 100644
--- a/src/napari_convpaint/utils.py
+++ b/src/napari_convpaint/utils.py
@@ -406,6 +406,10 @@ def rescale_class_labels(label_img, output_shape):
rescaled_label_img : np.ndarray
Rescaled label image.
"""
+ if label_img.shape == tuple(output_shape):
+ # Identity resize: skimage doesn't short-circuit, so skip it (keeping
+ # this function's uint8 output contract).
+ return label_img if label_img.dtype == np.uint8 else label_img.astype(np.uint8)
rescaled_label_img = skimage.transform.resize(label_img, output_shape, order=0, mode='reflect', preserve_range=True).astype(np.uint8)
return rescaled_label_img
@@ -426,6 +430,8 @@ def rescale_outputs(output_img, output_shape, order=0):
rescaled_output : np.ndarray
Rescaled class probability or feature image.
"""
+ if output_img.shape == tuple(output_shape):
+ return output_img # identity resize: skimage doesn't short-circuit, so skip it
rescaled_output = skimage.transform.resize(output_img, output_shape, order=order, mode='reflect', preserve_range=True)
return rescaled_output
@@ -548,6 +554,10 @@ def pad_to_shape(feat, target_shape):
pad_before = diff // 2
pad_after = diff - pad_before # ensures bottom/right get the extra pixel if diff is odd
pad.append((pad_before, pad_after))
+ if isinstance(feat, torch.Tensor):
+ # torch.nn.functional.pad orders pads (left, right, top, bottom) for the
+ # last two dims; same symmetric convention as the numpy branch.
+ return torch.nn.functional.pad(feat, (pad[3][0], pad[3][1], pad[2][0], pad[2][1]))
return np.pad(feat, pad, mode='constant')
def align_up(val, alignment):