Skip to content

training_setup_screen

anomaly_match_ui.screens.training_setup_screen

Training setup screen — choose source folder, labels CSV, and model checkpoint.

TrainingSetupScreen(app)

Bases: BaseScreen

Setup screen for configuring a training run.

Provides file choosers for the source folder, label CSV, and optional metadata CSV, plus normalisation configuration.

Parameters:

Name Type Description Default
app AnomalyMatchApp

The parent application instance used for navigation.

required
Source code in anomaly_match_ui/screens/training_setup_screen.py
def __init__(self, app: AnomalyMatchApp) -> None:
    super().__init__(app)
    self._cfg = BackendInterface.get_config()
    self._source_ok = False
    self._labels_ok = False
    self._source_count: int | None = None
    self._counting = False
    self._counting_status: str = "Counting images..."
    self._image_files: list[str] = []
    self._label_message: str = ""
    self._label_validation_message: str = ""
    self._detected_source_type: DataSourceType = DataSourceType.IMAGE_FOLDER
    self._found_source_ids: list[str] = []
    self._suppress_norm_preview: bool = False
    # Monotonic counter — incremented each time a new background job
    # starts.  Any running job whose generation doesn't match bails out.
    self._job_generation: int = 0
    # Cache the most recent label-validation result so norm-only changes
    # don't re-run expensive ID-matching against the data source.
    # Keyed on (source_dir, source_type, label_path, label_mtime).
    self._validation_cache: dict | None = None
    # Signature of the inputs the last on_enter auto-validation ran against
    # (source path, label file, normalisation).  Lets a pure back-navigation
    # — e.g. returning from the image-detail screen with nothing changed —
    # reuse the held count + preview instead of re-running the full
    # count → validate → preview job.  Reset on explicit source/label edits.
    self._auto_validate_signature: tuple | None = None
    # Mirror Cutana's INFO-level progress messages into the preview
    # spinner caption so the user sees what's happening while Cutana
    # extracts cutouts.  Installed/removed on screen enter/leave.
    self._progress_tap: LogProgressTap | None = None
    # Last norm config seen by ``_on_norm_change`` — deduplicates the
    # observer cascade when ``update_from_config`` (or
    # ``_add_channel``/``_remove_channel``) assigns several widgets
    # in a row, each firing ``observe``.  Without this, a single
    # screen-enter could spawn ten or more background jobs.
    self._last_norm_config: dict | None = None
    # In-memory store for label edits made on this screen.  Holds
    # only the user's setup-time toggles; existing CSV labels are
    # *not* mirrored here — round-tripping happens at render time
    # via the labeled_map argument to ``_render_preview_items``.
    # On Start Training we merge ``_setup_labels`` into the chosen
    # CSV via ``BackendInterface.merge_gallery_labels`` so the
    # training subprocess sees the up-to-date label set (#427).
    self._setup_labels: dict[str, LabelState] = {}
    # Decoded preview samples for the current source, paged by the
    # gallery.  Tuples of ``(name, png_bytes, csv_label_str)`` so the
    # gallery can render the existing label badge on each cell.
    self._preview_samples: list[tuple[str, bytes, str]] = []
    # Absolute offset of the first cell currently rendered into the
    # gallery.  Tracked in addition to the gallery's ``_current_page``
    # so slider changes (which alter the effective page size) can
    # land the user back near the same cell instead of the same
    # *page index* under a stale chunk size (#432).
    self._preview_first_offset: int = 0
    # Numpy-array versions of the preview samples, keyed by
    # filename — needed for the detail screen, which expects raw
    # HWC uint8 arrays (PreviewWidget.display_standalone) rather
    # than the PNG bytes the gallery cell displays.
    self._preview_raw_images: dict[str, np.ndarray] = {}
    # Normalisation-independent raw float32 cutouts behind the current
    # Cutana preview, in gallery order.  Held so a normalisation change
    # re-decodes them in-process instead of re-reading the labeled cache
    # over NFS and spawning a fresh count+validate+read job (#501).  None
    # whenever the preview isn't fully cache-backed (image folder, Zarr, or
    # a Cutana cache miss) — those fall back to the full reload path.
    self._preview_raw_cutouts: list[tuple[str, np.ndarray]] | None = None
    # Label map captured alongside the held raws so the in-memory re-decode
    # (#501) can render label badges without re-reading the CSV over NFS —
    # the label map can't change during a normalisation tweak (editing the
    # CSV chooser triggers a separate full reload).
    self._preview_labeled_map: dict[str, str] = {}
    # Serialises the nested setattr sequence in
    # ``_apply_norm_overrides_to_cfg`` so the UI-thread invariant
    # survives a future refactor that lets a worker thread call it.
    # ipykernel already serialises widget callbacks today, but that's
    # convention not contract — without this lock a debounced observer
    # firing from a background timer would silently re-introduce the
    # torn-write race against ``self._cfg.normalisation``.
    self._norm_apply_lock = threading.Lock()
    # Debounce timer + guard for coalescing rapid normalisation edits (e.g.
    # a zoom-out slider drag) into a single preview refresh.
    self._norm_refresh_timer: threading.Timer | None = None
    self._norm_refresh_lock = threading.Lock()
    # Set when a debounced edit touched an EXTRACTION_AFFECTING_NORM_FIELD.
    # Those change the pixels Cutana extracts, so the held raws are stale and
    # the in-memory re-decode below cannot show the change (#592:
    # the zoom-out slider appeared to do nothing on a warm cache).
    self._norm_refresh_needs_extraction = False
    # Single-flight guard for preview jobs.  ``_job_in_flight`` is True while
    # ANY preview job runs (initial load, source/label change, norm refresh).
    # A norm change arriving while one is in flight sets ``_refresh_pending``
    # and the in-flight job's completion re-runs the refresh once — by then
    # the raws are held, so it takes the fast in-memory path.  This stops a
    # norm change from spawning a *concurrent* cache read on top of a still-
    # running initial load (the raws aren't held until that load finishes).
    self._job_in_flight = False
    self._refresh_pending = False
    self._log_sink = ScreenLogSink("TrainingSetupScreen")

build()

Build and return the training setup layout.

Returns:

Type Description
Widget

The root VBox widget for the setup screen.

Source code in anomaly_match_ui/screens/training_setup_screen.py
def build(self) -> widgets.Widget:
    """Build and return the training setup layout.

    Returns:
        The root VBox widget for the setup screen.
    """
    self._out = widgets.Output(
        layout=widgets.Layout(
            border="1px solid white",
            height="120px",
            background_color=BG_COLOR,
            overflow="auto",
        ),
    )
    header = HeaderWidget(self.app, self._out, show_back_button=True)

    # ── Source folder chooser ─────────────────────────────────
    # Resolve the source-chooser starting dir once so auto-validation
    # scans the same folder the chooser displays.  When persistence
    # picks a different dir than cfg.data_dir, _auto_validate_from_config
    # would otherwise walk cfg.data_dir (the test-data default) while
    # the chooser shows the persisted dir, and the user gets
    # mismatched "Source folder = q1_data / Found 19 images" reports.
    self._source_initial_path = initial_dir(self._cfg.data_dir, purpose="source_chooser")
    source_has_default = bool(
        self._source_initial_path and os.path.isdir(self._source_initial_path)
    )
    self._source_chooser = AMFileChooser(
        path=self._source_initial_path,
        show_only_dirs=True,
        select_default=source_has_default,
        title="Source folder",
        purpose="source_chooser",
    )
    self._source_chooser.register_callback(self._on_source_change)

    # ── Label CSV chooser ─────────────────────────────────────
    label_dir, label_name = initial_file(
        self._cfg.label_file,
        purpose="label_chooser",
        is_shipped_default=BackendInterface.is_shipped_default_path(
            "label_file", self._cfg.label_file
        ),
    )
    self._label_chooser = AMFileChooser(
        path=label_dir,
        filename=label_name,
        select_default=bool(label_name),
        filter_pattern="*.csv",
        title="Labelled data CSV",
        purpose="label_chooser",
    )
    self._label_chooser.register_callback(self._on_label_change)

    # ── Metadata CSV chooser (optional) ───────────────────────
    meta_dir, meta_name = initial_file(
        self._cfg.metadata_file,
        purpose="metadata_chooser",
        is_shipped_default=BackendInterface.is_shipped_default_path(
            "metadata_file", self._cfg.metadata_file
        ),
    )
    self._meta_chooser = AMFileChooser(
        path=meta_dir,
        filename=meta_name,
        select_default=bool(meta_name),
        filter_pattern="*.csv",
        title="Metadata CSV (optional)",
        purpose="metadata_chooser",
    )

    # ── Preview grid ──────────────────────────────────────────
    # Keep PreviewGrid for the spinner / error-message panel — its
    # status / message HTML is exactly what we want for the
    # "Loading preview…" + "Select a label CSV" affordances.
    # The actual cell rendering moves to GalleryWidget so cells get
    # pagination + per-cell label toggles (#427).
    self._preview_grid = PreviewGrid()
    # Hide the static cell grid that ships with PreviewGrid — we
    # render cells through GalleryWidget instead.
    self._preview_grid._grid.layout.display = "none"

    def _setup_cell_factory(**kwargs: object) -> LabelableThumbnailCell:
        # The factory is called with on_magnify / on_star kwargs by
        # GalleryWidget; override on_magnify with our own handler so
        # the magnifying-glass button on a setup-screen cell opens
        # the image detail view (the gallery's default on_magnify
        # is None at setup time).
        kwargs["on_magnify"] = self._on_setup_magnify
        return LabelableThumbnailCell(on_label_change=self._on_setup_label_change, **kwargs)

    self._preview_gallery = GalleryWidget(
        setup_mode=True,
        cell_factory=_setup_cell_factory,
    )
    self._preview_gallery.set_on_page_change(self._on_preview_page_change)
    # Slider widens cells → fewer fit in two rows → effective page
    # size shrinks → re-render the current page so paging walks the
    # new chunk size and the user doesn't skip cells (#432).
    self._preview_gallery.set_on_scale_change(self._on_preview_scale_change)
    # Hide until we have items.
    self._preview_gallery.widget.layout.display = "none"

    # Wrap the three PreviewGrid status-affording methods so any
    # spinner / message / clear also hides the cell gallery — the
    # two panels share the same screen real estate and only one
    # should be visible at a time.
    _orig_show_spinner = self._preview_grid.show_spinner
    _orig_show_message = self._preview_grid.show_message
    _orig_clear = self._preview_grid.clear

    def _show_spinner(*args: object, **kwargs: object) -> None:
        self._preview_gallery.widget.layout.display = "none"
        _orig_show_spinner(*args, **kwargs)

    def _show_message(*args: object, **kwargs: object) -> None:
        self._preview_gallery.widget.layout.display = "none"
        _orig_show_message(*args, **kwargs)

    def _clear() -> None:
        self._preview_gallery.widget.layout.display = "none"
        _orig_clear()

    self._preview_grid.show_spinner = _show_spinner  # type: ignore[method-assign]
    self._preview_grid.show_message = _show_message  # type: ignore[method-assign]
    self._preview_grid.clear = _clear  # type: ignore[method-assign]

    # ── Normalisation config ─────────────────────────────────
    norm_cfg = self._cfg.normalisation
    self._norm_widget = NormalisationConfigWidget(
        n_channels=norm_cfg.n_output_channels,
        image_size=list(norm_cfg.image_size),
    )
    # Restore before registering the observer so the widget assignments
    # can't spawn a preview job, and before ``on_enter`` reads
    # ``cfg.normalisation`` back into the widget.  Same rationale as
    # ``initial_dir``: what the user last ran with is more current than a
    # notebook's cfg defaults, which would otherwise clobber it on every
    # fresh kernel.  Nothing is reading cfg yet at build time, so pushing
    # the restored values across is safe here.
    if self._restore_persisted_normalisation():
        self._apply_norm_overrides_to_cfg()
    self._norm_widget.register_on_change(self._on_norm_change)
    # Seed the dedup snapshot with the widget's initial state so the
    # observer cascade from build-time widget assignments is not
    # mistaken for a real user change.
    self._last_norm_config = serialisable_config(self._norm_widget.get_normalisation_config())

    # ── Training iterations ───────────────────────────────────
    self._iter_slider = widgets.IntSlider(
        value=self._cfg.num_train_iter,
        min=50,
        max=600,
        step=10,
        description="Iterations:",
        style={"description_width": "initial", "handle_color": "white"},
        layout=widgets.Layout(width="320px"),
    )

    # ── Source-size stratification (Cutana only) ──────────────
    # Euclid Q1/DR1 sources skew heavily small.  When the unlabelled source
    # is a Cutana catalogue (which carries each source's diameter), sample
    # the unlabelled training pool ~uniform in size so tiny sources don't
    # dominate.  Disabled for image/Zarr sources, which expose no per-source
    # size — ``_refresh_validation`` enables it once a Cutana source is
    # detected.
    self._stratify_checkbox = widgets.Checkbox(
        value=self._cfg.cutana_stratify_source_size,
        description="Stratify source size per tile (Cutana)",
        indent=False,
        disabled=True,
        style={"description_width": "initial"},
        layout=widgets.Layout(width="360px"),
    )

    # ── Validation panel + start button ───────────────────────
    self._validation_html = widgets.HTML(
        value=self._render_validation(),
        layout=widgets.Layout(padding="8px 0"),
    )

    self._start_btn = widgets.Button(
        description="Start Training",
        button_style="success",
        disabled=True,
        layout=widgets.Layout(width="200px", height="40px"),
        style={"font_size": "14px"},
    )
    self._start_btn.on_click(lambda _: self._on_start())

    # ── Layout ────────────────────────────────────────────────
    # Left: file choosers (constrained so preview gets space)
    left_panel = VBox(
        [self._source_chooser, self._label_chooser, self._meta_chooser],
        layout=widgets.Layout(
            background_color=BG_COLOR,
            padding="12px",
            gap="4px",
            flex="0 1 auto",
            max_width="420px",
        ),
    )

    # Centre: image preview (flex so it grows with window width)
    centre_panel = VBox(
        [self._preview_grid.widget, self._preview_gallery.widget],
        layout=widgets.Layout(
            background_color=BG_COLOR,
            padding="8px",
            flex="1 1 auto",
            min_width="140px",
            border_left=f"1px solid {ESA_BLUE_DEEP}",
            border_right=f"1px solid {ESA_BLUE_DEEP}",
        ),
    )

    # Right: normalisation config + iterations + validation + start
    right_panel = VBox(
        [
            self._norm_widget.norm_section,
            self._norm_widget.channel_section,
            self._iter_slider,
            self._stratify_checkbox,
            self._validation_html,
            self._start_btn,
        ],
        layout=widgets.Layout(
            background_color=BG_COLOR,
            padding="8px 12px",
            gap="8px",
            width="380px",
            flex="0 0 auto",
        ),
    )

    form = HBox(
        [left_panel, centre_panel, right_panel],
        layout=widgets.Layout(
            background_color=BG_COLOR,
        ),
    )

    return VBox(
        [
            header.widget,
            form,
            self._out,
        ],
        layout=widgets.Layout(
            background_color=BG_COLOR,
            border=f"1px solid {ESA_BLUE_DEEP}",
            width="100%",
        ),
    )

on_enter()

Set up terminal output and auto-validate from config.

Source code in anomaly_match_ui/screens/training_setup_screen.py
def on_enter(self) -> None:
    """Set up terminal output and auto-validate from config."""
    BackendInterface.set_terminal_output(self._out)
    self._log_sink.attach(self._out)
    self._install_progress_tap()
    self._auto_validate_from_config()
    # update_from_config() writes many widget values in sequence; each
    # assignment fires the value-observer callback, and without
    # suppression that would spawn a separate background job per field
    # (setup-screen gens jumped from 1 to 14 on a single re-enter).
    self._suppress_norm_preview = True
    try:
        self._norm_widget.update_from_config(dict(self._cfg.normalisation))
    finally:
        self._suppress_norm_preview = False
    # Resync the dedup snapshot so the next genuine user change
    # registers as a diff against the post-restore widget state.
    self._last_norm_config = serialisable_config(self._norm_widget.get_normalisation_config())

on_leave()

Remove the progress tap and widget sink when navigating away.

Source code in anomaly_match_ui/screens/training_setup_screen.py
def on_leave(self) -> None:
    """Remove the progress tap and widget sink when navigating away."""
    with self._norm_refresh_lock:
        if self._norm_refresh_timer is not None:
            self._norm_refresh_timer.cancel()
            self._norm_refresh_timer = None
    self._uninstall_progress_tap()
    self._log_sink.detach()