diff --git a/.gitattributes b/.gitattributes index 4cbaeb3..054a10b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -14,5 +14,5 @@ # committing it without the exclusion would force every `git clone` to # download 2.65 GB and spend the repo's LFS bandwidth quota. Fetch it via the # "BiRefNet setup" page in the app, or: -# git lfs pull --include="weights/birefnet_leaf.pth" +# git lfs pull -X "" -I "weights/birefnet_leaf.pth" weights/*.pth filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f623461..748d624 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,3 +26,18 @@ jobs: python -m pip install pytest - name: Run tests run: python -m pytest tests/ -q + + pre-cleanup-integration: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install image test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install --no-deps -e . + python -m pip install pytest numpy opencv-python-headless + - name: Run pre-cleanup integration tests + run: python -m pytest tests/test_pre_cleanup_exports.py -q -rs diff --git a/.lfsconfig b/.lfsconfig index b6799f7..639c41b 100644 --- a/.lfsconfig +++ b/.lfsconfig @@ -12,9 +12,9 @@ # Opt in via the "BiRefNet setup" page in the MATS app, `mats fetch-weights # --only birefnet --source lfs`, or by hand: # -# git lfs pull --include="weights/birefnet_leaf.pth" +# git lfs pull -X "" -I "weights/birefnet_leaf.pth" # -# Note that a bare `git lfs pull` will NOT fetch it -- only -I/--include -# overrides fetchexclude. That's deliberate: no accidental 2.65 GB pulls. -# [lfs] -# fetchexclude = weights/birefnet_leaf.pth +# A bare pull and an include-only pull both retain fetchexclude. `-X ""` +# clears that exclusion for this invocation; `-I` limits the pull to BiRefNet. +[lfs] + fetchexclude = weights/birefnet_leaf.pth diff --git a/AGENTS.md b/AGENTS.md index 374b870..7375bfb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -75,8 +75,9 @@ other system libraries (QR codes are decoded with OpenCV). > fresh checkout fails to detect markers, and it looks like a model problem > rather than a setup problem. Verify with `ls -l weights/rf_detr_marker.pth` > (~134 MB, not ~134 bytes) or `mats doctor`; repair with -> `git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth"` — the `--exclude` matters: a bare `git lfs pull` can also -> fetch the 2.65 GB BiRefNet checkpoint. +> `git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth"` — the `--exclude` keeps the repair to RF-DETR. With +> `.lfsconfig` active, a bare `git lfs pull` also leaves the 2.65 GB BiRefNet checkpoint out; fetch it only with +> `git lfs pull -X "" -I "weights/birefnet_leaf.pth"` (or `mats fetch-weights --only birefnet --source lfs`). ```bash git lfs install # one-time, per machine, BEFORE cloning @@ -128,6 +129,23 @@ Currently true, and worth knowing because users ask: `--input-dir`). - Defaults that surprise people: `--mask-method threshold`, `--threshold-level auto`, `--csv-schema full`, `--results-unit cm`, `--output-mode masks`. +- `--threshold-level` takes a preset (`auto`, `low`, `medium`, `high`) or an + integer cutoff `1`–`255`; the GUI's **custom** level is the same thing. +- `--mask-method both` (two checked methods in the GUI) measures each image with + Otsu and BiRefNet and writes one CSV and failure log per method + (`*_threshold.csv`, `*_birefnet.csv`); a single-method run keeps the + unsuffixed names and identical output. +- `--measure-pre-cleanup` is not "every raw pixel": it clears a `--clean-margin` + band along the target-box edge (default 1% of the shorter side, where the + printed box outline lands), keeps the largest remaining object (the leaf), + and drops pieces that touch the band or lie beyond `--stray-gap` (default + `0.25` × the leaf's bounding-box diagonal) — see `mask_cleanup.clean_raw_mask`. + `--clean-size` (pixels, default `0` = off, pre-cleanup only) then removes + specks and fills holes below that inscribed radius — the app's **Clean size**, + via `mask_cleanup.raw_measurement_mask`; Adjust's Overwrite saves it for Otsu + specimens. `--clean-margin` is the edge margin, not the clean size. + Clean image and Remove flashfill use the same margin. `--export pre-cleanup` + still writes the untouched raw mask. The default cleaned path is unaffected. - `mats fetch-weights` with no flag fetches **RF-DETR only**; BiRefNet needs `--only birefnet` or `--all`. `--source {auto,hf,lfs}` picks the channel — `lfs` is the one to use on networks that block huggingface.co. After a @@ -202,8 +220,8 @@ Rules that keep this repo working: module-level import of torch, rfdetr, or streamlit in an imported path breaks CI even when it works locally. - **`paths.py` stays import-light** — standard library only. Models load lazily - inside `core.py`; `samples.py`, `dimensions.py`, and `scaling.py` follow the - same contract. + inside `core.py`; `samples.py`, `dimensions.py`, `scaling.py`, and + `mask_settings.py` follow the same contract. - **The CLI and the GUI share one execution path** (`run_leaf_morpho_batch`). Never fork pipeline logic between them — divergence would mean the two interfaces report different measurements. diff --git a/CHANGELOG.md b/CHANGELOG.md index 3cb78d2..84cd2f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,35 @@ All notable changes to MATs are documented here. This project adheres to ## [Unreleased] ### Changed +- Adjust now lists specimens in a searchable measurement table with sortable + columns. Clicking a row (its **View** button) previews that specimen; for Otsu + runs, ticking **Marked for Adjustment** marks it, and marks persist across + searches. Its duplicate + measurement cards and saved-image viewer have been removed, leaving the + controls and live preview. +- **Remove flashfill from this preview** now sits under **Clean size** inside + **Explore and adjust output**. +- Adjust saves marked specimens with its own **Overwrite all marked specimens + (N)** button, below **Overwrite this specimen** (the specimen in View), + replacing the "Apply these adjustments to all marked specimens" checkbox. +- Analyze now makes the Classic thresholding/BiRefNet comparison switch explicit + and keeps each method's table selection separate, so the specimen viewer + always follows the active method's table. +- Setup now places threshold level directly under Classic thresholding and + labels the default method without the Otsu parenthetical. +- The Analyze measurement table and selected-specimen viewer occupy separate + full-width blocks. Interactive controls now live in the Adjust tab. +- Analyze shows only the specimen selected in the measurement table. Adjust + carries that selection into its preview and overwrite controls. +- The Analyze sample viewer now shows the raw or cleaned mask used for the + completed run, even when that image was not selected for export. +- The app's segmentation choice is now two checkboxes. **Threshold level** + appears only while Otsu is checked, and a single **Pre-cleanup masks** export + covers every checked method. +- The workbench now uses Setup, Analyze, Adjust, and Export, with Diagnostics + in the sidebar. Image exports are selected before a run; downloads include + only files recorded for that run. +- Target boxes and cleaned masks remain default exports but can be disabled. - Documentation now matches the shipped weights-delivery behavior: Git LFS is listed as an install prerequisite (a clone without it yields a 134-byte pointer stub, not the model), RF-DETR is documented as arriving *with* the @@ -13,6 +42,49 @@ All notable changes to MATs are documented here. This project adheres to described as the repair path it is. ### Added +- An interactive threshold preview for the selected Otsu sample in Adjust. + Dragging updates its raw mask, and a color panel of the masked leaf beside + it, immediately; releasing applies the existing + cleanup when cleaned measurements are selected. A chosen cutoff can be used + for the next run or overwrite the selected specimen's mask, measurement row, + and dependent images with one button. Marked specimens can receive the same + settings in one bulk save, each using its own calibration. +- A **Clean size** slider in Adjust's Explore and adjust box, for Otsu and + BiRefNet specimens. Above 0 it live-previews Clean image, which drops small + white specks and fills small enclosed holes without flash-filling the leaf. + For Otsu specimens, **Overwrite this specimen** and **Overwrite all marked + specimens** save that mask, re-measure it, and record the clean size; + BiRefNet specimens preview it only. +- A run-wide clean size for pre-cleanup measurements: `--clean-size PX` in the + CLI, **Clean size (px)** in Setup (default 0, off). It removes specks and + fills enclosed holes whose inscribed radius is below that many pixels before + measuring, exactly as the Adjust slider previews, and is recorded in each + results CSV's `.meta.json`. It requires `--measure-pre-cleanup`. +- Optional measurements from pre-cleanup binary masks in the app and CLI + (`--measure-pre-cleanup`), with per-CSV metadata recording the source. + A band along the target-box edge (`--clean-margin`, **Edge margin** in the + app; default 1% of the box's shorter side) is cleared first, so the + template's printed box outline can never outweigh and replace a small leaf. + The largest remaining object then anchors the measurement: pieces touching + that band, or farther from the leaf than `--stray-gap` (**Stray-piece + distance**; default 0.25 × the leaf's bounding-box diagonal), are dropped so + printed box lines and distant debris no longer stretch width and length. + Clean image (a clean size above 0), Remove flashfill, and the pre-cleanup + threshold explorer apply + the same cleanup; the default cleaned measurements are unchanged. In Setup + these settings are grayed out unless pre-cleanup measurement is checked; each + specimen's explorer can adjust them for its own preview and overwrite. +- Measure with Otsu and BiRefNet in one run: `--mask-method both`, or check both + segmentation methods in the app. Markers are detected once per image; each + method writes its own results CSV and failure log (`_threshold`/`_birefnet` + suffixes) in the single-method schema, and the Analyze view can switch + between them. Single-method runs write the same files as before. +- Separate pre-cleanup binary mask exports for threshold/Otsu and BiRefNet, + including both methods in one run while measurements use the selected method. +- Optional overlay, cutout, and measurement-axis exports in the CLI and app. +- A custom grayscale threshold: `--threshold-level` accepts an integer cutoff + `1`–`255`, and the app's **custom** threshold level shows a slider with the + low/medium/high presets marked. - A **Robust QR setup** sidebar page that explains the optional pyzbar and QReader fallbacks, reports their usable status, and keeps Conda optional. - A **Help** page in the app (sidebar) with a quick start, a photography guide diff --git a/README.md b/README.md index 013c6a1..0c27a61 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# MATS — Morphometric Analysis Toolbox +# MATS — Morphometric Analysis Toolbox for Segmentation Measure leaf **area, length, and width** in real-world units from a photo of leaves laid on a printed calibration template. @@ -68,8 +68,7 @@ ls -l weights/rf_detr_marker.pth **Already cloned without Git LFS?** No need to start over — install Git LFS as above, then repair the checkout in place. The `--exclude` keeps this to the -~134 MB RF-DETR file; a bare `git lfs pull` can also fetch the 2.65 GB BiRefNet -checkpoint: +~134 MB RF-DETR file, avoiding the optional 2.65 GB BiRefNet checkpoint: ```bash git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth" @@ -148,12 +147,14 @@ checkpoint or install the pyzbar/QReader robust-QR fallbacks. The required ~134 MB RF-DETR marker checkpoint is different: it is mandatory for every run, so it ships **in the clone** via Git LFS and needs no separate -download step. If it is ever missing — a clone made without Git LFS, or an -install outside a Git checkout — MATS fetches it once on first use and prints -`Fetching weights/rf_detr_marker.pth via Git LFS ...` while it does. Set -`MATS_NO_AUTO_FETCH=1` to turn that off and require pre-staged weights instead -(recommended on HPC login nodes). The app never does this silently: a missing -RF-DETR checkpoint is a blocking Preflight error. +download step. If it is missing from a Git checkout — for example, after cloning +without Git LFS — MATS can fetch it once on first use and prints +`Fetching weights/rf_detr_marker.pth via Git LFS ...` while it does. An install +outside a Git checkout must use a pre-staged checkpoint (or a separately +configured Hugging Face source). Set `MATS_NO_AUTO_FETCH=1` to turn automatic +fetching off and require pre-staged weights instead (recommended on HPC login +nodes). The app never does this silently: a missing RF-DETR checkpoint is a +blocking Preflight error. This keeps the initial network and disk footprint predictable, avoids native `zbar` failures on managed machines, and works better on HPC systems and @@ -235,10 +236,14 @@ Common options (full reference in [docs/cli.md](docs/cli.md)): | `-r, --results_path` | Measurement CSV path | `./leaf_morpho_results.csv` | | `--sheet-dimensions` | Finished Template Creator sheet size, `x` | read from QR | | `-t, --template_dimensions` | Legacy/custom marker-centre calibration area | unused | -| `--mask-method` | `birefnet` (accurate, GPU) or `threshold` (fast) | `threshold` | -| `--threshold-level` | `auto` (Otsu) / `low` / `medium` / `high` | `auto` | +| `--mask-method` | `birefnet` (accurate, GPU), `threshold` (fast), or `both` | `threshold` | +| `--threshold-level` | `auto` (Otsu) / `low` / `medium` / `high`, or a custom cutoff `1`–`255` | `auto` | | `--csv-schema` | `full` (area/width/length + per-axis pixels-per-selected-unit) or `compact` | `full` | | `--results-unit` | Measurement-output unit: `mm`, `cm`, or `in` | `cm` | +| `--measure-pre-cleanup` | Measure from raw binary masks before cleanup, minus the edge margin and stray pieces | off | +| `--clean-margin` | With `--measure-pre-cleanup`: band cleared along the target-box edge, as a percent of its shorter side | `1` | +| `--stray-gap` | With `--measure-pre-cleanup`: how far a piece may lie from the leaf, as a fraction of its bounding-box diagonal | `0.25` | +| `--clean-size` | With `--measure-pre-cleanup`: remove specks and fill enclosed holes smaller than this radius in pixels (the app's **Clean size**) | `0` (off) | | `-w, --workers` | Parallel workers (threshold path only) | auto | | `--save-axes` | Also save length/width overlay images for QC | off | @@ -247,7 +252,9 @@ no GPU, no extra download, and good for clean, high-contrast backgrounds where a leaf sits on plain white. `birefnet` is more accurate on cluttered or low-contrast backgrounds and uses a GPU when available (CPU works but is slow), at the cost of the ~2.65 GB checkpoint — fetch it once with -`mats fetch-weights --only birefnet`. +`mats fetch-weights --only birefnet`. To compare them, `--mask-method both` +(or checking both methods in the app) measures every image with each method +and writes one results CSV per method. --- @@ -258,6 +265,29 @@ Per image, in the output folder: - `{sample_id}_target_box.jpg` — the perspective-corrected observation box - `{sample_id}_mask.png` — the leaf segmentation mask +Both image exports are on by default and can be disabled with `--no-target-boxes` +or `--no-masks`. Pre-cropped target-box inputs are not copied. Optional +`--export pre-cleanup` saves each measurement method's binary mask before cleanup; +`--pre-cleanup-methods both` runs Otsu/threshold and BiRefNet and saves their +pre-cleanup masks separately. Only a `--mask-method` method determines +measurements. BiRefNet exports require its checkpoint to be installed locally. +With `--mask-method both`, each method's masks and QC images end in +`_threshold` or `_birefnet` (for example `{sample_id}_mask_birefnet.png`). +`--export overlay`, `--export cutout`, and `--export axes` add QC images; repeat +the flag to request multiple kinds. The app offers the same choices in Setup. +Use `--measure-pre-cleanup` (or **Measure from pre-cleanup masks** in the app) +to calculate measurements from the raw binary segmentation. A thin band along +the target-box edge, where the template's printed box outline lands, is cleared +first (`--clean-margin`). The largest remaining object is taken as the leaf; +other pieces that touch that band, or lie farther from the leaf than +`--stray-gap` times its bounding-box diagonal, are dropped. Area counts the remaining foreground pixels, and width and length span +their extent, so specks near the leaf can still affect the result; `--clean-size` +(for example `--clean-size 3`) removes specks and fills enclosed holes whose +inscribed radius is below that many pixels, without flash-filling the leaf. This choice +is independent of `--export pre-cleanup`, which saves the mask with every piece. +Each results CSV has a `.meta.json` companion that records the measurement source +and, for pre-cleanup runs, these settings. + Plus a measurements CSV. Choose `mm`, `cm` (the default), or `in` with `--results-unit` in the CLI or the **Result units** control in the app. The selection changes results, dashboard labels, and unit-bearing CSV column names; @@ -279,7 +309,11 @@ it does not change calibration math. Two schemas: millimeters or inches selected, `cm` is replaced consistently in the measurement column names. -A `leaf_morpho_failures.csv` records per-image warnings and failures. +A `leaf_morpho_failures.csv` records per-image warnings and failures. With +`--mask-method both` (or both methods checked in the app), each method writes its +own results CSV and failure log with a method suffix, for example +`leaf_morpho_results_birefnet.csv` and `leaf_morpho_failures_birefnet.csv`, in +the same schema as a single-method run. > **Migration note:** earlier versions reported three isotropic scale > conventions (`*_meanscale`, `*_widthscale`, `*_heightscale`). Old CSVs remain diff --git a/deploy/ondemand/mats/README.md b/deploy/ondemand/mats/README.md index 36bd255..e5cb7a6 100644 --- a/deploy/ondemand/mats/README.md +++ b/deploy/ondemand/mats/README.md @@ -13,12 +13,16 @@ compute node and exposes it through the Open OnDemand reverse proxy. ``` Then set `CONDA_ENV` in `template/script.sh.erb` to that env name (`mats`). (A plain virtualenv works too — the default install needs no system libs.) -2. **The model checkpoints.** Fetch them once, ideally to a shared location: +2. **The model checkpoints.** Pre-stage them once in a shared location. From a + Git checkout with Git LFS installed, materialize the files, then copy them to + the shared directory: ```bash export MATS_WEIGHTS_DIR=/shared/models/mats - mats fetch-weights --all # RF-DETR + BiRefNet (the GUI defaults to Otsu, but a - # GPU-backed OOD app is the typical BiRefNet use case) - mats doctor # confirm they resolve + mkdir -p "$MATS_WEIGHTS_DIR" + mats fetch-weights --all # writes RF-DETR + BiRefNet to the checkout's weights/ + cp weights/rf_detr_marker.pth "$MATS_WEIGHTS_DIR/" + cp weights/birefnet_leaf.pth "$MATS_WEIGHTS_DIR/" + mats doctor # confirm they resolve ``` Point the same `MATS_WEIGHTS_DIR` at that path in `template/script.sh.erb`. 3. *(Optional)* Enhanced QR reading (`pip install -e ".[qr]"`) adds the `pyzbar` diff --git a/docs/cli.md b/docs/cli.md index 5e17aa1..fb9f549 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -27,12 +27,95 @@ mats run -i ./images -o ./out -r results.csv --sheet-dimensions 12x12in | `--sheet-dimensions` | Finished Template Creator sheet size as `x`, e.g. `12x12in` or `30x30cm`; MATS derives calibration using Creator margins. | read from QR | | `-t, --template_dimensions, --template-dimensions` | Legacy/custom marker-centre calibration area. Retained for existing scripts and non-Creator sheets. | unused | | `--output-mode` | `masks` (segment leaves) or `target-boxes` (only save corrected boxes). | `masks` | -| `--mask-method` | `birefnet` (accurate, GPU) or `threshold` (fast). | `threshold` | -| `--threshold-level` | For `threshold`: `auto` (Otsu), `low` (100), `medium` (125), `high` (150). | `auto` | +| `--mask-method` | `birefnet` (accurate, GPU), `threshold` (fast), or `both` (measure every image with each method; see below). | `threshold` | +| `--threshold-level` | For `threshold` and threshold pre-cleanup exports: `auto` (Otsu), `low` (100), `medium` (125), `high` (150), or a custom integer cutoff `1`–`255` (e.g. `--threshold-level 140`). Grayscale pixels at or below the cutoff count as leaf. | `auto` | | `--csv-schema` | `full` (area/width/length + per-axis pixels-per-selected-unit) or `compact`. | `full` | | `--results-unit` | CSV measurement unit: `mm`, `cm`, or `in`. | `cm` | +| `--measure-pre-cleanup` | Measure from the raw binary segmentation before gap closing and hole filling, after clearing the edge margin and dropping stray pieces (see below). | off (cleaned mask) | +| `--clean-margin` | With `--measure-pre-cleanup`: width of the band cleared along every target-box edge, as a percent of the box's shorter side, `0`–`10`. `0` clears nothing. | `1` | +| `--stray-gap` | With `--measure-pre-cleanup`: drop pieces whose nearest pixel is farther from the leaf than this fraction of the leaf's bounding-box diagonal, `0`–`10`. `0` keeps only the leaf. | `0.25` | +| `--clean-size` | With `--measure-pre-cleanup`: after the margin and stray pieces are cleared, remove white specks and fill enclosed holes whose inscribed radius is below this many pixels, `0`–`50`. The app's **Clean size**. | `0` (off) | | `-w, --workers` | Parallel workers. Only the CPU `threshold` path over pre-made target boxes parallelizes; model-backed runs use one worker. | auto | | `--save-axes` | Also write per-image length/width overlay images for QC. | off | +| `--export` | Repeatable: `pre-cleanup`, `overlay`, `cutout`, `axes`. Overlays, cutouts, and axes use the selected measurement mask. | none | +| `--pre-cleanup-methods` | With `--export pre-cleanup`: `selected` (every `--mask-method` method), `threshold`, `birefnet`, or `both`. Requesting BiRefNet runs it for every image and requires its local checkpoint. | `selected` | +| `--no-target-boxes` | Do not save newly rectified target boxes. Existing target-box inputs are never copied. | off | +| `--no-masks` | Do not save cleaned masks. This does not change the measurement source. | off | +| `--no-failure-log` | Do not write `leaf_morpho_failures.csv`. | off | +| `--dataset-format` | Also create a training dataset ZIP: `png` image/mask pairs, `yolo-seg`, `yolo-detect`, or `coco` segmentation. | off | +| `--dataset-split TRAIN VAL TEST` | Train/validation/test percentages; must total 100. | `70 20 10` | +| `--dataset-seed` | Seed for reproducible split assignments. | `42` | +| `--dataset-method` | Labeling method when `--mask-method both` is selected. | first selected method | +| `--dataset-mask-source` | `measured` or `raw` pre-cleanup mask. | `measured` | +| `--dataset-groups CSV` | UTF-8 CSV with `sample_id,group_id`; keep photos of a group in one split. | no grouping | + +For example, create a YOLO segmentation dataset alongside the measurement CSV: + +```bash +mats run -i ./images -o ./out -r results.csv --sheet-dimensions 12x12in \ + --dataset-format yolo-seg --dataset-split 70 20 10 --dataset-seed 42 +``` + +The ZIP is written to `-o` as `mats_training__.zip` (with a +numeric suffix if that filename already exists). It includes a `manifest.json` +with actual split counts, excluded pairs, and conversion notes. Target-box +images and measured masks are paired even if `--no-target-boxes` or `--no-masks` +was used; temporary previews are removed after the ZIP is built. Labels are +generated by MATS and should be reviewed before training. YOLO polygons cannot +preserve mask holes; COCO run-length masks can. A group CSV may shift the exact +split percentages because a group always stays together. + +Pre-cleanup masks are binary segmentations before gap closing, hole filling, +and removal of smaller objects. BiRefNet masks are already thresholded, not +probability maps. `--pre-cleanup-methods both` writes +`{sample_id}_mask_precleanup_threshold.png` and +`{sample_id}_mask_precleanup_birefnet.png`. Only a `--mask-method` method +determines area, width, and length in the CSV. With `--output-mode target-boxes`, +segmentation exports are ignored. The results CSV is always written. + +`--measure-pre-cleanup` uses the selected method's raw binary mask for measurements. +First it clears a band `--clean-margin` percent of the box's shorter side wide along +every edge of the target box. The template's printed box outline runs through the +marker centres, so after perspective correction it lies on that edge; clearing it +first means the outline can never outweigh, and replace, a small leaf. The +measurement is then anchored on the leaf, the largest remaining object. Any other +piece that touches the cleared band (the rest of a printed line, marker remnants, +shadows at the sheet edge) is dropped, and so is any piece farther from the leaf +than `--stray-gap` times the leaf's bounding-box diagonal. A thin piece in the band +is never taken as the leaf. Area counts every remaining foreground pixel, +including specks near the leaf, while width and length span the remaining +foreground extent. Holes remain excluded from area. Lay leaves inside the printed +box: any part of a leaf within the margin is cleared too. Raise `--stray-gap` when +a leaf's parts lie apart, such as separated leaflets; lower it to drop specks +closer to the leaf. + +`--clean-size` then cleans what is left by size. Each remaining piece and each +enclosed hole is sized by its inscribed radius, the distance from its deepest +pixel to its edge. White specks with a radius below the clean size are removed, +and holes with a radius below it are filled; the leaf is always kept, and nothing +is flash-filled, so a hole at least that large stays excluded from area. It is the +**Clean size** slider in the app, and gives exactly the mask that slider previews. +For example, measure raw masks at the medium threshold, removing specks and +filling holes under 3 px: + +```bash +mats run -i ./images -o ./out -r results.csv --sheet-dimensions 12x12in \ + --measure-pre-cleanup --threshold-level medium --clean-size 3 +``` + +Without this flag, measurements use the cleaned mask as before. This setting is +independent of `--export pre-cleanup`, which writes the raw mask exactly as +segmented, stray pieces included. Each results CSV has a `.meta.json` companion +recording its measurement source, segmentation method, unit, and schema, plus +the clean margin, stray gap, and clean size for pre-cleanup runs. + +`--mask-method both` detects markers once per image and then measures it with +Otsu and with BiRefNet. Each method gets its own results CSV and failure log, +named from `-r` with a method suffix — `leaf_morpho_results_threshold.csv` and +`leaf_morpho_results_birefnet.csv` — in the same schema as a single-method run. +Each method's cleaned masks, overlays, cutouts, and axes also end in +`_threshold` or `_birefnet` (for example `{sample_id}_mask_birefnet.png`); +target boxes are shared. A run with one method keeps the unsuffixed names. ### Interactive vs non-interactive @@ -57,7 +140,9 @@ mats fetch-weights --only birefnet --source lfs # explicitly fetch just BiRefNet mats fetch-weights --force # re-download even if present ``` -Downloads to `~/.cache/mats/weights` (or `$MATS_WEIGHTS_DIR`). See +Git LFS downloads write to the Git checkout's `weights/` directory. +`MATS_WEIGHTS_DIR` is for pre-staged local or shared checkpoints; a configured +Hugging Face source may use it as its download destination. See [weights.md](weights.md). ## `mats doctor` diff --git a/docs/faq.md b/docs/faq.md index 545b86f..d87d2db 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -37,8 +37,13 @@ Repair an existing clone without re-cloning: git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth" ``` -The `--exclude` keeps the repair to the ~134 MB RF-DETR file; a bare -`git lfs pull` can also fetch the 2.65 GB BiRefNet checkpoint. +The `--exclude` keeps the repair to the ~134 MB RF-DETR file. With this +repository's `lfs.fetchexclude`, a bare `git lfs pull` also leaves the 2.65 GB +BiRefNet checkpoint out; fetch it explicitly only when needed: + +```bash +git lfs pull -X "" -I "weights/birefnet_leaf.pth" +``` **Do I need to download the model first?** No. The clone brings RF-DETR with it, so there is no separate download step — run `mats doctor` and you're done. (If @@ -120,7 +125,9 @@ you want when a batch mixes several template sizes. | Best for | Clean, high-contrast backgrounds (a leaf on plain white) | Cluttered or low-contrast backgrounds | Start with the default. Switch with `--mask-method birefnet` only if the masks -disappoint you. +disappoint you. To compare the two on your own photos, use `--mask-method both` +(or check both methods in the app): every image is measured with each method, +and each method gets its own results CSV. **How many workers?** `-w/--workers` applies to the threshold path. One worker uses CUDA/MPS when available; two or more switch to parallel CPU processing and @@ -167,6 +174,14 @@ Full detail, checksums, and the complete resolution order: [weights.md](weights. **What comes out?** Per image, a perspective-corrected `{sample_id}_target_box.jpg` and a `{sample_id}_mask.png`, plus one measurements CSV and a `leaf_morpho_failures.csv` listing anything that warned or failed. +Target-box and cleaned-mask files are optional, checked by default in the app; +the failure log is also optional. Pre-cleanup masks, overlays, cutouts, and +measurement axes can be requested in Setup or with CLI export flags. The +pre-cleanup mask keeps holes and smaller objects. You may export separate Otsu +and BiRefNet pre-cleanup masks in one run; only the measurement methods +determine the CSV values. A BiRefNet export requires a locally available model. +Measuring with both methods writes one results CSV and one failure log per +method, suffixed `_threshold` and `_birefnet`. **Which columns?** `--csv-schema full` (the default) gives `sample_id`, `leaf_area_cm2`, `width_cm`, `length_cm`, `px_per_cm_width`, `px_per_cm_height`, @@ -202,7 +217,7 @@ Run `mats doctor` first; it reports most of these. | No markers detected, on a fresh clone | Check for a Git LFS placeholder first: `ls -l weights/rf_detr_marker.pth` should be ~134 MB. If it's ~134 bytes, run `git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth"` | | "QR code not read" | Pass the size yourself: `--sheet-dimensions 12x12in`. To add decoders: `pip install -e ".[qr]"` (plus the native `zbar` for pyzbar) | | No markers detected | Get all four markers in frame; print at 100 % scale in the template's marker colour | -| Masks include the background | Try `--threshold-level low/medium/high`, or `--mask-method birefnet` | +| Masks include the background | Try `--threshold-level low/medium/high` or a custom cutoff such as `--threshold-level 140`, or `--mask-method birefnet` | | CUDA out of memory | Only with BiRefNet — process fewer images at a time, or use the default `threshold` | | Very slow run | Use `threshold` and raise `-w/--workers` | | Blank page in Open OnDemand | Reverse-proxy `baseUrlPath` mismatch — see [deploy/ondemand/mats/README.md](../deploy/ondemand/mats/README.md) | diff --git a/docs/gui.md b/docs/gui.md index 8d4c141..47e5078 100644 --- a/docs/gui.md +++ b/docs/gui.md @@ -9,14 +9,16 @@ mats app Streamlit opens in your browser (default http://localhost:8501). The sidebar provides the image source and output destination. The MATS Analysis Workbench has a prominent, sticky **WORKSPACE NAVIGATION** bar with large, button-like -controls for **Analyze**, **Results**, and **Diagnostics**. Its filled active section stays visible while the workbench -scrolls, including after a preflight shortcut opens Diagnostics. Analyze presents a -numbered workflow for scale, segmentation, output, and preflight; its highlighted -launch surface follows Preflight in the normal page flow. The separate sidebar -pages are **Template Creator** (making printable templates), **BiRefNet setup** +controls for **Setup**, **Analyze**, **Adjust**, and **Export**. Its filled active section stays visible while the workbench +scrolls. Setup holds +scale, segmentation, measurement output, and image output options. Analyze holds +Preflight, the launch button, progress, and results. Adjust holds specimen previews +and saving controls. Export holds saved files and downloads. The separate sidebar +pages are **Diagnostics** (compute status and preflight), **Template Creator** (making printable templates), **BiRefNet setup** (optional model install and hardware diagnostics), **CPU Options**, **Robust QR setup** (optional QR fallbacks), and **Help** (packaged sample images, a settings guide, and the results-CSV glossary). +After a run, **Go to Export** in the sidebar opens the Export tab. ## Home — measuring leaves @@ -25,7 +27,7 @@ guide, and the results-CSV glossary). the output folder there as well. The picker works on Windows, macOS, and Linux where a desktop folder-dialog backend is available. Accepted: `.jpg .jpeg .png .tif .tiff .bmp`. -2. In **Analyze**, enter the finished **printed sheet size** — outer-sheet +2. In **Setup**, enter the finished **printed sheet size** — outer-sheet **width** and **height** — and choose **in** or **cm** (defaults to 12x12 in). MATS uses the Template Creator's fixed margins to derive the marker-centre calibration area. Older/custom templates with different margins remain @@ -35,16 +37,108 @@ guide, and the results-CSV glossary). QR code instead. OpenCV handles clear codes with no extra setup. For glare, skew, or blur, **Robust QR setup** explains the optional `mats-morpho[qr]` fallbacks and the optional native `zbar` library. -3. Choose a **segmentation method**: *Classic thresholding (Otsu)* (fast, default; best - on clean backgrounds) or *BiRefNet* (more accurate on cluttered backgrounds; - uses a GPU when available; needs its optional ~2.65 GB local checkpoint). -4. Choose **output options**: choose whether measurements should be shown and exported in - **mm**, **cm**, or **in** (this is separate from the printed-sheet calibration unit), - then pick the **Full research schema** CSV for area/width/length plus per-axis +3. Choose one or both **segmentation methods**: *Classic thresholding* (fast, + default; best on clean backgrounds) and/or *BiRefNet* (more accurate on cluttered + backgrounds; uses a GPU when available; needs its optional ~2.65 GB local + checkpoint). Check both to measure every image with each method and compare + them: each method then gets its own results CSV + (`leaf_morpho_results_threshold.csv`, `leaf_morpho_results_birefnet.csv`), + failure log, and masks, while marker detection runs once per image. + **Threshold level** appears immediately below Classic thresholding, before + BiRefNet, and sets its cutoff: `auto` (Otsu, the default), the + `low`/`medium`/`high` presets, or **custom**, which shows a 1–255 slider with + the presets marked on it. +4. Choose **output options**: optionally check **Measure from pre-cleanup masks** + to use the raw binary segmentation for area, width, and length. Raw mode + first clears **Edge margin**, a band along every edge of the target box + (default 1% of its shorter side), where the template's printed box outline + lands. It then keeps the leaf (the largest object) and drops pieces that + touch that band or lie farther from the leaf than **Stray-piece distance**, + a fraction of the leaf's bounding-box diagonal (default 0.25; 0 keeps only + the leaf). Specks near the leaf still count toward area, width, and length, + and holes stay excluded from area, unless **Clean size** (px, default 0) is + above 0: it then removes white specks and fills enclosed holes whose + inscribed radius is below that many pixels, without flash-filling the leaf. + Lay leaves inside the printed box, since any part of a leaf within the + margin is cleared too. All three settings are grayed out unless the box is + checked. The default uses the cleaned mask. + This choice is separate from exporting pre-cleanup mask images, which keep + every piece. Choose result + units in **mm**, **cm**, or **in** (separate from the printed-sheet calibration + unit), then pick the **Full research schema** CSV for area/width/length plus per-axis pixels-per-selected-unit and a `scale_aspect_ratio` QC column, or **Compact** for a trimmed export. In **Variable dimensions** mode, the Full schema also records each installed QR decoder's outcome (`success`, - `failed`, or `unused`). Optionally write a failures log. + `failed`, or `unused`). + The failure log is also checked by default in **Measurement Output**. + In **4 · Image output options**, choose target-box and cleaned-mask images + (checked by default), or optional pre-cleanup masks, overlays, cutouts, + and measurement axes. Pre-cleanup masks are binary segmentations before + closing, hole filling, and removal of smaller objects; one is written for + each checked segmentation method. With both methods checked, masks, overlays, + cutouts, and axes end in `_threshold` or `_birefnet` + (for example `{id}_mask_birefnet.png`). Existing target-box inputs are not + copied. After a run with both methods, **Measurement table and specimen + view** switches the summary, table, charts, and selected specimen between + Classic thresholding and BiRefNet. Each method has its own table selection; + choose a row to see that method's specimen mask. + Select one row in the Analyze measurement table to inspect that specimen; + **Adjust selected specimen** opens its controls in Adjust. Adjust lists + specimens in a measurement table (click a column header to sort); search + specimen names to narrow it, and click a row, or its **View** button, to open + its live preview. Adjust contains the controls and live preview, + while saved images and measurement cards are shown in Analyze. The Analyze sample viewer shows + the mask used for that run's measurements: + raw with the edge margin cleared and stray pieces dropped when **Measure + from pre-cleanup masks** was checked, cleaned otherwise. In a pre-cleanup + run, the threshold explorer also clears the margin while you drag and shows + that measured mask when you release. **Edge margin** and **Stray-piece + distance** also appear in each specimen's Adjust controls, starting from the run's + values. There they adjust just that specimen's preview, and are enabled + only while flash fill is off: in pre-cleanup runs, with a **Clean size** + above 0, or (the margin alone) with **Remove flashfill**. Overwriting a specimen saves + its mask with the values shown and records them in the `.meta.json`. + It can show the mask and target box even if their exports were turned off; + those extra previews last only for the current app session. + **Remove flashfill from this preview**, under **Clean size** in + **Explore and adjust output**, leaves enclosed holes unfilled while + retaining the other cleanup for the selected specimen. It is available for + cleaned-mask runs and remains a preview until Overwrite is pressed. + For a selected Otsu sample, **Explore and adjust output** lets you + drag a cutoff and see the raw mask change immediately, with the masked leaf + in its original colors beside it so you can see which parts of the leaf the + cutoff keeps or loses. Releasing the slider + applies MATS cleanup when the run used cleaned masks, honoring the flashfill + checkbox. The preview does not change files by itself. + **Overwrite this specimen** replaces only the mask and CSV row of the + specimen selected in the table's **View** column. Existing overlays, cutouts, and measurement axes + for that sample are regenerated; every other sample remains unchanged. + To save the same cutoff and cleanup settings across specimens, tick their + **Marked for Adjustment** boxes in the Adjust table (tick again to unmark). + Marks persist across searches. **Mark all matching** marks every search match; + **Mark all** with an empty search marks every available specimen. + Then press **Overwrite all marked specimens (N)**, below **Overwrite this + specimen**, to save the settings to every specimen ticked in **Marked for + Adjustment**. + Each specimen keeps its own scale; all + marked specimens are validated before any saved output changes. + **Reset to saved threshold** restores the sample's current saved cutoff; + **Use threshold for next run** sets a custom cutoff in Setup for the next + analysis. BiRefNet specimens have no threshold, reset, next-run, or + overwrite controls. The **Clean size** slider, in the same box for both + methods, starts at the specimen's saved clean size, else the run's (0 + unless set in Setup); 0 turns Clean image off, and its **?** explains it. + Above 0 it previews Clean image, a gentler alternative to MATS cleanup: + it clears the specimen's edge margin, drops pieces that touch it or lie + beyond the stray-piece distance, then drops disconnected white specks and + fills enclosed holes smaller than the clean size. It always keeps the leaf + and never flash-fills. For an Otsu specimen, either Overwrite button saves + the mask shown, re-measures it under the run's measurement source, and + records the clean size with the cutoff in the `.meta.json`; in a + cleaned-mask run, Clean image then replaces MATS cleanup (and Remove + flashfill) in the saved mask. BiRefNet adjustments remain preview-only. + Each results CSV has a `.meta.json` companion recording which measurement + source was used and any per-sample threshold adjustments. 5. **CPU Options** retains the worker controls. The app detects the CPU workers assigned to it (including HPC scheduler limits). One worker uses CUDA/MPS when available. Selecting two or more workers enables parallel CPU processing and disables CUDA/MPS for that @@ -52,17 +146,37 @@ guide, and the results-CSV glossary). 25% or less of the CPU allocation, yellow through 50%, and red through 75%. Counts above 75% require a one-run **Break the glass** acknowledgement. 6. The compact **Preflight** in Analyze shows readiness without overwhelming - the workspace. The **Diagnostics** tab has compute status, a compact + the workspace. The **Diagnostics** sidebar page has compute status, a compact overview of every check, and a collapsible detailed report. They show green, yellow, and red checks for weights, BiRefNet compute availability, printed sheet calibration, QR-reader availability, and input images. OpenCV and pyzbar decode a known test QR during QR-mode Preflight; QReader is import- checked without initialization so Preflight cannot trigger a model download. A missing optional BiRefNet checkpoint is yellow and links to **BiRefNet setup**; - it only blocks a run when BiRefNet segmentation is selected. -7. After the run, MATS opens **Results** for measurement cards and charts, a - selectable measurements table beside a linked target-box / mask specimen - inspector, CSV download, and a ZIP of all outputs. + it blocks a run when BiRefNet is needed for measurements or an export. +7. After the run, **Analyze** shows measurement cards and charts, a + full-width selectable measurements table followed by a full-width linked + target-box / mask specimen inspector. **Export** lists the files saved for + that run and its output folder. Choose a method (or both), select file types + for a ZIP, review its file list and size, and set its download filename. + All available file types are selected by default. + Results CSVs can also be downloaded individually. File types not saved during + the run cannot be added to the ZIP; enable image output options in Setup + before the next run if needed. + The **Training dataset** section creates a separate ZIP from the current + run's aligned target-box images and masks, including saved specimen + adjustments. Choose one segmentation method, the measured mask or an + available raw pre-cleanup mask, and image/PNG-mask, YOLO segmentation, YOLO + detection, or COCO segmentation format. Set train/validation/test percentages + (default 70/20/10; they must total 100) and a random seed. An optional UTF-8 + CSV with `sample_id,group_id` columns keeps repeated photos of the same + plant or specimen in one split; group sizes can shift the exact percentages. + The ZIP contains `manifest.json` with actual counts, exclusions, and any + conversion notes. MATS generates these labels from its masks, so review + them before training. YOLO segmentation polygons cannot retain holes in + masks; COCO segmentation uses run-length masks that retain them. This export + uses current-session previews when regular image output was turned off, so + prepare it before ending the app session. Large batches (>200 images) ask for confirmation and run synchronously — keep the browser tab open until they finish. diff --git a/docs/hpc.md b/docs/hpc.md index 69087b1..d8e6e03 100644 --- a/docs/hpc.md +++ b/docs/hpc.md @@ -12,15 +12,21 @@ conda env create -f environment.yml conda activate mats pip install -e ".[app]" export MATS_WEIGHTS_DIR=/project//mats_weights # shared, readable -mats fetch-weights --all # populate both checkpoints once, from a data-transfer node +mkdir -p "$MATS_WEIGHTS_DIR" +# Run from a Git checkout with Git LFS installed. It materializes both files +# in that checkout's weights/ directory; copy the verified files to /project. +mats fetch-weights --all +cp weights/rf_detr_marker.pth "$MATS_WEIGHTS_DIR/" +cp weights/birefnet_leaf.pth "$MATS_WEIGHTS_DIR/" mats doctor ``` On USDA **SCINet** (Ceres/Atlas), a `/project` directory is a mounted filesystem shared across the project, so every job reads the weights in place — no per-user -copy. Fetch them once to that path and point `MATS_WEIGHTS_DIR` at it for all -users. External collaborators without SCINet accounts can pull the same directory -via a **Globus guest collection** (they need a free Globus login). +copy. Materialize the weights once in a Git checkout, copy the verified files to +that path, and point `MATS_WEIGHTS_DIR` at it for all users. External collaborators +without SCINet accounts can pull the same directory via a **Globus guest collection** +(they need a free Globus login). Set `MATS_NO_AUTO_FETCH=1` in your jobs so a misconfigured path fails fast with a clear error instead of triggering a 2.65 GB download on a login or compute node diff --git a/docs/weights.md b/docs/weights.md index 6ac7f88..c0b1e08 100644 --- a/docs/weights.md +++ b/docs/weights.md @@ -58,9 +58,11 @@ mats doctor # show resolved paths, channels, and so ``` After a clone made with Git LFS installed, bare `mats fetch-weights` is a no-op -that prints "already present" — RF-DETR arrived with the checkout. The command -exists to repair a checkout made *without* Git LFS, and to populate a shared -`MATS_WEIGHTS_DIR`. +that prints "already present" — RF-DETR arrived with the checkout. The Git LFS +channel repairs a checkout made *without* Git LFS and always writes to that +checkout's `weights/` directory. `MATS_WEIGHTS_DIR` controls where MATS looks +for pre-staged files; provision a shared directory by copying verified +checkpoints there rather than expecting a Git LFS fetch to populate it. BiRefNet is never downloaded automatically. If it is absent when selected, MATs reports the missing local checkpoint and leaves Otsu fully usable. @@ -78,7 +80,7 @@ only option that needs no per-user download at all: ```bash export MATS_WEIGHTS_DIR=/project//mats_weights -mats fetch-weights --all # populates it once (from a data-transfer node) +# Pre-stage the verified checkpoint files in this directory. mats doctor # confirm it resolves ``` @@ -96,10 +98,12 @@ excluded from the default clone and from a bare `git lfs pull` by `.lfsconfig` place, not the 2.65 GB file. Pull it explicitly: ```bash -git lfs pull --include="weights/birefnet_leaf.pth" +git lfs pull -X "" -I "weights/birefnet_leaf.pth" ``` or `mats fetch-weights --only birefnet --source lfs`, or the setup page. +The empty `-X` value clears `.lfsconfig`'s exclusion for this invocation, and +`-I` limits the pull to the BiRefNet checkpoint. This exclusion exists because committing BiRefNet without it would force *every* `git clone` to download 2.65 GB and spend the repository's Git LFS @@ -111,8 +115,9 @@ If your Git LFS version predates the exclusion behavior (needs the `.lfsconfig` fetchexclude to be read from the repo index/HEAD during the initial clone — true for modern Git LFS), a clone could pull BiRefNet anyway. `GIT_LFS_SKIP_SMUDGE=1 git clone ...` is a guaranteed way to skip *all* LFS -content on clone if you want to be certain, then `git lfs pull --include=...` -each file you actually need. +content on clone if you want to be certain. Afterward, fetch RF-DETR with +`git lfs pull --exclude="weights/birefnet_leaf.pth"`; add BiRefNet later with +`git lfs pull -X "" -I "weights/birefnet_leaf.pth"` if needed. ## Manual / air-gapped diff --git a/src/mats/app/Home.py b/src/mats/app/Home.py index 6c508c4..dd71e91 100644 --- a/src/mats/app/Home.py +++ b/src/mats/app/Home.py @@ -1,6 +1,7 @@ import csv import tempfile import traceback +import uuid import zipfile from pathlib import Path @@ -20,8 +21,17 @@ ) from mats.app.folder_picker import FolderPickerError, choose_folder from mats.app.runtime_paths import current_python, display_path +from mats.dataset_export import pairs_from_manifest from mats.qr_runtime import qr_preflight_status, qr_runtime_status from mats.scaling import DEFAULT_RESULTS_UNIT, QR_TRACE_FIELDNAMES, RESULT_UNITS +from mats.mask_settings import ( + CLEAN_MARGIN_DEFAULT, + CLEAN_MARGIN_MAX, + CLEAN_SIZE_DEFAULT, + CLEAN_SIZE_MAX, + STRAY_GAP_DEFAULT, + STRAY_GAP_MAX, +) from mats.template_layout import ( TemplateLayoutError, build_template_layout, @@ -30,6 +40,14 @@ minimum_template_edge, round_to_increment, ) +from mats.thresholds import ( + CUSTOM_THRESHOLD_LEVEL, + THRESHOLD_LEVEL_OPTIONS, + THRESHOLD_LEVELS, + THRESHOLD_MAX, + THRESHOLD_MIN, + threshold_value_for, +) APP_DIR = Path(__file__).resolve().parent @@ -42,13 +60,44 @@ OUTPUT_DIR_KEY = "output_directory" FOLDER_PICKER_ERROR_KEY = "folder_picker_error" INPUT_SOURCE_KEY = "input_source" -SEGMENTATION_METHOD_KEY = "segmentation_method" +SEGMENT_THRESHOLD_KEY = "segment_threshold" +SEGMENT_BIREFNET_KEY = "segment_birefnet" +# Measurement methods in output order, with their Setup checkbox keys. +SEGMENTATION_METHOD_KEYS = {"threshold": SEGMENT_THRESHOLD_KEY, "birefnet": SEGMENT_BIREFNET_KEY} +SEGMENTATION_METHOD_LABELS = {"threshold": "Classic thresholding", "birefnet": "BiRefNet"} +RESULTS_METHOD_KEY = "results_method" THRESHOLD_LEVEL_KEY = "threshold_level" +THRESHOLD_CUSTOM_VALUE_KEY = "threshold_custom_value" +THRESHOLD_MARK_LABELS = {"low": "low", "medium": "med", "high": "high"} RESULTS_SCHEMA_KEY = "results_schema" RESULTS_UNIT_KEY = "results_unit" +MEASURE_PRE_CLEANUP_KEY = "measure_pre_cleanup" +STRAY_GAP_KEY = "stray_gap" +CLEAN_MARGIN_KEY = "clean_margin" +CLEAN_SIZE_KEY = "clean_size" WRITE_FAILURES_KEY = "write_failures" +EXPORT_TARGET_BOXES_KEY = "export_target_boxes" +EXPORT_MASKS_KEY = "export_cleaned_masks" +EXPORT_PRE_CLEANUP_KEY = "export_pre_cleanup" +EXPORT_OVERLAY_KEY = "export_overlay" +EXPORT_CUTOUT_KEY = "export_cutout" +EXPORT_AXES_KEY = "export_axes" WORKSPACE_TAB_KEY = "analysis_workspace_tab" PENDING_WORKSPACE_TAB_KEY = "pending_analysis_workspace_tab" +EXPORT_RUN_ID_KEY = "export_run_id" +EXPORT_METHODS_KEY = "export_selected_methods" +EXPORT_ZIP_NAME_KEY = "export_zip_name" +EXPORT_FILE_TYPES = ( + ("results_csv", "Results CSVs"), + ("results_metadata", "Measurement metadata"), + ("failure_log", "Failure logs"), + ("target_box", "Target boxes"), + ("mask", "Cleaned masks"), + ("pre_cleanup", "Pre-cleanup masks"), + ("overlay", "Overlays"), + ("cutout", "Cutouts"), + ("axes", "Measurement axes"), +) RESULTS_DASHBOARD_MAX_ROWS = 5_000 RESULTS_TABLE_MAX_ROWS = 1_000 RESULTS_UNIT_LABELS = { @@ -215,6 +264,41 @@ transition: none; } } +/* Preset cutoffs marked beneath the custom-threshold slider. */ +.mats-threshold-marks { + color: #5F6C66; + container-type: inline-size; + font-size: 0.72rem; + line-height: 1.15; + margin: -0.4rem 0.5rem 0; +} +.mats-threshold-track { + height: 2.1rem; + position: relative; +} +.mats-threshold-mark { + position: absolute; + text-align: center; + transform: translateX(-50%); + white-space: nowrap; +} +.mats-threshold-mark::before { + background: #9AB0A0; + content: ""; + display: block; + height: 0.35rem; + margin: 0 auto 0.1rem; + width: 1px; +} +/* In a narrow column the presets sit ~17 px apart: drop "med" to a second row. */ +@container (max-width: 300px) { + .mats-threshold-track { + height: 3.5rem; + } + .mats-threshold-mark-medium::before { + height: 1.75rem; + } +} """ @@ -257,8 +341,6 @@ def save_uploaded_images(uploaded_files, destination): # Guardrails so a few-hundred-image run cannot exhaust memory or overwhelm the page. ZIP_SIZE_WARN_BYTES = 2 * 1024 ** 3 # 2 GB -PREVIEW_HARD_MAX = 24 -PREVIEW_AUTO_HIDE_THRESHOLD = 50 PREVIEW_IMAGE_WIDTH = 280 LARGE_BATCH_THRESHOLD = 200 OVERLAY_TINT_COLOR = (255, 0, 255) # BGR magenta -- reads clearly against green foliage @@ -337,6 +419,11 @@ def gather_output_files(output_dir, results_path, include_overlay=False, include return files +def files_from_manifest(artifacts): + """Select only paths this run successfully wrote, preserving run order.""" + return [Path(item["path"]) for item in artifacts if Path(item["path"]).is_file()] + + def estimate_zip_inputs(files): total_bytes = 0 for path in files: @@ -377,6 +464,25 @@ def read_results_dataframe(results_path, modified_ns, limit): return pd.read_csv(results_path, nrows=limit) +def _scale_axes_by_sample(summary): + """Retain canonical calibration for per-specimen output adjustments.""" + axes_by_method = {} + for method, outcome in summary["by_method"].items(): + axes_by_sample = {} + for row in outcome["result_rows"]: + try: + axes = ( + float(row["px_per_cm_width"]), + float(row["px_per_cm_height"]), + ) + except (KeyError, TypeError, ValueError): + continue + if all(np.isfinite(value) and value > 0 for value in axes): + axes_by_sample[str(row["sample_id"])] = axes + axes_by_method[method] = axes_by_sample + return axes_by_method + + def normalize_measurements(frame, results_unit=DEFAULT_RESULTS_UNIT): """Return numeric, consistently named measurements from either CSV schema.""" if results_unit not in RESULT_UNITS: @@ -404,6 +510,37 @@ def normalize_measurements(frame, results_unit=DEFAULT_RESULTS_UNIT): return measurements.dropna(subset=("leaf_area", "width", "length")) +def _unit_symbols(run): + """Return the length and area unit labels for a run's results.""" + results_unit = run.get("results_unit", DEFAULT_RESULTS_UNIT) + if results_unit not in RESULT_UNITS: + results_unit = DEFAULT_RESULTS_UNIT + unit_symbol = RESULTS_UNIT_SYMBOLS[results_unit] + return unit_symbol, f"{unit_symbol}²" + + +def _measurement_columns(unit_symbol="cm", area_symbol="cm²"): + """(column, label, decimals) shared by the Analyze and Adjust tables; None is text.""" + return [ + ("sample_id", "Sample", None), + ("leaf_area", f"Leaf area ({area_symbol})", 2), + ("width", f"Leaf width ({unit_symbol})", 2), + ("length", f"Leaf length ({unit_symbol})", 2), + ("scale_aspect_ratio", "Scale axis ratio", 3), + ] + + +def _measurement_column_config(unit_symbol="cm", area_symbol="cm²"): + """Streamlit column config for the Analyze measurement table.""" + return { + column: ( + st.column_config.TextColumn(label, pinned=True) if decimals is None + else st.column_config.NumberColumn(label, format=f"%.{decimals}f") + ) + for column, label, decimals in _measurement_columns(unit_symbol, area_symbol) + } + + def summarize_measurements(measurements): """Return dashboard-ready robust summary statistics for leaf measurements.""" if measurements.empty: @@ -546,6 +683,22 @@ def _layout_conversion_label(layout): ) +def _threshold_marks_html(): + """Mark the low/medium/high preset cutoffs beneath the custom-threshold slider.""" + span = THRESHOLD_MAX - THRESHOLD_MIN + marks = "".join( + f'' + f"{THRESHOLD_MARK_LABELS[name]}
{value}
" + for name, value in THRESHOLD_LEVELS.items() + if value is not None + ) + return ( + '
' + f'
{marks}
' + ) + + def _choose_folder(widget_key, title): """Open a local desktop picker and copy the result into a path widget.""" try: @@ -561,16 +714,16 @@ def _choose_folder(widget_key, title): st.session_state[widget_key] = selected -def _open_diagnostics(): - """Open the detailed preflight tab on the next Streamlit rerun.""" - st.session_state[WORKSPACE_TAB_KEY] = "Diagnostics" - - def _switch_workspace_tab(tab_name): """Open a named workspace tab through its stateful Streamlit key.""" st.session_state[WORKSPACE_TAB_KEY] = tab_name +def _open_export(): + """Open the Export workspace from the sidebar.""" + st.session_state[WORKSPACE_TAB_KEY] = "Export" + + def apply_pending_workspace_tab(): """Apply deferred tab navigation before the tab widget is instantiated.""" tab_name = st.session_state.pop(PENDING_WORKSPACE_TAB_KEY, None) @@ -583,10 +736,10 @@ def render_workspace_navigation(): with st.container(key="workspace_navigation_intro"): st.markdown("**WORKSPACE NAVIGATION**") st.caption( - "Move between analysis setup, measurement results, and system diagnostics." + "Set up, analyze, adjust specimens, then export saved results." ) return st.tabs( - ["Analyze", "Results", "Diagnostics"], + ["Setup", "Analyze", "Adjust", "Export"], key=WORKSPACE_TAB_KEY, on_change="rerun", ) @@ -603,9 +756,9 @@ def render_workspace_context(config, execution_plan): st.badge( config["segmentation_label"], icon=( - ":material/contrast:" - if config["mask_method"] == "threshold" - else ":material/auto_awesome:" + ":material/auto_awesome:" + if config["needs_birefnet"] + else ":material/contrast:" ), color="green", ) @@ -674,7 +827,15 @@ def render_file_sidebar(): if folder_picker_error: st.warning(folder_picker_error, icon=":material/folder_off:") - st.caption("Scale, segmentation, and export settings are in Analyze.") + st.caption("Choose image outputs in Setup. Download completed results in Export.") + st.button( + "Go to Export", + key="open_export", + icon=":material/download:", + width="stretch", + disabled=not st.session_state.get("last_run"), + on_click=_open_export, + ) return input_source, uploaded_files, input_dir, output_dir @@ -817,33 +978,124 @@ def render_analysis_settings(lm): with segmentation_column.container(border=True): st.markdown("**2 · Segmentation**") - st.caption("Choose the mask method best suited to the image background.") - segmentation_label = st.radio( - "Select a segmentation method", - ["Classic thresholding (Otsu)", "BiRefNet"], - captions=[ - "Fast and dependable for clean, well-lit backgrounds.", - "Better on cluttered backgrounds; needs the optional local checkpoint " - "and benefits from a GPU.", - ], - key=SEGMENTATION_METHOD_KEY, + st.caption( + "Choose the mask method best suited to the image background. Check both " + "to measure every image with each method and compare them." + ) + use_threshold = st.checkbox( + SEGMENTATION_METHOD_LABELS["threshold"], + key=SEGMENT_THRESHOLD_KEY, persist_state="session", - width="stretch", ) - if segmentation_label == "Classic thresholding (Otsu)": - st.selectbox( + st.caption("Fast and dependable for clean, well-lit backgrounds.") + # The threshold level only affects Otsu, so it is hidden without it. + if use_threshold: + threshold_level = st.selectbox( "Threshold level", - list(lm.THRESHOLD_LEVELS.keys()), + list(THRESHOLD_LEVEL_OPTIONS), key=THRESHOLD_LEVEL_KEY, - help="auto = Otsu (adapts per image). low/medium/high = fixed 100/125/150.", + help=( + "auto = Otsu (adapts per image). custom = choose the cutoff yourself. " + "low/medium/high = fixed 100/125/150." + ), persist_state="page", ) - else: - st.caption("BiRefNet uses its calibrated confidence threshold automatically.") + if threshold_level == CUSTOM_THRESHOLD_LEVEL: + st.slider( + "Custom threshold", + min_value=THRESHOLD_MIN, + max_value=THRESHOLD_MAX, + step=1, + key=THRESHOLD_CUSTOM_VALUE_KEY, + help=( + "Grayscale cutoff (0 = black, 255 = white): pixels at or below it " + "count as leaf. Raise it to include paler leaf tissue; lower it to " + "exclude shadows or a dark background." + ), + persist_state="page", + ) + st.html(_threshold_marks_html()) + use_birefnet = st.checkbox( + SEGMENTATION_METHOD_LABELS["birefnet"], + key=SEGMENT_BIREFNET_KEY, + persist_state="session", + ) + st.caption( + "Better on cluttered backgrounds; needs the optional local checkpoint " + "and benefits from a GPU." + ) + if not use_threshold and not use_birefnet: + st.warning("Choose at least one segmentation method.", icon=":material/warning:") with st.container(border=True): st.markdown("**3 · Measurement Output**") - st.caption("Choose result units, the export schema, and whether to keep a per-image failure report.") + st.caption("Choose the measurement mask, result units, and CSV schema.") + measure_pre_cleanup = st.checkbox( + "Measure from pre-cleanup masks", + key=MEASURE_PRE_CLEANUP_KEY, + help=( + "Measure the raw binary segmentation before gap closing and hole " + "filling. The edge margin is cleared, the largest object is the " + "leaf, and pieces touching the margin or beyond the stray-piece " + "distance are dropped. Specks near the leaf still count toward " + "area, width, and length." + ), + persist_state="page", + ) + st.caption( + "Pre-cleanup settings for this run. Each specimen's explorer in Adjust can " + "change them for its own preview and, for Classic thresholding, overwrite." + ) + margin_column, gap_column, size_column = st.columns(3) + with margin_column: + st.number_input( + "Edge margin (% of box)", + min_value=0.0, + max_value=CLEAN_MARGIN_MAX, + step=0.25, + format="%.2f", + key=CLEAN_MARGIN_KEY, + disabled=not measure_pre_cleanup, + help=( + "Clears a band this percent of the target box's shorter side along " + "every edge, where the template's printed box outline lands. Lay " + "leaves inside the box so none of the leaf falls in the band; 0 " + "clears nothing." + ), + persist_state="page", + ) + with gap_column: + st.number_input( + "Stray-piece distance (× leaf size)", + min_value=0.0, + max_value=STRAY_GAP_MAX, + step=0.05, + format="%.2f", + key=STRAY_GAP_KEY, + disabled=not measure_pre_cleanup, + help=( + "Pieces farther from the leaf than this fraction of its " + "bounding-box diagonal are dropped; 0 keeps only the leaf. Raise " + "it when a leaf's parts lie apart, such as separated leaflets." + ), + persist_state="page", + ) + with size_column: + st.number_input( + "Clean size (px)", + min_value=0, + max_value=CLEAN_SIZE_MAX, + step=1, + key=CLEAN_SIZE_KEY, + disabled=not measure_pre_cleanup, + help=( + "After the edge margin and stray pieces are cleared, removes white " + "specks and fills enclosed holes whose inscribed radius is below " + "this many pixels. The leaf is always kept and nothing is " + "flash-filled; 0 turns it off." + ), + persist_state="page", + ) st.segmented_control( "Result units", options=list(RESULT_UNITS), @@ -856,42 +1108,56 @@ def render_analysis_settings(lm): ), persist_state="page", ) - schema_column, failures_column = st.columns((1.35, 0.65), vertical_alignment="bottom") - with schema_column: - st.selectbox( - "Results CSV schema", - [ - "Full research schema (area/width/length + px-per-cm)", - "Compact (sample_id, area, width, length)", - ], - key=RESULTS_SCHEMA_KEY, - help=( - "Full adds per-axis pixels-per-unit columns and scale_aspect_ratio " - "for QC. Compact is the trimmed UI export." - ), - persist_state="page", - ) - with failures_column: - st.toggle( - "Write Failure Log", - key=WRITE_FAILURES_KEY, - help="Per-image failure/warning report used by the failure-taxonomy analysis.", - persist_state="page", - ) + st.selectbox( + "Results CSV schema", + [ + "Full research schema (area/width/length + px-per-cm)", + "Compact (sample_id, area, width, length)", + ], + key=RESULTS_SCHEMA_KEY, + help="Full includes independent axis scales and their ratio for QC.", + persist_state="page", + ) + st.checkbox("Failure log · leaf_morpho_failures.csv", key=WRITE_FAILURES_KEY, + persist_state="page") st.caption( - "Results, visual summaries, image previews, and downloads appear in the " - "Results tab after a run." + "The results CSV is always written. Results appear in Analyze after a run." ) + with st.container(border=True): + st.markdown("**4 · Image output options**") + st.caption("Choose image files to write during this run. QC images use the selected measurement mask.") + st.checkbox("Target boxes · {id}_target_box.jpg", key=EXPORT_TARGET_BOXES_KEY, + persist_state="page") + st.checkbox("Cleaned leaf masks · {id}_mask.png", key=EXPORT_MASKS_KEY, + persist_state="page") + st.checkbox("Pre-cleanup masks · {id}_mask_precleanup_{method}.png", + key=EXPORT_PRE_CLEANUP_KEY, persist_state="page", + help="One per selected segmentation method. Pre-cleanup masks keep " + "holes, specks, and small objects.") + st.checkbox("Overlays · {id}_overlay.jpg", key=EXPORT_OVERLAY_KEY, persist_state="page") + st.checkbox("Cutouts · {id}_cutout.jpg", key=EXPORT_CUTOUT_KEY, persist_state="page") + st.checkbox("Measurement axes · {id}_measurement_axes.jpg", key=EXPORT_AXES_KEY, + persist_state="page") + if st.session_state[SEGMENT_THRESHOLD_KEY] and st.session_state[SEGMENT_BIREFNET_KEY]: + st.caption( + "With both methods selected, each method writes its own masks, overlays, " + "cutouts, axes, results CSV, and failure log, with names ending in " + "_threshold or _birefnet (for example {id}_mask_birefnet.png)." + ) + def current_analysis_config(lm): """Resolve the persisted controls into the values needed for one batch.""" - segmentation_label = st.session_state[SEGMENTATION_METHOD_KEY] - mask_method = ( - "threshold" - if segmentation_label == "Classic thresholding (Otsu)" - else "birefnet" + mask_methods = tuple( + method for method, key in SEGMENTATION_METHOD_KEYS.items() if st.session_state[key] ) + if len(mask_methods) == 1: + segmentation_label = SEGMENTATION_METHOD_LABELS[mask_methods[0]] + elif mask_methods: + segmentation_label = "Classic thresholding + BiRefNet" + else: + segmentation_label = "No segmentation method" threshold_level = st.session_state[THRESHOLD_LEVEL_KEY] use_qr = st.session_state[QR_MODE_KEY] use_legacy_dimensions = st.session_state[USE_LEGACY_DIMENSIONS_KEY] @@ -917,15 +1183,35 @@ def current_analysis_config(lm): return { "segmentation_label": segmentation_label, - "mask_method": mask_method, - "threshold_value": ( - lm.THRESHOLD_LEVELS[threshold_level] - if mask_method == "threshold" - else lm.BIREFNET_THRESHOLD + "mask_methods": mask_methods, + "needs_birefnet": "birefnet" in mask_methods, + "threshold_value": threshold_value_for( + threshold_level, st.session_state[THRESHOLD_CUSTOM_VALUE_KEY] ), "compact_csv": st.session_state[RESULTS_SCHEMA_KEY].startswith("Compact"), "results_unit": st.session_state[RESULTS_UNIT_KEY], + "measurement_source": ( + "pre-cleanup" if st.session_state[MEASURE_PRE_CLEANUP_KEY] else "cleaned" + ), + "stray_gap": float(st.session_state[STRAY_GAP_KEY]), + "clean_margin": float(st.session_state[CLEAN_MARGIN_KEY]), + # The clean size shapes only pre-cleanup measurements; the pipeline + # rejects it otherwise, so a grayed-out value never reaches a run. + "clean_size": ( + int(st.session_state[CLEAN_SIZE_KEY]) + if st.session_state[MEASURE_PRE_CLEANUP_KEY] else CLEAN_SIZE_DEFAULT + ), "write_failures": st.session_state[WRITE_FAILURES_KEY], + "export_options": { + "target_boxes": st.session_state[EXPORT_TARGET_BOXES_KEY], + "cleaned_masks": st.session_state[EXPORT_MASKS_KEY], + "pre_cleanup_methods": ( + mask_methods if st.session_state[EXPORT_PRE_CLEANUP_KEY] else () + ), + "overlay": st.session_state[EXPORT_OVERLAY_KEY], + "cutout": st.session_state[EXPORT_CUTOUT_KEY], + "axes": st.session_state[EXPORT_AXES_KEY], + }, "use_qr": use_qr, "use_legacy_dimensions": use_legacy_dimensions, "sheet_dimensions": sheet_dimensions, @@ -971,7 +1257,20 @@ def build_preflight_checks( None, )) - if config["mask_method"] == "birefnet": + checks.append(( + "Segmentation method", + "success" if config["mask_methods"] else "error", + not config["mask_methods"], + ( + config["segmentation_label"] + if config["mask_methods"] + else "No segmentation method is selected. Check Otsu, BiRefNet, or both in Setup." + ), + None, + )) + + needs_birefnet = config["needs_birefnet"] + if needs_birefnet: birefnet_runtime = birefnet_runtime_status() runtime_level = "success" if birefnet_runtime.ready else "error" checks.append(( @@ -987,7 +1286,7 @@ def build_preflight_checks( checks.append(( "BiRefNet checkpoint", birefnet_level, - config["mask_method"] == "birefnet" and birefnet_status.state != "ready", + needs_birefnet and birefnet_status.state != "ready", ( birefnet_status.detail if birefnet_status.state != "ready" @@ -1054,7 +1353,7 @@ def build_preflight_checks( None, )) - if config["mask_method"] == "birefnet" and execution_plan.execution_device == "cpu": + if needs_birefnet and execution_plan.execution_device == "cpu": birefnet_cpu_detail = ( "GPU is available but disabled by an advanced CPU-only override." if compute_settings.accelerator_available @@ -1072,7 +1371,7 @@ def build_preflight_checks( checks.append(( "BiRefNet compute", device_report.severity, - config["mask_method"] == "birefnet" and device_report.severity == "error", + needs_birefnet and device_report.severity == "error", device_report.detail, None, )) @@ -1187,7 +1486,7 @@ def render_preflight_summary(checks): blockers = [check for check in checks if check[2]] attention = [check for check in checks if check[1] != "success" or check[2]] with st.container(border=True): - st.markdown("**4 · Preflight**") + st.markdown("**5 · Preflight**") if blockers: st.badge( f"{len(blockers)} blocking item(s)", @@ -1197,11 +1496,10 @@ def render_preflight_summary(checks): st.caption( "Needs attention: " + ", ".join(check[0] for check in blockers) + "." ) - st.button( - "Open Diagnostics", - key="open_diagnostics", + st.page_link( + "pages/0_Diagnostics.py", + label="Open Diagnostics", icon=":material/troubleshoot:", - on_click=_open_diagnostics, width="content", ) elif attention: @@ -1234,7 +1532,7 @@ def render_compute_status(compute_settings, execution_plan, config): color="green" if execution_plan.uses_gpu else "orange", ) st.caption(execution_plan.detail) - if config["mask_method"] == "birefnet" and not execution_plan.uses_gpu: + if config["needs_birefnet"] and not execution_plan.uses_gpu: gpu_notice = ( "GPU available but disabled. BiRefNet is running on CPU because an " "advanced CPU-only override is active." @@ -1379,12 +1677,11 @@ def execute_leaf_analysis( ): """Run the batch and return its summary plus output artifacts for this session.""" output_path.mkdir(parents=True, exist_ok=True) - run_sample_ids = [] - succeeded_so_far = 0 + preview_cache = tempfile.TemporaryDirectory(prefix="mats_preview_") if worker_risk.requires_break_glass: st.session_state.pop("break_glass_available_workers", None) if ( - config["mask_method"] == "birefnet" + config["needs_birefnet"] and execution_plan.execution_device == "cpu" and execution_plan.workers > 1 ): @@ -1399,17 +1696,6 @@ def execute_leaf_analysis( counts_box = st.empty() def update_progress(status): - nonlocal succeeded_so_far - if status["succeeded"] > succeeded_so_far: - current_image = status["current_image"] - sample_id = ( - lm.target_box_sample_id(current_image) - if lm.is_target_box_image(current_image) - else Path(current_image).stem - ) - run_sample_ids.append(sample_id) - succeeded_so_far = status["succeeded"] - total = max(status["total"], 1) progress_bar.progress(status["processed"] / total) counts_box.write( @@ -1426,7 +1712,7 @@ def update_progress(status): str(results_path), template_dimensions=config["template_dimensions"], output_mode="masks", - mask_method=config["mask_method"], + mask_method=config["mask_methods"], threshold_value=config["threshold_value"], workers=int(execution_plan.workers), execution_device=execution_plan.execution_device, @@ -1437,13 +1723,32 @@ def update_progress(status): write_failures=config["write_failures"], compact_csv=config["compact_csv"], results_unit=config["results_unit"], - save_measurement_axes=False, + save_measurement_axes=config["export_options"]["axes"], + export_options={**config["export_options"], "preview_dir": preview_cache.name}, + measurement_source=config["measurement_source"], + stray_gap=config["stray_gap"], + clean_margin=config["clean_margin"], + clean_size=config["clean_size"], ) except ValueError as exc: + preview_cache.cleanup() st.error(f"Run prevented by worker safety checks: {exc}") - return None, [] - - return summary, collect_output_pairs(output_path, run_sample_ids) + return None, {} + + input_images = image_paths if input_source == "Local folder" else () + run_pairs = { + method: pairs_from_manifest( + summary["artifacts"], input_images, method, + measurement_source=summary["measurement_source"], + preview_artifacts=summary["preview_artifacts"], + ) + for method in summary["methods"] + } + previous_cache = st.session_state.get("preview_cache") + st.session_state["preview_cache"] = preview_cache + if previous_cache is not None: + previous_cache.cleanup() + return summary, run_pairs def main(): @@ -1454,7 +1759,7 @@ def main(): ) branding.apply_logo() st.html(_WORKBENCH_STYLES) - st.session_state.setdefault("viewer_pairs", []) + st.session_state.setdefault("viewer_pairs", {}) st.session_state.setdefault(UNIT_KEY, "in") st.session_state.setdefault(PREVIOUS_UNIT_KEY, "in") st.session_state.setdefault(WIDTH_KEY, 12.0) @@ -1463,14 +1768,29 @@ def main(): st.session_state.setdefault(QR_MODE_KEY, False) st.session_state.setdefault(USE_LEGACY_DIMENSIONS_KEY, False) st.session_state.setdefault(LEGACY_DIMENSIONS_TEXT_KEY, "10x9.5in") - st.session_state.setdefault(SEGMENTATION_METHOD_KEY, "Classic thresholding (Otsu)") + st.session_state.setdefault(SEGMENT_THRESHOLD_KEY, True) + st.session_state.setdefault(SEGMENT_BIREFNET_KEY, False) st.session_state.setdefault(THRESHOLD_LEVEL_KEY, "auto") + st.session_state.setdefault(THRESHOLD_CUSTOM_VALUE_KEY, THRESHOLD_LEVELS["medium"]) + # Streamlit discards a widget's value once it stops rendering. Re-saving + # keeps the threshold choice while Otsu is unchecked or the level isn't custom. + for key in (THRESHOLD_LEVEL_KEY, THRESHOLD_CUSTOM_VALUE_KEY): + st.session_state[key] = st.session_state[key] st.session_state.setdefault( RESULTS_SCHEMA_KEY, "Full research schema (area/width/length + px-per-cm)", ) st.session_state.setdefault(RESULTS_UNIT_KEY, DEFAULT_RESULTS_UNIT) + st.session_state.setdefault(MEASURE_PRE_CLEANUP_KEY, False) + st.session_state.setdefault(STRAY_GAP_KEY, STRAY_GAP_DEFAULT) + st.session_state.setdefault(CLEAN_MARGIN_KEY, CLEAN_MARGIN_DEFAULT) + st.session_state.setdefault(CLEAN_SIZE_KEY, CLEAN_SIZE_DEFAULT) st.session_state.setdefault(WRITE_FAILURES_KEY, True) + for key, default in ((EXPORT_TARGET_BOXES_KEY, True), (EXPORT_MASKS_KEY, True), + (EXPORT_PRE_CLEANUP_KEY, False), + (EXPORT_OVERLAY_KEY, False), (EXPORT_CUTOUT_KEY, False), + (EXPORT_AXES_KEY, False)): + st.session_state.setdefault(key, default) st.session_state.setdefault(INPUT_DIR_KEY, str(DEFAULT_INPUT_DIR)) st.session_state.setdefault(OUTPUT_DIR_KEY, str(DEFAULT_OUTPUT_DIR)) st.session_state.setdefault(INPUT_SOURCE_KEY, "Local folder") @@ -1484,8 +1804,10 @@ def main(): with st.expander("About This Workspace", icon=":material/info:"): st.markdown( "Choose input and output locations in the sidebar, then configure scale, " - "segmentation, and export options in Analyze. Results contains measurement " - "visuals and downloads; Diagnostics contains compute and preflight details. " + "segmentation, and output options in Setup. Analyze contains the launch and " + "measurement visuals; Adjust contains specimen previews and saving; " + "Export contains saved files and downloads. Diagnostics in the sidebar " + "contains compute and preflight details. " "Template Creator, BiRefNet Setup, CPU Options, Robust QR Setup, and Help " "remain available in the app navigation." ) @@ -1517,10 +1839,10 @@ def main(): input_source, uploaded_files, input_dir, output_dir = render_file_sidebar() apply_pending_workspace_tab() - analyze_tab, results_tab, diagnostics_tab = render_workspace_navigation() + setup_tab, analyze_tab, adjust_tab, export_tab = render_workspace_navigation() - if analyze_tab.open: - with analyze_tab: + if setup_tab.open: + with setup_tab: render_analysis_settings(lm) config = current_analysis_config(lm) @@ -1529,7 +1851,7 @@ def main(): birefnet_parallel_is_unlocked = birefnet_parallel_unlocked(available_workers) execution_plan = resolve_execution_plan( compute_settings, - config["mask_method"], + "birefnet" if config["needs_birefnet"] else "threshold", birefnet_parallel_allowed=birefnet_parallel_is_unlocked, ) worker_risk = lm.worker_risk_report(execution_plan.workers, available_workers) @@ -1547,8 +1869,11 @@ def main(): worker_risk, break_glass_is_unlocked, ) + st.session_state["diagnostics_context"] = config + st.session_state["diagnostics_inputs"] = ( + input_source, uploaded_files, input_dir + ) - run_clicked = False if analyze_tab.open: with analyze_tab: render_workspace_context(config, execution_plan) @@ -1560,49 +1885,54 @@ def main(): execution_plan, output_path, ) - elif results_tab.open: - with results_tab: + if run_clicked: + summary, run_pairs = execute_leaf_analysis( + lm, config, image_paths, uploaded_files, input_source, + output_path, results_path, execution_plan, worker_risk, + break_glass_is_unlocked, birefnet_parallel_is_unlocked, + ) + if summary is not None: + st.session_state["last_run"] = { + "run_id": uuid.uuid4().hex, + "succeeded": summary["succeeded"], + "failed": summary["failed"], + "total": summary["total"], + "workers": summary["workers"], + "worker_reason": summary["worker_reason"], + "execution_device": summary["execution_device"], + "output_path": str(output_path), + "mask_methods": summary["methods"], + "by_method": { + method: { + "succeeded": outcome["succeeded"], + "failed": outcome["failed"], + "results_path": outcome["results_path"], + "failure_rows": outcome["failure_rows"][:200], + "failure_overflow": max(0, len(outcome["failure_rows"]) - 200), + } + for method, outcome in summary["by_method"].items() + }, + "pre_cleanup_methods": config["export_options"]["pre_cleanup_methods"], + "results_unit": config["results_unit"], + "measurement_source": summary["measurement_source"], + "stray_gap": summary["stray_gap"], + "clean_margin": summary["clean_margin"], + "clean_size": summary["clean_size"], + "threshold_value": config["threshold_value"], + "scale_axes_by_sample": _scale_axes_by_sample(summary), + "artifacts": summary["artifacts"], + "export_options": dict(config["export_options"]), + } + st.session_state["viewer_pairs"] = run_pairs + _clear_export_zip_cache() + st.rerun() render_results(lm) - elif diagnostics_tab.open: - with diagnostics_tab: - render_diagnostics(compute_settings, execution_plan, config, checks) - - if run_clicked: - summary, run_pairs = execute_leaf_analysis( - lm, - config, - image_paths, - uploaded_files, - input_source, - output_path, - results_path, - execution_plan, - worker_risk, - break_glass_is_unlocked, - birefnet_parallel_is_unlocked, - ) - if summary is not None: - st.session_state["last_run"] = { - "succeeded": summary["succeeded"], - "failed": summary["failed"], - "total": summary["total"], - "workers": summary["workers"], - "worker_reason": summary["worker_reason"], - "execution_device": summary["execution_device"], - "failure_rows": summary["failure_rows"][:200], - "failure_overflow": max(0, len(summary["failure_rows"]) - 200), - "results_path": str(results_path), - "output_path": str(output_path), - "mask_method": config["mask_method"], - "results_unit": config["results_unit"], - } - st.session_state["viewer_pairs"] = merge_viewer_pairs( - st.session_state["viewer_pairs"], - run_pairs, - ) - st.session_state.pop("export_zip_path", None) - st.session_state[PENDING_WORKSPACE_TAB_KEY] = "Results" - st.rerun() + elif export_tab.open: + with export_tab: + render_export_section() + elif adjust_tab.open: + with adjust_tab: + render_adjust_workspace() def render_results(lm): @@ -1624,19 +1954,52 @@ def render_results(lm): ) return - results_path = Path(run["results_path"]) - output_path = Path(run["output_path"]) - results_unit = run.get("results_unit", DEFAULT_RESULTS_UNIT) - if results_unit not in RESULT_UNITS: - results_unit = DEFAULT_RESULTS_UNIT - unit_symbol = RESULTS_UNIT_SYMBOLS[results_unit] - area_symbol = f"{unit_symbol}²" + unit_symbol, area_symbol = _unit_symbols(run) + methods = tuple(run["mask_methods"]) st.subheader("Results", anchor=False) - st.success( - f"Completed {run['succeeded']} measurement(s); {run['failed']} failed " - f"(of {run['total']} input image(s)).", - icon=":material/check_circle:", + st.caption( + "Measurements from " + + ("pre-cleanup masks" if run.get("measurement_source") == "pre-cleanup" + else "cleaned masks") + + "." ) + if len(methods) > 1: + per_method = "; ".join( + f"{SEGMENTATION_METHOD_LABELS[method]}: {run['by_method'][method]['succeeded']} " + f"measured, {run['by_method'][method]['failed']} failed" + for method in methods + ) + st.success( + f"Completed {run['total']} input image(s) with each method. {per_method}.", + icon=":material/check_circle:", + ) + if st.session_state.get(RESULTS_METHOD_KEY) not in methods: + st.session_state[RESULTS_METHOD_KEY] = methods[0] + with st.container(border=True): + st.markdown("**Compare measurement methods**") + method = st.segmented_control( + "Measurement table and specimen view", + options=list(methods), + format_func=SEGMENTATION_METHOD_LABELS.get, + required=True, + key=RESULTS_METHOD_KEY, + help="Each method has its own measurements, masks, and failure log.", + ) + st.caption( + "The summary, charts, measurement table, and selected specimen below " + "follow this choice." + ) + else: + method = methods[0] + outcome = run["by_method"][method] + results_path = Path(outcome["results_path"]) + output_pairs = st.session_state.get("viewer_pairs", {}).get(method, []) + if len(methods) == 1: + st.success( + f"Completed {outcome['succeeded']} measurement(s); {outcome['failed']} failed " + f"(of {run['total']} input image(s)).", + icon=":material/check_circle:", + ) device_label = { "cpu": "CPU only", "hybrid": "GPU RF-DETR + CPU Otsu fan-out", @@ -1645,21 +2008,19 @@ def render_results(lm): f"Workers used: {run['workers']} ({run['worker_reason']}); compute: {device_label}." ) - if run["failure_rows"]: - with st.expander(f"Processing warnings and failures ({run['failed']})"): - for row in run["failure_rows"]: + if outcome["failure_rows"]: + with st.expander(f"Processing warnings and failures ({outcome['failed']})"): + for row in outcome["failure_rows"]: st.write(f"{row['sample_id']}: {row['status']}") - if run["failure_overflow"]: - st.write(f"...and {run['failure_overflow']} more (see the failures CSV).") + if outcome["failure_overflow"]: + st.write(f"...and {outcome['failure_overflow']} more (see the failures CSV).") if not results_path.is_file(): st.warning( - f"The results CSV is not available at {display_path(results_path)}. The output preview " - "and exports may still be available below.", + f"The results CSV is not available at {display_path(results_path)}. " + "Other saved files may still be available in Export.", icon=":material/folder_off:", ) - render_output_preview(st.session_state.get("viewer_pairs", [])) - render_export_section(results_path, output_path, unit_symbol) return total_rows = count_csv_rows(results_path) @@ -1673,7 +2034,7 @@ def render_results(lm): st.error(f"Could not read the results CSV: {exc}") return - measurements = normalize_measurements(frame, results_unit) + measurements = normalize_measurements(frame, run.get("results_unit", DEFAULT_RESULTS_UNIT)) if measurements.empty: st.warning( "The CSV contains no complete area, width, and length measurements to visualize.", @@ -1684,8 +2045,8 @@ def render_results(lm): with st.container(horizontal=True): st.metric( "Successful measurements", - f"{run['succeeded']:,}", - delta=f"{run['failed']} failed", + f"{outcome['succeeded']:,}", + delta=f"{outcome['failed']} failed", delta_color="inverse", border=True, ) @@ -1787,6 +2148,67 @@ def render_results(lm): "horizontal and vertical axes." ) + table_columns = ["sample_id", "leaf_area", "width", "length"] + if "scale_aspect_ratio" in measurements: + table_columns.append("scale_aspect_ratio") + table_data = measurements.loc[:, table_columns].head(RESULTS_TABLE_MAX_ROWS).copy() + selected_key = f"selected_specimen_{run.get('run_id', '')}_{method}" + sample_ids = table_data["sample_id"].astype(str).tolist() + previous_sample_id = st.session_state.get(selected_key) + selection_default = ( + {"selection": {"rows": [sample_ids.index(previous_sample_id)]}} + if previous_sample_id in sample_ids else None + ) + with st.container(border=True): + st.markdown("**Measurement Table**") + st.caption( + f"{SEGMENTATION_METHOD_LABELS[method]} measurements. Select a row to " + "inspect its matching target-box image and segmentation mask." + ) + table_event = st.dataframe( + table_data, + column_config=_measurement_column_config(unit_symbol, area_symbol), + hide_index=True, + height=440, + width="stretch", + key=f"measurement_table_{run.get('run_id', '')}_{method}", + on_select="rerun", + selection_mode="single-row", + selection_default=selection_default, + ) + if total_rows > len(table_data): + st.caption( + f"Showing the first {len(table_data):,} of {total_rows:,} rows. " + "Download the CSV for the full dataset." + ) + + selected_row = None + if table_event.selection.rows: + selected_row = table_data.iloc[table_event.selection.rows[0]] + selected_sample_id = ( + str(selected_row["sample_id"]) if selected_row is not None else None + ) + if selected_sample_id is None: + st.session_state.pop(selected_key, None) + else: + st.session_state[selected_key] = selected_sample_id + render_specimen_inspector( + selected_sample_id, + selected_row, + output_pairs, + unit_symbol, + area_symbol, + run=run, + method=method, + ) + if selected_sample_id is not None: + st.button( + "Adjust selected specimen", + icon=":material/tune:", + on_click=_switch_workspace_tab, + args=("Adjust",), + ) + qr_trace_columns = [column for column in QR_TRACE_FIELDNAMES if column in frame.columns] if qr_trace_columns: with st.container(border=True): @@ -1810,70 +2232,162 @@ def render_results(lm): height=260, ) - table_columns = ["sample_id", "leaf_area", "width", "length"] - if "scale_aspect_ratio" in measurements: - table_columns.append("scale_aspect_ratio") - table_data = measurements.loc[:, table_columns].head(RESULTS_TABLE_MAX_ROWS).copy() - output_pairs = st.session_state.get("viewer_pairs", []) - table_column, inspector_column = st.columns((1.15, 0.85), vertical_alignment="top") - with table_column.container(border=True): - st.markdown("**Measurement Table**") - st.caption("Select a row to inspect the matching target-box image and segmentation mask.") - table_event = st.dataframe( - table_data, - column_config={ - "sample_id": st.column_config.TextColumn("Sample", pinned=True), - "leaf_area": st.column_config.NumberColumn( - f"Leaf area ({area_symbol})", - format="%.2f", - ), - "width": st.column_config.NumberColumn( - f"Leaf width ({unit_symbol})", format="%.2f" - ), - "length": st.column_config.NumberColumn( - f"Leaf length ({unit_symbol})", format="%.2f" - ), - "scale_aspect_ratio": st.column_config.NumberColumn( - "Scale axis ratio", - format="%.3f", - ), - }, - hide_index=True, - height=440, - key="measurement_table", - on_select="rerun", - selection_mode="single-row", - ) - if total_rows > len(table_data): - st.caption( - f"Showing the first {len(table_data):,} of {total_rows:,} rows. " - "Download the CSV for the full dataset." - ) - selected_row = None - if table_event.selection.rows: - selected_row = table_data.iloc[table_event.selection.rows[0]] - elif not table_data.empty: - selected_row = table_data.iloc[0] - selected_sample_id = ( - str(selected_row["sample_id"]) if selected_row is not None else None +def render_adjust_workspace(): + """Inspect one specimen and optionally save its settings to marked peers.""" + run = st.session_state.get("last_run") + st.subheader("Adjust", anchor=False) + if not run: + st.info("Run an analysis to adjust its specimen masks and measurements.") + st.button("Open Analyze", on_click=_switch_workspace_tab, args=("Analyze",)) + return + methods = tuple(run["mask_methods"]) + if st.session_state.get(RESULTS_METHOD_KEY) not in methods: + st.session_state[RESULTS_METHOD_KEY] = methods[0] + if len(methods) > 1: + method = st.segmented_control( + "Segmentation method", list(methods), + format_func=SEGMENTATION_METHOD_LABELS.get, required=True, + key=RESULTS_METHOD_KEY, + ) + else: + method = methods[0] + st.caption(SEGMENTATION_METHOD_LABELS[method]) + results_path = Path(run["by_method"][method]["results_path"]) + if not results_path.is_file(): + st.warning("This method's results CSV is unavailable.") + return + try: + frame = read_results_dataframe( + str(results_path), results_path.stat().st_mtime_ns, None ) - with inspector_column: - render_specimen_inspector( - selected_sample_id, - selected_row, - output_pairs, - unit_symbol, - area_symbol, + except (OSError, pd.errors.EmptyDataError, pd.errors.ParserError, ValueError) as exc: + st.error(f"Could not read results: {exc}") + return + measurements = normalize_measurements(frame, run.get("results_unit", DEFAULT_RESULTS_UNIT)) + pairs = st.session_state.get("viewer_pairs", {}).get(method, []) + pairs_by_id = {str(pair["sample_id"]): pair for pair in pairs} + choices = [ + str(sample_id) for sample_id in measurements["sample_id"] + if str(sample_id) in pairs_by_id + and (method != "threshold" or pairs_by_id[str(sample_id)].get("target_box")) + ] + if not choices: + st.info("No measured specimens with previews are available in this session.") + return + unit_symbol, area_symbol = _unit_symbols(run) + sample_id, marked = render_adjust_browser( + choices, run.get("run_id", ""), method, measurements=measurements, + unit_symbol=unit_symbol, area_symbol=area_symbol, + ) + if method != "threshold": + st.caption("BiRefNet cleanup controls currently preview changes only.") + render_adjust_controls( + pairs_by_id[sample_id], run, method, + marked_pairs=[pairs_by_id[item] for item in marked], + ) + + +def _record_table_view(table_key, selected_key): + sample_id = _component_value(table_key, "view") + if sample_id is not None: + st.session_state[selected_key] = str(sample_id) + + +def _record_table_mark(table_key, marks_key): + """Apply one Marked for Adjustment tick; setting, not toggling, keeps repeats harmless.""" + change = _component_value(table_key, "mark") + if not change: + return + marked = set(st.session_state.get(marks_key, [])) + if change["marked"]: + marked.add(str(change["id"])) + else: + marked.discard(str(change["id"])) + st.session_state[marks_key] = sorted(marked) + + +def _set_marked_specimens(key, sample_ids): + st.session_state[key] = list(sample_ids) + + +def render_adjust_browser( + choices, run_id, method, *, measurements=None, unit_symbol="cm", area_symbol="cm²", +): + """Search specimens in a measurement table; marks persist across searches.""" + scope = f"{run_id}_{method}" + selected_key = f"selected_specimen_{scope}" + marks_key = f"adjust_marked_{scope}" + search_key = f"adjust_search_{scope}" + if st.session_state.get(selected_key) not in choices: + st.session_state[selected_key] = choices[0] + valid = set(choices) + marked = set(st.session_state.get(marks_key, [])) & valid + st.session_state[marks_key] = sorted(marked) + query = st.text_input("Search specimen names", key=search_key).strip().casefold() + filtered = [sample_id for sample_id in choices if query in sample_id.casefold()] + if method == "threshold": + with st.container(horizontal=True): + st.button( + "Mark all matching" if query else "Mark all", + key=f"adjust_mark_all_{scope}", disabled=not filtered, + on_click=_set_marked_specimens, + args=(marks_key, sorted(marked | set(filtered))), ) - render_export_section(results_path, output_path, unit_symbol) + st.button( + "Clear marks", key=f"adjust_clear_marks_{scope}", disabled=not marked, + on_click=_set_marked_specimens, args=(marks_key, []), + ) + st.caption(f"{len(marked):,} marked across all searches.") + if not filtered: + st.info("No specimen names match this search. Clear the search to see all specimens.") + return st.session_state[selected_key], sorted(marked) + from mats.app.specimen_table import show_specimen_table + + if measurements is None: + table = pd.DataFrame({"sample_id": filtered}) + else: + table = measurements.assign(sample_id=measurements["sample_id"].astype(str)) + table = table.loc[table["sample_id"].isin(set(filtered))] + columns = [ + {"key": column, "label": label, "decimals": decimals} + for column, label, decimals in _measurement_columns(unit_symbol, area_symbol) + if column in table + ] + table = table.loc[:, [column["key"] for column in columns]] + rows = table.astype(object).where(table.notna(), None).to_dict("records") + markable = method == "threshold" + st.caption( + f"{len(filtered):,} matching specimens. Click a row to view it" + + ("; tick Marked for Adjustment to include it in the marked overwrite." if markable else ".") + ) + table_key = f"adjust_table_{scope}" + show_specimen_table( + key=table_key, columns=columns, rows=rows, + view=st.session_state[selected_key], marked=sorted(marked), markable=markable, + on_view_change=lambda: _record_table_view(table_key, selected_key), + on_mark_change=lambda: _record_table_mark(table_key, marks_key), + ) + return st.session_state[selected_key], sorted(marked) - with st.expander("Browse Output Previews", icon=":material/photo_library:"): - render_output_preview(st.session_state.get("viewer_pairs", [])) + +def render_adjust_controls(pair, run, method, *, marked_pairs=None): + """Render only adjustment controls and the live mask/color preview.""" + st.caption(f"Sample: {pair['sample_id']}") + with st.container(border=True): + if method == "threshold" and pair.get("target_box"): + render_threshold_explorer(pair, run, marked_pairs=marked_pairs) + elif method == "birefnet" and _raw_mask_path(pair): + render_mask_explorer(pair, run, method) + else: + st.info("The original segmentation mask is unavailable for adjustment.") -def render_specimen_inspector(sample_id, measurement, pairs, unit_symbol="cm", area_symbol="cm²"): - """Render the selected specimen's measurements beside its generated artifacts.""" +def render_specimen_inspector( + sample_id, measurement, pairs, unit_symbol="cm", area_symbol="cm²", + run=None, method=None, +): + """Render the selected specimen's measurements and artifacts below the table.""" with st.container(border=True): st.markdown("**Selected Specimen**") if sample_id is None: @@ -1900,7 +2414,404 @@ def render_specimen_inspector(sample_id, measurement, pairs, unit_symbol="cm", a render_output_pair(pair, width="stretch") -def render_output_pair(pair, width=PREVIEW_IMAGE_WIDTH): + +def _component_value(component_key, name): + component_state = st.session_state.get(component_key) + value = getattr(component_state, name, None) + if value is None and isinstance(component_state, dict): + value = component_state.get(name) + return value + + +def _record_threshold_preview(component_key, state_key): + cutoff = _component_value(component_key, "cutoff") + if cutoff is not None and 0 <= int(cutoff) <= THRESHOLD_MAX: + st.session_state.setdefault("threshold_preview_cutoffs", {})[state_key] = int(cutoff) + + +def _record_clean_radius(component_key, state_key): + from mats.mask_cleanup import CLEAN_RADIUS_MAX + + radius = _component_value(component_key, "clean_radius") + if radius is not None and 0 <= int(radius) <= CLEAN_RADIUS_MAX: + st.session_state.setdefault("clean_preview_radii", {})[state_key] = int(radius) + + +CLEAN_SIZE_HELP = ( + "0 turns Clean image off. Above 0, Clean image replaces MATS cleanup: the edge " + "margin is cleared, pieces touching it or far from the leaf are removed, and " + "white specks and enclosed black holes with an inscribed radius below the " + "clean size (px) are removed or filled. The leaf is always kept and nothing is " + "flash-filled." +) +CLEAN_SIZE_SAVE_HELP = CLEAN_SIZE_HELP + " Overwrite saves and measures the mask shown." +CLEAN_SIZE_PREVIEW_HELP = CLEAN_SIZE_HELP + " BiRefNet specimens preview it only." +CLEAN_SIZE_ON_CAPTION = ( + "Clean size is above 0, so this preview shows Clean image instead of MATS cleanup." +) + + +def _clean_radius(run, method, sample_id): + """This specimen's clean size: slider edits, else saved, else the run's; 0 is off.""" + saved = run.get("threshold_adjustments", {}).get(sample_id, {}) if method == "threshold" else {} + state_key = f"{run.get('run_id', '')}:{method}:{sample_id}" + return st.session_state.setdefault("clean_preview_radii", {}).get( + state_key, saved.get("clean_size", run.get("clean_size", CLEAN_SIZE_DEFAULT)) + ) + + +REMOVE_FILL_RAW_NOTE = "This run measured raw masks, so hole filling is already off." + + +def _record_remove_fill(component_key, state_key): + remove_fill = _component_value(component_key, "remove_fill") + if remove_fill is not None: + st.session_state.setdefault("remove_fill_previews", {})[state_key] = bool(remove_fill) + + +def _remove_fill(run, state_key): + """Whether Remove flashfill is on for this specimen; only cleaned runs fill holes.""" + if run.get("measurement_source", "cleaned") != "cleaned": + return False + return st.session_state.setdefault("remove_fill_previews", {}).get(state_key, False) + + +def _remove_fill_note(run): + """Why Remove flashfill is unavailable for this run, or None when it is available.""" + return None if run.get("measurement_source", "cleaned") == "cleaned" else REMOVE_FILL_RAW_NOTE + + +def _run_stray_gap(run): + """The stray-piece gap this run measured with; the default for older runs.""" + return run.get("stray_gap", STRAY_GAP_DEFAULT) + + +def _run_clean_margin(run): + """The edge margin this run measured with; the default for older runs.""" + return run.get("clean_margin", CLEAN_MARGIN_DEFAULT) + + +def _specimen_cleanup_keys(run, method, sample_id): + base = f"{run.get('run_id', '')}:{method}:{sample_id}" + return f"specimen_clean_margin_{base}", f"specimen_stray_gap_{base}" + + +def _specimen_cleanup(run, method, sample_id): + """This specimen's edge margin and stray gap: explorer edits, else saved, else the run's.""" + saved = run.get("threshold_adjustments", {}).get(sample_id, {}) if method == "threshold" else {} + margin_key, gap_key = _specimen_cleanup_keys(run, method, sample_id) + margin = st.session_state.get(margin_key, saved.get("clean_margin", _run_clean_margin(run))) + gap = st.session_state.get(gap_key, saved.get("stray_gap", _run_stray_gap(run))) + return float(margin), float(gap) + + +def _specimen_cleanup_inputs(run, method, sample_id, *, margin_active, gap_active): + """Render this specimen's edge-margin and stray-gap inputs; return their values. + + Each input is enabled only while the view it drives has flash fill off. + """ + margin_key, gap_key = _specimen_cleanup_keys(run, method, sample_id) + margin, gap = _specimen_cleanup(run, method, sample_id) + st.session_state[margin_key], st.session_state[gap_key] = margin, gap + margin_column, gap_column = st.columns(2) + with margin_column: + st.number_input( + "Edge margin (% of box)", + min_value=0.0, + max_value=CLEAN_MARGIN_MAX, + step=0.25, + format="%.2f", + key=margin_key, + disabled=not margin_active, + help=( + "Band cleared along every edge of this specimen's target box, where " + "the printed box outline lands. Applies to pre-cleanup runs, a clean " + "size above 0, and Remove flashfill." + ), + ) + with gap_column: + st.number_input( + "Stray-piece distance (× leaf size)", + min_value=0.0, + max_value=STRAY_GAP_MAX, + step=0.05, + format="%.2f", + key=gap_key, + disabled=not gap_active, + help=( + "Pieces farther from the leaf than this fraction of its bounding-box " + "diagonal are dropped; 0 keeps only the leaf. Applies to pre-cleanup " + "runs and a clean size above 0." + ), + ) + return float(st.session_state[margin_key]), float(st.session_state[gap_key]) + + +def _raw_mask_path(pair): + """The specimen's mask before MATS cleanup, when this session has it.""" + return pair.get("raw_mask") + + +def _reset_threshold_preview(state_key): + st.session_state.setdefault("threshold_preview_cutoffs", {}).pop(state_key, None) + + +def _use_preview_threshold(cutoff): + st.session_state[THRESHOLD_LEVEL_KEY] = CUSTOM_THRESHOLD_LEVEL + st.session_state[THRESHOLD_CUSTOM_VALUE_KEY] = cutoff + + +@st.fragment +def render_threshold_explorer(pair, run, *, marked_pairs=None): + """Preview an Otsu cutoff and save it to this specimen or marked specimens.""" + from mats.app.threshold_preview import ( + clean_levels_for_threshold, cleaned_sample, color_sample, grayscale_sample, + pre_cleanup_sample, show_threshold_preview, + ) + + target_path = pair["target_box"] + try: + grayscale_image, otsu_cutoff = grayscale_sample(str(target_path)) + except ValueError as exc: + st.info(str(exc)) + return + try: + color_image = color_sample(str(target_path)) + except ValueError: + color_image = None # The mask panel still works without the color panel. + + saved_adjustment = run.get("threshold_adjustments", {}).get(pair["sample_id"], {}) + saved_cutoff = saved_adjustment.get("cutoff", run.get("threshold_value")) + if saved_cutoff is None: + saved_cutoff = otsu_cutoff + state_key = f"{run.get('run_id', '')}:{pair['sample_id']}" + component_key = f"threshold_preview_{state_key}" + cutoff = st.session_state.setdefault("threshold_preview_cutoffs", {}).get( + state_key, saved_cutoff + ) + clean_key = f"{run.get('run_id', '')}:threshold:{pair['sample_id']}" + remove_fill = _remove_fill(run, clean_key) + + st.markdown("**Explore and adjust output**") + clean_radius = _clean_radius(run, "threshold", pair["sample_id"]) + clean_on = clean_radius > 0 + if clean_on: + st.caption( + CLEAN_SIZE_ON_CAPTION + " Overwrite saves it." + + (" Remove flashfill doesn't apply while it is above 0." if remove_fill else "") + ) + elif run.get("measurement_source") == "cleaned": + st.caption( + "Drag to preview a new cutoff; the masked leaf beside the mask shows which " + "parts of the leaf it keeps. Release to apply MATS cleanup to this " + "sample. Saved outputs change only when you press Overwrite below." + ) + else: + st.caption( + "Drag to preview a new raw mask for this sample; the masked leaf beside it " + "shows which parts of the leaf it keeps. Release to see the mask this run " + "measures, with the edge margin cleared and stray pieces dropped. Nothing " + "saved changes until you press Overwrite below." + ) + pre_cleanup = run.get("measurement_source") == "pre-cleanup" + # Views without flash fill clear the edge margin, including while dragging. + flash_fill_off = clean_on or pre_cleanup or remove_fill + margin, gap = _specimen_cleanup_inputs( + run, "threshold", pair["sample_id"], + margin_active=flash_fill_off, gap_active=clean_on or pre_cleanup, + ) + # Both images go to the browser, so the clean-size slider is live from 0. + try: + clean_levels_image = clean_levels_for_threshold(str(target_path), cutoff, gap, margin) + if pre_cleanup: + cleaned_image = pre_cleanup_sample(str(target_path), cutoff, margin, gap) + else: + cleaned_image = cleaned_sample( + str(target_path), cutoff, fill_holes=not remove_fill, clean_margin=margin, + ) + except ValueError as exc: + st.info(str(exc)) + return + show_threshold_preview( + key=component_key, + grayscale_image=grayscale_image, + cutoff=cutoff, + measurement_source=run.get("measurement_source", "cleaned"), + color_image=color_image, + cleaned_image=cleaned_image, + cleaned_cutoff=cutoff, + remove_fill=remove_fill, + clean_levels_image=clean_levels_image, + clean_cutoff=cutoff, + clean_radius=clean_radius, + clean_help=CLEAN_SIZE_SAVE_HELP, + fill_toggle=True, + fill_note=_remove_fill_note(run), + live_margin=margin if flash_fill_off else 0, + on_cutoff_change=lambda: _record_threshold_preview(component_key, state_key), + on_clean_radius_change=lambda: _record_clean_radius(component_key, clean_key), + on_remove_fill_change=lambda: _record_remove_fill(component_key, clean_key), + ) + left, right = st.columns(2) + with left: + st.button( + "Reset to saved threshold", key=f"reset_{state_key}", + on_click=_reset_threshold_preview, args=(state_key,), width="stretch", + ) + with right: + st.button( + "Use threshold for next run", key=f"use_{state_key}", + disabled=not THRESHOLD_MIN <= cutoff <= THRESHOLD_MAX, + on_click=_use_preview_threshold, args=(cutoff,), width="stretch", + help=( + "Auto selected cutoff 0. Drag to at least 1 to use a custom threshold." + if cutoff == 0 else + "Sets Setup to this custom threshold; run analysis again to update the CSV." + ), + ) + + notice_key = f"threshold_adjustment_notice_{state_key}" + notice = st.session_state.pop(notice_key, None) + if notice: + st.success(notice, icon=":material/check_circle:") + marked_pairs = marked_pairs or [] + st.caption( + "Saving replaces the selected mask, CSV measurements, and any saved " + "overlays, cutouts, and measurement axes." + ) + unsavable = not THRESHOLD_MIN <= cutoff <= THRESHOLD_MAX + overwrite_this = st.button( + "Overwrite this specimen", + key=f"apply_adjustment_{state_key}", + type="primary", + icon=":material/save:", + disabled=unsavable, + help="Saves these settings, including the clean size, to the specimen selected in View.", + width="stretch", + ) + overwrite_marked = st.button( + f"Overwrite all marked specimens ({len(marked_pairs)})", + key=f"apply_marked_adjustment_{state_key}", + icon=":material/done_all:", + disabled=unsavable or not marked_pairs, + help=( + "Saves these settings, including the clean size, to every specimen " + "checked in Marked; each keeps its own calibration." + ), + width="stretch", + ) + if overwrite_this or overwrite_marked: + from mats.app.output_adjustment import ( + apply_threshold_adjustment, apply_threshold_adjustments, + ) + + bulk = overwrite_marked + count = len(marked_pairs) if bulk else 1 + try: + if bulk: + apply_threshold_adjustments( + run, marked_pairs, int(cutoff), remove_fill=remove_fill, + clean_margin=margin, stray_gap=gap, clean_size=clean_radius, + ) + else: + apply_threshold_adjustment( + run, pair, int(cutoff), remove_fill=remove_fill, + clean_margin=margin, stray_gap=gap, clean_size=clean_radius, + ) + except (OSError, ValueError) as exc: + st.error(f"Could not save adjustments: {exc}") + else: + read_results_dataframe.clear() + _clear_export_zip_cache() + st.session_state[notice_key] = ( + f"Updated {count} specimen(s) at threshold {int(cutoff)}" + + (f", clean size {clean_radius} px." if clean_on else ".") + ) + st.rerun() + + +@st.fragment +def render_mask_explorer(pair, run, method): + """Preview cleanup on one specimen's raw mask; nothing is saved.""" + from mats.app.threshold_preview import ( + clean_levels_for_mask, color_sample, mask_sample, pre_cleanup_mask_sample, + show_threshold_preview, unfilled_sample, + ) + + state_key = f"{run.get('run_id', '')}:{method}:{pair['sample_id']}" + component_key = f"mask_preview_{state_key}" + remove_fill = _remove_fill(run, state_key) + pre_cleanup = run.get("measurement_source") == "pre-cleanup" + st.markdown("**Explore and adjust output**") + clean_radius = _clean_radius(run, method, pair["sample_id"]) + clean_on = clean_radius > 0 + if clean_on: + st.caption(CLEAN_SIZE_ON_CAPTION + " Nothing is saved.") + elif pre_cleanup: + st.caption( + "Showing this specimen's pre-cleanup measurement mask. Adjust the edge " + "margin and stray-piece distance to preview it; nothing is saved." + ) + elif remove_fill: + st.caption( + "Showing the mask with flashfill removed. Adjust the edge margin to " + "preview it; nothing is saved." + ) + else: + st.caption( + "Showing the saved measurement mask. Drag the clean size above 0 to " + "preview removing small specks and filling small holes in the raw mask; " + "nothing is saved." + ) + margin, gap = _specimen_cleanup_inputs( + run, method, pair["sample_id"], + margin_active=clean_on or pre_cleanup or remove_fill, + gap_active=clean_on or pre_cleanup, + ) + mask_status = "Saved measurement mask" + # Both images go to the browser, so the clean-size slider is live from 0. + try: + raw_path = Path(_raw_mask_path(pair)) + clean_levels_image = clean_levels_for_mask( + str(raw_path), raw_path.stat().st_mtime_ns, gap, margin, + ) + if pre_cleanup: + mask_image = pre_cleanup_mask_sample( + str(raw_path), raw_path.stat().st_mtime_ns, margin, gap, + ) + mask_status = "Pre-cleanup mask: edge margin cleared and stray pieces dropped" + elif remove_fill: + mask_image = unfilled_sample(str(raw_path), margin) + mask_status = "Mask with hole filling removed and the edge margin cleared" + else: + shown_path = Path(pair.get("mask") or raw_path) + mask_image = mask_sample(str(shown_path), shown_path.stat().st_mtime_ns) + except (OSError, ValueError) as exc: + st.info(str(exc)) + return + color_image = None + if pair.get("target_box"): + try: + color_image = color_sample(str(pair["target_box"])) + except ValueError: + pass # The mask panel still works without the color panel. + show_threshold_preview( + key=component_key, + mask_image=mask_image, + mask_status=mask_status, + color_image=color_image, + clean_levels_image=clean_levels_image, + clean_radius=clean_radius, + clean_help=CLEAN_SIZE_PREVIEW_HELP, + remove_fill=remove_fill, + fill_toggle=True, + fill_note=_remove_fill_note(run), + on_clean_radius_change=lambda: _record_clean_radius(component_key, state_key), + on_remove_fill_change=lambda: _record_remove_fill(component_key, state_key), + ) + + +def render_output_pair(pair, width=PREVIEW_IMAGE_WIDTH, mask_override=None): st.caption(f"Sample: {pair['sample_id']}") left, right = st.columns(2) with left: @@ -1913,179 +2824,395 @@ def render_output_pair(pair, width=PREVIEW_IMAGE_WIDTH): else: st.caption("Target box: not available") with right: - st.image( - str(pair["mask"]), - caption="Leaf segmentation mask", - width=width, - ) + if mask_override is not None: + st.image(mask_override, caption="Preview with flashfill removed", width=width) + elif pair["mask"] is not None: + caption = ( + "Pre-cleanup measurement mask" + if pair.get("mask_source") == "pre-cleanup" + else "Cleaned measurement mask" + ) + st.image(str(pair["mask"]), caption=caption, width=width) + else: + st.caption("Measurement mask preview unavailable for this sample") -def render_output_preview(pairs, selected_sample_id=None): - if not pairs: - return +def _clear_export_zip_cache(): + """Discard the prepared ZIP when its run or selection changes.""" + zip_path = st.session_state.pop("export_zip_path", None) + st.session_state.pop("export_zip_signature", None) + if zip_path: + try: + Path(zip_path).unlink(missing_ok=True) + except OSError: + pass + dataset_path = st.session_state.pop("dataset_zip_path", None) + st.session_state.pop("dataset_zip_signature", None) + st.session_state.pop("dataset_manifest", None) + if dataset_path: + try: + Path(dataset_path).unlink(missing_ok=True) + except OSError: + pass - if selected_sample_id is not None: - selected_pairs = [ - pair for pair in pairs - if str(pair["sample_id"]) == selected_sample_id - ] - st.subheader("Selected output", anchor=False) - if not selected_pairs: - st.info( - "No matched target-box and mask pair is available for the selected row.", - icon=":material/image_not_supported:", - ) - return - render_output_pair(selected_pairs[0]) - return - st.subheader("Output preview", anchor=False) - show_default = len(pairs) <= PREVIEW_AUTO_HIDE_THRESHOLD - show_preview = st.checkbox( - f"Show image preview ({len(pairs)} output(s))", - value=show_default, - key="show_output_preview", - help="Renders matched target-box / mask pairs. Hidden by default for large batches.", - ) - if not show_preview: - return +def select_export_files(artifacts, methods, kinds): + """Select existing files from this run, including shared files just once.""" + selected_methods = set(methods) + if not selected_methods: + return [] + selected_kinds = set(kinds) + files = [] + seen = set() + for item in artifacts: + if item.get("kind") not in selected_kinds: + continue + method = item.get("method") + if method is not None and method not in selected_methods: + continue + path = Path(item["path"]) + if path not in seen and path.is_file(): + files.append(path) + seen.add(path) + return files - max_count = min(PREVIEW_HARD_MAX, len(pairs)) - if len(pairs) == 1: - count = 1 - else: - count = st.slider( - "Pairs to display", - min_value=1, - max_value=max_count, - value=min(6, max_count), - key="output_preview_count", - ) - st.caption(f"Showing {count} of {len(pairs)} output pair(s).") - for pair in pairs[:count]: - render_output_pair(pair) +def zip_download_name(raw_name): + """Keep the requested download name a single ZIP filename.""" + name = str(raw_name or "").replace("\\", "/").rsplit("/", 1)[-1].strip() + if name in {"", ".", ".."}: + name = "leaf_morpho_outputs" + return name if name.lower().endswith(".zip") else f"{name}.zip" -def _clear_export_zip_cache(): - """Invalidate a previously prepared ZIP when the export selection changes.""" - st.session_state.pop("export_zip_path", None) + +def _export_run_token(run): + return run.get("run_id") or "|".join( + str(run.get("by_method", {}).get(method, {}).get("results_path", "")) + for method in run.get("mask_methods", ()) + ) + + +def _export_signature(run_token, files, download_name): + """Identify the exact ZIP contents, including files changed on disk.""" + return ( + run_token, + tuple((str(path), path.stat().st_size, path.stat().st_mtime_ns) for path in files), + download_name, + ) -def render_export_section(results_path, output_path, unit_symbol): - """Render a clearly defined Export section: CSV-only vs. the full ZIP bundle.""" - output_path = Path(output_path) +def render_export_section(): + """Choose and download files recorded for the completed analysis run.""" st.subheader("Export", anchor=False) - st.caption("Download the outputs from this run. Pick exactly what you need.") - csv_column, zip_column = st.columns(2, vertical_alignment="top") - - with csv_column.container(border=True): - st.markdown("**Measurements only**") - st.caption("Just the results CSV — sample IDs, area, width, length. No images.") - if results_path.is_file(): - st.download_button( - f"Download results CSV ({unit_symbol})", - data=results_path.read_bytes(), - file_name=results_path.name, - mime="text/csv", - icon=":material/download:", - key="download_csv_only", - ) - else: - st.caption("Not available yet.") + run = st.session_state.get("last_run") + if not run: + st.info("Run an analysis to see and download its saved files.", icon=":material/folder_open:") + st.button("Open Analyze", on_click=_switch_workspace_tab, args=("Analyze",), + icon=":material/science:") + return - base_files = gather_output_files(output_path, results_path) - with zip_column.container(border=True): - st.markdown("**Full export (ZIP)**") - st.caption( - "Results CSV + failure log + segmentation masks + specimen photos, " - "bundled together." + methods = tuple(run["mask_methods"]) + artifacts = tuple(run.get("artifacts", ())) + available = tuple(item for item in artifacts if Path(item["path"]).is_file()) + missing_count = len(artifacts) - len(available) + all_files = select_export_files(available, methods, (kind for kind, _ in EXPORT_FILE_TYPES)) + _, all_bytes = estimate_zip_inputs(all_files) + output_path = Path(run["output_path"]) + st.caption(f"Completed run · {len(all_files):,} saved file(s) · {human_bytes(all_bytes)}") + st.caption(f"Output folder: `{display_path(output_path)}`") + if missing_count: + st.warning( + f"{missing_count} file(s) recorded for this run are no longer available on disk. " + "They will be left out of downloads.", + icon=":material/folder_off:", ) - if not base_files: - st.caption("No output files found yet.") - return + if not all_files: + st.warning("No saved files from this run are available for download.") + return - pair_count = len(list(output_path.glob("*_mask.png"))) - include_overlay = st.checkbox( - "Include overlay images (mask highlighted on photo)", - key="export_include_overlay", - help=( - "One extra JPG per specimen: the photo with the detected leaf " - "region tinted and outlined, for visually checking segmentation " - "accuracy." - ), - on_change=_clear_export_zip_cache, + run_token = _export_run_token(run) + if st.session_state.get(EXPORT_RUN_ID_KEY) != run_token: + _clear_export_zip_cache() + st.session_state[EXPORT_RUN_ID_KEY] = run_token + st.session_state[EXPORT_METHODS_KEY] = list(methods) + st.session_state[EXPORT_ZIP_NAME_KEY] = "leaf_morpho_outputs.zip" + available_kinds = {item["kind"] for item in available} + for kind, _ in EXPORT_FILE_TYPES: + st.session_state[f"export_include_{kind}"] = kind in available_kinds + + if len(methods) > 1: + chosen_methods = st.multiselect( + "Methods to download", + options=list(methods), + format_func=SEGMENTATION_METHOD_LABELS.get, + key=EXPORT_METHODS_KEY, persist_state="page", ) - include_cutout = st.checkbox( - "Include specimen cutouts (background removed)", - key="export_include_cutout", - help=( - "One extra JPG per specimen: just the leaf pixels, with the " - "background blacked out." - ), - on_change=_clear_export_zip_cache, - persist_state="page", - ) - enabled_extra_count = int(include_overlay) + int(include_cutout) + else: + chosen_methods = list(methods) + st.caption(f"Method: {SEGMENTATION_METHOD_LABELS[methods[0]]}") - count, total_bytes = estimate_zip_inputs(base_files) - if enabled_extra_count and pair_count: - target_box_paths = list(output_path.glob("*_target_box.jpg")) - avg_target_box_bytes = ( - sum(path.stat().st_size for path in target_box_paths) / len(target_box_paths) - if target_box_paths - else 0 + st.markdown("**Files to include in ZIP**") + st.caption("These choices package existing files. Image output settings for the next run are in Setup.") + left, right = st.columns(2) + selected_kinds = [] + for index, (kind, label) in enumerate(EXPORT_FILE_TYPES): + count = sum( + item["kind"] == kind and + (item.get("method") is None or item["method"] in chosen_methods) + for item in available + ) + ever_saved = any(item["kind"] == kind for item in available) + with (left if index < 5 else right): + checked = st.checkbox( + f"{label} · {count} file(s)", + key=f"export_include_{kind}", + disabled=not ever_saved, + persist_state="page", ) - count += pair_count * enabled_extra_count - total_bytes += int(pair_count * avg_target_box_bytes * enabled_extra_count) + if checked and count: + selected_kinds.append(kind) + st.caption( + "Unavailable types were not saved in this run. Image types can be enabled " + "in Setup before the next analysis." + ) + + selected_files = select_export_files(available, chosen_methods, selected_kinds) + count, total_bytes = estimate_zip_inputs(selected_files) + st.markdown("**Download**") + csv_files = select_export_files(available, chosen_methods, ("results_csv",)) + with st.container(border=True): + st.markdown("**Measurement CSVs**") + st.caption("Download a CSV directly, regardless of the ZIP file choices above.") + if csv_files: + for path in csv_files: + with open(path, "rb") as csv_file: + st.download_button( + f"Download {path.name}", + data=csv_file, + file_name=path.name, + mime="text/csv", + icon=":material/download:", + key=f"download_csv_{path.name}", + ) + else: + st.caption("No results CSV is available for the selected method(s).") - st.caption(f"{count} file(s), ~{human_bytes(total_bytes)} uncompressed.") + with st.container(border=True): + st.markdown("**Selected files (ZIP)**") + raw_name = st.text_input( + "ZIP filename", + key=EXPORT_ZIP_NAME_KEY, + max_chars=120, + persist_state="page", + ) + download_name = zip_download_name(raw_name) + st.caption(f"{count:,} file(s) · {human_bytes(total_bytes)} uncompressed") + with st.expander(f"View files in ZIP ({count:,})"): + for path in selected_files[:200]: + st.write(path.name) + if count > 200: + st.caption(f"Showing the first 200 of {count:,} files.") if total_bytes > ZIP_SIZE_WARN_BYTES: st.warning( f"Outputs total ~{human_bytes(total_bytes)}. Building a ZIP this large can be " f"slow and memory-heavy. Consider collecting files directly from " - f"`{display_path(output_path)}` " - "instead." + f"`{display_path(output_path)}` instead." ) - if st.button("Prepare ZIP for download", icon=":material/folder_zip:"): - large_batch = enabled_extra_count and pair_count > LARGE_BATCH_THRESHOLD - if large_batch: - with st.spinner("Generating overlay/cutout images..."): - generate_export_overlays(output_path, include_overlay, include_cutout) + try: + signature = _export_signature(run_token, selected_files, download_name) + except OSError: + st.warning("A selected file changed or disappeared. Refresh Export and try again.") + _clear_export_zip_cache() + return + if st.session_state.get("export_zip_signature") != signature: + _clear_export_zip_cache() + if st.button("Prepare ZIP for download", icon=":material/folder_zip:", + key="prepare_zip_export", disabled=not selected_files): + dest = None + try: with st.spinner("Building ZIP..."): - files = gather_output_files( - output_path, results_path, include_overlay, include_cutout - ) - dest = Path(tempfile.gettempdir()) / "leaf_morpho_outputs.zip" - write_output_zip(files, dest) + with tempfile.NamedTemporaryFile( + prefix="mats_outputs_", suffix=".zip", delete=False + ) as temp_zip: + dest = Path(temp_zip.name) + write_output_zip(selected_files, dest) + except OSError as exc: + if dest is not None: + dest.unlink(missing_ok=True) + st.error(f"Could not prepare the ZIP: {exc}") else: - with st.spinner("Building ZIP..."): - generate_export_overlays(output_path, include_overlay, include_cutout) - files = gather_output_files( - output_path, results_path, include_overlay, include_cutout - ) - dest = Path(tempfile.gettempdir()) / "leaf_morpho_outputs.zip" - write_output_zip(files, dest) - st.session_state["export_zip_path"] = str(dest) + st.session_state["export_zip_path"] = str(dest) + st.session_state["export_zip_signature"] = signature zip_path = st.session_state.get("export_zip_path") if zip_path and Path(zip_path).is_file(): - contents = ["target boxes", "masks", "CSV"] - if include_overlay: - contents.append("overlays") - if include_cutout: - contents.append("cutouts") - with open(zip_path, "rb") as zf: + with open(zip_path, "rb") as zip_file: st.download_button( - f"Download ZIP ({', '.join(contents)})", - data=zf, - file_name="leaf_morpho_outputs.zip", + f"Download ZIP ({count:,} files)", + data=zip_file, + file_name=download_name, mime="application/zip", icon=":material/download:", key="download_zip_export", ) + render_training_dataset_export(run, methods) + + +def render_training_dataset_export(run, methods): + """Package the current run's aligned image and mask pairs for training.""" + from mats.dataset_export import read_group_csv, split_counts, write_dataset_zip + + with st.container(border=True): + st.markdown("**Training dataset (ZIP)**") + st.caption( + "Use target-box images and their masks from this run, including saved " + "specimen adjustments. Labels are generated by MATS; review them before training." + ) + method = st.selectbox( + "Labeling method", methods, format_func=SEGMENTATION_METHOD_LABELS.get, + key="dataset_method", + ) + pairs = st.session_state.get("viewer_pairs", {}).get(method, ()) + mask_sources = ["measured"] + if any(pair.get("raw_mask") for pair in pairs): + mask_sources.append("raw") + mask_source = st.selectbox( + "Masks to use", mask_sources, + format_func=lambda source: ( + "Measured mask" if source == "measured" else "Raw pre-cleanup mask" + ), + key="dataset_mask_source", + ) + format_labels = { + "png": "Image + PNG mask", + "yolo-seg": "YOLO segmentation", + "yolo-detect": "YOLO detection", + "coco": "COCO segmentation", + } + dataset_format = st.selectbox( + "Dataset format", tuple(format_labels), format_func=format_labels.get, + key="dataset_format", + ) + cols = st.columns(3) + percentages = tuple(int(cols[index].number_input( + label, min_value=0, max_value=100, value=default, step=1, + key=f"dataset_percentage_{split}", + )) for index, (split, label, default) in enumerate(( + ("train", "Train %", 70), ("val", "Validation %", 20), + ("test", "Test %", 10), + ))) + seed = int(st.number_input( + "Random seed", min_value=0, max_value=2_147_483_647, value=42, + step=1, key="dataset_seed", + )) + group_file = st.file_uploader( + "Optional group CSV (sample_id,group_id)", type="csv", + key="dataset_group_csv", + help="Keep repeated photos of the same plant or specimen in one split.", + ) + group_content = group_file.getvalue() if group_file is not None else b"" + try: + groups = read_group_csv(group_content) + split_counts(0, percentages) + error = None + except ValueError as exc: + groups = {} + error = str(exc) + st.error(error) + candidates = sum( + bool(pair.get("target_box") and + pair.get("mask" if mask_source == "measured" else "raw_mask") and + Path(pair["target_box"]).is_file() and + Path(pair["mask" if mask_source == "measured" else "raw_mask"]).is_file()) + for pair in pairs + ) + if not error: + estimated = split_counts(candidates, percentages) + st.caption( + f"{candidates} candidate pair(s) · expected split: " + f"{estimated['train']} train / {estimated['val']} validation / " + f"{estimated['test']} test. Validation may exclude unreadable or mismatched files." + ) + if candidates and (estimated["train"] == 0 or estimated["val"] == 0): + st.warning( + "This dataset may have no training or validation images. " + "Add more specimens or adjust the split before training." + ) + if groups: + st.caption("Group assignments may change the exact split counts.") + if dataset_format == "yolo-seg": + st.caption( + "YOLO polygons cannot represent holes; any affected specimens " + "will be listed in the dataset manifest." + ) + source_paths = tuple( + path for pair in pairs + for path in (pair.get("target_box"), pair.get( + "mask" if mask_source == "measured" else "raw_mask")) + if path and Path(path).is_file() + ) + try: + signature = ( + run.get("run_id"), method, mask_source, dataset_format, + percentages, seed, group_content, + tuple((str(path), Path(path).stat().st_size, Path(path).stat().st_mtime_ns) + for path in source_paths), + ) + except OSError: + signature = None + st.warning("A pair changed on disk. Refresh Export and try again.") + if st.session_state.get("dataset_zip_signature") != signature: + old_path = st.session_state.pop("dataset_zip_path", None) + st.session_state.pop("dataset_zip_signature", None) + st.session_state.pop("dataset_manifest", None) + if old_path: + Path(old_path).unlink(missing_ok=True) + if st.button( + "Prepare training dataset", icon=":material/folder_zip:", + key="prepare_training_dataset", + disabled=bool(error) or not candidates or signature is None, + ): + dest = None + try: + with st.spinner("Building training dataset..."): + with tempfile.NamedTemporaryFile( + prefix="mats_dataset_", suffix=".zip", delete=False, + ) as temp_zip: + dest = Path(temp_zip.name) + manifest = write_dataset_zip( + pairs, dest, dataset_format=dataset_format, + percentages=percentages, seed=seed, method=method, + mask_source=mask_source, groups=groups, + ) + except (OSError, ValueError, zipfile.BadZipFile) as exc: + if dest is not None: + dest.unlink(missing_ok=True) + st.error(f"Could not prepare training dataset: {exc}") + else: + st.session_state["dataset_zip_path"] = str(dest) + st.session_state["dataset_zip_signature"] = signature + st.session_state["dataset_manifest"] = manifest + dataset_path = st.session_state.get("dataset_zip_path") + manifest = st.session_state.get("dataset_manifest") + if dataset_path and manifest and Path(dataset_path).is_file(): + counts = manifest["counts"] + st.success( + f"Ready: {counts['train']} train / {counts['val']} validation / " + f"{counts['test']} test; {len(manifest['excluded'])} excluded." + ) + if manifest["conversion_notes"]: + st.warning(f"{len(manifest['conversion_notes'])} conversion note(s) in manifest.json.") + with open(dataset_path, "rb") as dataset_file: + st.download_button( + "Download training dataset", data=dataset_file, + file_name=f"mats_training_{method}_{dataset_format}.zip", + mime="application/zip", icon=":material/download:", + key="download_training_dataset", + ) + if __name__ == "__main__": main() diff --git a/src/mats/app/compute.py b/src/mats/app/compute.py index 76bb953..87a645b 100644 --- a/src/mats/app/compute.py +++ b/src/mats/app/compute.py @@ -121,12 +121,8 @@ def hybrid_fanout_disabled() -> bool: def selected_mask_method() -> str: - """Return the segmentation method selected on Home, defaulting to Otsu.""" - return ( - "birefnet" - if st.session_state.get("segmentation_method") == "BiRefNet" - else "threshold" - ) + """Return the heaviest segmentation method checked on Home, defaulting to Otsu.""" + return "birefnet" if st.session_state.get("segment_birefnet", False) else "threshold" def break_glass_unlocked(available_workers: int) -> bool: diff --git a/src/mats/app/output_adjustment.py b/src/mats/app/output_adjustment.py new file mode 100644 index 0000000..9c1a402 --- /dev/null +++ b/src/mats/app/output_adjustment.py @@ -0,0 +1,411 @@ +"""Re-measure one threshold specimen and replace its saved output atomically.""" + +import csv +import copy +import io +import json +import os +import shutil +import tempfile +from pathlib import Path + +import cv2 + +from mats.mask_cleanup import raw_measurement_mask +from mats.scaling import ( + NA_VALUE, + compact_measurement_row, + converted_measurement_row, + length_conversion_factor, + result_measurement_fieldnames, + validate_results_unit, +) +from mats.mask_settings import ( + CLEAN_MARGIN_DEFAULT, + CLEAN_SIZE_DEFAULT, + STRAY_GAP_DEFAULT, + checked_clean_margin, + checked_clean_size, + checked_stray_gap, +) + + +def _results_csv(run): + path = Path(run["by_method"]["threshold"]["results_path"]) + if not path.is_file(): + raise ValueError("The Classic thresholding results CSV is unavailable.") + with path.open(newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + fields = reader.fieldnames or [] + rows = list(reader) + unit = validate_results_unit(run.get("results_unit", "cm")) + full_area = result_measurement_fieldnames(unit)[0] + compact_area = result_measurement_fieldnames(unit, compact=True)[0] + if full_area in fields: + schema = "full" + elif compact_area in fields: + schema = "compact" + else: + raise ValueError("The results CSV has no recognized measurement columns.") + return path, fields, rows, unit, schema + + +def _sample_row(rows, sample_id): + matches = [index for index, row in enumerate(rows) if row.get("sample_id") == sample_id] + if len(matches) != 1: + raise ValueError("The selected specimen must have exactly one results row.") + return matches[0] + + +def _scale_axes(run, sample_id, csv_row, unit): + saved = run.get("scale_axes_by_sample", {}).get("threshold", {}).get(sample_id) + if saved is not None: + axes = tuple(float(value) for value in saved) + else: + factor = length_conversion_factor(unit) + try: + axes = ( + float(csv_row[f"px_per_{unit}_width"]) * factor, + float(csv_row[f"px_per_{unit}_height"]) * factor, + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + "Calibration for this specimen is unavailable; rerun analysis before adjusting it." + ) from exc + if len(axes) != 2 or not all(0 < value < float("inf") for value in axes): + raise ValueError("Calibration for this specimen is invalid.") + return axes + + +def measure_threshold_adjustment( + run, pair, cutoff, *, remove_fill=False, clean_margin=None, stray_gap=None, + clean_size=None, +): + """Compute the exact saved-mask measurement for one selected specimen. + + ``clean_margin``, ``stray_gap``, and ``clean_size`` default to the run's + values. The margin and gap apply to pre-cleanup runs and a clean size above + 0, and the margin also to Remove flashfill. A clean size above 0 measures + Clean image, the mask its preview shows; in a cleaned run it also replaces + MATS cleanup, and Remove flashfill, in the saved mask. + """ + from mats import core + + if not isinstance(cutoff, int) or not 0 <= cutoff <= 255: + raise ValueError("Threshold cutoff must be an integer from 0 to 255.") + sample_id = str(pair["sample_id"]) + if not sample_id or Path(sample_id).name != sample_id: + raise ValueError("Invalid specimen ID.") + source = run.get("measurement_source", "cleaned") + if source not in {"cleaned", "pre-cleanup"}: + raise ValueError("Unknown measurement-mask source.") + target_path = pair.get("target_box") + target = cv2.imread(str(target_path), cv2.IMREAD_COLOR) if target_path else None + if target is None: + raise ValueError("The selected specimen's target-box image is unavailable.") + + results_path, fields, rows, unit, schema = _results_csv(run) + row_index = _sample_row(rows, sample_id) + if rows[row_index].get(result_measurement_fieldnames(unit, schema == "compact")[0]) in { + NA_VALUE, "", None, + }: + raise ValueError("Only successfully measured specimens can be adjusted.") + axes = _scale_axes(run, sample_id, rows[row_index], unit) + margin = checked_clean_margin( + run.get("clean_margin", CLEAN_MARGIN_DEFAULT) if clean_margin is None else clean_margin + ) + gap = checked_stray_gap( + run.get("stray_gap", STRAY_GAP_DEFAULT) if stray_gap is None else stray_gap + ) + size = checked_clean_size( + run.get("clean_size", CLEAN_SIZE_DEFAULT) if clean_size is None else clean_size + ) + pre_cleanup = source == "pre-cleanup" + raw = core.threshold_mask(target, cutoff) + unfilled = raw_measurement_mask(raw, margin, gap, size) if pre_cleanup or size else None + # A pre-cleanup run keeps its cleaned-mask export and measures the unfilled + # mask; a cleaned run saves and measures one mask, Clean image when on. + fill_holes = not remove_fill and (pre_cleanup or not size) + if size and not pre_cleanup: + cleaned = unfilled + elif remove_fill: + cleaned = core.unfilled_leaf_mask(raw, margin) + else: + cleaned = core.clean_leaf_mask(raw.copy()) + measurement_mask = unfilled if pre_cleanup else cleaned + measurement = core.measurement_row_from_mask( + sample_id, measurement_mask, *axes, measurement_source=source + ) + if measurement["leaf_area_cm2"] == NA_VALUE: + raise ValueError("This threshold produces no measurable leaf mask.") + converted = ( + compact_measurement_row(measurement, unit) + if schema == "compact" else converted_measurement_row(measurement, unit) + ) + return { + "sample_id": sample_id, + "results_path": results_path, + "fields": fields, + "rows": rows, + "row_index": row_index, + "unit": unit, + "schema": schema, + "target": target, + "raw": raw, + "cleaned": cleaned, + "measurement_mask": measurement_mask, + "converted": converted, + "clean_margin": margin, + "stray_gap": gap, + "clean_size": size, + "fill_holes": fill_holes, + } + + +def _encoded_image(path, image): + suffix = path.suffix.lower() + if suffix not in {".png", ".jpg", ".jpeg"}: + raise ValueError(f"Unsupported saved image format: {suffix}") + ok, encoded = cv2.imencode(suffix, image) + if not ok: + raise ValueError(f"Could not encode {path.name}") + return encoded.tobytes() + + +def _commit_files(changes): + """Stage every file first; restore earlier files if a later replacement fails.""" + staged = [] + replaced = [] + try: + for destination, content in changes.items(): + destination.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=".mats_adjust_", dir=destination.parent) + with os.fdopen(fd, "wb") as handle: + handle.write(content) + staged.append((destination, Path(temporary))) + for destination, temporary in staged: + backup = None + if destination.exists(): + fd, backup_name = tempfile.mkstemp( + prefix=".mats_adjust_backup_", dir=destination.parent + ) + os.close(fd) + backup = Path(backup_name) + shutil.copy2(destination, backup) + os.replace(temporary, destination) + replaced.append((destination, backup)) + except Exception: + for destination, backup in reversed(replaced): + if backup is None: + destination.unlink(missing_ok=True) + else: + os.replace(backup, destination) + raise + finally: + for _, temporary in staged: + temporary.unlink(missing_ok=True) + for _, backup in replaced: + if backup is not None: + backup.unlink(missing_ok=True) + + +def apply_threshold_adjustment( + run, pair, cutoff, *, remove_fill=False, clean_margin=None, stray_gap=None, + clean_size=None, +): + """Overwrite only this threshold specimen's measurements and dependent images.""" + from mats import core + + prepared = measure_threshold_adjustment( + run, pair, cutoff, remove_fill=remove_fill, + clean_margin=clean_margin, stray_gap=stray_gap, clean_size=clean_size, + ) + sample_id = prepared["sample_id"] + output_dir = Path(run["output_path"]) + results_path = prepared["results_path"] + if not results_path.resolve().is_relative_to(output_dir.resolve()): + raise ValueError("The results CSV is outside this run's output folder.") + suffix = "_threshold" if len(run["mask_methods"]) > 1 else "" + artifacts = run.setdefault("artifacts", []) + + def artifact_path(kind): + return next(( + Path(item["path"]) for item in artifacts + if item.get("sample_id") == sample_id and item.get("kind") == kind + and item.get("method") in {None, "threshold"} + ), None) + + pre_cleanup = run.get("measurement_source", "cleaned") == "pre-cleanup" + mask_path = artifact_path("mask") or output_dir / f"{sample_id}_mask{suffix}.png" + raw_path = artifact_path("pre_cleanup") + if pre_cleanup and raw_path is None: + raw_path = output_dir / f"{sample_id}_mask_precleanup_threshold.png" + changes = {mask_path: _encoded_image(mask_path, prepared["cleaned"])} + if raw_path is not None: + changes[raw_path] = _encoded_image(raw_path, prepared["raw"]) + if pair.get("raw_mask"): + preview_raw_path = Path(pair["raw_mask"]) + changes[preview_raw_path] = _encoded_image(preview_raw_path, prepared["raw"]) + # A pre-cleanup run measured its raw mask minus stray pieces (and, at a clean + # size, small specks and holes); the preview mask holds that, while the + # pre-cleanup export stays the raw mask. + measured_path = None + if pre_cleanup and pair.get("mask") and Path(pair["mask"]) != raw_path: + measured_path = Path(pair["mask"]) + changes[measured_path] = _encoded_image(measured_path, prepared["measurement_mask"]) + + dependent = { + "overlay": lambda: core.build_overlay_image(prepared["target"], prepared["measurement_mask"]), + "cutout": lambda: core.build_cutout_image(prepared["target"], prepared["measurement_mask"]), + "axes": lambda: core.draw_measurement_axes( + prepared["target"], prepared["measurement_mask"], + run.get("measurement_source", "cleaned"), + ), + } + for kind, build in dependent.items(): + path = artifact_path(kind) + if path is not None: + image = build() + if image is None: + raise ValueError(f"Could not regenerate {kind} for this specimen.") + changes[path] = _encoded_image(path, image) + + converted = prepared["converted"] + fields_to_update = ( + result_measurement_fieldnames(prepared["unit"], prepared["schema"] == "compact") + ) + if prepared["schema"] == "full": + fields_to_update += [ + f"px_per_{prepared['unit']}_width", + f"px_per_{prepared['unit']}_height", + "scale_aspect_ratio", + ] + prepared["rows"][prepared["row_index"]].update({ + field: str(converted[field]) for field in fields_to_update + }) + csv_buffer = io.StringIO(newline="") + writer = csv.DictWriter(csv_buffer, fieldnames=prepared["fields"]) + writer.writeheader() + writer.writerows(prepared["rows"]) + changes[results_path] = csv_buffer.getvalue().encode("utf-8") + + metadata_path = Path(f"{results_path}.meta.json") + if metadata_path.exists(): + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError("Could not read the run's measurement metadata.") from exc + else: + metadata = { + "measurement_source": run.get("measurement_source", "cleaned"), + "mask_method": "threshold", + "results_unit": prepared["unit"], + "csv_schema": prepared["schema"], + } + if pre_cleanup: + metadata["clean_margin"] = run.get("clean_margin", CLEAN_MARGIN_DEFAULT) + metadata["stray_gap"] = run.get("stray_gap", STRAY_GAP_DEFAULT) + metadata["clean_size"] = run.get("clean_size", CLEAN_SIZE_DEFAULT) + # Record the cleanup settings only where they shaped the saved mask. + clean_on = bool(prepared["clean_size"]) + adjustment = {"cutoff": cutoff, "fill_holes": prepared["fill_holes"]} + if pre_cleanup or remove_fill or clean_on: + adjustment["clean_margin"] = prepared["clean_margin"] + if pre_cleanup or clean_on: + adjustment["stray_gap"] = prepared["stray_gap"] + adjustment["clean_size"] = prepared["clean_size"] + metadata.setdefault("threshold_adjustments", {})[sample_id] = adjustment + changes[metadata_path] = (json.dumps(metadata, indent=2) + "\n").encode("utf-8") + + _commit_files(changes) + if artifact_path("mask") is None: + artifacts.append({ + "path": str(mask_path), "sample_id": sample_id, + "kind": "mask", "method": "threshold", + }) + if raw_path is not None and artifact_path("pre_cleanup") is None: + artifacts.append({ + "path": str(raw_path), "sample_id": sample_id, + "kind": "pre_cleanup", "method": "threshold", + }) + if not any(item.get("kind") == "results_metadata" and item.get("method") == "threshold" + for item in artifacts): + artifacts.append({ + "path": str(metadata_path), "sample_id": None, + "kind": "results_metadata", "method": "threshold", + }) + if pre_cleanup: + pair["raw_mask"] = str(raw_path) + pair["mask"] = None if measured_path is None else str(measured_path) + pair["mask_source"] = "pre-cleanup" + else: + pair["mask"] = str(mask_path) + run.setdefault("threshold_adjustments", {})[sample_id] = dict(adjustment) + return converted + + +def apply_threshold_adjustments( + run, pairs, cutoff, *, remove_fill=False, clean_margin=None, stray_gap=None, + clean_size=None, +): + """Apply one set of controls to marked specimens, restoring all on failure.""" + pairs = list(pairs) + sample_ids = [str(pair["sample_id"]) for pair in pairs] + if not pairs or len(set(sample_ids)) != len(sample_ids): + raise ValueError("Mark one or more distinct specimens before saving.") + # Validate all measurements and target images before touching any output. + for pair in pairs: + measure_threshold_adjustment( + run, pair, cutoff, remove_fill=remove_fill, + clean_margin=clean_margin, stray_gap=stray_gap, clean_size=clean_size, + ) + + output_dir = Path(run["output_path"]) + results_path = Path(run["by_method"]["threshold"]["results_path"]) + suffix = "_threshold" if len(run["mask_methods"]) > 1 else "" + candidates = {results_path, Path(f"{results_path}.meta.json")} + for item in run.get("artifacts", []): + if (item.get("sample_id") in sample_ids + and item.get("method") in {None, "threshold"} + and item.get("kind") in {"mask", "pre_cleanup", "overlay", "cutout", "axes"}): + candidates.add(Path(item["path"])) + for pair in pairs: + sample_id = str(pair["sample_id"]) + candidates.add(output_dir / f"{sample_id}_mask{suffix}.png") + if run.get("measurement_source") == "pre-cleanup": + candidates.add(output_dir / f"{sample_id}_mask_precleanup_threshold.png") + for key in ("mask", "raw_mask"): + if pair.get(key): + candidates.add(Path(pair[key])) + + original_artifacts = copy.deepcopy(run.get("artifacts", [])) + original_adjustments = copy.deepcopy(run.get("threshold_adjustments", {})) + original_pairs = [dict(pair) for pair in pairs] + with tempfile.TemporaryDirectory(prefix="mats_adjust_batch_") as backup_dir: + backups = {} + for index, path in enumerate(candidates): + backup = Path(backup_dir) / str(index) + if path.exists(): + shutil.copy2(path, backup) + backups[path] = backup + else: + backups[path] = None + try: + for pair in pairs: + apply_threshold_adjustment( + run, pair, cutoff, remove_fill=remove_fill, + clean_margin=clean_margin, stray_gap=stray_gap, clean_size=clean_size, + ) + except Exception: + for path, backup in backups.items(): + if backup is None: + path.unlink(missing_ok=True) + else: + shutil.copy2(backup, path) + run["artifacts"] = original_artifacts + run["threshold_adjustments"] = original_adjustments + for pair, original in zip(pairs, original_pairs): + pair.clear() + pair.update(original) + raise + return sample_ids diff --git a/src/mats/app/pages/0_Diagnostics.py b/src/mats/app/pages/0_Diagnostics.py new file mode 100644 index 0000000..1abb36b --- /dev/null +++ b/src/mats/app/pages/0_Diagnostics.py @@ -0,0 +1,40 @@ +"""Compute status and the latest Home preflight report.""" + +import streamlit as st + +from mats.app import branding +from mats.app.Home import ( + build_preflight_checks, load_pipeline_module, render_diagnostics, +) +from mats.app.compute import ( + birefnet_parallel_unlocked, break_glass_unlocked, get_compute_settings, + resolve_execution_plan, +) + + +st.set_page_config(page_title="MATS — Diagnostics", page_icon=branding.page_icon(), layout="wide") +branding.apply_logo() + +config = st.session_state.get("diagnostics_context") +inputs = st.session_state.get("diagnostics_inputs") +if config is None or inputs is None: + st.title("Diagnostics") + st.info("Open Home to check the current analysis settings and input images.") + try: + st.page_link("Home.py", label="Open Home", icon=":material/home:") + except (KeyError, ValueError): + st.caption("Open Home from the sidebar.") +else: + lm = load_pipeline_module() + compute = get_compute_settings(lm) + workers = compute.available_workers + plan = resolve_execution_plan( + compute, + "birefnet" if config["needs_birefnet"] else "threshold", + birefnet_parallel_allowed=birefnet_parallel_unlocked(workers), + ) + risk = lm.worker_risk_report(plan.workers, workers) + checks, _, _ = build_preflight_checks( + lm, config, *inputs, compute, plan, risk, break_glass_unlocked(workers), + ) + render_diagnostics(compute, plan, config, checks) diff --git a/src/mats/app/pages/2_BiRefNet_Setup.py b/src/mats/app/pages/2_BiRefNet_Setup.py index 9b16c69..3671d3e 100644 --- a/src/mats/app/pages/2_BiRefNet_Setup.py +++ b/src/mats/app/pages/2_BiRefNet_Setup.py @@ -126,7 +126,7 @@ def update(phase, completed, total): "place, no download for anyone who can mount it.\n" "- Or set `BIREFNET_CHECKPOINT` to an explicit checkpoint path.\n" "- Or, from a terminal in a Git checkout: " - "`git lfs pull --include=\"weights/birefnet_leaf.pth\"`.\n" + "`git lfs pull -X \"\" -I \"weights/birefnet_leaf.pth\"`.\n" "- On air-gapped systems, pre-stage the checkpoint and verify its SHA-256 before " "launching MATS." ) diff --git a/src/mats/app/pages/3_CPU_Options.py b/src/mats/app/pages/3_CPU_Options.py index 327462e..e7226dc 100644 --- a/src/mats/app/pages/3_CPU_Options.py +++ b/src/mats/app/pages/3_CPU_Options.py @@ -58,7 +58,13 @@ def _worker_changed() -> None: st.stop() mask_method = selected_mask_method() -method_label = "BiRefNet" if mask_method == "birefnet" else "Classic thresholding (Otsu)" +if mask_method == "birefnet": + method_label = ( + "Classic thresholding + BiRefNet" + if st.session_state.get("segment_threshold", True) else "BiRefNet" + ) +else: + method_label = "Classic thresholding" settings = get_compute_settings(lm) available_workers = settings.available_workers birefnet_unlocked = birefnet_parallel_unlocked(available_workers) diff --git a/src/mats/app/pages/5_Help.py b/src/mats/app/pages/5_Help.py index 881dcdc..1011a20 100644 --- a/src/mats/app/pages/5_Help.py +++ b/src/mats/app/pages/5_Help.py @@ -65,12 +65,13 @@ def _samples_archive(): "2. **Photograph** one specimen, flat, inside the box, with all four corner " "markers in frame.\n" "3. **Pick your folders** in the sidebar: images in, results out.\n" - "4. **Set the scale** in Analyze — enter the finished printed-sheet width " + "4. **Set the scale** in Setup — enter the finished printed-sheet width " "and height. MATS derives the calibrated marker-centre area using the same " "margins as Template Creator. Tick *Variable " "dimensions* to read it from each image's QR code instead; if reads are " "unreliable, **Robust QR setup** adds sturdier decoders.\n" - "5. **Run**, then read **Results** and download the CSV." + "5. Choose image outputs in Setup, run and read results in Analyze, " + "adjust specimens in Adjust, then download saved files in Export." ) _page_link("pages/1_Template_Creator.py", "Open Template Creator", ":material/grid_on:") _page_link("pages/4_Robust_QR_Setup.py", "Open Robust QR setup", ":material/qr_code_scanner:") @@ -260,7 +261,7 @@ def _samples_archive(): "— still legible despite the blur. Because this is a legacy sheet, re-run " "it with the compatibility calibration option " f"(`-t {samples.QR_FAILURE_SAMPLE['calibration_dimensions']}`, or untick " - "*Variable dimensions* and use **Older or custom template?** in Analyze)." + "*Variable dimensions* and use **Older or custom template?** in Setup)." ) # ------------------------------------------------------- segmentation method @@ -268,11 +269,13 @@ def _samples_archive(): otsu_column, birefnet_column = st.columns(2) with otsu_column: with st.container(border=True): - st.markdown("**Classic thresholding (Otsu)** — the default") + st.markdown("**Classic thresholding** — the default") st.markdown( "- Fast, no GPU, no extra download.\n" "- Separates leaf from background by brightness.\n" "- Right choice for the **bench** samples above.\n" + "- If Otsu picks the wrong cutoff, choose **custom** under Threshold " + "level in Setup and drag the slider.\n" ) with birefnet_column: with st.container(border=True): @@ -286,8 +289,11 @@ def _samples_archive(): ) _page_link("pages/2_BiRefNet_Setup.py", "Open BiRefNet setup", ":material/download:") st.caption( - "Start with Otsu. Switch only when the mask shown in Results is visibly " - "wrong — that is the signal, not the file size or the leaf species." + "Start with Otsu. Switch only when the mask shown in Analyze is visibly " + "wrong — that is the signal, not the file size or the leaf species. To " + "compare them on your own photos, check both methods in Setup: each image is " + "measured with each method, and each method gets its own results CSV " + "(`leaf_morpho_results_threshold.csv`, `leaf_morpho_results_birefnet.csv`)." ) # ---------------------------------------------------------------- csv glossary @@ -295,7 +301,7 @@ def _samples_archive(): with st.container(border=True): st.markdown("**Full research schema** (default)") st.caption( - "Choose result units in Analyze before running: `mm`, `cm` (the default), " + "Choose result units in Setup before running: `mm`, `cm` (the default), " "or `in`. The selected unit appears in every measurement and pixels-per-unit " "column name; the table below shows the default centimeter names." ) @@ -304,7 +310,7 @@ def _samples_archive(): "|---|---|\n" "| `sample_id` | Input filename without its extension. Matches " "`{sample_id}_target_box.jpg` and `{sample_id}_mask.png` in the output " - "folder. |\n" + "folder when those exports are selected. |\n" "| `leaf_area_cm2` | Segmented leaf area — mask pixel count divided by " "`px_per_cm_width` x `px_per_cm_height`. |\n" "| `width_cm` | Horizontal extent of the leaf's bounding box, divided by " @@ -357,10 +363,53 @@ def _samples_archive(): ) st.caption( "Unmeasurable values are written as the literal `NA`. Also written per " - "image: `{sample_id}_target_box.jpg` (perspective-corrected box) and " - "`{sample_id}_mask.png` (segmentation mask). A failures log, when " + "image by default: `{sample_id}_target_box.jpg` (newly perspective-corrected " + "box) and `{sample_id}_mask.png` (cleaned mask). Setup can " + "also export `{sample_id}_mask_precleanup_{method}.png` for each checked " + "method: binary masks before cleanup. With both methods checked, the " + "cleaned masks, results CSV, and failure log are written once per method, " + "with names ending in `_threshold` or `_birefnet`. A failures log, when " "enabled, lists `sample_id, input_image, stage, failure_mode, status`." ) + st.caption( + "Measure from pre-cleanup masks in Setup to calculate area and extents after " + "the edge margin and stray pieces are removed. A Clean size " + "above 0 there also removes small specks and fills small holes before " + "measuring. Each results CSV has a `.meta.json` companion recording these " + "choices." + ) + st.caption( + "Select one measurement-table row to inspect that specimen. The Analyze " + "viewer shows the mask that produced its measurements, including when " + "that mask was not exported. In Adjust, check Remove flashfill under " + "Clean size for the selected cleaned-mask sample. The preview changes saved outputs " + "only after you press Overwrite. Extra previews last for this " + "app session." + ) + st.caption( + "For a classic thresholding sample, use Adjust to drag " + "the threshold preview slider; the masked leaf beside the mask shows, in " + "color, which parts of the leaf the cutoff keeps. Releasing it applies " + "cleanup to the preview when the run measured cleaned masks. Overwrite " + "this specimen replaces that sample's saved mask and CSV row. To use the " + "same settings for several specimens, tick their Marked for Adjustment " + "boxes in the table. Marks stay selected across searches. " + "Then press Overwrite all marked specimens, below Overwrite this " + "specimen, to save the settings to every marked specimen. Existing " + "overlays, cutouts, and axes are regenerated. Clean " + "image drops small specks and fills small holes with a clean-size slider; " + "Overwrite saves the clean size with the threshold and re-measures the " + "specimen. BiRefNet samples can preview Clean image, but it never changes " + "their saved files." + ) + st.caption( + "Export lists this run's saved files. Select Otsu, BiRefNet, or both; " + "choose which files go in the ZIP; or download a results CSV directly. " + "Changing ZIP choices does not change measurements or create new images. " + "Training dataset builds a separate image-and-mask ZIP with configurable " + "train/validation/test splits (70/20/10 by default) in PNG-mask, YOLO, " + "or COCO format. Review the generated labels before model training." + ) # -------------------------------------------------------------- troubleshooting st.subheader("Troubleshooting", anchor=False) diff --git a/src/mats/app/specimen_table.py b/src/mats/app/specimen_table.py new file mode 100644 index 0000000..fd4b44b --- /dev/null +++ b/src/mats/app/specimen_table.py @@ -0,0 +1,260 @@ +"""Adjust's specimen table: choose the specimen to view and mark specimens for adjustment.""" + +import streamlit as st + +_TABLE_HTML = """ +
+
+ + + +
+
+
+""" + +_TABLE_CSS = """ +.mats-specimen-table { + color: var(--st-text-color); font-family: var(--st-font); font-size: .875rem; +} +.mats-table-scroll { + position: relative; overflow: auto; + border: 1px solid var(--st-dataframe-border-color, var(--st-border-color)); + border-radius: var(--st-base-radius); +} +.mats-specimen-table table { width: 100%; border-collapse: separate; border-spacing: 0; } +.mats-specimen-table th { + position: sticky; top: 0; z-index: 1; padding: 0; text-align: left; font-weight: 400; + vertical-align: bottom; + color: color-mix(in srgb, var(--st-text-color) 65%, transparent); + background: var(--st-dataframe-header-background-color, var(--st-secondary-background-color)); + border-bottom: 1px solid var(--st-dataframe-border-color, var(--st-border-color)); +} +.mats-specimen-table th button { + all: unset; box-sizing: border-box; width: 100%; padding: .45rem .6rem; cursor: pointer; +} +.mats-specimen-table th button:focus-visible { outline: 2px solid var(--st-primary-color); } +.mats-specimen-table td { + padding: .4rem .6rem; white-space: nowrap; + border-bottom: 1px solid var(--st-dataframe-border-color, var(--st-border-color)); +} +.mats-specimen-table tbody tr:last-child td { border-bottom: 0; } +.mats-specimen-table .mats-num { text-align: right; font-variant-numeric: tabular-nums; } +.mats-specimen-table th.mats-num button { text-align: right; } +.mats-specimen-table .mats-col-view, +.mats-specimen-table .mats-col-mark { width: 1%; text-align: center; } +.mats-specimen-table th.mats-col-view, +.mats-specimen-table th.mats-col-mark { padding: .45rem .6rem; } +/* Wrap the long header rather than push the table wider than Adjust. */ +.mats-specimen-table th.mats-col-mark { min-width: 5.5rem; } +.mats-specimen-table tbody tr { cursor: pointer; } +.mats-specimen-table tbody tr:hover td { + background: color-mix(in srgb, var(--st-primary-color) 6%, transparent); +} +.mats-specimen-table tbody tr.mats-viewed td { + background: color-mix(in srgb, var(--st-primary-color) 12%, transparent); +} +.mats-specimen-table input { margin: 0; cursor: pointer; accent-color: var(--st-primary-color); } +/* The whole Marked cell toggles its checkbox. */ +.mats-specimen-table .mats-mark-hit { + display: flex; align-items: center; justify-content: center; + margin: -.4rem -.6rem; padding: .4rem .6rem; cursor: pointer; +} +.mats-specimen-table .mats-sort { margin-left: .3rem; font-size: .7em; } +""" + +_TABLE_JS = """ +export default function(component) { + const { data, parentElement, setTriggerValue } = component; + const scroller = parentElement.querySelector('.mats-table-scroll'); + const head = parentElement.querySelector('thead tr'); + const body = parentElement.querySelector('tbody'); + if (!scroller || !head || !body) return; + + // Sort order and scroll survive data updates because the table DOM does. + const state = parentElement.__matsTable || (parentElement.__matsTable = { + sort: null, scrolled: false, group: `mats-view-${Math.random().toString(36).slice(2)}`, + }); + const columns = data.columns || []; + const markable = Boolean(data.markable); + const marked = new Set(data.marked || []); + let view = data.view; + scroller.style.maxHeight = `${Number(data.height) || 300}px`; + + const numeric = (column) => column.decimals !== null && column.decimals !== undefined; + const format = (value, column) => { + if (value === null || value === undefined || value === '') return ''; + return numeric(column) ? Number(value).toFixed(column.decimals) : String(value); + }; + + function sortedRows() { + const rows = (data.rows || []).slice(); + const sort = state.sort; + const column = sort && columns.find((item) => item.key === sort.key); + if (!column) return rows; + rows.sort((a, b) => { + const x = a[column.key]; + const y = b[column.key]; + if (x === null || x === undefined) return 1; // blanks last either way + if (y === null || y === undefined) return -1; + const order = numeric(column) + ? x - y + : String(x).localeCompare(String(y), undefined, { numeric: true }); + return sort.descending ? -order : order; + }); + return rows; + } + + function headerCell(text, className) { + const th = document.createElement('th'); + th.scope = 'col'; + if (className) th.className = className; + th.textContent = text; + return th; + } + + function renderHead() { + const cells = [headerCell('View', 'mats-col-view')]; + for (const column of columns) { + const active = Boolean(state.sort) && state.sort.key === column.key; + const th = headerCell('', numeric(column) ? 'mats-num' : ''); + th.setAttribute('aria-sort', active + ? (state.sort.descending ? 'descending' : 'ascending') : 'none'); + const button = document.createElement('button'); + button.type = 'button'; + button.dataset.sort = column.key; + button.textContent = column.label; + const arrow = document.createElement('span'); + arrow.className = 'mats-sort'; + arrow.textContent = active ? (state.sort.descending ? '▼' : '▲') : ''; + button.append(arrow); + th.append(button); + cells.push(th); + } + if (markable) cells.push(headerCell('Marked for Adjustment', 'mats-col-mark')); + head.replaceChildren(...cells); + } + + function renderBody() { + const fragment = document.createDocumentFragment(); + for (const row of sortedRows()) { + const id = String(row.sample_id); + const tr = document.createElement('tr'); + tr.dataset.id = id; + tr.classList.toggle('mats-viewed', id === view); + const viewCell = document.createElement('td'); + viewCell.className = 'mats-col-view'; + const radio = document.createElement('input'); + radio.type = 'radio'; + radio.name = state.group; + radio.checked = id === view; + radio.setAttribute('aria-label', `View ${id}`); + viewCell.append(radio); + tr.append(viewCell); + for (const column of columns) { + const td = document.createElement('td'); + if (numeric(column)) td.className = 'mats-num'; + td.textContent = format(row[column.key], column); + tr.append(td); + } + if (markable) { + const td = document.createElement('td'); + td.className = 'mats-col-mark'; + const label = document.createElement('label'); + label.className = 'mats-mark-hit'; + const box = document.createElement('input'); + box.type = 'checkbox'; + box.checked = marked.has(id); + box.setAttribute('aria-label', `Mark ${id} for adjustment`); + label.append(box); + td.append(label); + tr.append(td); + } + fragment.append(tr); + } + body.replaceChildren(fragment); + } + + function render() { + renderHead(); + renderBody(); + } + + function choose(id) { + view = id; + for (const tr of body.rows) { + const on = tr.dataset.id === id; + tr.classList.toggle('mats-viewed', on); + tr.querySelector('input[type=radio]').checked = on; + } + setTriggerValue('view', id); + } + + head.onclick = (event) => { + const button = event.target.closest('button[data-sort]'); + if (!button) return; + const key = button.dataset.sort; + const again = Boolean(state.sort) && state.sort.key === key && !state.sort.descending; + state.sort = { key, descending: again }; + render(); + }; + body.onclick = (event) => { + if (event.target.closest('.mats-col-mark')) return; // marking never changes the view + const tr = event.target.closest('tr'); + if (tr && tr.dataset.id !== view) choose(tr.dataset.id); + }; + body.onchange = (event) => { + const box = event.target; + if (box.type !== 'checkbox') return; + const id = box.closest('tr').dataset.id; + if (box.checked) marked.add(id); else marked.delete(id); + setTriggerValue('mark', { id, marked: box.checked }); + }; + + render(); + // On first mount, bring the viewed specimen into sight. + if (!state.scrolled) { + state.scrolled = true; + const viewed = body.querySelector('tr.mats-viewed'); + if (viewed && viewed.offsetTop + viewed.offsetHeight > scroller.clientHeight) { + scroller.scrollTop = viewed.offsetTop - head.offsetHeight; + } + } + + return () => { + head.onclick = null; + body.onclick = null; + body.onchange = null; + }; +} +""" + + +def show_specimen_table( + *, key, columns, rows, view, marked=(), markable=False, height=300, + on_view_change=None, on_mark_change=None, +): + """Mount the specimen table. + + ``columns`` are ``{"key", "label", "decimals"}`` dicts (``decimals`` None for + text) and ``rows`` carry those keys, including ``sample_id``. Clicking a row + emits ``view`` (its sample id); with ``markable``, ticking **Marked for + Adjustment** emits ``mark`` as ``{"id", "marked"}``. + """ + # Register in the active Streamlit runtime, as threshold_preview does. + component = st.components.v2.component( + "mats_specimen_table", html=_TABLE_HTML, css=_TABLE_CSS, js=_TABLE_JS, + ) + return component( + key=key, + data={ + "columns": list(columns), + "rows": list(rows), + "view": view, + "marked": list(marked), + "markable": markable, + "height": height, + }, + on_view_change=on_view_change, + on_mark_change=on_mark_change, + ) diff --git a/src/mats/app/threshold_preview.py b/src/mats/app/threshold_preview.py new file mode 100644 index 0000000..e7371b7 --- /dev/null +++ b/src/mats/app/threshold_preview.py @@ -0,0 +1,543 @@ +"""Interactive, sample-only threshold and clean-image preview for the Analyze workbench.""" + +import base64 +from pathlib import Path + +import cv2 +import streamlit as st + +from mats.mask_cleanup import ( + CLEAN_RADIUS_DEFAULT, CLEAN_RADIUS_MAX, clean_levels, clean_raw_mask, +) +from mats.mask_settings import CLEAN_MARGIN_DEFAULT, STRAY_GAP_DEFAULT + +_PREVIEW_HTML = """ +
+ + + +
+
+ +
Mask
+
+
+ +
Masked leaf
+
+
+

+
+""" + +_PREVIEW_CSS = """ +.mats-threshold-preview { width: 100%; } +.mats-threshold-preview label { display: block; margin-bottom: .35rem; font-weight: 600; } +.mats-threshold-preview input { width: 100%; accent-color: var(--st-primary-color); } +.mats-threshold-presets { position: relative; height: 1.9rem; margin: -.15rem .45rem 0; } +.mats-threshold-presets span { + position: absolute; top: 0; transform: translateX(-50%); white-space: nowrap; + color: var(--st-text-color); font-size: .7rem; text-align: center; +} +/* A tick from the bar down to each preset label. */ +.mats-threshold-presets span::before { + content: ""; display: block; width: 2px; height: .55rem; margin: 0 auto .1rem; + border-radius: 1px; background: var(--st-text-color); opacity: .6; +} +.mats-threshold-panels { + display: grid; grid-template-columns: repeat(auto-fit, minmax(12rem, 1fr)); + gap: .75rem; margin-top: .65rem; +} +.mats-threshold-panels figure { margin: 0; min-width: 0; } +.mats-threshold-preview [hidden] { display: none; } +.mats-clean-control { margin-top: .5rem; } +.mats-clean-label { display: flex; align-items: center; gap: .35rem; margin-bottom: .35rem; } +.mats-threshold-preview .mats-clean-label label { margin-bottom: 0; } +.mats-fill-control { margin-top: .5rem; } +.mats-threshold-preview label.mats-fill-label { + display: inline-flex; align-items: center; gap: .45rem; margin: 0; + font-weight: 400; cursor: pointer; +} +.mats-threshold-preview input.mats-remove-fill { width: auto; margin: 0; } +.mats-threshold-preview label.mats-fill-label:has(input:disabled) { cursor: not-allowed; opacity: .6; } +.mats-fill-note { margin: .25rem 0 0; color: var(--st-text-color); font-size: .8rem; opacity: .75; } +.mats-help { + position: relative; display: inline-flex; color: var(--st-text-color); + opacity: .6; cursor: help; border-radius: 50%; +} +.mats-help:hover, .mats-help:focus-visible { opacity: 1; } +.mats-help-tip { + position: absolute; left: 0; bottom: calc(100% + .4rem); z-index: 10; + width: max-content; max-width: 20rem; padding: .5rem .65rem; + border-radius: var(--st-base-radius); background: var(--st-text-color); + color: var(--st-background-color); font-size: .8rem; font-weight: 400; + line-height: 1.4; visibility: hidden; opacity: 0; pointer-events: none; + transition: opacity .12s; +} +.mats-help:hover .mats-help-tip, .mats-help:focus-visible .mats-help-tip { + visibility: visible; opacity: 1; +} +.mats-threshold-panels canvas { + display: block; width: 100%; height: auto; max-height: 32rem; + object-fit: contain; background: var(--st-secondary-background-color); +} +.mats-threshold-panels figcaption { + margin-top: .25rem; color: var(--st-text-color); font-size: .8rem; +} +.mats-threshold-status { margin: .4rem 0 0; color: var(--st-text-color); font-size: .85rem; } +""" + +_PREVIEW_JS = """ +export default function(component) { + const { data, parentElement, setTriggerValue } = component; + const find = (selector) => parentElement.querySelector(selector); + const thresholdControl = find('.mats-threshold-control'); + const slider = find('input.mats-threshold-range'); + const valueLabel = find('output.mats-threshold-value'); + const cleanControl = find('.mats-clean-control'); + const cleanSlider = find('input.mats-clean-range'); + const cleanLabel = find('output.mats-clean-value'); + const cleanHelp = find('.mats-clean-control .mats-help-tip'); + const fillControl = find('.mats-fill-control'); + const fillBox = find('input.mats-remove-fill'); + const fillNote = find('.mats-fill-note'); + const maskCanvas = find('canvas.mats-threshold-mask'); + const leafCanvas = find('canvas.mats-threshold-leaf'); + const leafPanel = find('.mats-threshold-leaf-panel'); + const status = find('.mats-threshold-status'); + if (![thresholdControl, slider, valueLabel, cleanControl, cleanSlider, cleanLabel, + cleanHelp, fillControl, fillBox, fillNote, maskCanvas, leafCanvas, leafPanel, + status].every(Boolean)) return; + + const maskContext = maskCanvas.getContext('2d'); + const leafContext = leafCanvas.getContext('2d'); + // A grayscale image means a live threshold; otherwise the mask is fixed. + const thresholded = Boolean(data.grayscale_image); + // Levels make the clean-size slider live; a clean size of 0 turns it off. + const cleaning = Boolean(data.clean_levels_image); + const cleanActive = () => cleaning && Number(cleanSlider.value) > 0; + let pixels = null; + let baseMask = null; + let levels = null; + let colors = null; + let frame = null; + let disposed = false; + let dragging = false; + thresholdControl.hidden = !thresholded; + cleanControl.hidden = !cleaning; + leafPanel.hidden = !data.color_image; + if (thresholded) { + slider.value = String(Number(data.cutoff)); + valueLabel.textContent = slider.value; + } + cleanSlider.max = String(Number(data.clean_max)); + cleanSlider.value = String(Number(data.clean_radius)); + cleanLabel.textContent = cleanSlider.value; + cleanHelp.textContent = data.clean_help || ''; + fillControl.hidden = !data.fill_toggle; + fillBox.checked = Boolean(data.remove_fill); + fillBox.disabled = Boolean(data.fill_note); + fillNote.textContent = data.fill_note || ''; + fillNote.hidden = !data.fill_note; + + function loadImage(src) { + return new Promise((resolve, reject) => { + const image = new Image(); + image.onload = () => resolve(image); + image.onerror = reject; + image.src = src; + }); + } + + // Nearest-neighbour sampling keeps the grayscale, color, mask, and level + // pixels aligned one-to-one on the preview canvases. + function readPixels(image) { + const offscreen = document.createElement('canvas'); + offscreen.width = maskCanvas.width; + offscreen.height = maskCanvas.height; + const offscreenContext = offscreen.getContext('2d', { willReadFrequently: true }); + offscreenContext.imageSmoothingEnabled = false; + offscreenContext.drawImage(image, 0, 0, offscreen.width, offscreen.height); + return offscreenContext.getImageData(0, 0, offscreen.width, offscreen.height).data; + } + + // Paint the binary mask, and beside it the sample's own colors wherever the + // mask is foreground; excluded pixels stay transparent. + function paint(isForeground) { + const mask = maskContext.createImageData(maskCanvas.width, maskCanvas.height); + const leaf = colors + ? leafContext.createImageData(leafCanvas.width, leafCanvas.height) + : null; + for (let i = 0; i < mask.data.length; i += 4) { + mask.data[i + 3] = 255; + if (!isForeground(i)) continue; + mask.data[i] = 255; + mask.data[i + 1] = 255; + mask.data[i + 2] = 255; + if (leaf) { + leaf.data[i] = colors[i]; + leaf.data[i + 1] = colors[i + 1]; + leaf.data[i + 2] = colors[i + 2]; + leaf.data[i + 3] = 255; + } + } + maskContext.putImageData(mask, 0, 0); + if (leaf) leafContext.putImageData(leaf, 0, 0); + } + + // The edge band that mats.mask_cleanup.clear_margin sets to background: + // live_margin percent of the shorter side, in canvas pixels. + function marginBand() { + const shorter = Math.min(maskCanvas.width, maskCanvas.height); + return Math.round(shorter * Number(data.live_margin || 0) / 100); + } + + function drawRaw() { + frame = null; + if (!pixels || disposed) return; + const cutoff = Number(slider.value); + const band = marginBand(); + const width = maskCanvas.width; + const height = maskCanvas.height; + paint((i) => { + if (pixels[i] > cutoff) return false; + if (!band) return true; + const x = (i / 4) % width; + const y = Math.floor(i / 4 / width); + return x >= band && y >= band && x < width - band && y < height - band; + }); + const edge = band ? ' with the edge margin cleared' : ''; + if (cleanActive()) { + status.textContent = `Live mask${edge} before cleaning; release to see the clean result`; + } else if (data.measurement_source === 'pre-cleanup') { + status.textContent = `Live raw mask${edge}; release to see the measured mask`; + } else { + status.textContent = `Live mask${edge} before cleanup; release to see the cleaned result`; + } + } + + // Levels hold each pixel's keep and fill radii for the mask they were + // measured on (mats.mask_cleanup); with a threshold, that is one cutoff. + function levelsCurrent() { + return Boolean(levels) && + (!thresholded || Number(slider.value) === Number(data.clean_cutoff)); + } + + function drawClean() { + frame = null; + if (disposed || !levelsCurrent()) return; + const radius = Number(cleanSlider.value); + paint((i) => radius < levels[i] || (levels[i + 1] <= radius && radius <= levels[i + 2])); + status.textContent = radius < 2 + ? 'Clean preview: edge margin and stray pieces removed; no specks or holes change at this size' + : `Clean preview: edge margin and stray pieces removed; specks removed and holes filled below ${radius} px`; + } + + function showSettled() { + if (disposed || dragging) return; + if (cleanActive()) { + if (levelsCurrent()) drawClean(); else drawRaw(); + return; + } + if (!thresholded) { + if (!baseMask) return; + paint((i) => baseMask[i] >= 128); + status.textContent = data.mask_status || 'Saved measurement mask'; + return; + } + if (!pixels) return; + const cutoff = Number(slider.value); + if (cutoff !== Number(data.cleaned_cutoff) || !data.cleaned_image) { + drawRaw(); + return; + } + loadImage(data.cleaned_image).then((cleaned) => { + if (disposed || dragging || cleanActive() || Number(slider.value) !== cutoff) return; + const settled = readPixels(cleaned); + paint((i) => settled[i] >= 128); + if (data.measurement_source === 'pre-cleanup') { + status.textContent = + 'Pre-cleanup measurement mask: edge margin cleared and stray pieces dropped'; + } else { + status.textContent = data.remove_fill + ? 'Preview with hole filling removed and the edge margin cleared' + : 'Cleaned preview from the MATS pipeline'; + } + }, () => {}); + } + + slider.oninput = () => { + dragging = true; + valueLabel.textContent = slider.value; + if (frame !== null) cancelAnimationFrame(frame); + frame = requestAnimationFrame(drawRaw); + }; + slider.onchange = () => { + dragging = false; + if (frame !== null) cancelAnimationFrame(frame); + drawRaw(); + if (cleanActive()) { + status.textContent = 'Measuring speck and hole sizes for this cutoff…'; + } else if (data.measurement_source === 'pre-cleanup') { + status.textContent = 'Dropping stray pieces for this cutoff…'; + } else { + status.textContent = 'Applying MATS cleanup to this sample…'; + } + setTriggerValue('cutoff', Number(slider.value)); + }; + cleanSlider.oninput = () => { + cleanLabel.textContent = cleanSlider.value; + if (frame !== null) cancelAnimationFrame(frame); + frame = null; + if (!cleanActive()) { + showSettled(); + return; + } + if (levelsCurrent()) frame = requestAnimationFrame(drawClean); + }; + cleanSlider.onchange = () => { + setTriggerValue('clean_radius', Number(cleanSlider.value)); + }; + fillBox.onchange = () => { + // The server recomputes the mask; hold the box until it answers. + fillBox.disabled = true; + status.textContent = fillBox.checked + ? 'Removing flashfill from this preview…' + : 'Restoring flashfill to this preview…'; + setTriggerValue('remove_fill', fillBox.checked); + }; + + const optional = (src) => (src ? loadImage(src) : Promise.resolve(null)); + Promise.all([ + optional(data.grayscale_image), + optional(data.mask_image), + optional(data.clean_levels_image), + optional(data.color_image).catch(() => null), + ]).then(([gray, mask, levelsImage, colorImage]) => { + if (disposed) return; + const source = gray || mask || levelsImage; + if (!source) { + status.textContent = 'Could not load this sample for preview.'; + return; + } + const scale = Math.min(1, 1400 / Math.max(source.naturalWidth, source.naturalHeight)); + const width = Math.max(1, Math.round(source.naturalWidth * scale)); + const height = Math.max(1, Math.round(source.naturalHeight * scale)); + maskCanvas.width = leafCanvas.width = width; + maskCanvas.height = leafCanvas.height = height; + pixels = gray ? readPixels(gray) : null; + baseMask = mask ? readPixels(mask) : null; + levels = levelsImage ? readPixels(levelsImage) : null; + colors = colorImage ? readPixels(colorImage) : null; + leafPanel.hidden = !colors; + showSettled(); + }, () => { status.textContent = 'Could not load this sample for preview.'; }); + + return () => { + disposed = true; + if (frame !== null) cancelAnimationFrame(frame); + slider.oninput = null; + slider.onchange = null; + cleanSlider.oninput = null; + cleanSlider.onchange = null; + fillBox.onchange = null; + }; +} +""" + + +def _png_data_url(image): + ok, encoded = cv2.imencode(".png", image) + if not ok: + raise ValueError("could not encode threshold preview") + return "data:image/png;base64," + base64.b64encode(encoded).decode("ascii") + + +@st.cache_data(max_entries=8, show_spinner=False) +def grayscale_sample(path): + """Return exact OpenCV grayscale pixels and the sample's Otsu cutoff.""" + target = cv2.imread(str(path), cv2.IMREAD_COLOR) + if target is None: + raise ValueError(f"Could not read sample target box: {Path(path).name}") + gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY) + otsu_cutoff, _ = cv2.threshold( + gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU + ) + return _png_data_url(gray), int(otsu_cutoff) + + +@st.cache_data(max_entries=8, show_spinner=False) +def color_sample(path): + """Return the sample's color target box for the masked-leaf preview panel.""" + target = cv2.imread(str(path), cv2.IMREAD_COLOR) + if target is None: + raise ValueError(f"Could not read sample target box: {Path(path).name}") + # Display only: JPEG keeps a full-resolution color payload small, and the + # mask itself always comes from the lossless grayscale image. + ok, encoded = cv2.imencode(".jpg", target, [cv2.IMWRITE_JPEG_QUALITY, 90]) + if not ok: + raise ValueError("could not encode color preview") + return "data:image/jpeg;base64," + base64.b64encode(encoded).decode("ascii") + + +@st.cache_data(max_entries=8, show_spinner=False) +def cleaned_sample(path, cutoff, *, fill_holes=True, clean_margin=CLEAN_MARGIN_DEFAULT): + """Run the production mask cleanup on just the selected threshold sample. + + ``fill_holes=False`` is Remove flashfill, which clears the edge margin first. + """ + from mats import core + + target = cv2.imread(str(path), cv2.IMREAD_COLOR) + if target is None: + raise ValueError(f"Could not read sample target box: {Path(path).name}") + raw = core.threshold_mask(target, cutoff) + if fill_holes: + return _png_data_url(core.clean_leaf_mask(raw.copy())) + return _png_data_url(core.unfilled_leaf_mask(raw, clean_margin)) + + +@st.cache_data(max_entries=8, show_spinner=False) +def unfilled_sample(raw_mask_path, clean_margin=CLEAN_MARGIN_DEFAULT): + """Keep the largest raw component, past the edge margin, without filling holes.""" + from mats import core + + raw = cv2.imread(str(raw_mask_path), cv2.IMREAD_GRAYSCALE) + if raw is None: + raise ValueError(f"Could not read raw mask: {Path(raw_mask_path).name}") + return _png_data_url(core.unfilled_leaf_mask(raw, clean_margin)) + + +@st.cache_data(max_entries=8, show_spinner=False) +def pre_cleanup_sample(path, cutoff, clean_margin=CLEAN_MARGIN_DEFAULT, + stray_gap=STRAY_GAP_DEFAULT): + """The pre-cleanup measurement mask for one cutoff (``clean_raw_mask``).""" + from mats import core + + target = cv2.imread(str(path), cv2.IMREAD_COLOR) + if target is None: + raise ValueError(f"Could not read sample target box: {Path(path).name}") + raw = core.threshold_mask(target, cutoff) + return _png_data_url(clean_raw_mask(raw, clean_margin, stray_gap)) + + +def _levels_data_url(levels): + # OpenCV writes BGR; reverse so the browser reads the channels in order. + return _png_data_url(cv2.cvtColor(levels, cv2.COLOR_RGB2BGR)) + + +@st.cache_data(max_entries=8, show_spinner=False) +def clean_levels_for_threshold(path, cutoff, stray_gap=STRAY_GAP_DEFAULT, + clean_margin=CLEAN_MARGIN_DEFAULT): + """Speck and hole sizes of the raw threshold mask, for the clean slider.""" + from mats import core + + target = cv2.imread(str(path), cv2.IMREAD_COLOR) + if target is None: + raise ValueError(f"Could not read sample target box: {Path(path).name}") + raw = core.threshold_mask(target, cutoff) + return _levels_data_url(clean_levels(raw, stray_gap, clean_margin)) + + +@st.cache_data(max_entries=8, show_spinner=False) +def clean_levels_for_mask(path, mtime_ns, stray_gap=STRAY_GAP_DEFAULT, + clean_margin=CLEAN_MARGIN_DEFAULT): + """Speck and hole sizes of a saved raw mask; ``mtime_ns`` keys the cache.""" + mask = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) + if mask is None: + raise ValueError(f"Could not read raw mask: {Path(path).name}") + return _levels_data_url(clean_levels(mask, stray_gap, clean_margin)) + + +@st.cache_data(max_entries=8, show_spinner=False) +def pre_cleanup_mask_sample(path, mtime_ns, clean_margin=CLEAN_MARGIN_DEFAULT, + stray_gap=STRAY_GAP_DEFAULT): + """A saved raw mask's pre-cleanup measurement mask; ``mtime_ns`` keys the cache.""" + mask = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) + if mask is None: + raise ValueError(f"Could not read raw mask: {Path(path).name}") + return _png_data_url(clean_raw_mask(mask, clean_margin, stray_gap)) + + +@st.cache_data(max_entries=8, show_spinner=False) +def mask_sample(path, mtime_ns): + """A saved mask for the fixed-mask preview; ``mtime_ns`` keys the cache.""" + mask = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) + if mask is None: + raise ValueError(f"Could not read mask: {Path(path).name}") + return _png_data_url(mask) + + +def show_threshold_preview( + *, key, grayscale_image=None, cutoff=None, measurement_source="cleaned", + mask_image=None, mask_status=None, color_image=None, cleaned_image=None, + cleaned_cutoff=None, remove_fill=False, clean_levels_image=None, clean_cutoff=None, + clean_radius=CLEAN_RADIUS_DEFAULT, clean_help=None, fill_toggle=False, fill_note=None, + live_margin=0, on_cutoff_change=None, on_clean_radius_change=None, + on_remove_fill_change=None, +): + """Mount the local drag preview; emit values only when a drag ends. + + With ``grayscale_image`` the threshold is live; otherwise ``mask_image`` is + shown as is, labelled ``mask_status``. ``clean_levels_image`` switches on the + clean-size slider; above 0 it previews Clean image, and ``clean_help`` fills + the help icon beside it. ``fill_toggle`` shows the Remove flashfill checkbox + under it, checked when ``remove_fill``; ``fill_note`` disables it and says why. + ``cleaned_image`` is the server's settled mask for ``cleaned_cutoff``, and + ``live_margin`` is the edge-margin percent cleared while dragging. + """ + # Register in the active Streamlit runtime. AppTest and the running app use + # separate registries, so a registration made at import time can be stale. + component = st.components.v2.component( + "mats_threshold_preview", html=_PREVIEW_HTML, css=_PREVIEW_CSS, js=_PREVIEW_JS, + ) + return component( + key=key, + data={ + "grayscale_image": grayscale_image, + "mask_image": mask_image, + "mask_status": mask_status, + "color_image": color_image, + "cutoff": cutoff, + "measurement_source": measurement_source, + "cleaned_image": cleaned_image, + "cleaned_cutoff": cleaned_cutoff, + "remove_fill": remove_fill, + "clean_levels_image": clean_levels_image, + "clean_cutoff": clean_cutoff, + "clean_radius": clean_radius, + "clean_help": clean_help, + "clean_max": CLEAN_RADIUS_MAX, + "fill_toggle": fill_toggle, + "fill_note": fill_note, + "live_margin": live_margin, + }, + on_cutoff_change=on_cutoff_change, + on_clean_radius_change=on_clean_radius_change, + on_remove_fill_change=on_remove_fill_change, + ) diff --git a/src/mats/cli.py b/src/mats/cli.py index 536c30c..df6f178 100644 --- a/src/mats/cli.py +++ b/src/mats/cli.py @@ -15,6 +15,19 @@ import os import sys +from .mask_settings import ( + CLEAN_MARGIN_DEFAULT, + CLEAN_MARGIN_MAX, + CLEAN_SIZE_DEFAULT, + CLEAN_SIZE_MAX, + STRAY_GAP_DEFAULT, + STRAY_GAP_MAX, + checked_clean_margin, + checked_clean_size, + checked_stray_gap, +) +from .thresholds import parse_threshold_level, threshold_value_for + _RUN_SUBCOMMAND = "run" _SUBCOMMANDS = {"run", "app", "fetch-weights", "doctor"} @@ -24,6 +37,42 @@ def _fail(message, code=2): raise SystemExit(code) +def _threshold_level_arg(text): + try: + return parse_threshold_level(text) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from None + + +def _stray_gap_arg(text): + try: + return checked_stray_gap(text) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from None + + +def _clean_margin_arg(text): + try: + return checked_clean_margin(text) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from None + + +def _clean_size_arg(text): + try: + return checked_clean_size(text) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from None + + +class _TrackProvided(argparse.Action): + """Keep an option's parsed value and whether it was explicitly supplied.""" + + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, values) + setattr(namespace, f"{self.dest}_provided", True) + + def build_parser(): """Build the top-level argument parser.""" parser = argparse.ArgumentParser( @@ -60,22 +109,80 @@ def build_parser(): ) run.add_argument('--output-mode', choices=('masks', 'target-boxes'), default='masks', help='Produce segmentation masks, or only perspective-corrected target boxes.') - run.add_argument('--mask-method', choices=('threshold', 'birefnet'), default='threshold', + run.add_argument('--mask-method', choices=('threshold', 'birefnet', 'both'), default='threshold', help='Mask method when --output-mode masks. threshold (Otsu) is fast, needs ' 'no GPU and no extra download; birefnet is more accurate on cluttered ' 'backgrounds but needs a locally installed ~2.65 GB checkpoint ' - '(`mats fetch-weights --only birefnet --source lfs`).') - run.add_argument('--threshold-level', choices=('auto', 'low', 'medium', 'high'), default='auto', - help="For --mask-method threshold: auto uses Otsu (recommended); " - "low=100, medium=125, high=150.") + '(`mats fetch-weights --only birefnet --source lfs`); both measures every ' + 'image with each method and writes one CSV per method, suffixed ' + '_threshold and _birefnet.') + run.add_argument('--threshold-level', type=_threshold_level_arg, default='auto', + metavar='{auto,low,medium,high,1-255}', + help="Cutoff for --mask-method threshold and threshold pre-cleanup exports: " + "auto uses Otsu (recommended); low=100, medium=125, high=150; or an " + "integer 1-255 for a custom cutoff. Grayscale pixels at or below the " + "cutoff are counted as leaf.") run.add_argument('--csv-schema', choices=('full', 'compact'), default='full', help='full = area/width/length plus per-axis pixels-per-unit and scale_aspect_ratio ' '(research schema); compact = sample_id, area, width, length.') run.add_argument('--results-unit', choices=('mm', 'cm', 'in'), default='cm', help='Unit for area, width, length, and pixels-per-unit columns in the CSV ' '(default: cm).') + run.add_argument('--measure-pre-cleanup', action='store_true', + help='Derive area, width, and length from the raw binary mask before ' + 'gap closing and hole filling. A band --clean-margin wide along the ' + 'target-box edge (where the printed box outline lands) is cleared, ' + 'the largest object is the leaf, and other pieces are dropped when ' + 'they touch that band or lie farther from the leaf than --stray-gap.') + run.add_argument('--clean-margin', type=_clean_margin_arg, action=_TrackProvided, + default=CLEAN_MARGIN_DEFAULT, + metavar='PERCENT', + help='With --measure-pre-cleanup: width of the band cleared along every ' + 'target-box edge, as a percent of the box\'s shorter side ' + f'(0-{CLEAN_MARGIN_MAX:g}; 0 clears nothing; ' + f'default {CLEAN_MARGIN_DEFAULT:g}).') + run.add_argument('--stray-gap', type=_stray_gap_arg, action=_TrackProvided, + default=STRAY_GAP_DEFAULT, + metavar='FRACTION', + help='With --measure-pre-cleanup: drop pieces whose nearest pixel is ' + 'farther from the leaf than this fraction of the leaf\'s ' + f'bounding-box diagonal (0-{STRAY_GAP_MAX:g}; 0 keeps only the leaf; ' + f'default {STRAY_GAP_DEFAULT:g}).') + run.add_argument('--clean-size', type=_clean_size_arg, action=_TrackProvided, + default=CLEAN_SIZE_DEFAULT, + metavar='PX', + help='With --measure-pre-cleanup: after the edge margin and stray pieces ' + 'are cleared, remove white specks and fill enclosed holes whose ' + 'inscribed radius is below this many pixels; the leaf is always kept ' + 'and nothing is flash-filled. This is the app\'s Clean size ' + f'(0-{CLEAN_SIZE_MAX}; 0 turns it off; default {CLEAN_SIZE_DEFAULT}).') run.add_argument('--save-axes', action='store_true', help='Also save per-image length/width measurement-axis overlays for QC.') + run.add_argument('--export', action='append', choices=('pre-cleanup', 'overlay', 'cutout', 'axes'), + default=[], help='Additional image export; repeat for multiple kinds.') + run.add_argument('--pre-cleanup-methods', choices=('selected', 'threshold', 'birefnet', 'both'), + default=None, help='Methods whose binary masks are saved before cleanup; ' + 'selected = every --mask-method method.') + run.add_argument('--no-target-boxes', action='store_true', + help='Do not save new perspective-corrected target-box images.') + run.add_argument('--no-masks', action='store_true', + help='Do not save the cleaned measurement masks.') + run.add_argument('--no-failure-log', action='store_true', + help='Do not write the failures/warnings CSV.') + run.add_argument('--dataset-format', choices=('png', 'yolo-seg', 'yolo-detect', 'coco'), + default=None, help='Also export an aligned training dataset ZIP to the output ' + 'folder: PNG mask pairs, YOLO segmentation, YOLO detection, or COCO segmentation.') + run.add_argument('--dataset-split', nargs=3, type=int, metavar=('TRAIN', 'VAL', 'TEST'), + default=None, help='Training/validation/test percentages; must total 100 ' + '(default: 70 20 10).') + run.add_argument('--dataset-seed', type=int, default=None, + help='Reproducible split seed (default: 42).') + run.add_argument('--dataset-method', choices=('threshold', 'birefnet'), default=None, + help='Labeling method when --mask-method both (default: first selected method).') + run.add_argument('--dataset-mask-source', choices=('measured', 'raw'), default=None, + help='Use the measured mask or raw pre-cleanup mask (default: measured).') + run.add_argument('--dataset-groups', default=None, metavar='CSV', + help='UTF-8 CSV with sample_id,group_id columns; keep each group in one split.') app = sub.add_parser('app', help='Launch the Streamlit GUI.') app.add_argument('extra', nargs=argparse.REMAINDER, @@ -112,18 +219,51 @@ def _normalize_argv(argv): return [_RUN_SUBCOMMAND] + argv +def _measured_methods(args): + """The methods this run measures with, in output order.""" + return ('threshold', 'birefnet') if args.mask_method == 'both' else (args.mask_method,) + + +def _pre_cleanup_methods(args): + """The methods whose pre-cleanup masks this run exports.""" + if 'pre-cleanup' not in args.export: + return () + if args.pre_cleanup_methods in (None, 'selected'): + return _measured_methods(args) + if args.pre_cleanup_methods == 'both': + return ('threshold', 'birefnet') + return (args.pre_cleanup_methods,) + + +def _uses_threshold(args): + """Whether this run thresholds for measurements or a pre-cleanup export.""" + return 'threshold' in _measured_methods(args) + _pre_cleanup_methods(args) + + def _print_run_banner(args, threshold_value): print(f"\nOutput mode: {args.output_mode}") print("Scale: independent per-axis pixels-per-cm (anisotropic)") if args.output_mode == "masks": print(f"Mask method: {args.mask_method}") - if args.mask_method == "threshold": + print(f"Measurement source: {'pre-cleanup' if args.measure_pre_cleanup else 'cleaned'}") + if args.measure_pre_cleanup: + print(f"Clean margin: {args.clean_margin:g}% of the target box's shorter side") + print(f"Stray gap: {args.stray_gap:g} x leaf bounding-box diagonal") + print(f"Clean size: {args.clean_size} px" if args.clean_size else "Clean size: off") + if _uses_threshold(args): if args.threshold_level == "auto": print("Threshold level: auto (Otsu's method)") + elif isinstance(args.threshold_level, int): + print(f"Threshold level: custom ({threshold_value})") else: print(f"Threshold level: {args.threshold_level} ({threshold_value})") else: print("Mask method and threshold level ignored because output mode is target-boxes.") + print(f"Additional exports: {', '.join(args.export) if args.export else 'none'}") + if args.output_mode == "target-boxes" and args.export: + print("Segmentation exports ignored because output mode is target-boxes.") + elif 'pre-cleanup' in args.export: + print(f"Pre-cleanup methods: {args.pre_cleanup_methods or 'selected'}") def _resolve_template_dims(args, lm): @@ -215,7 +355,9 @@ def _cb(info): def _require_local_birefnet_for_run(args): """Stop once, before a batch, when optional BiRefNet is unavailable.""" - if args.output_mode != "masks" or args.mask_method != "birefnet": + if args.output_mode != "masks" or ( + 'birefnet' not in _measured_methods(args) + _pre_cleanup_methods(args) + ): return from . import weights from .birefnet_runtime import require_birefnet_dependencies @@ -230,47 +372,188 @@ def _require_local_birefnet_for_run(args): def _cmd_run(args): from . import core as lm + dataset_options = ( + args.dataset_split, args.dataset_seed, args.dataset_method, + args.dataset_mask_source, args.dataset_groups, + ) + if not args.dataset_format and any(option is not None for option in dataset_options): + _fail('--dataset-format is required when using other --dataset-* options') + percentages = tuple(args.dataset_split or (70, 20, 10)) + dataset_seed = 42 if args.dataset_seed is None else args.dataset_seed + dataset_mask_source = args.dataset_mask_source or 'measured' + dataset_groups = {} + if args.dataset_format: + from .dataset_export import read_group_csv, split_counts + + if args.output_mode != 'masks': + _fail('--dataset-format requires --output-mode masks') + if args.dataset_method and args.dataset_method not in _measured_methods(args): + _fail('--dataset-method must be selected by --mask-method') + if dataset_seed < 0: + _fail('--dataset-seed must be zero or greater') + try: + split_counts(0, percentages) + if args.dataset_groups: + with open(args.dataset_groups, 'rb') as group_file: + dataset_groups = read_group_csv(group_file.read()) + except (OSError, ValueError) as exc: + _fail(f'invalid dataset settings: {exc}') + + if args.output_mode != 'masks' and args.measure_pre_cleanup: + _fail('--measure-pre-cleanup requires --output-mode masks') + if args.pre_cleanup_methods is not None and 'pre-cleanup' not in args.export: + _fail('--pre-cleanup-methods requires --export pre-cleanup') + if not args.measure_pre_cleanup and ( + getattr(args, 'stray_gap_provided', False) or args.stray_gap != STRAY_GAP_DEFAULT + ): + _fail('--stray-gap requires --measure-pre-cleanup') + if not args.measure_pre_cleanup and ( + getattr(args, 'clean_margin_provided', False) or args.clean_margin != CLEAN_MARGIN_DEFAULT + ): + _fail('--clean-margin requires --measure-pre-cleanup') + if not args.measure_pre_cleanup and ( + getattr(args, 'clean_size_provided', False) or args.clean_size != CLEAN_SIZE_DEFAULT + ): + _fail('--clean-size requires --measure-pre-cleanup') _require_local_birefnet_for_run(args) template_dims = _resolve_template_dims(args, lm) - threshold_value = lm.THRESHOLD_LEVELS[args.threshold_level] + threshold_value = threshold_value_for(args.threshold_level) _print_run_banner(args, threshold_value) results_path = args.results_path or os.path.join(os.getcwd(), "leaf_morpho_results.csv") - print("\nMeasurement CSV will be written to:", results_path) + if args.output_mode == "masks" and len(_measured_methods(args)) > 1: + print("\nMeasurement CSVs will be written to:") + for method in _measured_methods(args): + print(f" {method}: {lm.method_suffixed_path(results_path, method)}") + else: + print("\nMeasurement CSV will be written to:", results_path) input_dir = _resolve_input_dir(args) input_images = lm.get_input_images(input_dir) if not input_images: _fail(f"No images found in {input_dir}") print(f"Found {len(input_images)} image(s).") + if dataset_groups: + sample_ids = { + lm.target_box_sample_id(path) if lm.is_target_box_image(path) + else os.path.splitext(os.path.basename(path))[0] + for path in input_images + } + unknown = sorted(set(dataset_groups) - sample_ids) + if unknown: + _fail('dataset group CSV contains unknown sample IDs: ' + ', '.join(unknown[:5])) output_dir = _resolve_output_dir(args) - - result = lm.run_leaf_morpho_batch( - input_images=input_images, - output_dir=output_dir, - results_path=results_path, - template_dimensions=template_dims, - output_mode=args.output_mode, - mask_method=args.mask_method, - threshold_value=threshold_value, - workers=args.workers, - write_failures=True, - compact_csv=(args.csv_schema == "compact"), - results_unit=args.results_unit, - save_measurement_axes=args.save_axes, - serialize_model_inference=False, - progress_callback=_make_progress_callback(), - ) + if args.dataset_format and output_dir is False: + _fail('--dataset-format requires an output directory') + export_options = { + 'target_boxes': not args.no_target_boxes, + 'cleaned_masks': not args.no_masks, + 'pre_cleanup_methods': _pre_cleanup_methods(args), + 'overlay': 'overlay' in args.export, + 'cutout': 'cutout' in args.export, + 'axes': args.save_axes or 'axes' in args.export, + } + + from contextlib import nullcontext + from tempfile import TemporaryDirectory + + preview_context = (TemporaryDirectory(prefix='mats_dataset_preview_') + if args.dataset_format else nullcontext(None)) + with preview_context as preview_dir: + if preview_dir: + export_options['preview_dir'] = preview_dir + result = lm.run_leaf_morpho_batch( + input_images=input_images, + output_dir=output_dir, + results_path=results_path, + template_dimensions=template_dims, + output_mode=args.output_mode, + mask_method=args.mask_method, + threshold_value=threshold_value, + workers=args.workers, + write_failures=not args.no_failure_log, + compact_csv=(args.csv_schema == "compact"), + results_unit=args.results_unit, + save_measurement_axes=args.save_axes, + serialize_model_inference=False, + progress_callback=_make_progress_callback(), + export_options=export_options, + measurement_source='pre-cleanup' if args.measure_pre_cleanup else 'cleaned', + stray_gap=args.stray_gap, + clean_margin=args.clean_margin, + clean_size=args.clean_size, + ) + if args.dataset_format: + _export_run_dataset( + result, input_images, output_dir, args.dataset_format, + args.dataset_method or result['methods'][0], percentages, + dataset_seed, dataset_mask_source, dataset_groups, + ) print(f"\nDone. {result['succeeded']} succeeded, {result['failed']} failed " f"({result['workers']} worker(s): {result['worker_reason']}).") + if len(result["methods"]) > 1: + for method in result["methods"]: + outcome = result["by_method"][method] + print(f"{method}: {outcome['succeeded']} succeeded, {outcome['failed']} failed.") + print(f" Measurement CSV written to: {outcome['results_path']}") + if outcome["failure_report_path"]: + print(f" Failure report written to: {outcome['failure_report_path']}") + return 0 print(f"Measurement CSV written to: {result['results_path']}") if result.get("failure_report_path"): print(f"Failure report written to: {result['failure_report_path']}") return 0 +def _export_run_dataset( + result, input_images, output_dir, dataset_format, method, percentages, + seed, mask_source, groups, +): + """Package the completed run while its private preview files still exist.""" + import tempfile + from pathlib import Path + + from .dataset_export import pairs_from_manifest, write_dataset_zip + + pairs = pairs_from_manifest( + result['artifacts'], input_images, method, + measurement_source=result['measurement_source'], + preview_artifacts=result['preview_artifacts'], + ) + pair_ids = {pair['sample_id'] for pair in pairs} + groups = {sample_id: group for sample_id, group in groups.items() + if sample_id in pair_ids} + base = Path(output_dir) / f'mats_training_{method}_{dataset_format}' + destination = base.with_suffix('.zip') + index = 2 + while destination.exists(): + destination = Path(f'{base}_{index}.zip') + index += 1 + with tempfile.NamedTemporaryFile( + dir=output_dir, prefix='.mats_dataset_', suffix='.zip', delete=False, + ) as temp_file: + temporary = Path(temp_file.name) + try: + manifest = write_dataset_zip( + pairs, temporary, dataset_format=dataset_format, + percentages=percentages, seed=seed, method=method, + mask_source=mask_source, groups=groups, + ) + os.replace(temporary, destination) + except (OSError, ValueError) as exc: + _fail(f'training dataset export failed: {exc}', code=1) + finally: + temporary.unlink(missing_ok=True) + counts = manifest['counts'] + print(f'Training dataset written to: {destination}') + print(f" {counts['train']} train / {counts['val']} validation / " + f"{counts['test']} test; {len(manifest['excluded'])} excluded.") + if manifest['conversion_notes']: + print(f" {len(manifest['conversion_notes'])} conversion note(s) in manifest.json.") + + def _cmd_app(args): from .app.launcher import launch return launch(args.extra or []) diff --git a/src/mats/core.py b/src/mats/core.py index ad22ce8..092ef90 100644 --- a/src/mats/core.py +++ b/src/mats/core.py @@ -4,6 +4,7 @@ import os import re import csv +import json import threading from itertools import combinations import concurrent.futures @@ -17,11 +18,6 @@ # Third-party libraries import numpy as np import cv2 -import torch -import torch.nn.functional as F -from PIL import Image -from rfdetr import RFDETRLarge -from tqdm import tqdm # type: ignore[import-not-found] # Optional "enhanced QR reading" backends. The default install decodes QR codes # with OpenCV only (no system libraries). Installing the optional extra -- @@ -62,6 +58,22 @@ validate_results_unit, ) +# Threshold presets live in a dependency-light module so the CLI parser can +# validate them without importing torch. Re-exported for `core.THRESHOLD_LEVELS`. +from .thresholds import THRESHOLD_LEVELS + +# Pre-cleanup measurements drop pieces touching the border or far from the leaf, +# then, at a clean size above 0, small specks and holes. +from .mask_cleanup import clear_margin, raw_measurement_mask +from .mask_settings import ( + CLEAN_MARGIN_DEFAULT, + CLEAN_SIZE_DEFAULT, + STRAY_GAP_DEFAULT, + checked_clean_margin, + checked_clean_size, + checked_stray_gap, +) + # Checkpoint resolution and constants live in mats.paths; re-exported for API # compatibility with callers that read core.RF_DETR_MARKER_CHECKPOINT etc. from .paths import ( @@ -80,12 +92,6 @@ BIREFNET_MEAN = [0.485, 0.456, 0.406] BIREFNET_STD = [0.229, 0.224, 0.225] TARGET_BOX_SUFFIX = "_target_box" -THRESHOLD_LEVELS = { - "auto": None, # Otsu's method: threshold computed per-image from histogram - "low": 100, - "medium": 125, - "high": 150, -} _MARKER_MODELS = {} _MARKER_MODEL_LOCK = threading.Lock() _MARKER_INFERENCE_LOCKS = {} @@ -99,6 +105,8 @@ def resolve_rfdetr_device(device_override=None): forced_device = os.environ.get("RF_DETR_DEVICE") if forced_device: return forced_device.lower() + import torch + if torch.cuda.is_available(): return "cuda" if torch.backends.mps.is_available() and torch.backends.mps.is_built(): @@ -114,6 +122,8 @@ def get_marker_model(device_override=None): if device in _MARKER_MODELS: return _MARKER_MODELS[device] from . import weights + from rfdetr import RFDETRLarge + checkpoint = weights.ensure_weight("rf-detr") # resolves or auto-fetches once model = RFDETRLarge( resolution=RF_DETR_MARKER_RESOLUTION, @@ -146,6 +156,8 @@ def pad_to_square_for_rfdetr( target_size=RF_DETR_MARKER_RESOLUTION, fill=RF_DETR_MARKER_PAD_COLOR, ): + from PIL import Image + orig_w, orig_h = pil_img.size scale = target_size / max(orig_w, orig_h) new_w = int(round(orig_w * scale)) @@ -180,6 +192,8 @@ def unpad_xyxy(xyxy, pad_info): def detect_marker_geometry(image_bgr, confidence=RF_DETR_MARKER_CONFIDENCE, device_override=None): + from PIL import Image + pil_img = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)) padded_img, pad_info = pad_to_square_for_rfdetr(pil_img) device = resolve_rfdetr_device(device_override) @@ -213,6 +227,8 @@ def detect_marker_centers(image_bgr, confidence=RF_DETR_MARKER_CONFIDENCE, devic def resolve_birefnet_device(device_override=None): + import torch + if device_override is not None: return torch.device(device_override) report = birefnet_device_report() @@ -240,6 +256,8 @@ def get_birefnet_model(device_override=None): require_birefnet_dependencies() checkpoint = weights.require_local_weight("birefnet") + import torch + model = create_birefnet_model() ckpt = torch.load( str(checkpoint), @@ -255,6 +273,9 @@ def get_birefnet_model(device_override=None): def _preprocess_birefnet_image(image_bgr, image_size=BIREFNET_IMAGE_SIZE): + import torch + from PIL import Image + image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) pil_img = Image.fromarray(image_rgb).convert("RGB") pil_img = pil_img.resize((image_size, image_size), Image.BILINEAR) @@ -265,7 +286,6 @@ def _preprocess_birefnet_image(image_bgr, image_size=BIREFNET_IMAGE_SIZE): return tensor.unsqueeze(0) -@torch.no_grad() def predict_birefnet_mask( image_bgr, image_size=BIREFNET_IMAGE_SIZE, @@ -273,21 +293,25 @@ def predict_birefnet_mask( device_override=None, ): """Predict a single-channel 0/255 leaf foreground mask for a BGR image.""" - model = get_birefnet_model(device_override) - device = resolve_birefnet_device(device_override) - orig_h, orig_w = image_bgr.shape[:2] - inp = _preprocess_birefnet_image(image_bgr, image_size=image_size).to(device) - - outputs = model(inp) - pred = outputs[-1] if isinstance(outputs, (list, tuple)) else outputs - pred = torch.sigmoid(pred) - pred = F.interpolate( - pred, - size=(orig_h, orig_w), - mode="bilinear", - align_corners=False, - ) - return (pred[0, 0].cpu().numpy() > threshold).astype(np.uint8) * 255 + import torch + import torch.nn.functional as F + + with torch.no_grad(): + model = get_birefnet_model(device_override) + device = resolve_birefnet_device(device_override) + orig_h, orig_w = image_bgr.shape[:2] + inp = _preprocess_birefnet_image(image_bgr, image_size=image_size).to(device) + + outputs = model(inp) + pred = outputs[-1] if isinstance(outputs, (list, tuple)) else outputs + pred = torch.sigmoid(pred) + pred = F.interpolate( + pred, + size=(orig_h, orig_w), + mode="bilinear", + align_corners=False, + ) + return (pred[0, 0].cpu().numpy() > threshold).astype(np.uint8) * 255 # QReader instantiation downloads a detector model, so defer it until an # enhanced decode is actually needed (and only if the optional extra is present). @@ -577,7 +601,7 @@ def white_out_marker_boxes(image_bgr, marker_boxes): return whitefilled_image -def clean_leaf_mask(binary_mask): +def clean_leaf_mask(binary_mask, *, fill_holes=True): # Re-join tiny breaks using morphological closing (dilate then erode). h, w = binary_mask.shape[:2] k = max(3, min(11, ((min(h, w) // 300) * 2) + 3)) @@ -590,12 +614,24 @@ def clean_leaf_mask(binary_mask): contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if len(contours) > 0: largest_contour = max(contours, key=cv2.contourArea) - binary_mask = np.zeros(binary_mask.shape, dtype=np.uint8) - cv2.drawContours(binary_mask, [largest_contour], -1, (255), thickness=cv2.FILLED) + largest_mask = np.zeros(binary_mask.shape, dtype=np.uint8) + cv2.drawContours(largest_mask, [largest_contour], -1, (255), thickness=cv2.FILLED) + binary_mask = ( + largest_mask if fill_holes else cv2.bitwise_and(binary_mask, largest_mask) + ) return binary_mask +def unfilled_leaf_mask(raw_mask, clean_margin=CLEAN_MARGIN_DEFAULT): + """MATS cleanup without flash fill, after clearing the template's edge margin. + + Used by Remove flashfill. The margin goes first so the printed box outline + can never be kept as, or joined to, the largest object. + """ + return clean_leaf_mask(clear_margin(raw_mask, clean_margin), fill_holes=False) + + def keep_largest_mask_component(binary_mask): contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: @@ -610,17 +646,66 @@ def keep_largest_mask_component(binary_mask): return largest_mask +MASK_METHODS = ("threshold", "birefnet") + + +def resolve_mask_methods(mask_method): + """Normalize a mask-method choice into an ordered tuple of methods. + + Accepts ``"threshold"``, ``"birefnet"``, ``"both"``, or an iterable of + method names. Duplicates are dropped; order follows the input, except that + ``"both"`` expands to ``MASK_METHODS`` order. + """ + if isinstance(mask_method, str): + methods = MASK_METHODS if mask_method == "both" else (mask_method,) + else: + methods = tuple(dict.fromkeys(mask_method or ())) + if not methods: + raise ValueError("choose at least one mask method") + unknown = [method for method in methods if method not in MASK_METHODS] + if unknown: + raise ValueError(f"Unknown segmentation method: {unknown[0]}") + return methods + + +def method_suffixed_path(path, method): + """Insert ``_{method}`` before the extension: results.csv -> results_threshold.csv.""" + root, ext = os.path.splitext(path) + return f"{root}_{method}{ext}" + + def create_leaf_mask(target_box, mask_method, threshold_value, device_override=None): + _, cleaned = segment_leaf(target_box, mask_method, threshold_value, device_override) + return cleaned + + +def segment_leaf(target_box, mask_method, threshold_value, device_override=None): + """Return the binary segmentation before and after measurement cleanup.""" if mask_method == "threshold": - binary_mask = threshold_mask(target_box, threshold_value) + raw = threshold_mask(target_box, threshold_value) + elif mask_method == "birefnet": + raw = predict_birefnet_mask(target_box, device_override=device_override) else: - binary_mask = predict_birefnet_mask(target_box, device_override=device_override) + raise ValueError(f"Unknown segmentation method: {mask_method}") + if not np.any(raw): + return None, None + return raw, clean_leaf_mask(raw.copy()) + + +def build_overlay_image(target_box, binary_mask, color=(255, 0, 255), alpha=0.4): + mask_bool = binary_mask.astype(bool) + tint = np.full_like(target_box, color, dtype=target_box.dtype) + blended = cv2.addWeighted(target_box, 1.0 - alpha, tint, alpha, 0) + overlay = target_box.copy() + overlay[mask_bool] = blended[mask_bool] + contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + thickness = max(2, min(6, min(target_box.shape[:2]) // 300)) + cv2.drawContours(overlay, contours, -1, color, thickness) + return overlay - if not np.any(binary_mask): - return None - binary_mask = clean_leaf_mask(binary_mask) - return binary_mask +def build_cutout_image(target_box, binary_mask): + return cv2.bitwise_and(target_box, target_box, mask=binary_mask) def leaf_bounding_rect(binary_mask): @@ -634,8 +719,20 @@ def leaf_bounding_rect(binary_mask): return cv2.boundingRect(largest_contour) -def draw_measurement_axes(target_box, binary_mask): - rect = leaf_bounding_rect(binary_mask) +def measurement_bounding_rect(binary_mask, measurement_source="cleaned"): + if measurement_source == "cleaned": + return leaf_bounding_rect(binary_mask) + if measurement_source != "pre-cleanup": + raise ValueError("measurement_source must be 'cleaned' or 'pre-cleanup'") + ys, xs = np.nonzero(binary_mask == 255) + if not len(xs): + return None + x, y = int(xs.min()), int(ys.min()) + return x, y, int(xs.max()) - x + 1, int(ys.max()) - y + 1 + + +def draw_measurement_axes(target_box, binary_mask, measurement_source="cleaned"): + rect = measurement_bounding_rect(binary_mask, measurement_source) if rect is None: return None @@ -672,7 +769,10 @@ def measurement_row_from_mask( binary_mask, px_per_cm_width, px_per_cm_height, + measurement_source="cleaned", ): + if measurement_source not in ("cleaned", "pre-cleanup"): + raise ValueError("measurement_source must be 'cleaned' or 'pre-cleanup'") if px_per_cm_width is None or px_per_cm_height is None: return measurement_na_row(sample_id, "SCALE: physical dimensions unavailable") if px_per_cm_width <= 0 or px_per_cm_height <= 0: @@ -682,7 +782,7 @@ def measurement_row_from_mask( if white_pixels == 0: return measurement_na_row(sample_id, "LEAF_MASK: leaf not detected") - rect = leaf_bounding_rect(binary_mask) + rect = measurement_bounding_rect(binary_mask, measurement_source) if rect is None: return measurement_na_row(sample_id, "LEAF_MASK: leaf contour not detected") @@ -784,13 +884,13 @@ def warning_report_rows(input_image, result): return rows -def write_failure_report(failure_rows, output_dir): +def write_failure_report(failure_rows, output_dir, file_name="leaf_morpho_failures.csv"): if not failure_rows: return None report_dir = output_dir if output_dir is not False else os.getcwd() os.makedirs(report_dir, exist_ok=True) - report_path = os.path.join(report_dir, "leaf_morpho_failures.csv") + report_path = os.path.join(report_dir, file_name) fieldnames = ["sample_id", "input_image", "stage", "failure_mode", "status"] with open(report_path, "w", newline="") as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) @@ -812,7 +912,19 @@ def leaf_morpho( save_measurement_axes=True, execution_device="auto", qr_backend_fields=None, + export_options=None, + measurement_source="cleaned", + stray_gap=STRAY_GAP_DEFAULT, + clean_margin=CLEAN_MARGIN_DEFAULT, + clean_size=CLEAN_SIZE_DEFAULT, ): + if measurement_source not in ("cleaned", "pre-cleanup"): + raise ValueError("measurement_source must be 'cleaned' or 'pre-cleanup'") + if output_mode != "masks" and measurement_source == "pre-cleanup": + raise ValueError("pre-cleanup measurements require output_mode='masks'") + stray_gap = checked_stray_gap(stray_gap) + clean_margin = checked_clean_margin(clean_margin) + clean_size = _checked_run_clean_size(clean_size, measurement_source) if execution_device not in {"auto", "cpu", "hybrid"}: raise ValueError("execution_device must be 'auto', 'cpu', or 'hybrid'") device_override = "cpu" if execution_device == "cpu" else None @@ -831,32 +943,177 @@ def leaf_morpho( else: qr_backend_fields = tuple(qr_backend_fields) qr_trace = {field: "not_reached" for field in qr_backend_fields} + methods = resolve_mask_methods(mask_method) + # With several measurement methods, each writes its own suffixed files and + # gets its own result; a single method keeps the original file names. + multi = len(methods) > 1 + options = export_options or {} + pre_cleanup_methods = tuple(options.get("pre_cleanup_methods", ())) + artifacts = [] + preview_artifacts = [] + warnings = [] + + def _save(kind, image, suffix, method=None, sink=None): + if output_dir is False: + return + path = os.path.join(output_dir, f"{file_name}{suffix}") + try: + if not cv2.imwrite(path, image): + raise OSError("image writer returned false") + except Exception as exc: + (warnings if sink is None else sink).append( + f"EXPORT: {kind} could not be written: {exc}" + ) + return + artifacts.append({"path": path, "sample_id": file_name, "kind": kind, "method": method}) + + def _save_preview(kind, image, suffix, method=None): + preview_dir = options.get("preview_dir") + if not preview_dir: + return + path = os.path.join(preview_dir, f"{file_name}{suffix}") + try: + if cv2.imwrite(path, image): + preview_artifacts.append({ + "path": path, "sample_id": file_name, "kind": kind, "method": method, + }) + except Exception: + # Preview availability must not change a specimen's measurement. + pass + + def _segment_target(target_box): + """Segment once per method; return its selected measurement mask.""" + segmented = {} + for method in dict.fromkeys(methods + pre_cleanup_methods): + measured = method in methods + # Method-specific warnings stay with that method's result in a + # multi-method run; otherwise they share the image's warning list. + method_warnings = [] if measured and multi else warnings + try: + raw, cleaned = segment_leaf(target_box, method, threshold_value, device_override) + except Exception as exc: + if measured: + segmented[method] = (None, str(exc), method_warnings) + else: + warnings.append(f"EXPORT: {method} pre-cleanup mask failed: {exc}") + continue + if raw is None: + if measured: + segmented[method] = (None, None, method_warnings) + else: + warnings.append(f"EXPORT: {method} detected no leaf") + continue + if method in pre_cleanup_methods: + _save("pre_cleanup", raw, f"_mask_precleanup_{method}.png", method, method_warnings) + if not measured: + continue + # The pre-cleanup export above stays the literal raw mask; the + # measurement clears the edge margin, drops stray pieces, and at a + # clean size above 0 removes small specks and fills small holes. + measurement_mask = ( + raw_measurement_mask(raw, clean_margin, stray_gap, clean_size) + if measurement_source == "pre-cleanup" else cleaned + ) + segmented[method] = (measurement_mask, None, method_warnings) + suffix = f"_{method}" if multi else "" + # The explorer re-cleans the raw mask with each specimen's settings. + _save_preview( + "preview_raw_mask", raw, f"_preview_raw_mask{suffix}.png", method + ) + if measurement_source == "pre-cleanup" or not options.get("cleaned_masks", True): + _save_preview( + "preview_mask", measurement_mask, f"_preview_mask{suffix}.png", method + ) + if options.get("cleaned_masks", True): + _save("mask", cleaned, f"_mask{suffix}.png", method, method_warnings) + if options.get("axes", save_measurement_axes): + try: + axes = draw_measurement_axes(target_box, measurement_mask, measurement_source) + if axes is not None: + _save("axes", axes, f"_measurement_axes{suffix}.jpg", method, method_warnings) + except Exception as exc: + method_warnings.append(f"EXPORT: axes failed: {exc}") + if options.get("overlay", False): + try: + _save("overlay", build_overlay_image(target_box, measurement_mask), + f"_overlay{suffix}.jpg", method, method_warnings) + except Exception as exc: + method_warnings.append(f"EXPORT: overlay failed: {exc}") + if options.get("cutout", False): + try: + _save("cutout", build_cutout_image(target_box, measurement_mask), + f"_cutout{suffix}.jpg", method, method_warnings) + except Exception as exc: + method_warnings.append(f"EXPORT: cutout failed: {exc}") + return segmented def _with_qr_trace(result_row): if not qr_trace: return result_row return {**result_row, **qr_trace} - def _fail(stage: str, msg: str): + def _combined_warnings(extra_warnings): + combined = list(warnings) + if extra_warnings is not None and extra_warnings is not warnings: + combined.extend(extra_warnings) + return combined + + def _fail(stage: str, msg: str, extra_warnings=None): status = f'{stage}: {msg}' return { 'sample_id': file_name, 'status': status, 'result_row': _with_qr_trace(measurement_na_row(file_name, status)), + 'artifacts': artifacts, + 'preview_artifacts': preview_artifacts, + 'warnings': _combined_warnings(extra_warnings), } - def _ok(result_row, warnings=None): + def _ok(result_row, extra_warnings=None): result = { 'sample_id': file_name, 'status': 'ok', 'result_row': _with_qr_trace(result_row), + 'artifacts': artifacts, + 'preview_artifacts': preview_artifacts, } - if warnings: - result['warnings'] = warnings + combined_warnings = _combined_warnings(extra_warnings) + if combined_warnings: + result['warnings'] = combined_warnings return result + def _measured(stage, segmented, scale_axes, scale_status): + """Measure each method's mask against the shared scale. + + One method returns its result directly; several return + ``{'method_results': {method: result}}`` for the batch to split. + """ + results = {} + for method in methods: + measurement_mask, error, method_warnings = segmented[method] + if measurement_mask is None: + results[method] = _fail(stage, error or "leaf not detected", method_warnings) + continue + try: + if scale_axes is None: + row = measurement_na_row(file_name, scale_status) + else: + row = measurement_row_from_mask( + file_name, measurement_mask, *scale_axes, measurement_source + ) + except Exception as exc: + print(f"[leaf_morpho] ERROR '{input_image}' at {stage} ({method}): {exc}") + results[method] = _fail(stage, f"unexpected error: {exc}", method_warnings) + continue + results[method] = _ok(row, method_warnings) + if not multi: + return results[methods[0]] + return { + 'sample_id': file_name, 'artifacts': artifacts, + 'preview_artifacts': preview_artifacts, 'method_results': results, + } + stage = "START" - warnings = [] scale_failure = None try: @@ -868,47 +1125,24 @@ def _ok(result_row, warnings=None): if is_target_box_input: target_box = image + _save_preview("preview_target_box", target_box, "_preview_target_box.png") if output_mode == "target-boxes": return _ok(measurement_na_row(file_name, "OUTPUT_MODE: target-boxes only")) stage = "LEAF_MASK" - binary_mask = create_leaf_mask( - target_box, - mask_method, - threshold_value, - device_override, - ) - if binary_mask is None: - return _fail(stage, "leaf not detected") - - if output_dir is not False: - output_path = os.path.join(output_dir, f"{file_name}_mask.png") - cv2.imwrite(output_path, binary_mask) - if save_measurement_axes: - save_measurement_axes_image(output_dir, file_name, target_box, binary_mask) + segmented = _segment_target(target_box) if template_dimensions is None: - return _ok( - measurement_na_row( - file_name, - "SCALE: target_box input requires --sheet-dimensions or legacy --template-dimensions", - ), - warnings, + scale_axes = None + scale_status = ( + "SCALE: target_box input requires --sheet-dimensions or legacy --template-dimensions" ) + else: + width, height, unit = template_dimensions + scale_axes = px_per_cm_axes(target_box.shape[:2], width, height, unit) + scale_status = "SCALE: invalid calibration dimensions" - width, height, unit = template_dimensions - scale_axes = px_per_cm_axes(target_box.shape[:2], width, height, unit) - if scale_axes is None: - return _ok(measurement_na_row(file_name, "SCALE: invalid calibration dimensions")) - - return _ok( - measurement_row_from_mask( - file_name, - binary_mask, - *scale_axes, - ), - warnings, - ) + return _measured(stage, segmented, scale_axes, scale_status) masked_img = image.copy() @@ -1040,49 +1274,25 @@ def _ok(result_row, warnings=None): if target_box is None or target_box.size == 0: return _fail(stage, "failed to compute target box") - if output_dir is not False: - # Save the expanded crop to the output directory - output_path = os.path.join(output_dir, f"{file_name}_target_box.jpg") - cv2.imwrite(output_path, target_box) + if "threshold" in methods or not options.get("target_boxes", True): + _save_preview("preview_target_box", target_box, "_preview_target_box.png") + + if options.get("target_boxes", True): + _save("target_box", target_box, "_target_box.jpg") if output_mode == "target-boxes": return _ok(measurement_na_row(file_name, "OUTPUT_MODE: target-boxes only"), warnings) stage = "LEAF_MASK" - binary_mask = create_leaf_mask( - target_box, - mask_method, - threshold_value, - device_override, - ) - if binary_mask is None: - return _fail(stage, "leaf not detected") - - # Save ONLY the refined (most accurate) mask. - # NOTE: This replaces older outputs: - # - *_binary_mask.jpg - # - *_masked.jpg - # - *_accurately_masked_leaf.jpg - if output_dir is not False: - output_path = os.path.join(output_dir, f"{file_name}_mask.png") - cv2.imwrite(output_path, binary_mask) - if save_measurement_axes: - save_measurement_axes_image(output_dir, file_name, target_box, binary_mask) + segmented = _segment_target(target_box) scale_axes = px_per_cm_axes(target_box.shape[:2], width, height, unit) - if scale_axes is None: - result_row = measurement_na_row( - file_name, - scale_failure or "SCALE: physical dimensions unavailable", - ) - else: - result_row = measurement_row_from_mask( - file_name, - binary_mask, - *scale_axes, - ) - - return _ok(result_row, warnings) + return _measured( + stage, + segmented, + scale_axes, + scale_failure or "SCALE: physical dimensions unavailable", + ) except Exception as e: print(f"[leaf_morpho] ERROR '{input_image}' at {stage}: {e}") @@ -1097,6 +1307,7 @@ def _process_batch_image( mask_method, threshold_value, save_measurement_axes, + export_options, execution_device, qr_backend_fields, ): @@ -1112,10 +1323,23 @@ def _process_batch_image( save_measurement_axes, execution_device, qr_backend_fields, + export_options=export_options, + measurement_source=export_options.get("measurement_source", "cleaned"), + stray_gap=export_options.get("stray_gap", STRAY_GAP_DEFAULT), + clean_margin=export_options.get("clean_margin", CLEAN_MARGIN_DEFAULT), + clean_size=export_options.get("clean_size", CLEAN_SIZE_DEFAULT), ) return input_image, result, None +def _checked_run_clean_size(clean_size, measurement_source): + """Validate a run's clean size; above 0 it needs pre-cleanup measurements.""" + clean_size = checked_clean_size(clean_size) + if clean_size and measurement_source != "pre-cleanup": + raise ValueError("clean_size requires measurement_source='pre-cleanup'") + return clean_size + + def run_leaf_morpho_batch( input_images, output_dir, @@ -1136,6 +1360,11 @@ def run_leaf_morpho_batch( break_glass_acknowledged=False, birefnet_parallel_acknowledged=False, results_unit="cm", + export_options=None, + measurement_source="cleaned", + stray_gap=STRAY_GAP_DEFAULT, + clean_margin=CLEAN_MARGIN_DEFAULT, + clean_size=CLEAN_SIZE_DEFAULT, ): """Run the leaf morphometrics pipeline with per-image error isolation. @@ -1146,12 +1375,58 @@ def run_leaf_morpho_batch( enabled, counts above 75% of the available CPU allocation require an explicit ``break_glass_acknowledged`` override. CPU-parallel BiRefNet additionally requires its own explicit acknowledgement. + + ``mask_method`` may name several methods (``"both"`` or a sequence). Each + image is then marker-detected once and measured once per method, and each + method gets its own results CSV and failure log, suffixed ``_{method}`` + (``results_threshold.csv``). ``by_method`` holds the per-method summary; + top-level ``succeeded``/``failed`` count images, where an image succeeds + only when every method measured it. + + ``measurement_source='pre-cleanup'`` measures each method's raw mask after + clearing a ``clean_margin`` percent band along the target-box edge and + dropping pieces that touch that band or lie more than ``stray_gap`` times the + leaf's bounding-box diagonal from it (``mask_cleanup.clean_raw_mask``). A + ``clean_size`` above 0 (pre-cleanup only) then removes white specks and fills + enclosed holes whose inscribed radius is below that many pixels + (``mask_cleanup.raw_measurement_mask``). """ input_images = list(input_images or []) + if measurement_source not in ("cleaned", "pre-cleanup"): + raise ValueError("measurement_source must be 'cleaned' or 'pre-cleanup'") + stray_gap = checked_stray_gap(stray_gap) + clean_margin = checked_clean_margin(clean_margin) + clean_size = _checked_run_clean_size(clean_size, measurement_source) + methods = resolve_mask_methods(mask_method) + if output_mode != "masks" and measurement_source == "pre-cleanup": + raise ValueError("pre-cleanup measurements require output_mode='masks'") + if output_mode != "masks": + # Target-box runs never segment, so they keep a single, unsuffixed CSV. + methods = methods[:1] + multi = len(methods) > 1 + options = dict(export_options or {}) + options["measurement_source"] = measurement_source + options["stray_gap"] = stray_gap + options["clean_margin"] = clean_margin + options["clean_size"] = clean_size + if options.get("preview_dir"): + os.makedirs(options["preview_dir"], exist_ok=True) + pre_cleanup = tuple(dict.fromkeys(options.get("pre_cleanup_methods", ()))) + if any(method not in MASK_METHODS for method in pre_cleanup): + raise ValueError("pre_cleanup_methods must contain only threshold or birefnet") + options["pre_cleanup_methods"] = pre_cleanup if output_mode == "masks" else () + needs_birefnet = output_mode == "masks" and ( + "birefnet" in methods or "birefnet" in options["pre_cleanup_methods"] + ) + sample_ids = [target_box_sample_id(p) if is_target_box_image(p) + else os.path.splitext(os.path.basename(p))[0] for p in input_images] + sample_keys = [sample_id.casefold() for sample_id in sample_ids] + if len(sample_keys) != len(set(sample_keys)): + raise ValueError("Input images contain duplicate sample IDs; use unique basenames") if execution_device not in {"auto", "cpu", "hybrid"}: raise ValueError("execution_device must be 'auto', 'cpu', or 'hybrid'") validate_results_unit(results_unit) - if execution_device == "hybrid" and mask_method != "threshold": + if execution_device == "hybrid" and needs_birefnet: raise ValueError("hybrid execution is currently supported for Otsu thresholding only") if output_dir is not False: os.makedirs(output_dir, exist_ok=True) @@ -1166,7 +1441,9 @@ def run_leaf_morpho_batch( ) if workers is None: - workers, worker_reason = default_worker_count(input_images, output_mode, mask_method) + workers, worker_reason = default_worker_count( + input_images, output_mode, "birefnet" if needs_birefnet else methods[0] + ) else: worker_reason = "user override" workers = max(1, int(workers)) @@ -1181,7 +1458,7 @@ def run_leaf_morpho_batch( f"{workers} workers uses {worker_risk.utilization:.0%} of the available CPU allocation. " "High-risk worker counts require break-glass acknowledgement." ) - if mask_method == "birefnet" and workers > 1 and not birefnet_parallel_acknowledged: + if needs_birefnet and workers > 1 and execution_device == "cpu" and not birefnet_parallel_acknowledged: raise ValueError( "CPU-parallel BiRefNet requires a separate break-glass acknowledgement." ) @@ -1190,7 +1467,7 @@ def run_leaf_morpho_batch( serialize_model_inference and execution_device == "auto" and workers > 1 - and (mask_method == "birefnet" or not all_target_boxes) + and (needs_birefnet or not all_target_boxes) ): workers = 1 worker_reason = f"{worker_reason}; serialized for model-backed UI inference" @@ -1198,49 +1475,88 @@ def run_leaf_morpho_batch( succeeded = 0 failed = 0 processed = 0 - failure_rows = [] - result_rows = [] + artifacts = [] + preview_artifacts = [] + # One tally per measurement method: its rows, failures, and output files. + tallies = { + method: { + "succeeded": 0, + "failed": 0, + "failure_rows": [], + "result_rows": [], + "results_path": ( + method_suffixed_path(results_path, method) + if multi and results_path else results_path + ), + } + for method in methods + } - def _record_result(input_image, result, exception=None): - nonlocal succeeded, failed, processed + def _write_results(tally): + if compact_csv: + write_compact_results_csv(tally["result_rows"], tally["results_path"], results_unit) + else: + write_results_csv( + tally["result_rows"], + tally["results_path"], + qr_backend_fields, + results_unit, + ) + + def _record_method(tally, input_image, result, exception): + """Record one method's outcome for an image; return whether it succeeded.""" if exception is not None: - failed += 1 - failure_rows.append(failure_report_row(input_image, exception=exception)) + tally["failed"] += 1 + tally["failure_rows"].append(failure_report_row(input_image, exception=exception)) sample_id = os.path.splitext(os.path.basename(input_image))[0] row = measurement_na_row(sample_id, f"EXCEPTION: {exception}") row.update({field: "not_reached" for field in qr_backend_fields}) - result_rows.append(row) - elif result and result.get("result_row"): - result_rows.append(result["result_row"]) + tally["result_rows"].append(row) + return False + if result and result.get("result_row"): + tally["result_rows"].append(result["result_row"]) if result.get("status") == "ok": - succeeded += 1 - failure_rows.extend(warning_report_rows(input_image, result)) - else: - failed += 1 - failure_rows.append(failure_report_row(input_image, result=result)) - failure_rows.extend(warning_report_rows(input_image, result)) + tally["succeeded"] += 1 + tally["failure_rows"].extend(warning_report_rows(input_image, result)) + return True + tally["failed"] += 1 + tally["failure_rows"].append(failure_report_row(input_image, result=result)) + tally["failure_rows"].extend(warning_report_rows(input_image, result)) + return False + tally["failed"] += 1 + tally["failure_rows"].append(failure_report_row( + input_image, + result={"status": "UNKNOWN: no result returned"}, + )) + sample_id = os.path.splitext(os.path.basename(input_image))[0] + row = measurement_na_row(sample_id, "UNKNOWN: no result returned") + row.update({field: "not_reached" for field in qr_backend_fields}) + tally["result_rows"].append(row) + return False + + def _record_result(input_image, result, exception=None): + nonlocal succeeded, failed, processed + if exception is None and result: + artifacts.extend(result.get("artifacts", ())) + preview_artifacts.extend(result.get("preview_artifacts", ())) + # A multi-method result carries one result per method; anything else + # (a single method, or a failure before segmentation) applies to all. + method_results = (result or {}).get("method_results") or { + method: result for method in methods + } + outcomes = [ + _record_method(tallies[method], input_image, method_results.get(method), exception) + for method in methods + ] + if all(outcomes): + succeeded += 1 else: failed += 1 - failure_rows.append(failure_report_row( - input_image, - result={"status": "UNKNOWN: no result returned"}, - )) - sample_id = os.path.splitext(os.path.basename(input_image))[0] - row = measurement_na_row(sample_id, "UNKNOWN: no result returned") - row.update({field: "not_reached" for field in qr_backend_fields}) - result_rows.append(row) processed += 1 if results_path and processed % csv_update_interval == 0: - if compact_csv: - write_compact_results_csv(result_rows, results_path, results_unit) - else: - write_results_csv( - result_rows, - results_path, - qr_backend_fields, - results_unit, - ) + for tally in tallies.values(): + _write_results(tally) if progress_callback is not None: progress_callback({ "processed": processed, @@ -1258,9 +1574,10 @@ def _record_result(input_image, result, exception=None): output_dir, template_dimensions, output_mode, - mask_method, + methods, threshold_value, save_measurement_axes, + options, execution_device, qr_backend_fields, ) @@ -1276,9 +1593,10 @@ def _record_result(input_image, result, exception=None): output_dir, template_dimensions, output_mode, - mask_method, + methods, threshold_value, save_measurement_axes, + options, execution_device, qr_backend_fields, ): input_image @@ -1294,35 +1612,79 @@ def _record_result(input_image, result, exception=None): _record_result(input_image, None, exc) if results_path: - if compact_csv: - write_compact_results_csv(result_rows, results_path, results_unit) - else: - write_results_csv( - result_rows, - results_path, - qr_backend_fields, - results_unit, + for method, tally in tallies.items(): + _write_results(tally) + artifacts.append({"path": tally["results_path"], "sample_id": None, + "kind": "results_csv", "method": method}) + metadata_path = f"{tally['results_path']}.meta.json" + metadata = { + "measurement_source": measurement_source, + "mask_method": method, + "results_unit": results_unit, + "csv_schema": "compact" if compact_csv else "full", + } + if measurement_source == "pre-cleanup": + metadata["clean_margin"] = clean_margin + metadata["stray_gap"] = stray_gap + metadata["clean_size"] = clean_size + with open(metadata_path, "w", encoding="utf-8") as metadata_file: + json.dump(metadata, metadata_file, indent=2) + metadata_file.write("\n") + artifacts.append({"path": metadata_path, "sample_id": None, + "kind": "results_metadata", "method": method}) + + for method, tally in tallies.items(): + tally["failure_report_path"] = None + if write_failures: + tally["failure_report_path"] = write_failure_report( + tally["failure_rows"], + output_dir, + (method_suffixed_path("leaf_morpho_failures.csv", method) + if multi else "leaf_morpho_failures.csv"), ) - - failure_report_path = None - if write_failures: - failure_report_path = write_failure_report(failure_rows, output_dir) - - return { + if tally["failure_report_path"]: + artifacts.append({"path": tally["failure_report_path"], "sample_id": None, + "kind": "failure_log", "method": method if multi else None}) + + by_method = { + method: { + "succeeded": tally["succeeded"], + "failed": tally["failed"], + "results_path": tally["results_path"], + "failure_report_path": tally["failure_report_path"], + "failure_rows": tally["failure_rows"], + "result_rows": tally["result_rows"], + "compact_rows": [ + compact_measurement_row(row, results_unit) for row in tally["result_rows"] + ], + } + for method, tally in tallies.items() + } + summary = { "succeeded": succeeded, "failed": failed, "processed": processed, "total": len(input_images), - "results_path": results_path, - "failure_report_path": failure_report_path, - "failure_rows": failure_rows, - "result_rows": result_rows, - "compact_rows": [ - compact_measurement_row(row, results_unit) for row in result_rows - ], + "methods": methods, + "by_method": by_method, + "artifacts": artifacts, + "preview_artifacts": preview_artifacts, "results_unit": results_unit, + "measurement_source": measurement_source, + "clean_margin": clean_margin, + "stray_gap": stray_gap, + "clean_size": clean_size, "qr_backend_fields": qr_backend_fields, "workers": workers, "worker_reason": worker_reason, "execution_device": execution_device, } + if not multi: + # Single-method callers read these top-level keys. They are left out of + # a multi-method summary so that reading one fails loudly. + summary.update( + (key, by_method[methods[0]][key]) + for key in ("results_path", "failure_report_path", "failure_rows", + "result_rows", "compact_rows") + ) + return summary diff --git a/src/mats/dataset_export.py b/src/mats/dataset_export.py new file mode 100644 index 0000000..89b94c5 --- /dev/null +++ b/src/mats/dataset_export.py @@ -0,0 +1,342 @@ +"""Build training datasets from app or CLI target-box and mask pairs. + +Only standard-library modules are imported at module load so this remains safe +for the package's dependency-light CLI and test environments. +""" + +from __future__ import annotations + +import csv +import io +import json +import random +import zipfile +from pathlib import Path + + +SPLITS = ("train", "val", "test") +FORMATS = ("png", "yolo-seg", "yolo-detect", "coco") + + +def pairs_from_manifest( + artifacts, input_images=(), method=None, measurement_source="cleaned", + preview_artifacts=(), +): + """Pair each method's measured mask with its corrected target-box image.""" + if measurement_source not in {"cleaned", "pre-cleanup"}: + raise ValueError("unknown measurement source") + by_sample = {} + raw_kinds = {"preview_raw_mask", "pre_cleanup"} + mask_kinds = ( + {"preview_mask"} if measurement_source == "pre-cleanup" else {"mask", "preview_mask"} + ) + for item in (*artifacts, *preview_artifacts): + sample_id = item.get("sample_id") + kind = item["kind"] + if sample_id is None or kind not in { + "target_box", "preview_target_box", *mask_kinds, *raw_kinds, + }: + continue + if method is not None and kind in {*mask_kinds, *raw_kinds}: + if item.get("method") != method: + continue + by_sample.setdefault(sample_id, {"sample_id": sample_id, "target_box": None, "mask": None}) + if kind in raw_kinds: + by_sample[sample_id]["raw_mask"] = item["path"] + elif kind in mask_kinds: + by_sample[sample_id]["mask"] = item["path"] + else: + by_sample[sample_id]["target_box"] = item["path"] + for path in input_images: + source = Path(path) + if source.stem.endswith("_target_box") and source.is_file(): + sample_id = source.stem[:-len("_target_box")] + if sample_id in by_sample and by_sample[sample_id]["target_box"] is None: + by_sample[sample_id]["target_box"] = str(source) + if measurement_source == "pre-cleanup": + for pair in by_sample.values(): + pair["mask_source"] = measurement_source + return list(by_sample.values()) + + +def split_counts(total, percentages): + """Largest-remainder counts; sum is always exactly ``total``.""" + if len(percentages) != 3 or any(not 0 <= value <= 100 for value in percentages): + raise ValueError("Train, validation, and test must each be between 0 and 100%.") + if sum(percentages) != 100: + raise ValueError("Train, validation, and test percentages must total 100%.") + exact = [total * value / 100 for value in percentages] + counts = [int(value) for value in exact] + order = sorted(range(3), key=lambda index: (-(exact[index] - counts[index]), index)) + for index in order[:total - sum(counts)]: + counts[index] += 1 + return dict(zip(SPLITS, counts)) + + +def read_group_csv(content): + """Parse optional sample_id,group_id CSV without guessing group boundaries.""" + if not content: + return {} + try: + rows = csv.DictReader(io.StringIO(content.decode("utf-8-sig"))) + if not rows.fieldnames or not {"sample_id", "group_id"} <= set(rows.fieldnames): + raise ValueError("Group CSV needs sample_id and group_id columns.") + groups = {} + for row in rows: + sample_id = (row.get("sample_id") or "").strip() + group_id = (row.get("group_id") or "").strip() + if not sample_id or not group_id: + raise ValueError("Every group CSV row needs a sample_id and group_id.") + if sample_id in groups and groups[sample_id] != group_id: + raise ValueError(f"Conflicting group IDs for {sample_id}.") + groups[sample_id] = group_id + return groups + except UnicodeDecodeError as exc: + raise ValueError("Group CSV must be UTF-8 text.") from exc + + +def validate_pairs(pairs, mask_source="measured", dataset_format="png"): + """Read each image once; return usable pairs and explicit exclusions.""" + import cv2 + import numpy as np + + if mask_source not in {"measured", "raw"}: + raise ValueError("Unknown mask source.") + seen = set() + valid, excluded = [], [] + for pair in pairs: + sample_id = str(pair.get("sample_id", "")) + image_path = pair.get("target_box") + mask_path = pair.get("mask" if mask_source == "measured" else "raw_mask") + reason = None + if not sample_id or sample_id in seen: + reason = "missing or duplicate sample ID" + elif not image_path or not mask_path: + reason = "image or selected mask is unavailable" + elif not Path(image_path).is_file() or not Path(mask_path).is_file(): + reason = "image or selected mask file is missing" + else: + image = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) + if image is None or mask is None: + reason = "image or mask could not be read" + elif image.shape[:2] != mask.shape[:2]: + reason = "image and mask dimensions differ" + elif not np.any(mask): + reason = "mask is empty" + elif dataset_format == "yolo-seg" and not _yolo_segments( + (mask > 0).astype("uint8") + )[0]: + reason = "mask has no usable polygon" + seen.add(sample_id) + if reason: + excluded.append({"sample_id": sample_id, "reason": reason}) + else: + valid.append({ + "sample_id": sample_id, "image": Path(image_path), "mask": Path(mask_path), + "width": image.shape[1], "height": image.shape[0], + }) + return valid, excluded + + +def assign_splits(pairs, percentages=(70, 20, 10), seed=42, groups=None): + """Deterministically place each group wholly in one split.""" + targets = split_counts(len(pairs), percentages) + groups = groups or {} + buckets = {} + for pair in pairs: + sample_id = pair["sample_id"] + group_id = (("group", groups[sample_id]) if sample_id in groups + else ("sample", sample_id)) + buckets.setdefault(group_id, []).append(pair) + items = sorted(buckets.items()) + random.Random(seed).shuffle(items) + counts = dict.fromkeys(SPLITS, 0) + assignments = {} + for _, members in items: + size = len(members) + choice = min(SPLITS, key=lambda split: ( + sum(abs(targets[name] - counts[name] - (size if name == split else 0)) + for name in SPLITS), + -(targets[split] - counts[split]), + SPLITS.index(split), + )) + counts[choice] += size + for pair in members: + assignments[pair["sample_id"]] = choice + return assignments, counts + + +def _encoded_png(mask): + import cv2 + + ok, data = cv2.imencode(".png", mask) + if not ok: + raise ValueError("Could not encode a mask as PNG.") + return data.tobytes() + + +def _binary_mask(path): + import cv2 + + mask = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) + if mask is None: + raise ValueError(f"Could not read mask: {path.name}") + return (mask > 0).astype("uint8") + + +def _yolo_segments(mask): + import cv2 + + contours, hierarchy = cv2.findContours(mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) + if hierarchy is None: + return [], False + height, width = mask.shape + lines = [] + holes = False + for contour, relation in zip(contours, hierarchy[0]): + if relation[3] != -1: + holes = True + continue + points = contour.reshape(-1, 2) + if len(points) < 3 or cv2.contourArea(contour) == 0: + continue + coords = " ".join(f"{x / width:.6f} {y / height:.6f}" for x, y in points) + lines.append(f"0 {coords}") + return lines, holes + + +def _bbox(mask): + import cv2 + + points = cv2.findNonZero(mask) + if points is None: + raise ValueError("Mask is empty.") + return cv2.boundingRect(points) + + +def _coco_rle(mask): + """Uncompressed COCO RLE in column-major order, preserving mask holes.""" + import numpy as np + + pixels = mask.T.reshape(-1) + boundaries = np.flatnonzero(pixels[1:] != pixels[:-1]) + 1 + counts = np.diff(np.concatenate(([0], boundaries, [pixels.size]))).tolist() + if pixels[0]: + counts.insert(0, 0) + return {"size": list(mask.shape), "counts": counts} + + +def write_dataset_zip( + pairs, dest_path, *, dataset_format="png", percentages=(70, 20, 10), seed=42, + method="threshold", mask_source="measured", groups=None, +): + """Validate and stream one method's dataset into a ZIP; return its manifest.""" + if dataset_format not in FORMATS: + raise ValueError("Unknown dataset format.") + groups = groups or {} + valid, excluded = validate_pairs(pairs, mask_source, dataset_format) + if not valid: + raise ValueError("No usable image and mask pairs were found for this method.") + unknown_groups = set(groups) - {pair["sample_id"] for pair in pairs} + if unknown_groups: + raise ValueError("Group CSV contains unknown sample IDs: " + + ", ".join(sorted(unknown_groups)[:5])) + assignments, counts = assign_splits(valid, percentages, seed, groups) + manifest = { + "format": dataset_format, "label_method": method, "mask_source": mask_source, + "label_origin": "MATS generated masks; review labels before model training", + "class_names": ["leaf"], "percentages": dict(zip(SPLITS, percentages)), + "seed": seed, "counts": counts, "samples": [], "excluded": excluded, + "conversion_notes": [], + } + annotations = {split: {"images": [], "annotations": [], "categories": [ + {"id": 1, "name": "leaf", "supercategory": "plant"} + ]} for split in SPLITS} + used_names = set() + next_annotation_id = 1 + try: + with zipfile.ZipFile(dest_path, "w", zipfile.ZIP_DEFLATED, allowZip64=True) as archive: + for split in SPLITS: + archive.writestr(f"images/{split}/", b"") + if dataset_format.startswith("yolo"): + archive.writestr(f"labels/{split}/", b"") + elif dataset_format == "png": + archive.writestr(f"masks/{split}/", b"") + for image_id, pair in enumerate(valid, 1): + sample_id = pair["sample_id"] + split = assignments[sample_id] + # Never put untrusted sample IDs directly into ZIP member paths. + stem = f"sample_{image_id:06d}" + image_ext = pair["image"].suffix.lower() + if image_ext not in {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp"}: + raise ValueError(f"Unsupported target-box image type: {image_ext}") + image_name = stem + image_ext + image_member = f"images/{split}/{image_name}" + if image_member in used_names: + raise ValueError("Duplicate image name in dataset.") + used_names.add(image_member) + archive.write(pair["image"], image_member) + mask = _binary_mask(pair["mask"]) + if mask.shape != (pair["height"], pair["width"]) or not mask.any(): + raise ValueError(f"Mask changed during export: {sample_id}") + entry = { + "sample_id": sample_id, "group_id": groups.get(sample_id, sample_id), + "split": split, "image": image_member, + "width": pair["width"], "height": pair["height"], + } + if dataset_format == "png": + member = f"masks/{split}/{stem}.png" + archive.writestr(member, _encoded_png(mask * 255)) + entry["label"] = member + elif dataset_format == "yolo-seg": + member = f"labels/{split}/{stem}.txt" + lines, has_holes = _yolo_segments(mask) + if not lines: + raise ValueError(f"Mask has no usable polygon: {sample_id}") + archive.writestr(member, "\n".join(lines) + "\n") + entry["label"] = member + if has_holes: + manifest["conversion_notes"].append( + f"{sample_id}: YOLO polygons cannot preserve mask holes." + ) + if len(lines) > 1: + manifest["conversion_notes"].append( + f"{sample_id}: disconnected regions became separate YOLO segments." + ) + elif dataset_format == "yolo-detect": + member = f"labels/{split}/{stem}.txt" + x, y, width, height = _bbox(mask) + archive.writestr(member, ( + f"0 {(x + width / 2) / pair['width']:.6f} " + f"{(y + height / 2) / pair['height']:.6f} " + f"{width / pair['width']:.6f} {height / pair['height']:.6f}\n" + )) + entry["label"] = member + else: + x, y, width, height = _bbox(mask) + annotations[split]["images"].append({ + "id": image_id, "file_name": image_name, + "width": pair["width"], "height": pair["height"], + }) + annotations[split]["annotations"].append({ + "id": next_annotation_id, "image_id": image_id, "category_id": 1, + "segmentation": _coco_rle(mask), "area": int(mask.sum()), + "bbox": [x, y, width, height], "iscrowd": 0, + }) + next_annotation_id += 1 + entry["label"] = f"annotations/instances_{split}.json" + manifest["samples"].append(entry) + if dataset_format.startswith("yolo"): + archive.writestr("data.yaml", ( + "train: images/train\nval: images/val\ntest: images/test\n" + "names:\n 0: leaf\n" + )) + if dataset_format == "coco": + for split, data in annotations.items(): + archive.writestr(f"annotations/instances_{split}.json", json.dumps(data)) + archive.writestr("manifest.json", json.dumps(manifest, indent=2)) + except Exception: + Path(dest_path).unlink(missing_ok=True) + raise + return manifest diff --git a/src/mats/mask_cleanup.py b/src/mats/mask_cleanup.py new file mode 100644 index 0000000..0f4637b --- /dev/null +++ b/src/mats/mask_cleanup.py @@ -0,0 +1,258 @@ +"""Mask cleanup for masks that are never flash-filled. + +:func:`clean_raw_mask` is the cleanup every such mask gets. It first clears the +edge margin (:func:`clear_margin`), where the template's printed box outline +lands after perspective correction, so those lines can never outrank the leaf. +It then keeps the leaf -- the largest white component -- and drops every other +piece that touches the cleared margin or lies far from the leaf +(:func:`drop_stray_pieces`). Pre-cleanup measurements use it on every image, and +the Clean image preview applies it before anything else. + +Clean image -- a clean size above 0, an alternative to +:func:`mats.core.clean_leaf_mask` -- then drops small white specks and fills small +black holes: only pieces smaller than the chosen radius change, and the leaf is +always kept. + +Sizes are distance-transform inscribed radii, measured once on the input mask. +:func:`clean_levels` encodes them per pixel, so choosing a radius is a per-pixel +comparison. The browser preview applies the same levels with the same rule, so a +live slider shows exactly what :func:`clean_specks_and_holes` returns. + +Imports numpy and OpenCV only -- never torch -- so the offline tests can import it +without :mod:`mats.core`. + +:func:`raw_measurement_mask` is the one entry point for a saved mask: a +pre-cleanup run (``--clean-size``) measures it, and an Adjust overwrite saves it. +Either way the mask is measured under the run's measurement source (the leaf's +bounding box for cleaned runs, the extent of all white pixels for pre-cleanup +runs), so every row of a CSV follows the same rule. +""" + +import math + +import cv2 +import numpy as np + +from .mask_settings import ( + CLEAN_MARGIN_DEFAULT, + CLEAN_SIZE_DEFAULT, + CLEAN_SIZE_MAX, + STRAY_GAP_DEFAULT, + checked_clean_margin, + checked_clean_size, + checked_stray_gap, +) + +# The preview's names for the clean size; 0 means Clean image is off. +CLEAN_RADIUS_DEFAULT = CLEAN_SIZE_DEFAULT +CLEAN_RADIUS_MAX = CLEAN_SIZE_MAX + +# Level channels. A pixel is foreground at radius r when +# r < KEEP_BELOW or FILL_FROM <= r <= FILL_UNTIL. +KEEP_BELOW, FILL_FROM, FILL_UNTIL = 0, 1, 2 +_NEVER_FILL = (255, 0) + + +def _max_per_label(labels, values, count): + maxima = np.zeros(count, dtype=np.float32) + np.maximum.at(maxima, labels.ravel(), values.ravel()) + return np.floor(maxima).astype(np.int32) + + +def _first_pixels(labels, count): + """Row and column of each label's first pixel in raster order.""" + first = np.full(count, labels.size, dtype=np.int64) + np.minimum.at(first, labels.ravel(), np.arange(labels.size)) + return np.divmod(first, labels.shape[1]) + + +def _enclosing(labels, count, other_labels): + """Label, in ``other_labels``, of the component enclosing each component. + + The pixel directly above a component's first raster pixel lies outside it + and belongs to the opposite-colour component immediately around it. Row-0 + components touch the border and get -1. + """ + rows, cols = _first_pixels(labels, count) + enclosing = np.full(count, -1, dtype=np.int64) + inside = rows > 0 + enclosing[inside] = other_labels[rows[inside] - 1, cols[inside]] + return enclosing + + +def margin_width(shape, margin=CLEAN_MARGIN_DEFAULT): + """Pixels cleared from each edge: ``margin`` percent of the shorter side.""" + return int(round(min(shape[:2]) * checked_clean_margin(margin) / 100)) + + +def clear_margin(mask, margin=CLEAN_MARGIN_DEFAULT): + """Set a band ``margin`` percent of the shorter side wide along every edge to 0. + + The template's printed box outline runs through the marker centres, so it + lies on the edge of the perspective-corrected target box. Returns a 0/255 + uint8 mask; ``margin=0`` changes nothing. + """ + cleared = np.where(np.asarray(mask) > 127, 255, 0).astype(np.uint8) + band = margin_width(cleared.shape, margin) + if band: + cleared[:band] = cleared[-band:] = 0 + cleared[:, :band] = cleared[:, -band:] = 0 + return cleared + + +def drop_stray_pieces(mask, max_gap=STRAY_GAP_DEFAULT, border=0): + """Keep the leaf and the white pieces near it; drop the rest. + + The leaf is the largest 8-connected white component, and it is always kept, + even where it crosses the image border. Any other piece is stray when it + comes within ``border`` px of the image edge (touches it, for ``border=0``), + or when its nearest pixel is more than ``max_gap`` times the leaf's + bounding-box diagonal from the leaf; ``max_gap=0`` keeps the leaf alone. + Holes are never filled. Returns a 0/255 uint8 mask. + + A piece at the edge whose inscribed radius is at most ``border`` px -- a + printed line reaching past a cleared margin -- is never taken as the leaf, + however large, unless no other piece is left. + """ + max_gap = checked_stray_gap(max_gap) + white = (np.asarray(mask) > 127).astype(np.uint8) + count, labels, stats, _ = cv2.connectedComponentsWithStats(white, connectivity=8) + if count <= 2: + return white * 255 + height, width = white.shape + x, y, w, h = (stats[:, index] for index in range(4)) + at_edge = ( + (x <= border) | (y <= border) + | (x + w >= width - border) | (y + h >= height - border) + ) + area = stats[:, cv2.CC_STAT_AREA].astype(np.int64) + if border: + radius = _max_per_label(labels, cv2.distanceTransform(white, cv2.DIST_L2, 5), count) + edge_line = at_edge & (radius <= border) + edge_line[0] = False + if not edge_line[1:].all(): + area[edge_line] = -1 + leaf = 1 + np.argmax(area[1:]) + + gap = cv2.distanceTransform( + (labels != leaf).astype(np.uint8), cv2.DIST_L2, cv2.DIST_MASK_PRECISE + ) + others = (labels > 0) & (labels != leaf) + nearest = np.full(count, np.inf, dtype=np.float32) + np.minimum.at(nearest, labels[others], gap[others]) + + stray = at_edge | (nearest > max_gap * math.hypot(w[leaf], h[leaf])) + stray[0] = True + stray[leaf] = False + return np.where(stray[labels], 0, 255).astype(np.uint8) + + +def clean_raw_mask(mask, margin=CLEAN_MARGIN_DEFAULT, max_gap=STRAY_GAP_DEFAULT): + """Clear the edge margin, then drop stray pieces: the mask measured without flash fill. + + A piece touching the cleared band counts as touching the border, so a printed + line that reaches past the margin is still dropped. Idempotent. + """ + cleared = clear_margin(mask, margin) + return drop_stray_pieces(cleared, max_gap, border=margin_width(cleared.shape, margin)) + + +def clean_levels(mask, max_gap=STRAY_GAP_DEFAULT, margin=CLEAN_MARGIN_DEFAULT): + """Per-pixel keep and fill radii for :func:`apply_clean_levels`. + + Returns an ``H x W x 3`` uint8 array with channels ``KEEP_BELOW``, + ``FILL_FROM`` and ``FILL_UNTIL``: + + - The edge margin and stray pieces (:func:`clean_raw_mask`) are gone at + every radius, and any hole inside a stray piece never fills. + - White pixels stay while ``r`` is below their component's inscribed radius + plus one. The largest white component is always kept. + - A black hole (4-connected, not touching the border) fills once ``r`` + exceeds its inscribed radius, for as long as its enclosing white component + is kept, so removing a thin ring never leaves a solid disk behind. + - A white island inside a hole fills with that hole, so a filled hole shows + no black spot where a smaller island was removed. + """ + white = (clean_raw_mask(mask, margin, max_gap) > 127).astype(np.uint8) + height, width = white.shape + levels = np.empty((height, width, 3), dtype=np.uint8) + levels[..., KEEP_BELOW] = 0 + levels[..., FILL_FROM], levels[..., FILL_UNTIL] = _NEVER_FILL + + white_count, white_labels, white_stats, _ = cv2.connectedComponentsWithStats( + white, connectivity=8 + ) + if white_count == 1: + return levels + keep_below = np.minimum( + _max_per_label(white_labels, cv2.distanceTransform(white, cv2.DIST_L2, 5), white_count) + + 1, + 254, + ) + keep_below[0] = 0 + keep_below[1 + np.argmax(white_stats[1:, cv2.CC_STAT_AREA])] = 255 + levels[..., KEEP_BELOW] = keep_below[white_labels] + + black = 1 - white + black_count, black_labels, black_stats, _ = cv2.connectedComponentsWithStats( + black, connectivity=4 + ) + x, y, w, h = (black_stats[:, index] for index in range(4)) + is_hole = (x > 0) & (y > 0) & (x + w < width) & (y + h < height) + is_hole[0] = False + if not is_hole.any(): + return levels + hole_radius = _max_per_label( + black_labels, cv2.distanceTransform(black, cv2.DIST_L2, 5), black_count + ) + fill_from = np.full(black_count, 255, dtype=np.int32) + fill_until = np.zeros(black_count, dtype=np.int32) + enclosing_white = _enclosing(black_labels, black_count, white_labels) + fill_from[is_hole] = np.minimum(hole_radius[is_hole] + 1, 255) + fill_until[is_hole] = np.maximum(keep_below[enclosing_white[is_hole]] - 1, 0) + + # Islands take the fill interval of the hole around them. + enclosing_black = _enclosing(white_labels, white_count, black_labels) + white_fill_from = np.full(white_count, 255, dtype=np.int32) + white_fill_until = np.zeros(white_count, dtype=np.int32) + island = enclosing_black >= 0 + island[island] = is_hole[enclosing_black[island]] + island[0] = False + white_fill_from[island] = fill_from[enclosing_black[island]] + white_fill_until[island] = fill_until[enclosing_black[island]] + + is_white = white.astype(bool) + levels[..., FILL_FROM] = np.where( + is_white, white_fill_from[white_labels], fill_from[black_labels] + ) + levels[..., FILL_UNTIL] = np.where( + is_white, white_fill_until[white_labels], fill_until[black_labels] + ) + return levels + + +def apply_clean_levels(levels, radius): + """Return the 0/255 mask that :func:`clean_levels` encodes at ``radius``.""" + radius = int(radius) + keep = levels[..., KEEP_BELOW] > radius + fill = (levels[..., FILL_FROM] <= radius) & (radius <= levels[..., FILL_UNTIL]) + return np.where(keep | fill, 255, 0).astype(np.uint8) + + +def clean_specks_and_holes(mask, radius, max_gap=STRAY_GAP_DEFAULT, margin=CLEAN_MARGIN_DEFAULT): + """Clear the margin and stray pieces, then specks and holes below ``radius`` px.""" + return apply_clean_levels(clean_levels(mask, max_gap, margin), radius) + + +def raw_measurement_mask(mask, margin=CLEAN_MARGIN_DEFAULT, max_gap=STRAY_GAP_DEFAULT, + clean_size=CLEAN_SIZE_DEFAULT): + """The mask measured without flash fill, at a clean size. + + A ``clean_size`` of 0 is :func:`clean_raw_mask`. Above 0 it is Clean image, + :func:`clean_specks_and_holes` at that radius -- the mask the clean-size + slider previews. + """ + clean_size = checked_clean_size(clean_size) + if clean_size: + return clean_specks_and_holes(mask, clean_size, max_gap, margin) + return clean_raw_mask(mask, margin, max_gap) diff --git a/src/mats/mask_settings.py b/src/mats/mask_settings.py new file mode 100644 index 0000000..f89927c --- /dev/null +++ b/src/mats/mask_settings.py @@ -0,0 +1,70 @@ +"""Settings for masks that are never flash-filled. + +The clean margin clears a band along the target-box edge, where the template's +printed box outline lands after perspective correction. The stray gap sets how +far a mask piece may sit from the leaf before it counts as stray. The clean size +is the radius, in pixels, below which white specks are removed and enclosed black +holes are filled (Clean image); 0 turns that off. + +Deliberately dependency-light: this module imports only the standard library so +that the CLI parser and the Streamlit app can validate the settings without +importing :mod:`mats.mask_cleanup` (numpy, OpenCV) or :mod:`mats.core`. Keep it +that way -- do not add heavy imports here. +""" + +import math +import numbers + +# Percent of the target box's shorter side, cleared along every edge. +CLEAN_MARGIN_DEFAULT = 1.0 +CLEAN_MARGIN_MAX = 10.0 + +# A fraction of the leaf's bounding-box diagonal. +STRAY_GAP_DEFAULT = 0.25 +STRAY_GAP_MAX = 10.0 + +# Inscribed radius in pixels; 0 means Clean image is off. +CLEAN_SIZE_DEFAULT = 0 +CLEAN_SIZE_MAX = 50 + + +def _checked_number(value, name, maximum): + message = f"{name} must be a number 0-{maximum:g} (got {value!r})" + if isinstance(value, bool): + raise ValueError(message) + try: + number = float(value) + except (TypeError, ValueError): + raise ValueError(message) from None + if not (math.isfinite(number) and 0 <= number <= maximum): + raise ValueError(message) + return number + + +def checked_clean_margin(value): + """Return ``value`` as a float from 0 to ``CLEAN_MARGIN_MAX``, else raise ``ValueError``.""" + return _checked_number(value, "clean margin", CLEAN_MARGIN_MAX) + + +def checked_stray_gap(value): + """Return ``value`` as a float from 0 to ``STRAY_GAP_MAX``, else raise ``ValueError``.""" + return _checked_number(value, "stray gap", STRAY_GAP_MAX) + + +def checked_clean_size(value): + """Return ``value`` as an int from 0 to ``CLEAN_SIZE_MAX`` px, else raise ``ValueError``. + + Accepts an integer or its decimal text (as the CLI passes it); a fractional + radius is rejected rather than rounded. + """ + message = f"clean size must be a whole number of pixels 0-{CLEAN_SIZE_MAX} (got {value!r})" + if isinstance(value, bool): + raise ValueError(message) + if isinstance(value, str): + try: + value = int(value.strip()) + except ValueError: + raise ValueError(message) from None + if not isinstance(value, numbers.Integral) or not 0 <= value <= CLEAN_SIZE_MAX: + raise ValueError(message) + return int(value) diff --git a/src/mats/scaling.py b/src/mats/scaling.py index 8222c7e..c6984b0 100644 --- a/src/mats/scaling.py +++ b/src/mats/scaling.py @@ -2,8 +2,8 @@ Deliberately dependency-light: this module imports only the standard library so the offline test suite can exercise the scale math without importing -:mod:`mats.core` (which pulls in torch, rfdetr, cv2 and transformers). Keep it -that way -- do not add heavy imports here. +:mod:`mats.core` (which imports NumPy and OpenCV). Keep it that way -- do not +add heavy imports here. The pipeline calibrates against the *template*, not the leaf: the warped target-box raster's pixel width/height are compared to the template's known diff --git a/src/mats/thresholds.py b/src/mats/thresholds.py new file mode 100644 index 0000000..15ab398 --- /dev/null +++ b/src/mats/thresholds.py @@ -0,0 +1,70 @@ +"""Grayscale threshold levels for classic (non-BiRefNet) segmentation. + +Deliberately dependency-light: this module imports only the standard library so +that the CLI parser and the Streamlit app can validate threshold choices without +importing :mod:`mats.core` (which pulls in torch, rfdetr, cv2 and transformers). +Keep it that way -- do not add heavy imports here. +""" + +THRESHOLD_LEVELS = { + "auto": None, # Otsu's method: threshold computed per-image from histogram + "low": 100, + "medium": 125, + "high": 150, +} +THRESHOLD_MIN = 1 +THRESHOLD_MAX = 255 +# Deliberately not a THRESHOLD_LEVELS key: None there means Otsu, so a lost +# custom value must fail loudly instead of silently running Otsu. +CUSTOM_THRESHOLD_LEVEL = "custom" +THRESHOLD_LEVEL_OPTIONS = ("auto", CUSTOM_THRESHOLD_LEVEL, "low", "medium", "high") + + +def _checked_cutoff(value): + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError( + f"threshold must be an integer {THRESHOLD_MIN}-{THRESHOLD_MAX} (got {value!r})" + ) + if not THRESHOLD_MIN <= value <= THRESHOLD_MAX: + raise ValueError( + f"threshold must be an integer {THRESHOLD_MIN}-{THRESHOLD_MAX} (got {value})" + ) + return value + + +def parse_threshold_level(text): + """Parse a threshold choice into a preset name or an integer cutoff. + + Preset names (``auto``, ``low``, ``medium``, ``high``) are matched + case-insensitively and returned lowercased; a plain integer between + ``THRESHOLD_MIN`` and ``THRESHOLD_MAX`` is returned as an ``int``. Anything + else -- including a bare ``custom``, which names no cutoff -- raises + ``ValueError``. + """ + level = str(text).strip().lower() + if level in THRESHOLD_LEVELS: + return level + if level == CUSTOM_THRESHOLD_LEVEL: + raise ValueError( + "custom needs a number: pass the cutoff as an integer " + f"{THRESHOLD_MIN}-{THRESHOLD_MAX}" + ) + if not (level.isascii() and level.isdigit()): + raise ValueError( + "threshold level must be auto, low, medium, high, or an integer " + f"{THRESHOLD_MIN}-{THRESHOLD_MAX} (got {text!r})" + ) + return _checked_cutoff(int(level)) + + +def threshold_value_for(level, custom_value=None): + """Return the cutoff passed to the pipeline: ``None`` (Otsu) or an ``int``. + + ``level`` is a preset name, an integer cutoff, or ``"custom"`` -- in which + case ``custom_value`` supplies the cutoff. + """ + if level == CUSTOM_THRESHOLD_LEVEL: + return _checked_cutoff(custom_value) + if isinstance(level, int) and not isinstance(level, bool): + return _checked_cutoff(level) + return THRESHOLD_LEVELS[level] diff --git a/src/mats/weights.py b/src/mats/weights.py index ec40196..d0602e6 100644 --- a/src/mats/weights.py +++ b/src/mats/weights.py @@ -1,16 +1,17 @@ """Install, verify and resolve the MATs model checkpoints. The checkpoints are large (RF-DETR ~134 MB, BiRefNet ~2.65 GB) and are -delivered through two independent channels, plus a shared-filesystem escape -hatch -- :mod:`mats.paths` resolves to whichever channel produces a real file: - -1. **Hugging Face Hub** -- the default public host. Free, no account needed, - but unreachable on some institutional networks (notably USDA's). -2. **Git LFS** -- By default, MATS will only pull the RF-DETR checkpoint file - when fetching weights via ``git-lfs pull``. To pull the larger BiRefNet file, - you can run ``mats fetch-weights --only birefnet --source lfs`` or just - ``git-lfs pull``. This channel exists because Hugging Face is not reachable - from every collaborator's network. +delivered through Git LFS or an optional Hugging Face configuration, plus a +shared-filesystem escape hatch -- :mod:`mats.paths` resolves to whichever +channel produces a real file: + +1. **Hugging Face Hub** -- available only when a MATS weights repository is + configured. It is free and needs no account, but is unreachable on some + institutional networks (notably USDA's). +2. **Git LFS** -- The default clone and pull exclude BiRefNet, so RF-DETR is + available without an automatic 2.65 GB download. Install BiRefNet explicitly + with ``mats fetch-weights --only birefnet --source lfs``. This channel exists + because Hugging Face is not reachable from every collaborator's network. 3. **A shared/mounted filesystem** (e.g. USDA SCINet ``/project``) -- point ``MATS_WEIGHTS_DIR`` at it and the weights are read in place, no download, for anyone who can mount it. @@ -205,7 +206,17 @@ def get_weight_status(name): checkout = _checkout_target(name) if checkout is not None and looks_like_lfs_pointer(checkout): - detail = "Not yet fetched via Git LFS -- fetch it via the app, `mats fetch-weights --only birefnet` or Hugging Face." + if name == "birefnet": + detail = ( + "Excluded from the default Git LFS clone. Fetch it via the app or " + "`mats fetch-weights --only birefnet --source lfs`." + ) + else: + detail = ( + "Git LFS left an RF-DETR pointer instead of the checkpoint. Run " + "`git lfs install && git lfs pull --exclude=\"weights/birefnet_leaf.pth\"` " + "or `mats fetch-weights --only rf-detr --source lfs`." + ) return WeightStatus(name, checkout, "missing", detail, 0, spec["size_bytes"], sources) target = _download_target(name) @@ -220,18 +231,36 @@ def _emit(progress_callback, phase, completed, total): progress_callback(phase, completed, total) -def _manual_instructions(): +def _lfs_pull_args(name): + """Return a Git LFS pull command that fetches only the intended weights.""" + birefnet_rel = f"weights/{_MANIFEST['birefnet']['filename']}" + if name == "birefnet": + return ["git", "lfs", "pull", "-X", "", "-I", birefnet_rel] + return ["git", "lfs", "pull", "--exclude", birefnet_rel] + + +def _lfs_manual_command(name): + """Return the shell form of the checkpoint-specific Git LFS repair.""" + if name == "birefnet": + return 'git lfs pull -X "" -I "weights/birefnet_leaf.pth"' + return 'git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth"' + + +def _manual_instructions(name): + spec = _MANIFEST[name] + override = ( + "RF_DETR_MARKER_CHECKPOINT" if name == "rf-detr" else "BIREFNET_CHECKPOINT" + ) print( "No automatic download source is available in this build.\n\n" - "Get the checkpoints one of these ways:\n" - f" - Download them and place them here:\n" - f" {WEIGHTS_DIR / RF_DETR_MARKER_FILENAME}\n" - f" {WEIGHTS_DIR / BIREFNET_FILENAME}\n" - " - Or set MATS_WEIGHTS_DIR to a directory that already contains them\n" + f"Get {spec['filename']} one of these ways:\n" + " - Download it and place it here:\n" + f" {WEIGHTS_DIR / spec['filename']}\n" + " - Or set MATS_WEIGHTS_DIR to a directory that already contains it\n" " (e.g. a shared SCINet /project path).\n" - " - Or set RF_DETR_MARKER_CHECKPOINT / BIREFNET_CHECKPOINT to specific files.\n" + f" - Or set {override} to the specific file.\n" " - Or, from a Git checkout with Git LFS installed:\n" - " git lfs pull\n\n" + f" {_lfs_manual_command(name)}\n\n" "See docs/weights.md.", file=sys.stderr, ) @@ -334,11 +363,9 @@ def _emit_lfs_progress(progress_path, progress_callback, fallback_total, last_do def _download_from_lfs(name, progress_callback=None): """Fetch one checkpoint via Git LFS. - For BiRefNet, runs a plain ``git lfs pull`` (no flags) so the large - checkpoint is fetched without affecting other files. For all other - checkpoints, runs ``git lfs pull --exclude weights/birefnet_leaf.pth`` - so the 2.65 GB BiRefNet file is never pulled as a side-effect of an - unrelated weight update. + BiRefNet clears the repository exclusion for one invocation and includes + only its checkpoint. All other checkpoints explicitly exclude BiRefNet so + it is never pulled as a side effect of an unrelated weight update. Writes into the checkout's weights/ directory -- that's where Git LFS smudges content, and it's tier 3 of paths.py's resolution order, so the @@ -364,16 +391,10 @@ def _download_from_lfs(name, progress_callback=None): rel_path = f"weights/{spec['filename']}" print(f"Fetching {rel_path} via Git LFS -> {target}") - # For BiRefNet use a plain `git lfs pull` (no flags) -- without a - # fetchexclude in .lfsconfig a bare pull fetches all LFS files, which is - # what we want for this explicit opt-in download. - # For everything else, exclude the large BiRefNet checkpoint so it is - # never pulled as an unintended side-effect. - birefnet_rel = f"weights/{_MANIFEST['birefnet']['filename']}" - if name == "birefnet": - lfs_cmd = ["git", "lfs", "pull"] - else: - lfs_cmd = ["git", "lfs", "pull", "--exclude", birefnet_rel] + # `-X ""` clears .lfsconfig's exclusion for the explicit BiRefNet request; + # `-I` keeps that pull scoped to BiRefNet. Other requests explicitly + # exclude the large optional checkpoint under either repository setting. + lfs_cmd = _lfs_pull_args(name) with tempfile.TemporaryDirectory() as tmp: progress_path = Path(tmp) / "progress" @@ -495,7 +516,7 @@ def fetch(only=None, force=False, source="auto"): print(f"error: {by_id[source].label} is unavailable: {by_id[source].reason}", file=sys.stderr) else: - _manual_instructions() + _manual_instructions(name) ok = False continue @@ -530,13 +551,13 @@ def ensure_weight(name): raise FileNotFoundError( f"{spec['filename']} not found and auto-fetch is disabled " f"({_AUTO_FETCH_DISABLED} is set). Pre-stage the weights, or run " - f"`mats fetch-weights --only {name}` after unsetting {_AUTO_FETCH_DISABLED} " - f"(from a Git checkout, `git lfs pull` fetches all weights including birefnet, or " - f"`git lfs pull --exclude weights/{_MANIFEST['birefnet']['filename']}` fetches all others)." + f"`mats fetch-weights --only {name}` after unsetting " + f"{_AUTO_FETCH_DISABLED}. From a Git checkout with Git LFS, run " + f"`{_lfs_manual_command(name)}`." ) - # A pointer stub for BiRefNet means the user hasn't run `git lfs pull` - # for it yet (it's large and opt-in) -- fetch it rather than failing. + # A pointer stub means Git LFS has not materialized this checkpoint yet. + # Fetch the requested checkpoint rather than handing the stub to a model. checkout = _checkout_target(name) if checkout is not None and looks_like_lfs_pointer(checkout) and _download_from_lfs(name): return checkout @@ -570,8 +591,9 @@ def require_local_weight(name): ) raise FileNotFoundError( f"{_MANIFEST[name]['filename']} is not installed locally. " - f"BiRefNet is optional; install it explicitly with " - f"`mats fetch-weights --only {name} --source lfs`, or place it at {status.path}." + f"BiRefNet is optional; from a Git checkout with Git LFS, install it " + f"explicitly with `mats fetch-weights --only {name} --source lfs`, or " + f"place it at {status.path}." ) diff --git a/tests/test_birefnet_runtime.py b/tests/test_birefnet_runtime.py index 91a9c1d..844e630 100644 --- a/tests/test_birefnet_runtime.py +++ b/tests/test_birefnet_runtime.py @@ -24,6 +24,7 @@ def fake_import(name): def test_local_loader_never_calls_auto_fetch_or_transformers(monkeypatch, tmp_path): pytest.importorskip("numpy") + torch = pytest.importorskip("torch") from mats import core, weights class FakeModel: @@ -53,8 +54,8 @@ def eval(self): monkeypatch.setitem(sys.modules, "mats.models.birefnet", fake_model_module) monkeypatch.setattr(weights, "require_local_weight", lambda name: tmp_path / "birefnet_leaf.pth") monkeypatch.setattr(weights, "ensure_weight", lambda name: pytest.fail("must not auto-fetch")) - monkeypatch.setattr(core.torch, "load", lambda *args, **kwargs: {"model_state_dict": {"local": 1}}) - monkeypatch.setattr(core, "resolve_birefnet_device", lambda *_: core.torch.device("cpu")) + monkeypatch.setattr(torch, "load", lambda *args, **kwargs: {"model_state_dict": {"local": 1}}) + monkeypatch.setattr(core, "resolve_birefnet_device", lambda *_: torch.device("cpu")) core._BIREFNET_MODELS.clear() model = core.get_birefnet_model() diff --git a/tests/test_cli_args.py b/tests/test_cli_args.py index 6782439..2d889fa 100644 --- a/tests/test_cli_args.py +++ b/tests/test_cli_args.py @@ -1,17 +1,23 @@ """CLI argument parsing -- no pipeline execution, so no torch required.""" +import sys from types import SimpleNamespace import pytest from mats.cli import ( + _cmd_run, + _measured_methods, _normalize_argv, + _pre_cleanup_methods, + _print_run_banner, _require_local_birefnet_for_run, _resolve_fetch_only, _resolve_template_dims, build_parser, ) from mats.dimensions import parse_template_dimensions +from mats.mask_settings import CLEAN_MARGIN_DEFAULT, CLEAN_SIZE_DEFAULT, STRAY_GAP_DEFAULT def test_default_subcommand_inserted(): @@ -35,9 +41,138 @@ def test_run_defaults(): assert ns.threshold_level == "auto" assert ns.csv_schema == "full" # research schema by default assert ns.results_unit == "cm" + assert not ns.measure_pre_cleanup assert ns.save_axes is False assert ns.sheet_dimensions is None assert ns.template_dimensions is None + assert ns.export == [] + assert ns.pre_cleanup_methods is None + assert not ns.no_target_boxes and not ns.no_masks and not ns.no_failure_log + assert ns.dataset_format is None + + +def test_dataset_options_parse_and_validate(tmp_path, monkeypatch, capsys): + args = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), + "--dataset-format", "yolo-seg", "--dataset-split", "60", "30", "10", + "--dataset-seed", "17", "--dataset-mask-source", "raw", + ]) + assert args.dataset_format == "yolo-seg" + assert args.dataset_split == [60, 30, 10] + assert args.dataset_seed == 17 + assert args.dataset_mask_source == "raw" + + calls = [] + _fake_core(monkeypatch, calls) + bad = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), + "--dataset-format", "png", "--dataset-split", "70", "20", "9", + ]) + with pytest.raises(SystemExit): + _cmd_run(bad) + assert "total 100" in capsys.readouterr().err + assert not calls + + +def test_cli_dataset_uses_shared_pairing_and_exporter(tmp_path, monkeypatch): + calls = [] + _fake_core(monkeypatch, calls) + fake_core = sys.modules["mats.core"] + image_path = tmp_path / "leaf_target_box.png" + mask_path = tmp_path / "leaf_mask.png" + image_path.write_bytes(b"image") + mask_path.write_bytes(b"mask") + + def run_leaf_morpho_batch(**kwargs): + calls.append(kwargs) + return { + "succeeded": 1, "failed": 0, "workers": 1, "worker_reason": "test", + "methods": ("threshold",), "results_path": kwargs["results_path"], + "failure_report_path": None, "measurement_source": "cleaned", + "artifacts": [{ + "path": str(mask_path), "sample_id": "leaf", "kind": "mask", + "method": "threshold", + }], + "preview_artifacts": [{ + "path": str(image_path), "sample_id": "leaf", "kind": "preview_target_box", + "method": None, + }], + } + + fake_core.run_leaf_morpho_batch = run_leaf_morpho_batch + exported = {} + + def fake_write(pairs, dest, **options): + exported["pairs"] = pairs + exported["options"] = options + dest.write_bytes(b"zip") + return {"counts": {"train": 1, "val": 0, "test": 0}, + "excluded": [], "conversion_notes": []} + + monkeypatch.setattr("mats.dataset_export.write_dataset_zip", fake_write) + args = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), + "-t", "10x10cm", "--dataset-format", "png", "--no-target-boxes", + "--no-masks", + ]) + assert _cmd_run(args) == 0 + assert exported["pairs"] == [{ + "sample_id": "leaf", "target_box": str(image_path), "mask": str(mask_path), + }] + assert exported["options"]["percentages"] == (70, 20, 10) + assert calls[0]["export_options"]["preview_dir"] + assert (tmp_path / "out" / "mats_training_threshold_png.zip").read_bytes() == b"zip" + + +def test_repeated_exports_and_both_methods_parse(): + ns = build_parser().parse_args([ + "run", "--export", "pre-cleanup", "--export", "overlay", + "--pre-cleanup-methods", "both", "--no-target-boxes", "--no-masks", + ]) + assert ns.export == ["pre-cleanup", "overlay"] + assert ns.pre_cleanup_methods == "both" + assert ns.no_target_boxes and ns.no_masks + + +def test_mask_method_both_measures_and_exports_each_method(): + parse = build_parser().parse_args + both = parse(["run", "-i", "x", "--mask-method", "both", "--export", "pre-cleanup"]) + assert both.mask_method == "both" + assert _measured_methods(both) == ("threshold", "birefnet") + assert _pre_cleanup_methods(both) == ("threshold", "birefnet") + + single = parse(["run", "-i", "x", "--mask-method", "birefnet", "--export", "pre-cleanup"]) + assert _measured_methods(single) == ("birefnet",) + assert _pre_cleanup_methods(single) == ("birefnet",) + assert _pre_cleanup_methods(parse(["run", "-i", "x", "--mask-method", "both"])) == () + + +@pytest.mark.parametrize("method", ["selected", "threshold"]) +def test_pre_cleanup_method_option_requires_export(method, tmp_path, monkeypatch, capsys): + calls = [] + _fake_core(monkeypatch, calls) + args = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), + "--pre-cleanup-methods", method, + ]) + with pytest.raises(SystemExit) as exc: + _cmd_run(args) + assert exc.value.code == 2 + assert "--pre-cleanup-methods requires --export pre-cleanup" in capsys.readouterr().err + assert not calls + + +def test_default_pre_cleanup_method_applies_when_export_requested(): + args = build_parser().parse_args([ + "run", "--mask-method", "both", "--export", "pre-cleanup", + ]) + assert args.pre_cleanup_methods is None + assert _pre_cleanup_methods(args) == ("threshold", "birefnet") + + +def test_unknown_export_rejected(): + with pytest.raises(SystemExit): + build_parser().parse_args(["run", "--export", "unknown"]) def test_underscore_and_hyphen_aliases_agree(): @@ -54,6 +189,151 @@ def test_compact_and_axes_opt_in(): assert ns.save_axes is True +def test_pre_cleanup_measurement_flag_parses_independently_of_export(): + ns = build_parser().parse_args(["run", "--measure-pre-cleanup"]) + assert ns.measure_pre_cleanup + assert ns.export == [] + + +def test_stray_gap_defaults_and_parses(): + parse = build_parser().parse_args + assert parse(["run", "-i", "x"]).stray_gap == STRAY_GAP_DEFAULT + assert parse(["run", "-i", "x", "--measure-pre-cleanup", "--stray-gap", "0"]).stray_gap == 0.0 + assert parse(["run", "-i", "x", "--stray-gap", "1.5"]).stray_gap == 1.5 + + +@pytest.mark.parametrize("value", ["-0.1", "10.5", "nan", "inf", "far"]) +def test_stray_gap_rejects_invalid_values(value, capsys): + with pytest.raises(SystemExit) as exc: + build_parser().parse_args(["run", "-i", "x", "--stray-gap", value]) + assert exc.value.code == 2 + assert "--stray-gap" in capsys.readouterr().err + + +def test_clean_margin_defaults_and_parses(): + parse = build_parser().parse_args + assert parse(["run", "-i", "x"]).clean_margin == CLEAN_MARGIN_DEFAULT + assert parse(["run", "-i", "x", "--clean-margin", "0"]).clean_margin == 0.0 + assert parse(["run", "-i", "x", "--clean-margin", "2.5"]).clean_margin == 2.5 + + +@pytest.mark.parametrize("value", ["-1", "10.5", "nan", "wide"]) +def test_clean_margin_rejects_invalid_values(value, capsys): + with pytest.raises(SystemExit) as exc: + build_parser().parse_args(["run", "-i", "x", "--clean-margin", value]) + assert exc.value.code == 2 + assert "--clean-margin" in capsys.readouterr().err + + +def test_clean_size_defaults_and_parses(): + parse = build_parser().parse_args + assert parse(["run", "-i", "x"]).clean_size == CLEAN_SIZE_DEFAULT == 0 + assert parse(["run", "-i", "x", "--measure-pre-cleanup", "--clean-size", "3"]).clean_size == 3 + assert parse(["run", "-i", "x", "--clean-size", "50"]).clean_size == 50 + + +@pytest.mark.parametrize("value", ["-1", "51", "2.5", "big"]) +def test_clean_size_rejects_invalid_values(value, capsys): + with pytest.raises(SystemExit) as exc: + build_parser().parse_args(["run", "-i", "x", "--clean-size", value]) + assert exc.value.code == 2 + assert "--clean-size" in capsys.readouterr().err + + +def test_banner_reports_clean_size_for_pre_cleanup(capsys): + parse = build_parser().parse_args + _print_run_banner(parse(["run", "-i", "x", "--measure-pre-cleanup", "--clean-size", "3"]), None) + assert "Clean size: 3 px" in capsys.readouterr().out + _print_run_banner(parse(["run", "-i", "x", "--measure-pre-cleanup"]), None) + assert "Clean size: off" in capsys.readouterr().out + _print_run_banner(parse(["run", "-i", "x"]), None) + assert "Clean size" not in capsys.readouterr().out + + +def test_banner_reports_margin_and_stray_gap_only_for_pre_cleanup(capsys): + parse = build_parser().parse_args + _print_run_banner(parse([ + "run", "-i", "x", "--measure-pre-cleanup", "--stray-gap", "0.5", "--clean-margin", "2", + ]), None) + out = capsys.readouterr().out + assert "Clean margin: 2% of the target box's shorter side" in out + assert "Stray gap: 0.5 x leaf bounding-box diagonal" in out + _print_run_banner(parse(["run", "-i", "x"]), None) + out = capsys.readouterr().out + assert "Stray gap" not in out and "Clean margin" not in out + + +def _fake_core(monkeypatch, calls): + """Stand in for mats.core so _cmd_run runs without torch.""" + import mats + + def run_leaf_morpho_batch(**kwargs): + calls.append(kwargs) + return { + "succeeded": 1, "failed": 0, "workers": 1, "worker_reason": "test", + "methods": ("threshold",), "results_path": kwargs["results_path"], + "failure_report_path": None, + } + + fake = SimpleNamespace( + parse_template_dimensions=parse_template_dimensions, + method_suffixed_path=lambda path, method: path, + get_input_images=lambda input_dir: ["leaf.jpg"], + run_leaf_morpho_batch=run_leaf_morpho_batch, + ) + monkeypatch.setitem(sys.modules, "mats.core", fake) + monkeypatch.setattr(mats, "core", fake, raising=False) + + +def test_stray_gap_reaches_the_pipeline(tmp_path, monkeypatch): + calls = [] + _fake_core(monkeypatch, calls) + args = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), "-t", "10x10cm", + "--measure-pre-cleanup", "--stray-gap", "0.6", "--clean-margin", "2", + ]) + assert _cmd_run(args) == 0 + assert calls[0]["measurement_source"] == "pre-cleanup" + assert calls[0]["stray_gap"] == 0.6 + assert calls[0]["clean_margin"] == 2.0 + assert calls[0]["clean_size"] == CLEAN_SIZE_DEFAULT + + +def test_clean_size_reaches_the_pipeline(tmp_path, monkeypatch): + calls = [] + _fake_core(monkeypatch, calls) + args = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), "-t", "10x10cm", + "--measure-pre-cleanup", "--threshold-level", "medium", "--clean-size", "3", + ]) + assert _cmd_run(args) == 0 + assert calls[0]["measurement_source"] == "pre-cleanup" + assert calls[0]["threshold_value"] == 125 + assert calls[0]["clean_size"] == 3 + + +@pytest.mark.parametrize( + "flag, value", [ + ("--stray-gap", "0.6"), ("--stray-gap", str(STRAY_GAP_DEFAULT)), + ("--clean-margin", "0.6"), ("--clean-margin", str(CLEAN_MARGIN_DEFAULT)), + ("--clean-size", "3"), ("--clean-size", str(CLEAN_SIZE_DEFAULT)), + ], +) +def test_cleanup_settings_require_pre_cleanup_measurement( + flag, value, tmp_path, monkeypatch, capsys, +): + calls = [] + _fake_core(monkeypatch, calls) + args = build_parser().parse_args([ + "run", "-i", str(tmp_path), "-o", str(tmp_path / "out"), "-t", "10x10cm", + flag, value, + ]) + with pytest.raises(SystemExit): + _cmd_run(args) + assert f"{flag} requires --measure-pre-cleanup" in capsys.readouterr().err + assert not calls + + def test_results_unit_accepts_supported_choices(): ns = build_parser().parse_args(["run", "-i", "x", "--results-unit", "in"]) assert ns.results_unit == "in" @@ -64,6 +344,39 @@ def test_invalid_choice_rejected(): build_parser().parse_args(["run", "-i", "x", "--mask-method", "nonsense"]) +def test_threshold_level_accepts_presets_and_custom_cutoffs(): + parse = build_parser().parse_args + assert parse(["run", "-i", "x", "--threshold-level", "177"]).threshold_level == 177 + assert parse(["run", "-i", "x", "--threshold-level", "HIGH"]).threshold_level == "high" + + +@pytest.mark.parametrize("value", ["0", "256", "custom", "12.5"]) +def test_threshold_level_rejects_invalid_cutoffs(value, capsys): + with pytest.raises(SystemExit) as exc: + build_parser().parse_args(["run", "-i", "x", "--threshold-level", value]) + assert exc.value.code == 2 + assert "--threshold-level" in capsys.readouterr().err + + +@pytest.mark.parametrize("extra", [ + [], + ["--mask-method", "birefnet", "--export", "pre-cleanup", "--pre-cleanup-methods", "threshold"], + ["--mask-method", "both"], +]) +def test_banner_reports_custom_threshold_whenever_thresholding_runs(extra, capsys): + args = build_parser().parse_args(["run", "-i", "x", "--threshold-level", "177", *extra]) + _print_run_banner(args, 177) + assert "Threshold level: custom (177)" in capsys.readouterr().out + + +def test_banner_omits_threshold_when_birefnet_alone_segments(capsys): + args = build_parser().parse_args( + ["run", "-i", "x", "--mask-method", "birefnet", "--threshold-level", "177"] + ) + _print_run_banner(args, 177) + assert "Threshold level" not in capsys.readouterr().out + + def test_sheet_and_legacy_dimensions_are_mutually_exclusive(): with pytest.raises(SystemExit): build_parser().parse_args( @@ -85,6 +398,17 @@ def test_legacy_template_dimensions_keep_their_historical_meaning(): assert _resolve_template_dims(args, lm) == (10.5, 9.5, "in") +@pytest.mark.parametrize("method", ["birefnet", "both"]) +def test_local_birefnet_preflight_runs_whenever_birefnet_measures(method, monkeypatch): + args = build_parser().parse_args(["run", "-i", "x", "--mask-method", method]) + required = [] + monkeypatch.setattr("mats.birefnet_runtime.require_birefnet_dependencies", lambda: None) + monkeypatch.setattr("mats.weights.require_local_weight", required.append) + + _require_local_birefnet_for_run(args) + assert required == ["birefnet"] + + def test_local_birefnet_preflight_skips_otsu(monkeypatch): args = build_parser().parse_args(["run", "-i", "x", "--mask-method", "threshold"]) monkeypatch.setattr("mats.weights.require_local_weight", lambda name: pytest.fail("not needed")) diff --git a/tests/test_dataset_export.py b/tests/test_dataset_export.py new file mode 100644 index 0000000..f1a91ae --- /dev/null +++ b/tests/test_dataset_export.py @@ -0,0 +1,107 @@ +"""Training exports preserve pair alignment and deterministic split assignments.""" + +import json +import zipfile + +import pytest + +np = pytest.importorskip("numpy") +cv2 = pytest.importorskip("cv2") + +from mats.dataset_export import ( + _coco_rle, assign_splits, read_group_csv, split_counts, write_dataset_zip, +) + + +def _pair(folder, sample_id, *, hole=False): + image = np.full((16, 20, 3), 220, dtype=np.uint8) + mask = np.zeros((16, 20), dtype=np.uint8) + mask[3:13, 4:15] = 255 + if hole: + mask[6:9, 8:11] = 0 + image_path = folder / f"{sample_id}.png" + mask_path = folder / f"{sample_id}_mask.png" + assert cv2.imwrite(str(image_path), image) + assert cv2.imwrite(str(mask_path), mask) + return {"sample_id": sample_id, "target_box": str(image_path), "mask": str(mask_path)} + + +def test_split_counts_and_grouped_assignments_are_reproducible(): + assert split_counts(10, (70, 20, 10)) == {"train": 7, "val": 2, "test": 1} + pairs = [{"sample_id": str(index)} for index in range(10)] + groups = {"0": "plant_a", "1": "plant_a"} + first, counts = assign_splits(pairs, seed=17, groups=groups) + second, _ = assign_splits(pairs, seed=17, groups=groups) + assert first == second + assert first["0"] == first["1"] + assert sum(counts.values()) == 10 + with pytest.raises(ValueError, match="total 100"): + split_counts(10, (70, 20, 9)) + + +def test_png_and_yolo_exports_keep_images_and_labels_aligned(tmp_path): + pairs = [_pair(tmp_path, f"leaf{index}") for index in range(10)] + for dataset_format in ("png", "yolo-seg", "yolo-detect"): + output = tmp_path / f"{dataset_format}.zip" + manifest = write_dataset_zip(pairs, output, dataset_format=dataset_format) + assert manifest["counts"] == {"train": 7, "val": 2, "test": 1} + with zipfile.ZipFile(output) as archive: + assert json.loads(archive.read("manifest.json"))["samples"] == manifest["samples"] + for sample in manifest["samples"]: + assert sample["image"] in archive.namelist() + assert sample["label"] in archive.namelist() + if dataset_format == "png": + decoded = cv2.imdecode( + np.frombuffer(archive.read(sample["label"]), dtype=np.uint8), 0 + ) + assert decoded.shape == (16, 20) + assert decoded[4, 5] == 255 and decoded[0, 0] == 0 + else: + label = archive.read(sample["label"]).decode().strip().split() + assert label[0] == "0" + assert all(0 <= float(value) <= 1 for value in label[1:]) + if dataset_format.startswith("yolo"): + assert "data.yaml" in archive.namelist() + + +def test_coco_rle_preserves_holes_and_invalid_pairs_are_reported(tmp_path): + good = _pair(tmp_path, "good", hole=True) + bad = _pair(tmp_path, "bad") + bad["mask"] = str(tmp_path / "missing.png") + output = tmp_path / "coco.zip" + manifest = write_dataset_zip([good, bad], output, dataset_format="coco") + assert manifest["counts"] == {"train": 1, "val": 0, "test": 0} + assert manifest["excluded"] == [{ + "sample_id": "bad", "reason": "image or selected mask file is missing", + }] + with zipfile.ZipFile(output) as archive: + annotations = json.loads(archive.read("annotations/instances_train.json")) + annotation = annotations["annotations"][0] + assert annotation["area"] == 101 + runs = annotation["segmentation"]["counts"] + values = [] + for index, length in enumerate(runs): + values.extend([index % 2] * length) + restored = np.array(values, dtype=np.uint8).reshape((20, 16)).T + assert restored[7, 9] == 0 and restored[4, 5] == 1 + + +def test_group_csv_rejects_conflicts(): + assert read_group_csv(b"sample_id,group_id\na,plant1\n") == {"a": "plant1"} + with pytest.raises(ValueError, match="Conflicting"): + read_group_csv(b"sample_id,group_id\na,plant1\na,plant2\n") + + +def test_yolo_seg_excludes_mask_without_a_polygon(tmp_path): + good = _pair(tmp_path, "good") + tiny = _pair(tmp_path, "tiny") + mask = np.zeros((16, 20), dtype=np.uint8) + mask[5, 6] = 255 + assert cv2.imwrite(tiny["mask"], mask) + manifest = write_dataset_zip( + [good, tiny], tmp_path / "small.zip", dataset_format="yolo-seg", + ) + assert len(manifest["samples"]) == 1 + assert manifest["excluded"] == [{ + "sample_id": "tiny", "reason": "mask has no usable polygon", + }] diff --git a/tests/test_dual_method.py b/tests/test_dual_method.py new file mode 100644 index 0000000..ecfa944 --- /dev/null +++ b/tests/test_dual_method.py @@ -0,0 +1,151 @@ +"""Measure with Otsu and BiRefNet in one run, without model downloads.""" + +import pytest + +np = pytest.importorskip("numpy") +cv2 = pytest.importorskip("cv2") +pytest.importorskip("torch") +pytest.importorskip("rfdetr") +from mats import core + + +def _input(tmp_path, name="leaf"): + image = np.full((100, 100, 3), 255, dtype=np.uint8) + image[20:80, 20:80] = 0 + image[43:57, 43:57] = 255 + image[5:10, 5:10] = 0 + path = tmp_path / f"{name}_target_box.png" + assert cv2.imwrite(str(path), image) + return path + + +def _fake_birefnet(calls, leaf=True): + def predict(image, device_override=None): + calls.append(image.shape) + mask = np.zeros(image.shape[:2], dtype=np.uint8) + if leaf: + mask[30:70, 30:70] = 255 + return mask + return predict + + +def _run(sources, output, mask_method, **kwargs): + kwargs.setdefault("workers", 1) + return core.run_leaf_morpho_batch( + [str(source) for source in sources], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), mask_method=mask_method, + compact_csv=False, write_failures=True, **kwargs, + ) + + +def test_resolve_mask_methods_accepts_both_and_sequences(): + assert core.resolve_mask_methods("threshold") == ("threshold",) + assert core.resolve_mask_methods("both") == ("threshold", "birefnet") + assert core.resolve_mask_methods(["birefnet", "threshold", "birefnet"]) == ( + "birefnet", "threshold" + ) + for bad in ((), "nonsense", ["threshold", "nonsense"]): + with pytest.raises(ValueError): + core.resolve_mask_methods(bad) + assert core.method_suffixed_path("/out/results.csv", "birefnet") == "/out/results_birefnet.csv" + + +def test_each_method_matches_its_single_method_run(tmp_path, monkeypatch): + source = _input(tmp_path) + calls = [] + monkeypatch.setattr(core, "predict_birefnet_mask", _fake_birefnet(calls)) + exports = {"pre_cleanup_methods": ("threshold", "birefnet"), "overlay": True, + "cutout": True, "axes": True} + threshold = tmp_path / "threshold" + birefnet = tmp_path / "birefnet" + both = tmp_path / "both" + _run([source], threshold, "threshold", export_options=exports) + _run([source], birefnet, "birefnet", export_options=exports) + calls.clear() + run = _run([source], both, "both", export_options=exports) + + assert calls == [(100, 100, 3)] + assert run["methods"] == ("threshold", "birefnet") + assert (run["succeeded"], run["failed"]) == (1, 0) + assert "results_path" not in run and "result_rows" not in run + for method, single in (("threshold", threshold), ("birefnet", birefnet)): + assert (both / f"results_{method}.csv").read_bytes() == (single / "results.csv").read_bytes() + assert run["by_method"][method]["results_path"] == str(both / f"results_{method}.csv") + for suffix in ("_mask.png", "_overlay.jpg", "_cutout.jpg", "_measurement_axes.jpg"): + stem, ext = suffix.rsplit(".", 1) + assert (both / f"leaf{stem}_{method}.{ext}").read_bytes() == ( + single / f"leaf{suffix}" + ).read_bytes() + assert (both / f"leaf_mask_precleanup_{method}.png").exists() + assert not (both / "leaf_mask.png").exists() + assert {item["method"] for item in run["artifacts"] if item["kind"] == "results_csv"} == { + "threshold", "birefnet" + } + + +def test_a_method_failure_only_fails_that_methods_outputs(tmp_path, monkeypatch): + source = _input(tmp_path) + monkeypatch.setattr(core, "predict_birefnet_mask", _fake_birefnet([], leaf=False)) + single = tmp_path / "single" + both = tmp_path / "both" + _run([source], single, "birefnet") + run = _run([source], both, "both") + + assert (run["succeeded"], run["failed"]) == (0, 1) + assert run["by_method"]["threshold"]["succeeded"] == 1 + assert run["by_method"]["birefnet"]["failed"] == 1 + assert "LEAF_MASK: leaf not detected" in (both / "results_birefnet.csv").read_text() + assert (both / "leaf_morpho_failures_birefnet.csv").read_bytes() == ( + single / "leaf_morpho_failures.csv" + ).read_bytes() + assert not (both / "leaf_morpho_failures_threshold.csv").exists() + + +def test_raw_measurement_source_applies_to_each_method(tmp_path, monkeypatch): + import json + + source = _input(tmp_path) + monkeypatch.setattr(core, "predict_birefnet_mask", _fake_birefnet([])) + run = _run([source], tmp_path / "raw", "both", measurement_source="pre-cleanup") + assert run["succeeded"] == 1 + for method in ("threshold", "birefnet"): + assert run["by_method"][method]["result_rows"][0]["leaf_area_cm2"] != "NA" + metadata = json.loads((tmp_path / "raw" / f"results_{method}.csv.meta.json").read_text()) + assert metadata["measurement_source"] == "pre-cleanup" + assert metadata["mask_method"] == method + + +def test_failure_before_segmentation_is_recorded_for_every_method(tmp_path): + unreadable = tmp_path / "broken.jpg" + unreadable.write_bytes(b"not an image") + run = _run([unreadable], tmp_path / "out", "both") + + assert (run["succeeded"], run["failed"]) == (0, 1) + for method in ("threshold", "birefnet"): + assert "READ_IMAGE" in (tmp_path / "out" / f"results_{method}.csv").read_text() + assert (tmp_path / "out" / f"leaf_morpho_failures_{method}.csv").exists() + + +def test_parallel_workers_record_every_image_for_every_method(tmp_path, monkeypatch): + sources = [_input(tmp_path, "leaf_a"), _input(tmp_path, "leaf_b")] + monkeypatch.setattr(core, "predict_birefnet_mask", _fake_birefnet([])) + run = _run(sources, tmp_path / "out", "both", workers=2, execution_device="cpu") + + assert (run["succeeded"], run["processed"]) == (2, 2) + for method in ("threshold", "birefnet"): + rows = run["by_method"][method]["result_rows"] + assert sorted(row["sample_id"] for row in rows) == ["leaf_a", "leaf_b"] + + +def test_target_box_mode_keeps_a_single_csv(tmp_path): + source = _input(tmp_path) + run = _run([source], tmp_path / "out", "both", output_mode="target-boxes") + + assert run["methods"] == ("threshold",) + assert run["results_path"] == str(tmp_path / "out" / "results.csv") + assert (tmp_path / "out" / "results.csv").exists() + + +def test_both_methods_refuse_hybrid_execution(tmp_path): + with pytest.raises(ValueError, match="Otsu thresholding only"): + _run([_input(tmp_path)], tmp_path / "out", "both", workers=2, execution_device="hybrid") diff --git a/tests/test_home_app.py b/tests/test_home_app.py index c583e00..fa39799 100644 --- a/tests/test_home_app.py +++ b/tests/test_home_app.py @@ -1,30 +1,47 @@ from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch +import json +import zipfile -import numpy as np import pytest +# Every heavy import is gated: CI installs with `--no-deps`, so a bare +# module-level `import numpy` here is a collection error, not a skip. +np = pytest.importorskip("numpy") pd = pytest.importorskip("pandas") pytest.importorskip("streamlit") cv2 = pytest.importorskip("cv2") from streamlit.testing.v1 import AppTest from mats.app.Home import ( PENDING_WORKSPACE_TAB_KEY, + THRESHOLD_CUSTOM_VALUE_KEY, + THRESHOLD_LEVEL_KEY, WORKSPACE_TAB_KEY, _WORKBENCH_STYLES, _resolve_sheet_layout, + _scale_axes_by_sample, + _threshold_marks_html, build_cutout_image, build_overlay_image, collect_output_pairs, gather_output_files, + files_from_manifest, + pairs_from_manifest, generate_export_overlays, merge_viewer_pairs, normalize_measurements, + select_export_files, summarize_measurements, + write_output_zip, + zip_download_name, ) from mats.dimensions import parse_template_dimensions +from mats.mask_settings import CLEAN_MARGIN_DEFAULT, CLEAN_SIZE_DEFAULT, STRAY_GAP_DEFAULT HOME_PAGE = Path(__file__).resolve().parents[1] / "src" / "mats" / "app" / "Home.py" +DIAGNOSTICS_PAGE = HOME_PAGE.parent / "pages" / "0_Diagnostics.py" CPU_OPTIONS_PAGE = HOME_PAGE.parent / "pages" / "3_CPU_Options.py" @@ -67,6 +84,120 @@ def test_collect_output_pairs_ignores_outputs_from_earlier_sessions(tmp_path): }] +def test_manifest_selects_only_current_run_outputs(tmp_path): + old_target, old_mask = _write_output_pair(tmp_path, "old") + new_target, new_mask = _write_output_pair(tmp_path, "new") + manifest = [ + {"path": str(new_target), "sample_id": "new", "kind": "target_box"}, + {"path": str(new_mask), "sample_id": "new", "kind": "mask"}, + ] + assert files_from_manifest(manifest) == [new_target, new_mask] + assert pairs_from_manifest(manifest) == [{ + "sample_id": "new", "target_box": str(new_target), "mask": str(new_mask) + }] + assert old_target not in files_from_manifest(manifest) + assert old_mask not in files_from_manifest(manifest) + + +def test_manifest_pairs_each_method_with_its_own_mask(): + manifest = [ + {"path": "/out/leaf_target_box.jpg", "sample_id": "leaf", "kind": "target_box", + "method": None}, + {"path": "/out/leaf_mask_threshold.png", "sample_id": "leaf", "kind": "mask", + "method": "threshold"}, + {"path": "/out/leaf_mask_birefnet.png", "sample_id": "leaf", "kind": "mask", + "method": "birefnet"}, + ] + for method in ("threshold", "birefnet"): + assert pairs_from_manifest(manifest, method=method) == [{ + "sample_id": "leaf", + "target_box": "/out/leaf_target_box.jpg", + "mask": f"/out/leaf_mask_{method}.png", + }] + + +def test_manifest_uses_raw_mask_for_raw_measurements_and_keeps_methods_separate(): + exported = [ + {"path": "/out/leaf_target_box.jpg", "sample_id": "leaf", "kind": "target_box"}, + {"path": "/out/leaf_mask_threshold.png", "sample_id": "leaf", "kind": "mask", + "method": "threshold"}, + {"path": "/out/leaf_mask_birefnet.png", "sample_id": "leaf", "kind": "mask", + "method": "birefnet"}, + ] + previews = [ + {"path": "/tmp/leaf_target_box.png", "sample_id": "leaf", + "kind": "preview_target_box"}, + {"path": "/tmp/leaf_raw_threshold.png", "sample_id": "leaf", + "kind": "preview_mask", "method": "threshold"}, + {"path": "/tmp/leaf_raw_birefnet.png", "sample_id": "leaf", + "kind": "preview_mask", "method": "birefnet"}, + ] + assert pairs_from_manifest( + exported, method="threshold", measurement_source="pre-cleanup", + preview_artifacts=previews, + ) == [{ + "sample_id": "leaf", "target_box": "/tmp/leaf_target_box.png", + "mask": "/tmp/leaf_raw_threshold.png", "mask_source": "pre-cleanup", + }] + assert pairs_from_manifest( + exported, method="birefnet", measurement_source="pre-cleanup", + preview_artifacts=previews, + )[0]["mask"] == "/tmp/leaf_raw_birefnet.png" + + +def test_pre_cleanup_runs_show_the_measured_mask_not_the_raw_export(): + # The pre-cleanup export keeps stray pieces; the measured mask drops them. + exported = [ + {"path": "/out/leaf_mask_precleanup_threshold.png", "sample_id": "leaf", + "kind": "pre_cleanup", "method": "threshold"}, + ] + previews = [ + {"path": "/tmp/leaf_measured.png", "sample_id": "leaf", + "kind": "preview_mask", "method": "threshold"}, + ] + assert pairs_from_manifest( + exported, method="threshold", measurement_source="pre-cleanup", + preview_artifacts=previews, + ) == [{ + "sample_id": "leaf", "target_box": None, + "raw_mask": "/out/leaf_mask_precleanup_threshold.png", + "mask": "/tmp/leaf_measured.png", "mask_source": "pre-cleanup", + }] + + +def test_manifest_uses_private_previews_when_image_exports_are_disabled(): + previews = [ + {"path": "/tmp/leaf_box.png", "sample_id": "leaf", "kind": "preview_target_box"}, + {"path": "/tmp/leaf_mask.png", "sample_id": "leaf", "kind": "preview_mask", + "method": "threshold"}, + ] + assert pairs_from_manifest( + [], method="threshold", preview_artifacts=previews, + ) == [{"sample_id": "leaf", "target_box": "/tmp/leaf_box.png", + "mask": "/tmp/leaf_mask.png"}] + + +def test_manifest_attaches_only_the_selected_methods_raw_preview(): + previews = [ + {"path": f"/tmp/leaf_raw_{method}.png", "sample_id": "leaf", + "kind": "preview_raw_mask", "method": method} + for method in ("threshold", "birefnet") + ] + previews += [ + {"path": f"/tmp/leaf_clean_{method}.png", "sample_id": "leaf", + "kind": "preview_mask", "method": method} + for method in ("threshold", "birefnet") + ] + for method in ("threshold", "birefnet"): + assert pairs_from_manifest( + [], method=method, preview_artifacts=previews, + ) == [{ + "sample_id": "leaf", "target_box": None, + "raw_mask": f"/tmp/leaf_raw_{method}.png", + "mask": f"/tmp/leaf_clean_{method}.png", + }] + + def test_merge_viewer_pairs_accumulates_current_session_without_duplicates(): first = { "sample_id": "first", @@ -176,34 +307,55 @@ def test_gather_output_files_includes_overlay_and_cutout_only_when_requested(tmp def test_home_page_renders_analyze_view_without_worker_control(): - app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.run(timeout=30) assert not app.exception - assert [item.label for item in app.number_input] == ["Sheet width", "Sheet height"] + assert not app.number_input subheaders = [item.value for item in app.subheader] - assert "Analysis setup" in subheaders assert "Launch analysis" in subheaders assert "**WORKSPACE NAVIGATION**" in [item.value for item in app.markdown] - assert "**4 · Preflight**" in [item.value for item in app.markdown] - assert app.segmented_control(key="results_unit").value == "cm" + assert "**5 · Preflight**" in [item.value for item in app.markdown] + assert app.session_state["results_unit"] == "cm" -def test_diagnostics_tab_renders_compute_status_without_worker_control(): +def test_diagnostics_moves_out_of_workspace_tabs(): app = AppTest.from_file(str(HOME_PAGE)) - app.session_state[WORKSPACE_TAB_KEY] = "Diagnostics" app.run(timeout=30) assert not app.exception - assert [item.value for item in app.subheader if item.value == "Compute status"] == ["Compute status"] + assert "diagnostics_context" in app.session_state + assert "Diagnostics" not in app.session_state[WORKSPACE_TAB_KEY] + + +def test_sidebar_diagnostics_offers_home_when_opened_without_context(): + app = AppTest.from_file(str(DIAGNOSTICS_PAGE)).run(timeout=30) + + assert not app.exception + assert any(item.value == "Diagnostics" for item in app.title) + + +def test_sidebar_diagnostics_renders_latest_preflight(): + home = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app = AppTest.from_file(str(DIAGNOSTICS_PAGE)) + app.session_state["diagnostics_context"] = home.session_state["diagnostics_context"] + app.session_state["diagnostics_inputs"] = home.session_state["diagnostics_inputs"] + with patch("streamlit.page_link"): + app.run(timeout=30) + + assert not app.exception + assert "Compute status" in [item.value for item in app.subheader] + assert "**Preflight Overview**" in [item.value for item in app.markdown] def test_pending_workspace_tab_is_applied_before_navigation(): app = AppTest.from_file(str(HOME_PAGE)) - app.session_state[PENDING_WORKSPACE_TAB_KEY] = "Results" + app.session_state[PENDING_WORKSPACE_TAB_KEY] = "Analyze" app.run(timeout=30) assert not app.exception - assert app.session_state[WORKSPACE_TAB_KEY] == "Results" + assert app.session_state[WORKSPACE_TAB_KEY] == "Analyze" assert PENDING_WORKSPACE_TAB_KEY not in app.session_state assert any(item.value == "Results" for item in app.subheader) @@ -214,14 +366,39 @@ def test_home_page_offers_input_and_output_folder_pickers(): assert not app.exception assert app.button(key="choose_input_folder").label == "Choose input folder" assert app.button(key="choose_output_folder").label == "Choose output folder" + assert app.button(key="open_export").disabled assert app.text_input(key="input_directory").value == str(Path.home()) assert app.text_input(key="output_directory").value == str( Path.home() / "mats_outputs" ) +def test_sidebar_export_shortcut_opens_export_after_a_run(tmp_path): + results_path = tmp_path / "leaf_morpho_results.csv" + results_path.write_text( + "sample_id,leaf_area_cm2,width_cm,length_cm\nleaf_1,12.5,2.5,7.0\n" + ) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + ) + app.run(timeout=30) + + assert not app.exception + assert app.session_state[WORKSPACE_TAB_KEY] == "Setup" + assert not app.button(key="open_export").disabled + + app.button(key="open_export").click().run(timeout=30) + + assert not app.exception + assert app.session_state[WORKSPACE_TAB_KEY] == "Export" + assert any(item.value == "Export" for item in app.subheader) + + def test_home_page_has_a_dedicated_launch_section(): - app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.run(timeout=30) assert not app.exception assert any(item.value == "Launch analysis" for item in app.subheader) @@ -237,20 +414,142 @@ def test_analyze_keeps_scale_and_segmentation_in_distinct_blocks(): app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) assert not app.exception - assert app.radio(key="segmentation_method").label == "Select a segmentation method" + assert app.checkbox(key="segment_threshold").label == "Classic thresholding" + assert app.checkbox(key="segment_threshold").value + assert app.checkbox(key="segment_birefnet").label == "BiRefNet" + assert not app.checkbox(key="segment_birefnet").value assert {"**1 · Scale**", "**2 · Segmentation**", "**3 · Measurement Output**"} <= { item.value for item in app.markdown } + assert "**4 · Image output options**" in {item.value for item in app.markdown} + assert app.checkbox(key="export_target_boxes").value + assert app.checkbox(key="export_cleaned_masks").value + assert not app.checkbox(key="measure_pre_cleanup").value + assert app.checkbox(key="write_failures").value + assert not app.checkbox(key="export_pre_cleanup").value + assert not app.checkbox(key="export_overlay").value -def test_blocking_preflight_button_opens_diagnostics(): +def test_raw_measurement_checkbox_persists_in_setup(): app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app.checkbox(key="measure_pre_cleanup").set_value(True).run(timeout=30) + assert not app.exception + assert app.checkbox(key="measure_pre_cleanup").value - app.button(key="open_diagnostics").click().run(timeout=30) +def test_cleanup_settings_are_grayed_out_unless_measuring_pre_cleanup(): + app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + margin, gap = app.number_input(key="clean_margin"), app.number_input(key="stray_gap") + size = app.number_input(key="clean_size") + assert (margin.value, gap.value, size.value) == ( + CLEAN_MARGIN_DEFAULT, STRAY_GAP_DEFAULT, CLEAN_SIZE_DEFAULT, + ) + assert margin.disabled and gap.disabled and size.disabled + + app.checkbox(key="measure_pre_cleanup").set_value(True).run(timeout=30) + assert not app.number_input(key="clean_margin").disabled + assert not app.number_input(key="stray_gap").disabled + assert not app.number_input(key="clean_size").disabled + app.number_input(key="clean_margin").set_value(2.0).run(timeout=30) + app.number_input(key="stray_gap").set_value(0.6).run(timeout=30) + app.number_input(key="clean_size").set_value(3).run(timeout=30) + assert not app.exception + assert app.number_input(key="clean_margin").value == 2.0 + assert app.number_input(key="stray_gap").value == 0.6 + assert app.number_input(key="clean_size").value == 3 + assert app.session_state["diagnostics_context"]["clean_size"] == 3 + + # Unchecked, the grayed-out clean size keeps its value but never reaches a run. + app.checkbox(key="measure_pre_cleanup").set_value(False).run(timeout=30) assert not app.exception - assert app.session_state[WORKSPACE_TAB_KEY] == "Diagnostics" - assert any(item.value == "Diagnostics" for item in app.subheader) + assert app.number_input(key="clean_size").disabled + assert app.session_state["diagnostics_context"]["clean_size"] == 0 + + +def _custom_threshold_sliders(app): + return [slider for slider in app.slider if slider.key == THRESHOLD_CUSTOM_VALUE_KEY] + + +def test_threshold_level_offers_custom_after_auto_without_a_slider(): + app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + + assert not app.exception + level = app.selectbox(key=THRESHOLD_LEVEL_KEY) + assert list(level.options) == ["auto", "custom", "low", "medium", "high"] + assert level.value == "auto" + assert _custom_threshold_sliders(app) == [] + + +@pytest.mark.parametrize("with_birefnet", [False, True]) +def test_custom_threshold_level_shows_slider_while_otsu_is_checked(with_birefnet): + app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app.checkbox(key="segment_birefnet").set_value(with_birefnet).run(timeout=30) + app.selectbox(key=THRESHOLD_LEVEL_KEY).set_value("custom").run(timeout=30) + + assert not app.exception + (slider,) = _custom_threshold_sliders(app) + assert (slider.min, slider.max, slider.value) == (1, 255, 125) + widget_keys = [getattr(node, "key", None) for node in app._tree] + assert widget_keys.index("segment_threshold") < widget_keys.index(THRESHOLD_LEVEL_KEY) + assert widget_keys.index(THRESHOLD_LEVEL_KEY) < widget_keys.index(THRESHOLD_CUSTOM_VALUE_KEY) + assert widget_keys.index(THRESHOLD_CUSTOM_VALUE_KEY) < widget_keys.index("segment_birefnet") + slider.set_value(177).run(timeout=30) + assert app.session_state[THRESHOLD_CUSTOM_VALUE_KEY] == 177 + + +def _threshold_level_boxes(app): + return [box for box in app.selectbox if box.key == THRESHOLD_LEVEL_KEY] + + +def test_unchecking_otsu_hides_the_threshold_level_and_keeps_the_choice(): + app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app.selectbox(key=THRESHOLD_LEVEL_KEY).set_value("custom").run(timeout=30) + _custom_threshold_sliders(app)[0].set_value(177).run(timeout=30) + app.checkbox(key="segment_birefnet").check().run(timeout=30) + app.checkbox(key="segment_threshold").uncheck().run(timeout=30) + + assert not app.exception + assert _threshold_level_boxes(app) == [] + assert _custom_threshold_sliders(app) == [] + + app.checkbox(key="segment_threshold").check().run(timeout=30) + assert app.selectbox(key=THRESHOLD_LEVEL_KEY).value == "custom" + assert _custom_threshold_sliders(app)[0].value == 177 + + +def test_no_segmentation_method_blocks_preflight(): + app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) + app.checkbox(key="segment_threshold").uncheck().run(timeout=30) + + assert not app.exception + assert "Choose at least one segmentation method." in [item.value for item in app.warning] + assert _threshold_level_boxes(app) == [] + + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.run(timeout=30) + assert app.button(key="run_leaf_morphometrics").disabled + assert any( + item.value.startswith("Needs attention:") and "Segmentation method" in item.value + for item in app.caption + ) + + +def test_threshold_marks_label_each_preset_at_its_track_position(): + html = _threshold_marks_html() + + for label, value, left in (("low", 100, "38.98%"), ("med", 125, "48.82%"), + ("high", 150, "58.66%")): + assert f'left: {left}">{label}
{value}' in html + + +def test_blocking_preflight_links_to_sidebar_diagnostics(): + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.run(timeout=30) + + assert not app.exception + assert app.session_state[WORKSPACE_TAB_KEY] == "Analyze" + assert "pages/0_Diagnostics.py" in HOME_PAGE.read_text() def test_measurement_normalization_supports_full_and_compact_schemas(): @@ -296,6 +595,31 @@ def test_measurement_normalization_supports_inch_columns(): assert measurements.loc[0, "length"] == 3.0 +def _last_run(results_paths, succeeded, failed, total, **extra): + """Build a session ``last_run`` for ``{method: results_path}``.""" + return { + "succeeded": succeeded, + "failed": failed, + "total": total, + "workers": 1, + "worker_reason": "test", + "execution_device": "cpu", + "output_path": str(Path(next(iter(results_paths.values()))).parent), + "mask_methods": tuple(results_paths), + "by_method": { + method: { + "succeeded": succeeded, + "failed": failed, + "results_path": str(path), + "failure_rows": [], + "failure_overflow": 0, + } + for method, path in results_paths.items() + }, + **extra, + } + + def test_results_tab_renders_measurement_dashboard(tmp_path): results_path = tmp_path / "leaf_morpho_results.csv" results_path.write_text( @@ -305,20 +629,10 @@ def test_results_tab_renders_measurement_dashboard(tmp_path): "leaf_2,8.0,2.0,4.0,100,101,0.99,0\n" ) app = AppTest.from_file(str(HOME_PAGE)) - app.session_state[WORKSPACE_TAB_KEY] = "Results" - app.session_state["last_run"] = { - "succeeded": 2, - "failed": 0, - "total": 2, - "workers": 1, - "worker_reason": "test", - "execution_device": "cpu", - "failure_rows": [], - "failure_overflow": 0, - "results_path": str(results_path), - "output_path": str(tmp_path), - "mask_method": "threshold", - } + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=2, failed=0, total=2, + ) app.run(timeout=30) assert not app.exception @@ -331,14 +645,764 @@ def test_results_tab_renders_measurement_dashboard(tmp_path): "Median axis-scale ratio", "Within 2% of 1.0", "Observed range", - "Leaf Area", - "Leaf Width", - "Leaf Length", ] assert {"**Measurement Table**", "**Selected Specimen**"} <= { item.value for item in app.markdown } + section_titles = [item.value for item in app.markdown] + table_index = section_titles.index("**Measurement Table**") + assert section_titles.index("**Leaf Width and Length**") < table_index + assert section_titles.index("**Leaf Area Distribution**") < table_index + assert section_titles.index("**Scale Quality Control**") < table_index assert "**QR decoder trace**" not in {item.value for item in app.markdown} + assert any("Select a row in the measurement table" in item.value for item in app.info) + + +def test_selected_threshold_sample_mounts_interactive_preview(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text( + "sample_id,area_cm2,width_cm,length_cm\n" + "leaf_1,4.0,2.0,2.0\nleaf_2,8.0,2.0,4.0\n" + ) + target = tmp_path / "leaf_1_preview_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 140 + image[10:30, 10:30] = 100 + assert cv2.imwrite(str(target), image) + mask = tmp_path / "leaf_1_mask.png" + original_mask = np.zeros((40, 40), dtype=np.uint8) + original_mask[10:30, 10:30] = 255 + assert cv2.imwrite(str(mask), original_mask) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=2, failed=0, total=2, + run_id="preview-test", threshold_value=125, measurement_source="cleaned", + results_unit="cm", scale_axes_by_sample={"threshold": {"leaf_1": (10, 10)}}, + artifacts=[{ + "path": str(mask), "sample_id": "leaf_1", + "kind": "mask", "method": "threshold", + }], + ) + app.session_state["threshold_preview_cutoffs"] = {"preview-test:leaf_1": 150} + app.session_state["viewer_pairs"] = { + "threshold": [{ + "sample_id": "leaf_1", "target_box": str(target), "mask": str(mask), + }] + } + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected): + app.run(timeout=30) + + assert not app.exception + assert "**Explore and adjust output**" in {item.value for item in app.markdown} + assert app.button(key="reset_preview-test:leaf_1").label == "Reset to saved threshold" + assert not app.button(key="use_preview-test:leaf_1").disabled + assert not app.button(key="apply_adjustment_preview-test:leaf_1").disabled + with patch("streamlit.dataframe", return_value=selected): + app.button(key="use_preview-test:leaf_1").click().run(timeout=30) + assert not app.exception + assert app.session_state[THRESHOLD_LEVEL_KEY] == "custom" + assert app.session_state[THRESHOLD_CUSTOM_VALUE_KEY] == 150 + + with patch("streamlit.dataframe", return_value=selected): + app.button(key="apply_adjustment_preview-test:leaf_1").click().run(timeout=30) + assert not app.exception + adjusted = pd.read_csv(results_path) + assert adjusted.loc[0, ["area_cm2", "width_cm", "length_cm"]].tolist() == [9.0, 3.0, 3.0] + assert adjusted.loc[1].to_dict() == { + "sample_id": "leaf_2", "area_cm2": 8.0, "width_cm": 2.0, "length_cm": 4.0, + } + saved_mask = cv2.imread(str(mask), cv2.IMREAD_GRAYSCALE) + assert saved_mask[7, 7] == 255 + assert app.session_state["last_run"]["threshold_adjustments"]["leaf_1"] == { + "cutoff": 150, "fill_holes": True, + } + + +def test_bulk_adjustment_updates_only_marked_specimens_and_preserves_each_scale(tmp_path): + from mats.app.output_adjustment import apply_threshold_adjustments + + results = tmp_path / "results.csv" + results.write_text( + "sample_id,area_cm2,width_cm,length_cm\n" + "leaf_1,1,1,1\nleaf_2,2,2,2\nleaf_3,3,3,3\n" + ) + pairs = [] + artifacts = [] + for sample_id in ("leaf_1", "leaf_2", "leaf_3"): + target = tmp_path / f"{sample_id}_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 100 + assert cv2.imwrite(str(target), image) + mask = tmp_path / f"{sample_id}_mask.png" + assert cv2.imwrite(str(mask), np.zeros((40, 40), dtype=np.uint8)) + pairs.append({"sample_id": sample_id, "target_box": str(target), "mask": str(mask)}) + artifacts.append({"path": str(mask), "sample_id": sample_id, + "kind": "mask", "method": "threshold"}) + run = _last_run( + {"threshold": results}, succeeded=3, failed=0, total=3, + results_unit="cm", measurement_source="cleaned", artifacts=artifacts, + scale_axes_by_sample={"threshold": { + "leaf_1": (10, 10), "leaf_2": (20, 20), "leaf_3": (10, 10), + }}, + ) + + assert apply_threshold_adjustments(run, pairs[:2], 125) == ["leaf_1", "leaf_2"] + + saved = pd.read_csv(results).set_index("sample_id") + assert saved.loc["leaf_1", "area_cm2"] == 9.0 + assert saved.loc["leaf_2", "area_cm2"] == 2.25 + assert saved.loc["leaf_3", "area_cm2"] == 3.0 + assert cv2.imread(pairs[0]["mask"], 0)[10, 10] == 255 + assert cv2.imread(pairs[1]["mask"], 0)[10, 10] == 255 + assert not cv2.imread(pairs[2]["mask"], 0).any() + + +def test_adjust_tab_overwrites_marked_specimens_from_their_own_button(tmp_path): + results = tmp_path / "results.csv" + results.write_text( + "sample_id,area_cm2,width_cm,length_cm\n" + "leaf_1,1,1,1\nleaf_2,2,2,2\n" + ) + pairs = [] + for sample_id in ("leaf_1", "leaf_2"): + target = tmp_path / f"{sample_id}_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 100 + assert cv2.imwrite(str(target), image) + mask = tmp_path / f"{sample_id}_mask.png" + assert cv2.imwrite(str(mask), np.zeros((40, 40), dtype=np.uint8)) + pairs.append({"sample_id": sample_id, "target_box": str(target), "mask": str(mask)}) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results}, succeeded=2, failed=0, total=2, + run_id="bulk-ui", results_unit="cm", measurement_source="cleaned", + threshold_value=125, + scale_axes_by_sample={"threshold": {"leaf_1": (10, 10), "leaf_2": (20, 20)}}, + ) + app.session_state["viewer_pairs"] = {"threshold": pairs} + app.run(timeout=30) + + assert not app.exception + assert not app.checkbox # no bulk switch any more + assert app.button(key="apply_adjustment_bulk-ui:leaf_1").label == "Overwrite this specimen" + marked_button = app.button(key="apply_marked_adjustment_bulk-ui:leaf_1") + assert marked_button.label == "Overwrite all marked specimens (0)" + assert marked_button.disabled + + # Overwrite this specimen saves only the one in View, even with others marked. + app.session_state["adjust_marked_bulk-ui_threshold"] = ["leaf_2"] + app.run(timeout=30) + assert app.button(key="apply_marked_adjustment_bulk-ui:leaf_1").label == ( + "Overwrite all marked specimens (1)" + ) + app.button(key="apply_adjustment_bulk-ui:leaf_1").click().run(timeout=30) + saved = pd.read_csv(results).set_index("sample_id") + assert saved.loc["leaf_1", "area_cm2"] == 9.0 + assert saved.loc["leaf_2", "area_cm2"] == 2.0 + + next(button for button in app.button if button.label == "Mark all").click().run(timeout=30) + marked_button = app.button(key="apply_marked_adjustment_bulk-ui:leaf_1") + assert marked_button.label == "Overwrite all marked specimens (2)" + assert not marked_button.disabled + assert app.button(key="apply_adjustment_bulk-ui:leaf_1").label == "Overwrite this specimen" + marked_button.click().run(timeout=30) + + assert not app.exception + saved = pd.read_csv(results).set_index("sample_id") + assert saved.loc["leaf_1", "area_cm2"] == 9.0 + assert saved.loc["leaf_2", "area_cm2"] == 2.25 + + +def test_bulk_adjustment_restores_previous_outputs_if_a_save_fails(tmp_path): + from mats.app import output_adjustment + + results = tmp_path / "results.csv" + original_csv = ( + "sample_id,area_cm2,width_cm,length_cm\n" + "leaf_1,1,1,1\nleaf_2,2,2,2\n" + ) + results.write_text(original_csv) + pairs = [] + artifacts = [] + for sample_id in ("leaf_1", "leaf_2"): + target = tmp_path / f"{sample_id}_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 100 + assert cv2.imwrite(str(target), image) + mask = tmp_path / f"{sample_id}_mask.png" + assert cv2.imwrite(str(mask), np.zeros((40, 40), dtype=np.uint8)) + pairs.append({"sample_id": sample_id, "target_box": str(target), "mask": str(mask)}) + artifacts.append({"path": str(mask), "sample_id": sample_id, + "kind": "mask", "method": "threshold"}) + run = _last_run( + {"threshold": results}, succeeded=2, failed=0, total=2, + results_unit="cm", measurement_source="cleaned", artifacts=artifacts, + scale_axes_by_sample={"threshold": {"leaf_1": (10, 10), "leaf_2": (10, 10)}}, + ) + actual = output_adjustment.apply_threshold_adjustment + + def fail_second(*args, **kwargs): + if args[1]["sample_id"] == "leaf_2": + raise OSError("simulated write failure") + return actual(*args, **kwargs) + + with patch.object(output_adjustment, "apply_threshold_adjustment", side_effect=fail_second): + with pytest.raises(OSError, match="simulated write failure"): + output_adjustment.apply_threshold_adjustments(run, pairs, 125) + + assert results.read_text() == original_csv + assert all(not cv2.imread(pair["mask"], 0).any() for pair in pairs) + assert not Path(f"{results}.meta.json").exists() + assert run["artifacts"] == artifacts[:2] + assert run["threshold_adjustments"] == {} + + +def test_scale_axes_by_sample_retains_successful_canonical_calibration(): + summary = { + "by_method": { + "threshold": {"result_rows": [ + {"sample_id": "leaf_1", "px_per_cm_width": 10, "px_per_cm_height": 12}, + {"sample_id": "failed", "px_per_cm_width": "NA", "px_per_cm_height": "NA"}, + ]}, + }, + } + + assert _scale_axes_by_sample(summary) == {"threshold": {"leaf_1": (10.0, 12.0)}} + + +def test_specimen_explorer_shows_only_selected_row_and_can_remove_flashfill(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text( + "sample_id,leaf_area_cm2,width_cm,length_cm\n" + "leaf_1,12.5,2.5,7.0\nleaf_2,8.0,2.0,4.0\n" + ) + raw = np.zeros((100, 100), dtype=np.uint8) + raw[20:80, 20:80] = 255 + raw[40:60, 40:60] = 0 + pairs = [] + for sample_id in ("leaf_1", "leaf_2"): + target = tmp_path / f"{sample_id}_target_box.png" + raw_path = tmp_path / f"{sample_id}_raw.png" + mask = tmp_path / f"{sample_id}_mask.png" + assert cv2.imwrite(str(target), cv2.cvtColor(255 - raw, cv2.COLOR_GRAY2BGR)) + assert cv2.imwrite(str(raw_path), raw) + assert cv2.imwrite(str(mask), raw) + pairs.append({ + "sample_id": sample_id, "target_box": str(target), + "raw_mask": str(raw_path), "mask": str(mask), + }) + + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=2, failed=0, total=2, + run_id="one-sample", threshold_value=125, measurement_source="cleaned", + ) + app.session_state["viewer_pairs"] = {"threshold": pairs} + app.session_state["selected_specimen_one-sample_threshold"] = "leaf_2" + mounted = [] + preview = patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ) + with preview: + app.run(timeout=30) + + assert not app.exception + assert "Sample: leaf_2" in {item.value for item in app.caption} + assert "Sample: leaf_1" not in {item.value for item in app.caption} + assert "Browse Output Previews" not in {item.label for item in app.expander} + # Remove flashfill sits in the explorer under Clean size, not above it. + assert not [box for box in app.checkbox if "flashfill" in (box.key or "")] + assert mounted[-1]["fill_toggle"] and mounted[-1]["fill_note"] is None + assert not mounted[-1]["remove_fill"] + + # Ticking it in the component records it for this specimen only. + app.session_state["remove_fill_previews"] = {"one-sample:threshold:leaf_2": True} + with preview: + app.run(timeout=30) + assert not app.exception + assert mounted[-1]["remove_fill"] + assert "Sample: leaf_1" not in {item.value for item in app.caption} + + # Clicking leaf_1's row in the table records it as the viewed specimen. + app.session_state["selected_specimen_one-sample_threshold"] = "leaf_1" + with preview: + app.run(timeout=30) + assert "Sample: leaf_1" in {item.value for item in app.caption} + assert "Sample: leaf_2" not in {item.value for item in app.caption} + assert not mounted[-1]["remove_fill"] + + +def test_remove_flashfill_component_value_is_recorded_per_specimen(): + app = AppTest.from_string(''' +import streamlit as st +from mats.app.Home import _record_remove_fill, _remove_fill +st.session_state["preview"] = {"remove_fill": True} +_record_remove_fill("preview", "run:threshold:leaf_1") +st.session_state["cleaned"] = _remove_fill({"measurement_source": "cleaned"}, "run:threshold:leaf_1") +st.session_state["raw"] = _remove_fill({"measurement_source": "pre-cleanup"}, "run:threshold:leaf_1") +''').run(timeout=30) + + assert not app.exception + assert app.session_state["remove_fill_previews"] == {"run:threshold:leaf_1": True} + assert app.session_state["cleaned"] is True + assert app.session_state["raw"] is False # raw-mask runs never fill holes + + +def test_analyze_selection_opens_same_specimen_in_adjust(tmp_path): + results = tmp_path / "results.csv" + results.write_text( + "sample_id,area_cm2,width_cm,length_cm\n" + "leaf_1,1,1,1\nleaf_2,2,2,2\n" + ) + pairs = [] + for sample_id in ("leaf_1", "leaf_2"): + target = tmp_path / f"{sample_id}_target_box.png" + mask = tmp_path / f"{sample_id}_mask.png" + assert cv2.imwrite(str(target), np.full((20, 20, 3), 255, dtype=np.uint8)) + assert cv2.imwrite(str(mask), np.zeros((20, 20), dtype=np.uint8)) + pairs.append({"sample_id": sample_id, "target_box": str(target), "mask": str(mask)}) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + {"threshold": results}, succeeded=2, failed=0, total=2, + run_id="carry-selection", results_unit="cm", measurement_source="cleaned", + threshold_value=125, + ) + app.session_state["viewer_pairs"] = {"threshold": pairs} + selected = SimpleNamespace(selection=SimpleNamespace(rows=[1])) + with patch("streamlit.dataframe", return_value=selected): + app.run(timeout=30) + next(button for button in app.button if button.label == "Adjust selected specimen").click().run(timeout=30) + + assert not app.exception + assert app.session_state[WORKSPACE_TAB_KEY] == "Adjust" + assert app.session_state["selected_specimen_carry-selection_threshold"] == "leaf_2" + assert "**Selected Specimen**" not in {item.value for item in app.markdown} + assert not app.metric + assert not app.selectbox + assert not app.multiselect + + +def _mount_capture(mounted): + return patch( + "mats.app.specimen_table.show_specimen_table", + side_effect=lambda **kwargs: mounted.append(kwargs), + ) + + +def test_adjust_browser_handles_hundreds_and_keeps_hidden_marks(): + app = AppTest.from_string(''' +from mats.app.Home import render_adjust_browser +import streamlit as st +choices = [f"leaf_{index:03d}" for index in range(500)] +st.session_state["browser_result"] = render_adjust_browser(choices, "large", "threshold") +''') + mounted = [] + with _mount_capture(mounted): + app.run(timeout=30) + assert not app.exception + # One scrollable table: no pages and no per-name widgets. + assert len(mounted[-1]["rows"]) == 500 + assert mounted[-1]["markable"] and mounted[-1]["view"] == "leaf_000" + assert not app.checkbox and not app.number_input and not app.selectbox + + # A tick made in the table survives searches that hide it. + app.session_state["adjust_marked_large_threshold"] = ["leaf_025"] + app.text_input(key="adjust_search_large_threshold").set_value("LEAF_49").run(timeout=30) + assert [row["sample_id"] for row in mounted[-1]["rows"]] == [ + f"leaf_{index}" for index in range(490, 500) + ] + assert mounted[-1]["marked"] == ["leaf_025"] + app.button(key="adjust_mark_all_large_threshold").click().run(timeout=30) + assert len(app.session_state["browser_result"][1]) == 11 + + shown = len(mounted) + app.text_input(key="adjust_search_large_threshold").set_value("no-match").run(timeout=30) + assert len(mounted) == shown # no table, just the notice + assert any("No specimen names match" in item.value for item in app.info) + assert len(app.session_state["browser_result"][1]) == 11 + + app.text_input(key="adjust_search_large_threshold").set_value("").run(timeout=30) + assert len(mounted[-1]["marked"]) == 11 + assert "11 marked across all searches." in {item.value for item in app.caption} + app.button(key="adjust_clear_marks_large_threshold").click().run(timeout=30) + assert app.session_state["browser_result"][1] == [] + assert mounted[-1]["marked"] == [] + assert not app.exception + + +def _adjust_table_script(method): + return f''' +import pandas as pd +import streamlit as st +from mats.app.Home import render_adjust_browser +measurements = pd.DataFrame({{ + "sample_id": ["leaf_04", "leaf_40", "leaf_41", "leaf_50"], + "leaf_area": [12.5, 8.0, 9.1, 3.0], + "width": [2.5, 2.0, 2.2, 1.0], + "length": [7.0, 4.0, 4.4, 3.0], + "scale_aspect_ratio": [1.0, None, 0.99, 1.0], +}}) +st.session_state["browser_result"] = render_adjust_browser( + measurements["sample_id"].tolist(), "tbl", "{method}", measurements=measurements, +) +''' + + +def test_adjust_table_shows_measurements_with_view_and_marked_columns(): + app = AppTest.from_string(_adjust_table_script("threshold")) + app.session_state["selected_specimen_tbl_threshold"] = "leaf_41" + app.session_state["adjust_marked_tbl_threshold"] = ["leaf_50"] + mounted = [] + with _mount_capture(mounted): + app.run(timeout=30) + assert not app.exception + table = mounted[-1] + assert [column["label"] for column in table["columns"]] == [ + "Sample", "Leaf area (cm²)", "Leaf width (cm)", "Leaf length (cm)", + "Scale axis ratio", + ] + assert [row["leaf_area"] for row in table["rows"]] == [12.5, 8.0, 9.1, 3.0] + assert table["rows"][1]["scale_aspect_ratio"] is None # blank, not NaN + assert table["view"] == "leaf_41" + assert table["markable"] and table["marked"] == ["leaf_50"] + # The old Mark/Unmark button is gone: the boxes are clickable now. + assert not [b for b in app.button if b.key.startswith("adjust_toggle_mark_")] + + app.text_input(key="adjust_search_tbl_threshold").set_value("LEAF_4").run(timeout=30) + assert [row["sample_id"] for row in mounted[-1]["rows"]] == ["leaf_40", "leaf_41"] + assert "1 marked across all searches." in {item.value for item in app.caption} + assert app.session_state["browser_result"] == ("leaf_41", ["leaf_50"]) + + from mats.app import specimen_table + assert "'Marked for Adjustment'" in specimen_table._TABLE_JS + assert "headerCell('View'" in specimen_table._TABLE_JS + + +def test_specimen_table_clicks_set_the_view_and_marks(): + app = AppTest.from_string(''' +import streamlit as st +from mats.app.Home import _record_table_mark, _record_table_view +st.session_state["marks"] = ["leaf_2"] +st.session_state["table"] = {"view": "leaf_3"} +_record_table_view("table", "selected") +st.session_state["table"] = {"mark": {"id": "leaf_1", "marked": True}} +_record_table_mark("table", "marks") +_record_table_mark("table", "marks") # a repeated tick changes nothing +st.session_state["after_tick"] = list(st.session_state["marks"]) +st.session_state["table"] = {"mark": {"id": "leaf_2", "marked": False}} +_record_table_mark("table", "marks") +''').run(timeout=30) + + assert not app.exception + assert app.session_state["selected"] == "leaf_3" + assert app.session_state["after_tick"] == ["leaf_1", "leaf_2"] + assert app.session_state["marks"] == ["leaf_1"] + + +def test_adjust_table_has_no_marking_for_birefnet(): + app = AppTest.from_string(_adjust_table_script("birefnet")) + mounted = [] + with _mount_capture(mounted): + app.run(timeout=30) + + assert not app.exception + assert not mounted[-1]["markable"] + assert not [ + button for button in app.button + if button.key.startswith(("adjust_mark_all_", "adjust_clear_marks_")) + ] + + +def test_clean_image_preview_is_saved_by_overwrite(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + target = tmp_path / "leaf_1_preview_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 100 + image[10:22, 10:22] = 255 # hole, inscribed radius 6 + image[2, 2] = 0 # 1 px speck beside the leaf + assert cv2.imwrite(str(target), image) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + run_id="clean-test", threshold_value=125, measurement_source="cleaned", + results_unit="cm", scale_axes_by_sample={"threshold": {"leaf_1": (10, 10)}}, + ) + app.session_state["viewer_pairs"] = { + "threshold": [{"sample_id": "leaf_1", "target_box": str(target), "mask": None}] + } + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + mounted = [] + with patch("streamlit.dataframe", return_value=selected), patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ): + app.run(timeout=30) + # No Clean image checkbox: the clean-size slider starts at 0, already live. + assert not any("clean_image" in (box.key or "") for box in app.checkbox) + assert mounted[-1]["clean_radius"] == 0 + assert mounted[-1]["clean_levels_image"] and mounted[-1]["cleaned_image"] + assert "0 turns Clean image off" in mounted[-1]["clean_help"] + assert "Overwrite saves" in mounted[-1]["clean_help"] + assert not app.button(key="apply_adjustment_clean-test:leaf_1").disabled + + # Releasing the slider above 0 records it, and Overwrite saves that mask. + app.session_state["clean_preview_radii"] = {"clean-test:threshold:leaf_1": 5} + app.run(timeout=30) + assert mounted[-1]["clean_radius"] == 5 + assert any("Overwrite saves it" in item.value for item in app.caption) + assert not app.button(key="apply_adjustment_clean-test:leaf_1").disabled + app.button(key="apply_adjustment_clean-test:leaf_1").click().run(timeout=30) + + assert not app.exception + assert any("clean size 5 px" in item.value for item in app.success) + saved = pd.read_csv(results_path).set_index("sample_id") + assert saved.loc["leaf_1", "area_cm2"] == pytest.approx((900 - 144) / 100) + mask = cv2.imread(str(tmp_path / "leaf_1_mask.png"), cv2.IMREAD_GRAYSCALE) + assert mask[15, 15] == 0 and mask[2, 2] == 0 and mask[30, 30] == 255 + assert app.session_state["last_run"]["threshold_adjustments"]["leaf_1"] == { + "cutoff": 125, "fill_holes": False, "clean_margin": CLEAN_MARGIN_DEFAULT, + "stray_gap": STRAY_GAP_DEFAULT, "clean_size": 5, + } + + +@pytest.mark.parametrize("saved, expected", [(None, 3), ({"cutoff": 125, "clean_size": 0}, 0)]) +def test_clean_size_slider_starts_at_the_saved_or_run_clean_size(tmp_path, saved, expected): + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + target = tmp_path / "leaf_1_preview_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 100 + assert cv2.imwrite(str(target), image) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + run_id="run-size", threshold_value=125, measurement_source="pre-cleanup", + clean_size=3, threshold_adjustments={"leaf_1": saved} if saved else {}, + ) + app.session_state["viewer_pairs"] = {"threshold": [{ + "sample_id": "leaf_1", "target_box": str(target), "mask": None, + "mask_source": "pre-cleanup", + }]} + mounted = [] + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected), patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ): + app.run(timeout=30) + + assert not app.exception + assert mounted[-1]["clean_radius"] == expected + + +def test_pre_cleanup_explorer_settles_on_the_measured_mask(tmp_path): + import base64 + + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + target = tmp_path / "leaf_1_preview_target_box.png" + image = np.full((100, 100, 3), 255, dtype=np.uint8) + image[0:2, 10:90] = 0 # printed outline on the edge + image[30:70, 30:70] = 100 + assert cv2.imwrite(str(target), image) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + run_id="raw-test", threshold_value=125, measurement_source="pre-cleanup", + clean_margin=2.0, stray_gap=0.25, + ) + app.session_state["viewer_pairs"] = {"threshold": [{ + "sample_id": "leaf_1", "target_box": str(target), "mask": None, + "mask_source": "pre-cleanup", + }]} + mounted = [] + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected), patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ): + app.run(timeout=30) + + assert not app.exception + preview = mounted[-1] + assert preview["fill_toggle"] and not preview["remove_fill"] + assert preview["fill_note"] == "This run measured raw masks, so hole filling is already off." + assert preview["live_margin"] == 2.0 + assert preview["cleaned_cutoff"] == 125 + settled = cv2.imdecode(np.frombuffer( + base64.b64decode(preview["cleaned_image"].split(",", 1)[1]), dtype=np.uint8, + ), cv2.IMREAD_GRAYSCALE) + assert not settled[:2].any() and settled[50, 50] == 255 + + # The specimen's own margin re-cleans the preview; 0 leaves the edge line in. + margin_key = "specimen_clean_margin_raw-test:threshold:leaf_1" + assert not app.number_input(key=margin_key).disabled + with patch("streamlit.dataframe", return_value=selected), patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ): + app.number_input(key=margin_key).set_value(0.0).run(timeout=30) + assert not app.exception + assert mounted[-1]["live_margin"] == 0 + unguarded = cv2.imdecode(np.frombuffer( + base64.b64decode(mounted[-1]["cleaned_image"].split(",", 1)[1]), dtype=np.uint8, + ), cv2.IMREAD_GRAYSCALE) + assert unguarded[0, 50] == 0 and unguarded[50, 50] == 255 # a stray edge strip + + +def test_specimen_cleanup_inputs_follow_flash_fill(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + target = tmp_path / "leaf_1_preview_target_box.png" + image = np.full((40, 40, 3), 255, dtype=np.uint8) + image[5:35, 5:35] = 100 + assert cv2.imwrite(str(target), image) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + run_id="flash", threshold_value=125, measurement_source="cleaned", + ) + app.session_state["viewer_pairs"] = { + "threshold": [{"sample_id": "leaf_1", "target_box": str(target), "mask": None}] + } + margin_key = "specimen_clean_margin_flash:threshold:leaf_1" + gap_key = "specimen_stray_gap_flash:threshold:leaf_1" + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected): + app.run(timeout=30) + assert app.number_input(key=margin_key).disabled # flash fill is on + assert app.number_input(key=gap_key).disabled + + app.session_state["remove_fill_previews"] = {"flash:threshold:leaf_1": True} + app.run(timeout=30) + assert not app.number_input(key=margin_key).disabled # margin only + assert app.number_input(key=gap_key).disabled + # The Remove flashfill preview above reruns with the new margin. + app.number_input(key=margin_key).set_value(5.0).run(timeout=30) + assert not app.exception + assert app.number_input(key=margin_key).value == 5.0 + + app.session_state["clean_preview_radii"] = {"flash:threshold:leaf_1": 3} + app.run(timeout=30) + assert not app.number_input(key=margin_key).disabled + assert not app.number_input(key=gap_key).disabled + assert not app.exception + + +def test_birefnet_pre_cleanup_explorer_recleans_the_raw_mask(tmp_path): + import base64 + + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + raw = np.zeros((100, 100), dtype=np.uint8) + raw[0:2, 10:90] = 255 # printed outline on the edge + raw[30:70, 30:70] = 255 + raw_path = tmp_path / "leaf_1_preview_raw_mask.png" + assert cv2.imwrite(str(raw_path), raw) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"birefnet": results_path}, succeeded=1, failed=0, total=1, + run_id="bir-raw", measurement_source="pre-cleanup", clean_margin=2.0, + ) + app.session_state["viewer_pairs"] = {"birefnet": [{ + "sample_id": "leaf_1", "target_box": None, "mask": None, + "raw_mask": str(raw_path), "mask_source": "pre-cleanup", + }]} + mounted = [] + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected), patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ): + app.run(timeout=30) + + assert not app.exception + shown = cv2.imdecode(np.frombuffer( + base64.b64decode(mounted[-1]["mask_image"].split(",", 1)[1]), dtype=np.uint8, + ), cv2.IMREAD_GRAYSCALE) + assert not shown[:2].any() and shown[50, 50] == 255 + assert "edge margin cleared" in mounted[-1]["mask_status"] + assert not app.number_input(key="specimen_stray_gap_bir-raw:birefnet:leaf_1").disabled + + +def test_birefnet_specimen_gets_a_preview_only_clean_image_explorer(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + raw = np.zeros((40, 40), dtype=np.uint8) + raw[5:35, 5:35] = 255 + raw[18:21, 18:21] = 0 + raw[1, 1] = 255 + raw_path = tmp_path / "leaf_1_preview_raw_mask.png" + mask_path = tmp_path / "leaf_1_mask.png" + target = tmp_path / "leaf_1_preview_target_box.png" + assert cv2.imwrite(str(raw_path), raw) + assert cv2.imwrite(str(mask_path), raw) + assert cv2.imwrite(str(target), cv2.cvtColor(255 - raw, cv2.COLOR_GRAY2BGR)) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"birefnet": results_path}, succeeded=1, failed=0, total=1, + run_id="bir-test", measurement_source="cleaned", + ) + app.session_state["viewer_pairs"] = {"birefnet": [{ + "sample_id": "leaf_1", "target_box": str(target), + "mask": str(mask_path), "raw_mask": str(raw_path), + }]} + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected): + app.run(timeout=30) + assert not app.exception + assert "**Explore and adjust output**" in {item.value for item in app.markdown} + assert any("Drag the clean size above 0" in item.value for item in app.caption) + app.session_state["clean_preview_radii"] = {"bir-test:birefnet:leaf_1": 4} + app.run(timeout=30) + + assert not app.exception + assert any("Clean size is above 0" in item.value for item in app.caption) + assert not any("apply_adjustment" in (button.key or "") for button in app.button) + + +def test_birefnet_clean_image_help_says_it_is_preview_only(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2,width_cm,length_cm\nleaf_1,4.0,2.0,2.0\n") + raw = np.zeros((40, 40), dtype=np.uint8) + raw[5:35, 5:35] = 255 + raw_path = tmp_path / "leaf_1_preview_raw_mask.png" + assert cv2.imwrite(str(raw_path), raw) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Adjust" + app.session_state["last_run"] = _last_run( + {"birefnet": results_path}, succeeded=1, failed=0, total=1, + run_id="bir-help", measurement_source="pre-cleanup", clean_size=4, + ) + app.session_state["viewer_pairs"] = {"birefnet": [{ + "sample_id": "leaf_1", "target_box": None, "mask": None, + "raw_mask": str(raw_path), "mask_source": "pre-cleanup", + }]} + mounted = [] + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected), patch( + "mats.app.threshold_preview.show_threshold_preview", + side_effect=lambda **kwargs: mounted.append(kwargs), + ): + app.run(timeout=30) + + assert not app.exception + assert mounted[-1]["clean_radius"] == 4 # the run's clean size + assert "preview it only" in mounted[-1]["clean_help"] + assert "Overwrite" not in mounted[-1]["clean_help"] def test_results_tab_uses_the_completed_run_unit(tmp_path): @@ -348,59 +1412,197 @@ def test_results_tab_uses_the_completed_run_unit(tmp_path): "leaf_1,2.5,1.5,3.0\n" ) app = AppTest.from_file(str(HOME_PAGE)) - app.session_state[WORKSPACE_TAB_KEY] = "Results" - app.session_state["last_run"] = { - "succeeded": 1, - "failed": 0, - "total": 1, - "workers": 1, - "worker_reason": "test", - "execution_device": "cpu", - "failure_rows": [], - "failure_overflow": 0, - "results_path": str(results_path), - "output_path": str(tmp_path), - "mask_method": "threshold", - "results_unit": "in", - } + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, results_unit="in", + ) app.run(timeout=30) assert not app.exception assert app.metric[1].value == "2.50 in²" assert app.metric[2].value == "1.50 in" - assert app.download_button[0].label == "Download results CSV (in)" + assert not app.download_button -def test_results_tab_has_a_clearly_defined_export_section(tmp_path): +def test_export_tab_lists_saved_results_and_downloads(tmp_path): results_path = tmp_path / "leaf_morpho_results.csv" results_path.write_text( "sample_id,leaf_area_cm2,width_cm,length_cm\nleaf_1,12.5,2.5,7.0\n" ) app = AppTest.from_file(str(HOME_PAGE)) - app.session_state[WORKSPACE_TAB_KEY] = "Results" - app.session_state["last_run"] = { - "succeeded": 1, - "failed": 0, - "total": 1, - "workers": 1, - "worker_reason": "test", - "execution_device": "cpu", - "failure_rows": [], - "failure_overflow": 0, - "results_path": str(results_path), - "output_path": str(tmp_path), - "mask_method": "threshold", - } + app.session_state[WORKSPACE_TAB_KEY] = "Export" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + artifacts=[{"path": str(results_path), "kind": "results_csv", "method": "threshold"}], + ) app.run(timeout=30) assert not app.exception assert any(item.value == "Export" for item in app.subheader) markdown_values = {item.value for item in app.markdown} - assert "**Measurements only**" in markdown_values - assert "**Full export (ZIP)**" in markdown_values - checkbox_labels = {item.label for item in app.checkbox} - assert "Include overlay images (mask highlighted on photo)" in checkbox_labels - assert "Include specimen cutouts (background removed)" in checkbox_labels + assert "**Files to include in ZIP**" in markdown_values + assert "**Measurement CSVs**" in markdown_values + assert app.download_button[0].label == f"Download {results_path.name}" + assert app.checkbox(key="export_include_results_csv").value + assert app.checkbox(key="export_include_overlay").disabled + + +def test_export_tab_explains_what_to_do_before_a_run(): + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Export" + app.run(timeout=30) + + assert not app.exception + assert any("Run an analysis" in item.value for item in app.info) + assert "prepare_zip_export" not in {button.key for button in app.button} + + +def test_export_tab_prepares_training_dataset_from_current_pairs(tmp_path): + results_path = tmp_path / "results.csv" + results_path.write_text("sample_id,area_cm2\nleaf,1\n") + image = np.full((16, 20, 3), 200, dtype=np.uint8) + mask = np.zeros((16, 20), dtype=np.uint8) + mask[3:13, 4:15] = 255 + image_path = tmp_path / "leaf_target_box.png" + mask_path = tmp_path / "leaf_mask.png" + assert cv2.imwrite(str(image_path), image) + assert cv2.imwrite(str(mask_path), mask) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Export" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + artifacts=[{"path": str(results_path), "kind": "results_csv", "method": "threshold"}], + ) + app.session_state["viewer_pairs"] = {"threshold": [{ + "sample_id": "leaf", "target_box": str(image_path), "mask": str(mask_path), + }]} + app.run(timeout=30) + assert not app.exception + assert not app.button(key="prepare_training_dataset").disabled + app.button(key="prepare_training_dataset").click().run(timeout=30) + assert not app.exception + prepared = Path(app.session_state["dataset_zip_path"]) + with zipfile.ZipFile(prepared) as archive: + manifest = json.loads(archive.read("manifest.json")) + assert manifest["counts"] == {"train": 1, "val": 0, "test": 0} + assert archive.read(manifest["samples"][0]["label"]) + assert any(button.key == "download_training_dataset" for button in app.download_button) + + +def test_export_selection_filters_methods_and_missing_files(tmp_path): + paths = { + name: tmp_path / name for name in ( + "leaf_target_box.jpg", "results_threshold.csv", "results_birefnet.csv", + "leaf_mask_threshold.png", "leaf_mask_birefnet.png", + ) + } + for path in paths.values(): + path.write_bytes(path.name.encode()) + artifacts = [ + {"path": str(paths["leaf_target_box.jpg"]), "kind": "target_box", "method": None}, + {"path": str(paths["results_threshold.csv"]), "kind": "results_csv", "method": "threshold"}, + {"path": str(paths["results_birefnet.csv"]), "kind": "results_csv", "method": "birefnet"}, + {"path": str(paths["leaf_mask_threshold.png"]), "kind": "mask", "method": "threshold"}, + {"path": str(paths["leaf_mask_birefnet.png"]), "kind": "mask", "method": "birefnet"}, + {"path": str(tmp_path / "missing_overlay.jpg"), "kind": "overlay", "method": "threshold"}, + ] + artifacts.append(artifacts[0]) # Shared target box appears only once in the ZIP. + + selected = select_export_files( + artifacts, ("birefnet",), ("target_box", "results_csv", "mask", "overlay") + ) + assert [path.name for path in selected] == [ + "leaf_target_box.jpg", "results_birefnet.csv", "leaf_mask_birefnet.png", + ] + assert select_export_files(artifacts, (), ("target_box",)) == [] + assert zip_download_name("folder\\leaf_results") == "leaf_results.zip" + + dest = tmp_path / "chosen.zip" + write_output_zip(selected, dest) + with zipfile.ZipFile(dest) as archive: + assert archive.namelist() == [path.name for path in selected] + + +def test_export_prepared_zip_is_cleared_when_selection_changes(tmp_path): + results_path = tmp_path / "results.csv" + mask_path = tmp_path / "leaf_mask.png" + results_path.write_text("sample_id,area_cm2\nleaf,1\n") + mask_path.write_bytes(b"mask") + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Export" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=0, total=1, + run_id="export-test", + artifacts=[ + {"path": str(results_path), "kind": "results_csv", "method": "threshold"}, + {"path": str(mask_path), "kind": "mask", "method": "threshold"}, + ], + ) + app.run(timeout=30) + assert not app.exception + app.button(key="prepare_zip_export").click().run(timeout=30) + assert not app.exception + prepared = Path(app.session_state["export_zip_path"]) + assert prepared.is_file() + with zipfile.ZipFile(prepared) as archive: + assert set(archive.namelist()) == {"results.csv", "leaf_mask.png"} + + app.text_input(key="export_zip_name").set_value("selected_leaves").run(timeout=30) + assert not app.exception + assert "export_zip_path" not in app.session_state + assert not prepared.exists() + + app.button(key="prepare_zip_export").click().run(timeout=30) + prepared = Path(app.session_state["export_zip_path"]) + assert app.download_button[-1].label == "Download ZIP (2 files)" + app.checkbox(key="export_include_mask").set_value(False).run(timeout=30) + assert not app.exception + assert "export_zip_path" not in app.session_state + assert not prepared.exists() + + app.session_state["last_run"] = { + **app.session_state["last_run"], "run_id": "next-run", + } + app.run(timeout=30) + assert not app.exception + assert app.checkbox(key="export_include_mask").value + assert app.text_input(key="export_zip_name").value == "leaf_morpho_outputs.zip" + + +def test_export_method_choice_filters_direct_csv_and_zip(tmp_path): + threshold_csv = tmp_path / "results_threshold.csv" + birefnet_csv = tmp_path / "results_birefnet.csv" + target_box = tmp_path / "leaf_target_box.jpg" + for path in (threshold_csv, birefnet_csv, target_box): + path.write_bytes(path.name.encode()) + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Export" + app.session_state["last_run"] = _last_run( + {"threshold": threshold_csv, "birefnet": birefnet_csv}, + succeeded=1, failed=0, total=1, run_id="dual-method-export", + artifacts=[ + {"path": str(target_box), "kind": "target_box", "method": None}, + {"path": str(threshold_csv), "kind": "results_csv", "method": "threshold"}, + {"path": str(birefnet_csv), "kind": "results_csv", "method": "birefnet"}, + ], + ) + app.run(timeout=30) + assert not app.exception + assert len(app.download_button) == 2 + + app.multiselect(key="export_selected_methods").set_value(["birefnet"]).run(timeout=30) + assert not app.exception + assert [button.label for button in app.download_button] == [ + f"Download {birefnet_csv.name}" + ] + app.button(key="prepare_zip_export").click().run(timeout=30) + assert not app.exception + prepared = Path(app.session_state["export_zip_path"]) + try: + with zipfile.ZipFile(prepared) as archive: + assert set(archive.namelist()) == {target_box.name, birefnet_csv.name} + finally: + prepared.unlink(missing_ok=True) def test_results_tab_shows_qr_trace_when_full_qr_columns_are_present(tmp_path): @@ -412,26 +1614,138 @@ def test_results_tab_shows_qr_trace_when_full_qr_columns_are_present(tmp_path): "leaf_2,NA,NA,NA,NA,NA,NA,QR_READ: QR not found/readable,failed,failed\n" ) app = AppTest.from_file(str(HOME_PAGE)) - app.session_state[WORKSPACE_TAB_KEY] = "Results" - app.session_state["last_run"] = { - "succeeded": 1, - "failed": 1, - "total": 2, - "workers": 1, - "worker_reason": "test", - "execution_device": "cpu", - "failure_rows": [], - "failure_overflow": 0, - "results_path": str(results_path), - "output_path": str(tmp_path), - "mask_method": "threshold", - } + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + {"threshold": results_path}, succeeded=1, failed=1, total=2, + ) app.run(timeout=30) assert not app.exception assert "**QR decoder trace**" in {item.value for item in app.markdown} +def test_results_tab_switches_between_method_results(tmp_path): + threshold_path = tmp_path / "leaf_morpho_results_threshold.csv" + birefnet_path = tmp_path / "leaf_morpho_results_birefnet.csv" + threshold_path.write_text("sample_id,leaf_area_cm2,width_cm,length_cm\nleaf_1,12.5,2.5,7.0\n") + birefnet_path.write_text("sample_id,leaf_area_cm2,width_cm,length_cm\nleaf_1,11.0,2.4,6.9\n") + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + {"threshold": threshold_path, "birefnet": birefnet_path}, + succeeded=1, failed=0, total=1, + ) + app.run(timeout=30) + + assert not app.exception + picker = app.segmented_control(key="results_method") + assert list(picker.options) == ["Classic thresholding", "BiRefNet"] + assert picker.value == "threshold" + assert app.metric[1].value == "12.50 cm²" + + picker.set_value("birefnet").run(timeout=30) + assert not app.exception + assert app.metric[1].value == "11.00 cm²" + + +def test_adjust_opens_with_the_method_selected_in_analyze(tmp_path): + paths = {} + pairs = {} + for method in ("threshold", "birefnet"): + results = tmp_path / f"results_{method}.csv" + results.write_text( + "sample_id,area_cm2,width_cm,length_cm\nleaf,4,2,2\n" + ) + paths[method] = results + target = tmp_path / f"target_{method}.png" + mask = tmp_path / f"mask_{method}.png" + assert cv2.imwrite(str(target), np.full((20, 20, 3), 255, dtype=np.uint8)) + assert cv2.imwrite(str(mask), np.zeros((20, 20), dtype=np.uint8)) + pairs[method] = [{"sample_id": "leaf", "target_box": str(target), "mask": str(mask)}] + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + paths, succeeded=1, failed=0, total=1, run_id="method-carry", + results_unit="cm", measurement_source="cleaned", threshold_value=125, + ) + app.session_state["viewer_pairs"] = pairs + selected = SimpleNamespace(selection=SimpleNamespace(rows=[0])) + with patch("streamlit.dataframe", return_value=selected): + app.run(timeout=30) + app.segmented_control(key="results_method").set_value("birefnet").run(timeout=30) + next(button for button in app.button if button.label == "Adjust selected specimen").click().run(timeout=30) + + assert not app.exception + assert app.segmented_control(key="results_method").value == "birefnet" + + +def test_method_switch_updates_table_and_selected_specimen_together(tmp_path): + paths = { + method: tmp_path / f"results_{method}.csv" + for method in ("threshold", "birefnet") + } + pairs = {} + for method, sample_id, area in ( + ("threshold", "classic_leaf", 12.5), + ("birefnet", "birefnet_leaf", 11.0), + ): + paths[method].write_text( + f"sample_id,leaf_area_cm2,width_cm,length_cm\n" + f"{sample_id},{area},2.5,7.0\n" + ) + mask = tmp_path / f"{sample_id}_mask.png" + assert cv2.imwrite(str(mask), np.full((12, 12), 255, dtype=np.uint8)) + pairs[method] = [{"sample_id": sample_id, "target_box": None, "mask": str(mask)}] + + app = AppTest.from_file(str(HOME_PAGE)) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.session_state["last_run"] = _last_run( + paths, succeeded=1, failed=0, total=1, run_id="method-switch-test", + ) + app.session_state["viewer_pairs"] = pairs + selected_rows = {"threshold": [0], "birefnet": []} + displayed_tables = [] + selection_defaults = [] + + def selected_table(data, **kwargs): + key = kwargs["key"] + method = key.rsplit("_", 1)[-1] + displayed_tables.append((key, data["sample_id"].tolist())) + selection_defaults.append((key, kwargs["selection_default"])) + return SimpleNamespace(selection=SimpleNamespace(rows=selected_rows[method])) + + with patch("streamlit.dataframe", side_effect=selected_table): + app.run(timeout=30) + assert "Sample: classic_leaf" in {item.value for item in app.caption} + assert displayed_tables[-1] == ( + "measurement_table_method-switch-test_threshold", ["classic_leaf"] + ) + + app.segmented_control(key="results_method").set_value("birefnet").run(timeout=30) + assert not app.exception + assert displayed_tables[-1] == ( + "measurement_table_method-switch-test_birefnet", ["birefnet_leaf"] + ) + assert "Sample: classic_leaf" not in {item.value for item in app.caption} + assert any("Select a row in the measurement table" in item.value for item in app.info) + + selected_rows["birefnet"] = [0] + app.run(timeout=30) + assert "Sample: birefnet_leaf" in {item.value for item in app.caption} + assert "Sample: classic_leaf" not in {item.value for item in app.caption} + + app.segmented_control(key="results_method").set_value("threshold").run(timeout=30) + assert displayed_tables[-1] == ( + "measurement_table_method-switch-test_threshold", ["classic_leaf"] + ) + assert selection_defaults[-1] == ( + "measurement_table_method-switch-test_threshold", + {"selection": {"rows": [0]}}, + ) + assert "Sample: classic_leaf" in {item.value for item in app.caption} + assert "Sample: birefnet_leaf" not in {item.value for item in app.caption} + + def test_cpu_options_page_renders_worker_control(): app = AppTest.from_file(str(CPU_OPTIONS_PAGE)).run(timeout=30) @@ -439,12 +1753,16 @@ def test_cpu_options_page_renders_worker_control(): assert [item.label for item in app.number_input] == ["CPU workers"] +def _sheet_size_inputs(app): + return [item for item in app.number_input if item.key in {"measure_width", "measure_height"}] + + def test_home_page_defaults_to_manual_dimensions(): app = AppTest.from_file(str(HOME_PAGE)).run(timeout=30) assert not app.exception assert app.checkbox(key="measure_use_qr").value is False - width_input, height_input = app.number_input + width_input, height_input = _sheet_size_inputs(app) assert width_input.value == 12.0 assert height_input.value == 12.0 assert not width_input.disabled @@ -478,7 +1796,7 @@ def test_legacy_calibration_is_kept_in_an_explicit_compatibility_control(): app.toggle(key="measure_use_legacy_calibration").set_value(True).run(timeout=30) assert not app.exception - assert not app.number_input + assert not _sheet_size_inputs(app) legacy_input = app.text_input(key="measure_legacy_dimensions_text") assert legacy_input.label == "Legacy calibration area" assert legacy_input.value == "10x9.5in" @@ -490,7 +1808,7 @@ def test_qr_mode_disables_manual_dimension_inputs(): app.checkbox(key="measure_use_qr").check().run(timeout=30) assert not app.exception - assert all(item.disabled for item in app.number_input) + assert all(item.disabled for item in _sheet_size_inputs(app)) dimensions_bar = [i for i in app.text_input if i.label == "Printed sheet size"][0] assert dimensions_bar.disabled assert dimensions_bar.value == "Variable dimensions — QR-derived" @@ -534,7 +1852,7 @@ def test_unit_conversion_snaps_selectors_to_nearest_half(): app.segmented_control(key="measure_unit").set_value("cm").run(timeout=30) - width_input, height_input = app.number_input + width_input, height_input = _sheet_size_inputs(app) assert width_input.value == pytest.approx(30.5) assert height_input.value == pytest.approx(30.5) dimensions_bar = [i for i in app.text_input if i.label == "Printed sheet size"][0] @@ -548,13 +1866,15 @@ def test_off_grid_sheet_size_is_rejected_with_creator_grid_guidance(): dimensions_bar.set_value("8.27x11.69in").run(timeout=30) assert not app.exception - width_input, height_input = app.number_input + width_input, height_input = _sheet_size_inputs(app) assert width_input.value == 12.0 assert height_input.value == 12.0 dimensions_bar = [i for i in app.text_input if i.label == "Printed sheet size"][0] assert dimensions_bar.value == "8.27x11.69in" - assert app.button(key="run_leaf_morphometrics").disabled assert any("0.5-unit increments" in item.value for item in app.caption) + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.run(timeout=30) + assert app.button(key="run_leaf_morphometrics").disabled def test_invalid_custom_dimensions_block_the_run(): @@ -564,5 +1884,7 @@ def test_invalid_custom_dimensions_block_the_run(): dimensions_bar.set_value("not-a-size").run(timeout=30) assert not app.exception + app.session_state[WORKSPACE_TAB_KEY] = "Analyze" + app.run(timeout=30) assert app.button(key="run_leaf_morphometrics").disabled assert any("Printed sheet size" in item.value for item in app.caption) diff --git a/tests/test_mask_cleanup.py b/tests/test_mask_cleanup.py new file mode 100644 index 0000000..13f3639 --- /dev/null +++ b/tests/test_mask_cleanup.py @@ -0,0 +1,252 @@ +"""Clean image: drop stray pieces, small specks and small holes, never the leaf itself.""" + +import pytest + +np = pytest.importorskip("numpy") +cv2 = pytest.importorskip("cv2") + +from mats.mask_cleanup import ( + CLEAN_RADIUS_MAX, apply_clean_levels, clean_levels, clean_raw_mask, + clean_specks_and_holes, clear_margin, drop_stray_pieces, margin_width, + raw_measurement_mask, +) +from mats.mask_settings import ( + CLEAN_MARGIN_MAX, CLEAN_SIZE_MAX, STRAY_GAP_MAX, checked_clean_size, +) + + +def _leaf(): + # Bounding box 141 x 141, so the default 0.25 gap limit is ~49.9 px. + mask = np.zeros((200, 200), dtype=np.uint8) + cv2.circle(mask, (100, 100), 70, 255, -1) + return mask + + +def test_small_radii_leave_the_mask_unchanged(): + mask = _leaf() + cv2.circle(mask, (80, 80), 3, 0, -1) + mask[25, 100] = 255 # 1 px speck beside the leaf + levels = clean_levels(mask) + assert np.array_equal(apply_clean_levels(levels, 0), mask) + assert np.array_equal(apply_clean_levels(levels, 1), mask) + + +def test_pieces_touching_the_border_are_stray(): + mask = _leaf() + mask[0:3, 60:140] = 255 # strip along the top edge + mask[190:200, 195:200] = 255 # corner fragment + kept = drop_stray_pieces(mask) + assert not kept[0:3].any() and not kept[190:, 195:].any() + assert np.array_equal(kept, _leaf()) + + +def test_pieces_far_from_the_leaf_are_stray_and_near_ones_stay(): + mask = _leaf() + cv2.circle(mask, (100, 20), 3, 255, -1) # gap ~7 px: near + cv2.circle(mask, (15, 15), 3, 255, -1) # gap ~47 px: under the limit + cv2.circle(mask, (185, 185), 5, 255, -1) # gap ~45 px: under the limit + far = np.zeros((400, 400), dtype=np.uint8) + far[:200, :200] = mask + cv2.circle(far, (300, 300), 5, 255, -1) # gap ~213 px: far + kept = drop_stray_pieces(far) + assert kept[20, 100] == 255 and kept[15, 15] == 255 and kept[185, 185] == 255 + assert kept[300, 300] == 0 + assert drop_stray_pieces(far, 0.1)[15, 15] == 0 # limit ~20 px + + +def test_the_leaf_stays_even_when_it_crosses_the_border(): + mask = np.zeros((100, 100), dtype=np.uint8) + mask[0:60, 20:80] = 255 # leaf runs off the top edge + mask[95:100, 0:10] = 255 # smaller border piece + kept = drop_stray_pieces(mask) + assert kept[0:60, 20:80].min() == 255 + assert not kept[95:, :10].any() + + +def test_zero_gap_keeps_only_the_leaf(): + mask = _leaf() + mask[25, 100] = 255 + assert np.array_equal(drop_stray_pieces(mask, 0), _leaf()) + + +def test_drop_stray_pieces_keeps_holes_and_is_idempotent(): + mask = _leaf() + cv2.circle(mask, (120, 110), 15, 0, -1) + mask[25, 100] = mask[2, 2] = 255 + kept = drop_stray_pieces(mask) + assert kept[110, 120] == 0 # hole is not filled + assert kept[25, 100] == 255 and kept[2, 2] == 0 + assert np.array_equal(drop_stray_pieces(kept), kept) + + +@pytest.mark.parametrize("gap", [-0.1, STRAY_GAP_MAX + 1, float("nan"), True, "far"]) +def test_invalid_gaps_raise(gap): + with pytest.raises(ValueError): + drop_stray_pieces(_leaf(), gap) + + +def test_clean_levels_drop_stray_pieces_at_every_radius(): + mask = _leaf() + mask[0:3, 60:140] = 255 + far = np.zeros((300, 300), dtype=np.uint8) + far[:200, :200] = mask + cv2.circle(far, (270, 270), 12, 255, 2) # far ring with a hole inside + levels = clean_levels(far) + for radius in (0, 5, CLEAN_RADIUS_MAX): + cleaned = apply_clean_levels(levels, radius) + assert not cleaned[0:3].any() + assert not cleaned[255:285, 255:285].any() # ring gone, its hole never filled + assert cleaned[100, 100] == 255 + assert np.array_equal(clean_specks_and_holes(far, 0, max_gap=0)[:200, :200], _leaf()) + + +def test_specks_go_by_size_and_the_leaf_always_stays(): + mask = _leaf() + mask[5, 5] = 255 # 1 px speck + cv2.circle(mask, (185, 20), 8, 255, -1) # ~8 px radius speck + levels = clean_levels(mask) + + small = apply_clean_levels(levels, 2) + assert small[5, 5] == 0 and small[20, 185] == 255 + large = apply_clean_levels(levels, 9) + assert large[20, 185] == 0 + assert apply_clean_levels(levels, CLEAN_RADIUS_MAX)[100, 100] == 255 + + +def test_small_enclosed_holes_fill_but_border_background_never_does(): + mask = _leaf() + cv2.circle(mask, (80, 80), 3, 0, -1) + cv2.circle(mask, (120, 110), 15, 0, -1) + cleaned = clean_specks_and_holes(mask, 6) + assert cleaned[80, 80] == 255 + assert cleaned[110, 120] == 0 + assert cleaned[0, 0] == 0 + assert clean_specks_and_holes(mask, CLEAN_RADIUS_MAX)[0, 0] == 0 + + +def test_removing_a_thin_ring_never_fills_its_hole(): + mask = _leaf() + cv2.circle(mask, (20, 185), 10, 255, 1) + for radius in range(2, CLEAN_RADIUS_MAX + 1, 4): + cleaned = clean_specks_and_holes(mask, radius) + assert cleaned[185, 20] == 0 and cleaned[185, 30] == 0 + + +def test_an_island_fills_with_its_hole(): + mask = _leaf() + cv2.circle(mask, (120, 110), 15, 0, -1) + cv2.circle(mask, (120, 110), 2, 255, -1) + levels = clean_levels(mask) + assert apply_clean_levels(levels, 4)[110, 120] == 0 # island removed as a speck + filled = apply_clean_levels(levels, CLEAN_RADIUS_MAX) + assert filled[100:121, 110:131].min() == 255 # no black spot remains + + +def test_empty_and_full_masks(): + empty = np.zeros((20, 20), dtype=np.uint8) + full = np.full((20, 20), 255, dtype=np.uint8) + assert not clean_specks_and_holes(empty, 10).any() + assert np.array_equal(clean_specks_and_holes(full, 10), full) + + +def _framed(size=400, line=4, corner=30, leaf=20): + """A small leaf inside the printed box outline, as the target box shows it. + + The outline lands on the box edge as four strips; the corners are white + where the marker boxes were whited out. Each strip outweighs the leaf. + """ + mask = np.zeros((size, size), dtype=np.uint8) + mask[:line, corner:-corner] = mask[-line:, corner:-corner] = 255 + mask[corner:-corner, :line] = mask[corner:-corner, -line:] = 255 + start = size // 2 - leaf // 2 + only_leaf = np.zeros_like(mask) + only_leaf[start:start + leaf, start:start + leaf] = 255 + return mask | only_leaf, only_leaf + + +def test_margin_is_a_percent_of_the_shorter_side(): + mask = np.full((200, 400), 255, dtype=np.uint8) + assert margin_width(mask.shape, 1) == 2 + cleared = clear_margin(mask, 2.5) # 5 px on every edge + assert not cleared[:5].any() and not cleared[-5:].any() + assert not cleared[:, :5].any() and not cleared[:, -5:].any() + assert cleared[5:-5, 5:-5].min() == 255 + assert np.array_equal(clear_margin(mask, 0), mask) + + +@pytest.mark.parametrize("margin", [-1, CLEAN_MARGIN_MAX + 1, float("nan"), True, "wide"]) +def test_invalid_margins_raise(margin): + with pytest.raises(ValueError): + clear_margin(_leaf(), margin) + + +def test_edge_lines_that_outweigh_the_leaf_never_replace_it(): + mask, only_leaf = _framed() + assert (mask[:4] > 0).sum() > (only_leaf > 0).sum() # a strip outweighs the leaf + # Without the margin a strip is the largest piece, and the leaf is dropped. + assert not drop_stray_pieces(mask)[200, 200] + assert np.array_equal(clean_raw_mask(mask), only_leaf) + + +def test_a_line_reaching_past_the_margin_is_still_dropped(): + mask, only_leaf = _framed(line=7) # the 1% margin clears only 4 px + assert np.array_equal(clean_raw_mask(mask), only_leaf) + + +def test_a_leaf_crossing_into_the_margin_loses_only_the_band(): + mask = np.zeros((400, 400), dtype=np.uint8) + mask[0:100, 150:250] = 255 # leaf runs off the top edge + mask[105:108, 200:203] = 255 # a nearby speck stays + kept = clean_raw_mask(mask) + assert not kept[:4].any() + assert kept[4:100, 150:250].min() == 255 + assert kept[106, 201] == 255 + + +def test_clean_raw_mask_is_idempotent_and_feeds_clean_levels(): + mask, only_leaf = _framed() + kept = clean_raw_mask(mask) + assert np.array_equal(clean_raw_mask(kept), kept) + assert np.array_equal(apply_clean_levels(clean_levels(mask), 0), only_leaf) + assert np.array_equal(clean_raw_mask(mask, margin=0), drop_stray_pieces(mask)) + + +def _speckled_leaf(): + mask = _leaf() + cv2.circle(mask, (80, 80), 3, 0, -1) # small hole in the leaf + mask[25, 100] = 255 # 1 px speck beside the leaf + return mask + + +def test_raw_measurement_mask_is_clean_raw_mask_at_clean_size_0(): + for mask in (_framed()[0], _speckled_leaf()): + assert np.array_equal(raw_measurement_mask(mask), clean_raw_mask(mask)) + assert np.array_equal( + raw_measurement_mask(mask, 0, 0.1, clean_size=0), clean_raw_mask(mask, 0, 0.1) + ) + + +def test_raw_measurement_mask_above_0_is_the_clean_image_preview(): + mask = _speckled_leaf() + cleaned = raw_measurement_mask(mask, clean_size=5) + assert np.array_equal(cleaned, clean_specks_and_holes(mask, 5)) + assert np.array_equal(cleaned, apply_clean_levels(clean_levels(mask), 5)) + assert cleaned[25, 100] == 0 and cleaned[80, 80] == 255 + assert np.array_equal( + raw_measurement_mask(mask, 2.0, 0.1, clean_size=5), + clean_specks_and_holes(mask, 5, max_gap=0.1, margin=2.0), + ) + + +@pytest.mark.parametrize("value", [0, 3, CLEAN_SIZE_MAX, "3", " 7 ", np.int64(4)]) +def test_clean_size_accepts_whole_pixels(value): + assert checked_clean_size(value) == int(value) + assert type(checked_clean_size(value)) is int + + +@pytest.mark.parametrize("value", [-1, CLEAN_SIZE_MAX + 1, 2.5, 3.0, True, "3.5", "big", None]) +def test_clean_size_rejects_other_values(value): + with pytest.raises(ValueError, match="clean size"): + checked_clean_size(value) + with pytest.raises(ValueError, match="clean size"): + raw_measurement_mask(_leaf(), clean_size=value) diff --git a/tests/test_parallel_workers.py b/tests/test_parallel_workers.py index e0930be..a0d6608 100644 --- a/tests/test_parallel_workers.py +++ b/tests/test_parallel_workers.py @@ -23,6 +23,21 @@ def _result_for(path): }, None +@pytest.mark.parametrize("input_images", [ + ["Sample.jpg", "sample.png"], + ["Sample_target_box.jpg", "sample.png"], +]) +def test_batch_rejects_case_insensitive_sample_id_collisions(input_images, tmp_path): + output_dir = tmp_path / "out" + with pytest.raises(ValueError, match="duplicate sample IDs"): + core.run_leaf_morpho_batch( + input_images, + str(output_dir), + str(output_dir / "results.csv"), + ) + assert not output_dir.exists() + + def test_parallel_cpu_policy_allows_model_backed_workers(monkeypatch, tmp_path): seen_devices = [] diff --git a/tests/test_pre_cleanup_exports.py b/tests/test_pre_cleanup_exports.py new file mode 100644 index 0000000..55d30de --- /dev/null +++ b/tests/test_pre_cleanup_exports.py @@ -0,0 +1,549 @@ +"""Exercise export behavior on a real Otsu target box without model downloads.""" + +import csv +import json +from pathlib import Path + +import pytest + +np = pytest.importorskip("numpy") +cv2 = pytest.importorskip("cv2") +from mats import core + + +def _input(tmp_path): + image = np.full((100, 100, 3), 255, dtype=np.uint8) + image[20:80, 20:80] = 0 + image[43:57, 43:57] = 255 + image[5:10, 5:10] = 0 + path = tmp_path / "leaf_target_box.png" + assert cv2.imwrite(str(path), image) + return path + + +def _run(source, output, **options): + return core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, + compact_csv=False, export_options=options, + ) + + +def test_pre_cleanup_and_optional_files_preserve_measurements(tmp_path): + source = _input(tmp_path) + source_bytes = source.read_bytes() + plain = tmp_path / "plain" + exported = tmp_path / "exported" + plain_run = _run(source, plain, target_boxes=False, cleaned_masks=False) + export_run = _run(source, exported, target_boxes=False, cleaned_masks=True, + pre_cleanup_methods=("threshold",), overlay=True, + cutout=True, axes=True) + + assert (plain / "results.csv").read_bytes() == (exported / "results.csv").read_bytes() + assert plain_run["result_rows"] == export_run["result_rows"] + raw = cv2.imread(str(exported / "leaf_mask_precleanup_threshold.png"), 0) + cleaned = cv2.imread(str(exported / "leaf_mask.png"), 0) + assert raw[50, 50] == 0 and raw[7, 7] == 255 + assert cleaned[50, 50] == 255 and cleaned[7, 7] == 0 + assert not (exported / "leaf_target_box.jpg").exists() + assert not (plain / "leaf_mask.png").exists() + assert source.read_bytes() == source_bytes + assert {item["kind"] for item in export_run["artifacts"]} == { + "results_csv", "results_metadata", "mask", "pre_cleanup", "overlay", "cutout", "axes" + } + + +def test_both_methods_export_once_and_do_not_change_otsu_measurement(tmp_path, monkeypatch): + source = _input(tmp_path) + baseline = _run(source, tmp_path / "baseline") + calls = [] + + def fake_birefnet(image, device_override=None): + calls.append(image.shape) + mask = np.zeros(image.shape[:2], dtype=np.uint8) + mask[30:70, 30:70] = 255 + return mask + + monkeypatch.setattr(core, "predict_birefnet_mask", fake_birefnet) + both = _run(source, tmp_path / "both", + pre_cleanup_methods=("threshold", "birefnet")) + assert calls == [(100, 100, 3)] + assert baseline["result_rows"] == both["result_rows"] + assert (tmp_path / "both" / "leaf_mask_precleanup_threshold.png").exists() + assert (tmp_path / "both" / "leaf_mask_precleanup_birefnet.png").exists() + + +def test_birefnet_measurements_are_independent_of_threshold_export(tmp_path, monkeypatch): + source = _input(tmp_path) + calls = [] + + def fake_birefnet(image, device_override=None): + calls.append(image.shape) + mask = np.zeros(image.shape[:2], dtype=np.uint8) + mask[30:70, 30:70] = 255 + return mask + + monkeypatch.setattr(core, "predict_birefnet_mask", fake_birefnet) + baseline = core.run_leaf_morpho_batch( + [str(source)], str(tmp_path / "birefnet"), str(tmp_path / "birefnet" / "results.csv"), + template_dimensions=(10, 10, "cm"), mask_method="birefnet", workers=1, + ) + both = core.run_leaf_morpho_batch( + [str(source)], str(tmp_path / "both"), str(tmp_path / "both" / "results.csv"), + template_dimensions=(10, 10, "cm"), mask_method="birefnet", workers=1, + export_options={"pre_cleanup_methods": ("threshold", "birefnet")}, + ) + assert calls == [(100, 100, 3), (100, 100, 3)] + assert baseline["result_rows"] == both["result_rows"] + + +def test_custom_cutoff_drives_the_threshold_pre_cleanup_mask(tmp_path): + # A pale (gray 160) leaf sits above the "high" preset but below a custom 177. + image = np.full((100, 100, 3), 255, dtype=np.uint8) + image[20:80, 20:80] = 160 + source = tmp_path / "pale_target_box.png" + assert cv2.imwrite(str(source), image) + + def run(output, threshold_value): + return core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), threshold_value=threshold_value, + workers=1, export_options={"pre_cleanup_methods": ("threshold",)}, + ) + + custom = run(tmp_path / "custom", 177) + raw = cv2.imread(str(tmp_path / "custom" / "pale_mask_precleanup_threshold.png"), 0) + _, expected = cv2.threshold( + cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), 177, 255, cv2.THRESH_BINARY_INV + ) + assert custom["succeeded"] == 1 + assert np.array_equal(raw, expected) + assert run(tmp_path / "high", core.THRESHOLD_LEVELS["high"])["succeeded"] == 0 + + +def test_secondary_method_failure_keeps_measurement(tmp_path, monkeypatch): + source = _input(tmp_path) + + def fail_birefnet(image, device_override=None): + raise RuntimeError("mock inference error") + + monkeypatch.setattr(core, "predict_birefnet_mask", fail_birefnet) + run = _run(source, tmp_path / "out", pre_cleanup_methods=("birefnet",)) + assert run["succeeded"] == 1 + assert run["result_rows"][0]["leaf_area_cm2"] != "NA" + assert any("mock inference error" in row["status"] for row in run["failure_rows"]) + + +def test_manifest_excludes_files_left_by_prior_run(tmp_path): + source = _input(tmp_path) + output = tmp_path / "reused" + _run(source, output, pre_cleanup_methods=("threshold",)) + later = _run(source, output, cleaned_masks=False) + assert (output / "leaf_mask_precleanup_threshold.png").exists() + assert {item["kind"] for item in later["artifacts"]} == {"results_csv", "results_metadata"} + + +def test_pre_cleanup_measurement_keeps_holes_and_nearby_specks(tmp_path): + source = _input(tmp_path) + cleaned = core.run_leaf_morpho_batch( + [str(source)], str(tmp_path / "cleaned"), str(tmp_path / "cleaned" / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, compact_csv=False, + export_options={"axes": True, "pre_cleanup_methods": ("threshold",)}, + ) + raw = core.run_leaf_morpho_batch( + [str(source)], str(tmp_path / "raw"), str(tmp_path / "raw" / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, compact_csv=False, + measurement_source="pre-cleanup", + export_options={"axes": True, "pre_cleanup_methods": ("threshold",)}, + ) + clean_row, raw_row = cleaned["result_rows"][0], raw["result_rows"][0] + assert raw_row["leaf_area_cm2"] < clean_row["leaf_area_cm2"] # hole stays empty + assert raw_row["width_cm"] > clean_row["width_cm"] # detached speck extends width + assert raw_row["length_cm"] > clean_row["length_cm"] + assert json.loads((tmp_path / "raw" / "results.csv.meta.json").read_text()) == { + "measurement_source": "pre-cleanup", + "mask_method": "threshold", + "results_unit": "cm", + "csv_schema": "full", + "clean_margin": 1.0, + "stray_gap": 0.25, + "clean_size": 0, + } + assert "stray_gap" not in json.loads( + (tmp_path / "cleaned" / "results.csv.meta.json").read_text() + ) + assert raw["measurement_source"] == "pre-cleanup" + assert raw["artifacts"][-1]["kind"] == "results_metadata" + assert not np.array_equal( + cv2.imread(str(tmp_path / "raw" / "leaf_measurement_axes.jpg")), + cv2.imread(str(tmp_path / "cleaned" / "leaf_measurement_axes.jpg")), + ) + + +def _stray_input(tmp_path): + """A leaf with a border strip, a far speck, and a speck just above it.""" + image = np.full((200, 200, 3), 255, dtype=np.uint8) + image[60:140, 60:140] = 0 # leaf: 80 px, so the 0.25 gap limit is ~28 px + image[0:4, 30:170] = 0 # printed edge along the top border + image[10:16, 10:16] = 0 # far speck, ~64 px from the leaf + image[45:50, 95:100] = 0 # near speck, ~10 px above the leaf + path = tmp_path / "leaf_target_box.png" + assert cv2.imwrite(str(path), image) + return path + + +def _pre_cleanup_run(source, output, **options): + return core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, compact_csv=False, + measurement_source="pre-cleanup", + export_options={"pre_cleanup_methods": ("threshold",), "axes": True, + "preview_dir": str(output / "previews")}, + **options, + ) + + +def _extent_px(row): + return (round(row["width_cm"] * row["px_per_cm_width"]), + round(row["length_cm"] * row["px_per_cm_height"])) + + +def test_pre_cleanup_measurement_drops_border_and_far_pieces(tmp_path): + source = _stray_input(tmp_path) + output = tmp_path / "out" + run = _pre_cleanup_run(source, output) + row = run["result_rows"][0] + # Width is the leaf alone; length reaches up to the near speck (rows 45-139). + assert _extent_px(row) == (80, 95) + assert round(row["leaf_area_cm2"] * row["px_per_cm_width"] * row["px_per_cm_height"]) == ( + 80 * 80 + 25 + ) + + raw = cv2.imread(str(output / "leaf_mask_precleanup_threshold.png"), 0) + assert raw[1, 100] == 255 and raw[12, 12] == 255 # the export stays raw + measured = cv2.imread(next( + item["path"] for item in run["preview_artifacts"] if item["kind"] == "preview_mask" + ), 0) + assert measured[1, 100] == 0 and measured[12, 12] == 0 + assert measured[47, 97] == 255 and measured[100, 100] == 255 + + leaf_only = _pre_cleanup_run(source, tmp_path / "leaf_only", stray_gap=0) + assert _extent_px(leaf_only["result_rows"][0]) == (80, 80) + + +def test_pre_cleanup_adjustment_measures_without_stray_pieces(tmp_path): + from mats.app.output_adjustment import apply_threshold_adjustment + + source = _stray_input(tmp_path) + output = tmp_path / "out" + summary = _pre_cleanup_run(source, output, threshold_value=125, stray_gap=0.5) + paths = { + item["kind"]: item["path"] + for item in (*summary["artifacts"], *summary["preview_artifacts"]) + if item["sample_id"] == "leaf" + } + run = { + "by_method": {"threshold": {"results_path": summary["results_path"]}}, + "results_unit": "cm", "measurement_source": "pre-cleanup", + "stray_gap": summary["stray_gap"], "output_path": str(output), + "mask_methods": ("threshold",), "artifacts": list(summary["artifacts"]), + } + pair = { + "sample_id": "leaf", "target_box": str(source), "mask": paths["preview_mask"], + "raw_mask": paths["pre_cleanup"], "mask_source": "pre-cleanup", + } + before = summary["result_rows"][0] + apply_threshold_adjustment(run, pair, 125) + + with open(summary["results_path"], newline="") as handle: + adjusted = next(csv.DictReader(handle)) + assert float(adjusted["width_cm"]) == pytest.approx(before["width_cm"]) + assert float(adjusted["length_cm"]) == pytest.approx(before["length_cm"]) + assert pair["mask"] == paths["preview_mask"] + assert pair["raw_mask"] == paths["pre_cleanup"] + assert cv2.imread(pair["mask"], 0)[1, 100] == 0 + assert cv2.imread(pair["raw_mask"], 0)[1, 100] == 255 + metadata = json.loads(Path(f"{summary['results_path']}.meta.json").read_text()) + assert metadata["stray_gap"] == 0.5 + + # The explorer's own settings drive the overwrite and are recorded with it. + apply_threshold_adjustment(run, pair, 125, stray_gap=0, clean_margin=2) + with open(summary["results_path"], newline="") as handle: + leaf_only = next(csv.DictReader(handle)) + width = float(leaf_only["width_cm"]) * float(leaf_only["px_per_cm_width"]) + length = float(leaf_only["length_cm"]) * float(leaf_only["px_per_cm_height"]) + assert (round(width), round(length)) == (80, 80) # the near speck is gone + metadata = json.loads(Path(f"{summary['results_path']}.meta.json").read_text()) + assert metadata["threshold_adjustments"]["leaf"] == { + "cutoff": 125, "fill_holes": True, "clean_margin": 2.0, "stray_gap": 0.0, + "clean_size": 0, + } + assert run["threshold_adjustments"]["leaf"]["stray_gap"] == 0.0 + + +def _speckled_input(tmp_path): + """A leaf with a small hole, a large hole, and a speck just above it.""" + image = np.full((200, 200, 3), 255, dtype=np.uint8) + image[60:140, 60:140] = 0 # leaf + image[99:102, 99:102] = 255 # small hole: inscribed radius 2 + image[70:90, 70:90] = 255 # large hole: inscribed radius 10 + image[45:50, 95:100] = 0 # near speck: inscribed radius 3 + path = tmp_path / "leaf_target_box.png" + assert cv2.imwrite(str(path), image) + return path + + +def _area_px(row): + return round(row["leaf_area_cm2"] * row["px_per_cm_width"] * row["px_per_cm_height"]) + + +def _preview_mask(summary): + return cv2.imread(next( + item["path"] for item in summary["preview_artifacts"] if item["kind"] == "preview_mask" + ), 0) + + +def test_clean_size_removes_specks_and_fills_small_holes_before_measuring(tmp_path): + from mats.mask_cleanup import clean_specks_and_holes + + source = _speckled_input(tmp_path) + plain = _pre_cleanup_run(source, tmp_path / "plain") + assert _extent_px(plain["result_rows"][0]) == (80, 95) # speck included + assert _area_px(plain["result_rows"][0]) == 80 * 80 - 400 - 9 + 25 + + output = tmp_path / "clean" + summary = _pre_cleanup_run(source, output, clean_size=5) + row = summary["result_rows"][0] + assert _extent_px(row) == (80, 80) # speck removed + assert _area_px(row) == 80 * 80 - 400 # small hole filled + raw = cv2.imread(str(output / "leaf_mask_precleanup_threshold.png"), 0) + assert raw[47, 97] == 255 and raw[100, 100] == 0 # the export stays raw + assert np.array_equal(_preview_mask(summary), clean_specks_and_holes(raw, 5)) + assert summary["clean_size"] == 5 + metadata = json.loads((output / "results.csv.meta.json").read_text()) + assert metadata["clean_size"] == 5 + # The cleaned-mask export is still MATS cleanup, whatever the clean size. + assert np.array_equal( + cv2.imread(str(output / "leaf_mask.png"), 0), + cv2.imread(str(tmp_path / "plain" / "leaf_mask.png"), 0), + ) + + +@pytest.mark.parametrize("clean_size", [3, "big", 51]) +def test_clean_size_needs_pre_cleanup_and_a_valid_radius(tmp_path, clean_size): + with pytest.raises(ValueError, match="clean"): + core.run_leaf_morpho_batch( + [str(_speckled_input(tmp_path))], str(tmp_path / "out"), + str(tmp_path / "out" / "results.csv"), template_dimensions=(10, 10, "cm"), + workers=1, clean_size=clean_size, + ) + assert not (tmp_path / "out" / "results.csv").exists() + + +def _adjustable(summary, source, output, measurement_source): + paths = { + item["kind"]: item["path"] + for item in (*summary["artifacts"], *summary["preview_artifacts"]) + if item["sample_id"] == "leaf" + } + run = { + "by_method": {"threshold": {"results_path": summary["results_path"]}}, + "results_unit": "cm", "measurement_source": measurement_source, + "clean_margin": summary["clean_margin"], "stray_gap": summary["stray_gap"], + "clean_size": summary["clean_size"], "output_path": str(output), + "mask_methods": ("threshold",), "artifacts": list(summary["artifacts"]), + } + pair = {"sample_id": "leaf", "target_box": str(source), "mask": paths.get("mask")} + if measurement_source == "pre-cleanup": + pair.update(mask=paths["preview_mask"], raw_mask=paths["pre_cleanup"], + mask_source="pre-cleanup") + return run, pair + + +def _saved_row(summary): + with open(summary["results_path"], newline="") as handle: + row = next(csv.DictReader(handle)) + return {key: float(value) for key, value in row.items() if key != "sample_id"} + + +def test_pre_cleanup_overwrite_saves_and_measures_the_clean_size(tmp_path): + from mats.app.output_adjustment import apply_threshold_adjustment + from mats.mask_cleanup import clean_specks_and_holes + + source = _speckled_input(tmp_path) + output = tmp_path / "out" + summary = _pre_cleanup_run(source, output, threshold_value=125) + run, pair = _adjustable(summary, source, output, "pre-cleanup") + + apply_threshold_adjustment(run, pair, 125, clean_size=5) + + row = _saved_row(summary) + assert _extent_px(row) == (80, 80) and _area_px(row) == 80 * 80 - 400 + raw = cv2.imread(pair["raw_mask"], 0) + assert raw[47, 97] == 255 # raw export untouched + assert np.array_equal(cv2.imread(pair["mask"], 0), clean_specks_and_holes(raw, 5)) + metadata = json.loads(Path(f"{summary['results_path']}.meta.json").read_text()) + assert metadata["clean_size"] == 0 # the run's own setting + assert metadata["threshold_adjustments"]["leaf"] == { + "cutoff": 125, "fill_holes": True, "clean_margin": 1.0, "stray_gap": 0.25, + "clean_size": 5, + } + assert run["threshold_adjustments"]["leaf"]["clean_size"] == 5 + + # Clean size 0 restores the run's own measurement. + apply_threshold_adjustment(run, pair, 125, clean_size=0) + assert _extent_px(_saved_row(summary)) == (80, 95) + + +def test_overwrite_defaults_to_the_run_clean_size(tmp_path): + from mats.app.output_adjustment import measure_threshold_adjustment + + source = _speckled_input(tmp_path) + output = tmp_path / "out" + summary = _pre_cleanup_run(source, output, threshold_value=125, clean_size=5) + run, pair = _adjustable(summary, source, output, "pre-cleanup") + prepared = measure_threshold_adjustment(run, pair, 125) + assert prepared["clean_size"] == 5 + assert np.array_equal(prepared["measurement_mask"], _preview_mask(summary)) + + +def test_cleaned_run_overwrite_replaces_mats_cleanup_with_clean_image(tmp_path): + from mats.app.output_adjustment import apply_threshold_adjustment + from mats.mask_cleanup import clean_specks_and_holes + + source = _speckled_input(tmp_path) + output = tmp_path / "out" + summary = core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, compact_csv=False, + threshold_value=125, export_options={"axes": True}, + ) + assert _area_px(summary["result_rows"][0]) == 80 * 80 # flash-filled + run, pair = _adjustable(summary, source, output, "cleaned") + + # Remove flashfill doesn't apply while the clean size is above 0. + apply_threshold_adjustment(run, pair, 125, clean_size=5, remove_fill=True) + + raw = core.threshold_mask(cv2.imread(str(source)), 125) + saved = cv2.imread(pair["mask"], 0) + assert np.array_equal(saved, clean_specks_and_holes(raw, 5)) + assert saved[80, 80] == 0 and saved[100, 100] == 255 and saved[47, 97] == 0 + row = _saved_row(summary) + assert _extent_px(row) == (80, 80) and _area_px(row) == 80 * 80 - 400 + metadata = json.loads(Path(f"{summary['results_path']}.meta.json").read_text()) + assert metadata["threshold_adjustments"]["leaf"] == { + "cutoff": 125, "fill_holes": False, "clean_margin": 1.0, "stray_gap": 0.25, + "clean_size": 5, + } + + +def _framed_input(tmp_path): + """A small leaf inside the printed box outline, as a target box shows it. + + The outline lands on the box edge as four strips, with white corners where + the marker boxes were whited out. Each strip outweighs the 20 x 20 leaf. + """ + image = np.full((400, 400, 3), 255, dtype=np.uint8) + image[:4, 30:-30] = image[-4:, 30:-30] = 0 + image[30:-30, :4] = image[30:-30, -4:] = 0 + image[190:210, 190:210] = 0 + path = tmp_path / "framed_target_box.png" + assert cv2.imwrite(str(path), image) + return path + + +def test_pre_cleanup_measurement_ignores_edge_lines_that_outweigh_the_leaf(tmp_path): + source = _framed_input(tmp_path) + output = tmp_path / "out" + row = _pre_cleanup_run(source, output)["result_rows"][0] + assert _extent_px(row) == (20, 20) + raw = cv2.imread(str(output / "framed_mask_precleanup_threshold.png"), 0) + assert raw[1, 200] == 255 # the export keeps the lines + metadata = json.loads((output / "results.csv.meta.json").read_text()) + assert metadata["clean_margin"] == 1.0 + + # Without the margin, a strip is the largest piece and replaces the leaf. + unguarded = _pre_cleanup_run(source, tmp_path / "no_margin", clean_margin=0) + assert _extent_px(unguarded["result_rows"][0]) != (20, 20) + + +def test_remove_flashfill_adjustment_clears_the_edge_margin(tmp_path): + from mats.app.output_adjustment import measure_threshold_adjustment + + source = _framed_input(tmp_path) + output = tmp_path / "out" + summary = core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, compact_csv=False, + threshold_value=125, + ) + run = { + "by_method": {"threshold": {"results_path": summary["results_path"]}}, + "results_unit": "cm", "measurement_source": "cleaned", + "output_path": str(output), "mask_methods": ("threshold",), + "artifacts": list(summary["artifacts"]), + } + prepared = measure_threshold_adjustment( + run, {"sample_id": "framed", "target_box": str(source)}, 125, remove_fill=True, + ) + only_leaf = np.zeros((400, 400), dtype=np.uint8) + only_leaf[190:210, 190:210] = 255 + assert np.array_equal(prepared["measurement_mask"], only_leaf) + + +def test_raw_preview_survives_removed_input_without_becoming_export(tmp_path): + source = _input(tmp_path) + output = tmp_path / "out" + preview_dir = tmp_path / "private_previews" + run = core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, + measurement_source="pre-cleanup", + export_options={ + "target_boxes": False, "cleaned_masks": False, + "preview_dir": str(preview_dir), + }, + ) + source.unlink() # Uploaded files are discarded after the GUI run. + previews = {item["kind"]: item for item in run["preview_artifacts"]} + assert {"preview_target_box", "preview_mask", "preview_raw_mask"} == set(previews) + assert cv2.imread(previews["preview_target_box"]["path"]) is not None + measured = cv2.imread(previews["preview_mask"]["path"], cv2.IMREAD_GRAYSCALE) + assert measured[50, 50] == 0 and measured[7, 7] == 255 + # The explorer re-cleans the true raw mask with each specimen's settings. + raw = cv2.imread(previews["preview_raw_mask"]["path"], cv2.IMREAD_GRAYSCALE) + assert np.array_equal(raw, core.threshold_mask(cv2.imread(previews["preview_target_box"]["path"]), None)) + assert {item["kind"] for item in run["artifacts"]} == { + "results_csv", "results_metadata", + } + + +def test_cleaned_preview_uses_measurement_mask_when_export_is_disabled(tmp_path): + source = _input(tmp_path) + output = tmp_path / "out" + run = core.run_leaf_morpho_batch( + [str(source)], str(output), str(output / "results.csv"), + template_dimensions=(10, 10, "cm"), workers=1, + export_options={"cleaned_masks": False, "preview_dir": str(tmp_path / "previews")}, + ) + preview = next( + item for item in run["preview_artifacts"] if item["kind"] == "preview_mask" + ) + cleaned = cv2.imread(preview["path"], cv2.IMREAD_GRAYSCALE) + assert cleaned[50, 50] == 255 and cleaned[7, 7] == 0 + raw_preview = next( + item for item in run["preview_artifacts"] if item["kind"] == "preview_raw_mask" + ) + raw = cv2.imread(raw_preview["path"], cv2.IMREAD_GRAYSCALE) + assert raw[50, 50] == 0 and raw[7, 7] == 255 + assert not (output / "leaf_mask.png").exists() + + +def test_secondary_birefnet_export_requires_supported_execution(tmp_path): + source = _input(tmp_path) + with pytest.raises(ValueError, match="Otsu thresholding only"): + core.run_leaf_morpho_batch( + [str(source)], str(tmp_path / "out"), str(tmp_path / "out" / "results.csv"), + workers=2, execution_device="hybrid", + export_options={"pre_cleanup_methods": ("birefnet",)}, + ) diff --git a/tests/test_threshold_preview.py b/tests/test_threshold_preview.py new file mode 100644 index 0000000..658e14e --- /dev/null +++ b/tests/test_threshold_preview.py @@ -0,0 +1,171 @@ +"""The interactive preview must agree with the production threshold path.""" + +import base64 + +import pytest + +np = pytest.importorskip("numpy") +cv2 = pytest.importorskip("cv2") +pytest.importorskip("streamlit") + +from mats.app.threshold_preview import ( + clean_levels_for_mask, cleaned_sample, color_sample, grayscale_sample, + pre_cleanup_sample, unfilled_sample, +) +from mats.mask_cleanup import clean_levels, clean_raw_mask + + +def _decode_png(url): + encoded = base64.b64decode(url.split(",", 1)[1]) + return cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8), cv2.IMREAD_GRAYSCALE) + + +def test_grayscale_preview_uses_exact_pipeline_pixels_and_otsu_cutoff(tmp_path): + image = np.full((40, 50, 3), (180, 210, 240), dtype=np.uint8) + image[8:32, 10:40] = (15, 25, 35) + path = tmp_path / "sample.png" + assert cv2.imwrite(str(path), image) + + data_url, cutoff = grayscale_sample(str(path)) + expected_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + expected_cutoff, _ = cv2.threshold( + expected_gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU + ) + assert np.array_equal(_decode_png(data_url), expected_gray) + assert cutoff == int(expected_cutoff) + + +def test_color_preview_matches_the_target_box_geometry_and_colors(tmp_path): + image = np.full((40, 50, 3), (180, 210, 240), dtype=np.uint8) + image[8:32, 10:40] = (40, 140, 60) + path = tmp_path / "sample.png" + assert cv2.imwrite(str(path), image) + + url = color_sample(str(path)) + assert url.startswith("data:image/jpeg;base64,") + decoded = cv2.imdecode( + np.frombuffer(base64.b64decode(url.split(",", 1)[1]), dtype=np.uint8), + cv2.IMREAD_COLOR, + ) + # Same pixel grid as the grayscale preview, so the two panels align. + assert decoded.shape == image.shape + assert np.abs(decoded.astype(int) - image.astype(int)).mean() < 3 + + +def test_clean_levels_reach_the_browser_in_channel_order(tmp_path): + mask = np.zeros((60, 60), dtype=np.uint8) + cv2.circle(mask, (30, 30), 20, 255, -1) + cv2.circle(mask, (30, 30), 3, 0, -1) + mask[2, 2] = 255 + path = tmp_path / "raw.png" + assert cv2.imwrite(str(path), mask) + + url = clean_levels_for_mask(str(path), path.stat().st_mtime_ns) + decoded = cv2.imdecode( + np.frombuffer(base64.b64decode(url.split(",", 1)[1]), dtype=np.uint8), + cv2.IMREAD_COLOR, + ) + # The browser reads PNG channels as RGB; OpenCV decodes them as BGR. + assert np.array_equal(cv2.cvtColor(decoded, cv2.COLOR_BGR2RGB), clean_levels(mask)) + + +def test_clean_levels_use_the_runs_stray_gap(tmp_path): + mask = np.zeros((60, 60), dtype=np.uint8) + cv2.circle(mask, (30, 30), 20, 255, -1) + mask[30, 4] = 255 # 6 px from the leaf + path = tmp_path / "raw.png" + assert cv2.imwrite(str(path), mask) + + def levels(stray_gap): + url = clean_levels_for_mask(str(path), path.stat().st_mtime_ns, stray_gap) + decoded = cv2.imdecode( + np.frombuffer(base64.b64decode(url.split(",", 1)[1]), dtype=np.uint8), + cv2.IMREAD_COLOR, + ) + return cv2.cvtColor(decoded, cv2.COLOR_BGR2RGB) + + assert np.array_equal(levels(0.25), clean_levels(mask, 0.25)) + assert np.array_equal(levels(0), clean_levels(mask, 0)) + assert not np.array_equal(levels(0), levels(0.25)) # separate cache entries + + +def test_settled_cleaned_preview_uses_production_mask_cleanup(tmp_path): + pytest.importorskip("torch") + pytest.importorskip("rfdetr") + from mats import core + + image = np.full((100, 100, 3), 255, dtype=np.uint8) + image[20:80, 20:80] = 0 + image[43:57, 43:57] = 255 + path = tmp_path / "sample.png" + assert cv2.imwrite(str(path), image) + + preview = _decode_png(cleaned_sample(str(path), 125)) + expected = core.clean_leaf_mask(core.threshold_mask(image, 125)) + assert np.array_equal(preview, expected) + + +def test_remove_flashfill_preserves_holes_but_keeps_other_cleanup(tmp_path): + pytest.importorskip("torch") + pytest.importorskip("rfdetr") + from mats import core + + raw = np.zeros((100, 100), dtype=np.uint8) + raw[20:80, 20:80] = 255 + raw[40:60, 40:60] = 0 + raw[4:9, 4:9] = 255 + path = tmp_path / "raw.png" + assert cv2.imwrite(str(path), raw) + + preview = _decode_png(unfilled_sample(str(path))) + cleaned = core.clean_leaf_mask(raw) + assert preview[50, 50] == 0 and cleaned[50, 50] == 255 + assert preview[6, 6] == 0 # The detached speck is still removed. + assert preview[25, 25] == 255 + + +def test_otsu_zero_cutoff_can_be_previewed(tmp_path): + image = np.full((20, 20, 3), 255, dtype=np.uint8) + image[5:15, 5:15] = 0 + path = tmp_path / "binary_sample.png" + assert cv2.imwrite(str(path), image) + _, cutoff = grayscale_sample(str(path)) + assert cutoff == 0 + + +def _framed_target(tmp_path): + image = np.full((400, 400, 3), 255, dtype=np.uint8) + image[:4, 30:-30] = image[-4:, 30:-30] = 0 # printed outline on the edge + image[30:-30, :4] = image[30:-30, -4:] = 0 + image[190:210, 190:210] = 0 # a leaf smaller than a strip + path = tmp_path / "framed.png" + assert cv2.imwrite(str(path), image) + return path, image + + +def test_pre_cleanup_preview_is_the_measured_mask(tmp_path): + pytest.importorskip("torch") + pytest.importorskip("rfdetr") + from mats import core + + path, image = _framed_target(tmp_path) + settled = _decode_png(pre_cleanup_sample(str(path), 125, 1.0, 0.25)) + expected = clean_raw_mask(core.threshold_mask(image, 125), 1.0, 0.25) + assert np.array_equal(settled, expected) + assert not settled[:4].any() and settled[200, 200] == 255 + + +def test_remove_flashfill_previews_clear_the_edge_margin(tmp_path): + pytest.importorskip("torch") + pytest.importorskip("rfdetr") + from mats import core + + path, image = _framed_target(tmp_path) + raw = core.threshold_mask(image, 125) + raw_path = tmp_path / "raw.png" + assert cv2.imwrite(str(raw_path), raw) + only_leaf = np.zeros((400, 400), dtype=np.uint8) + only_leaf[190:210, 190:210] = 255 + assert np.array_equal(_decode_png(unfilled_sample(str(raw_path), 1.0)), only_leaf) + unfilled = cleaned_sample(str(path), 125, fill_holes=False, clean_margin=1.0) + assert np.array_equal(_decode_png(unfilled), only_leaf) diff --git a/tests/test_thresholds.py b/tests/test_thresholds.py new file mode 100644 index 0000000..5dfe3e9 --- /dev/null +++ b/tests/test_thresholds.py @@ -0,0 +1,58 @@ +"""Threshold-level parsing -- standard library only, so it runs in CI.""" + +import pytest + +from mats.thresholds import ( + CUSTOM_THRESHOLD_LEVEL, + THRESHOLD_LEVEL_OPTIONS, + THRESHOLD_LEVELS, + parse_threshold_level, + threshold_value_for, +) + + +def test_presets_are_unchanged(): + # Changing these changes every preset run's measurements. + assert THRESHOLD_LEVELS == {"auto": None, "low": 100, "medium": 125, "high": 150} + + +def test_custom_is_an_option_but_never_a_preset(): + # None means Otsu in THRESHOLD_LEVELS; a "custom" key would risk silently running Otsu. + assert CUSTOM_THRESHOLD_LEVEL not in THRESHOLD_LEVELS + assert THRESHOLD_LEVEL_OPTIONS == ("auto", "custom", "low", "medium", "high") + + +@pytest.mark.parametrize("text, expected", [ + ("auto", "auto"), + ("HIGH", "high"), + (" Medium ", "medium"), + ("1", 1), + ("177", 177), + ("255", 255), +]) +def test_parse_accepts_presets_and_in_range_integers(text, expected): + assert parse_threshold_level(text) == expected + + +@pytest.mark.parametrize("text", ["0", "256", "-5", "12.5", "max", "1_77", "", "custom"]) +def test_parse_rejects_everything_else(text): + with pytest.raises(ValueError): + parse_threshold_level(text) + + +def test_bare_custom_explains_that_a_number_is_needed(): + with pytest.raises(ValueError, match="integer 1-255"): + parse_threshold_level("custom") + + +def test_threshold_value_for_presets_integers_and_custom(): + assert threshold_value_for("auto") is None + assert threshold_value_for("medium") == 125 + assert threshold_value_for(177) == 177 + assert threshold_value_for("custom", 140) == 140 + + +@pytest.mark.parametrize("custom_value", [None, 0, 256, 140.0, True]) +def test_custom_value_must_be_an_in_range_integer(custom_value): + with pytest.raises(ValueError): + threshold_value_for("custom", custom_value) diff --git a/tests/test_weights.py b/tests/test_weights.py index fd6ff01..bd6452e 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -76,12 +76,38 @@ def test_pointer_is_not_counted_present(monkeypatch, tmp_path): assert weights._is_present(tmp_path / "rf_detr_marker.pth") is False -def test_fetch_without_any_source_prints_manual(monkeypatch, tmp_path, capsys): +@pytest.mark.parametrize( + ("name", "filename", "override", "command", "forbidden"), + ( + ( + "rf-detr", + "rf_detr_marker.pth", + "RF_DETR_MARKER_CHECKPOINT", + 'git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth"', + 'git lfs pull -X "" -I', + ), + ( + "birefnet", + "birefnet_leaf.pth", + "BIREFNET_CHECKPOINT", + 'git lfs pull -X "" -I "weights/birefnet_leaf.pth"', + "git lfs install &&", + ), + ), +) +def test_fetch_without_any_source_prints_checkpoint_instructions( + monkeypatch, tmp_path, capsys, name, filename, override, command, forbidden +): weights = _fresh_weights(monkeypatch, tmp_path, MATS_WEIGHTS_DIR=str(tmp_path)) assert weights._HF_REPO_ID is None - code = weights.fetch() + code = weights.fetch(only=name) assert code == 1 - assert "No automatic download source" in capsys.readouterr().err + err = capsys.readouterr().err + assert "No automatic download source" in err + assert filename in err + assert override in err + assert command in err + assert forbidden not in err def test_ensure_weight_returns_present_file(monkeypatch, tmp_path): @@ -109,12 +135,34 @@ def test_require_local_weight_missing_never_attempts_fetch(monkeypatch, tmp_path weights.require_local_weight("birefnet") -def test_ensure_weight_honors_no_auto_fetch(monkeypatch, tmp_path): +@pytest.mark.parametrize( + ("name", "command", "forbidden"), + ( + ( + "rf-detr", + 'git lfs install && git lfs pull --exclude="weights/birefnet_leaf.pth"', + "--only birefnet", + ), + ( + "birefnet", + 'git lfs pull -X "" -I "weights/birefnet_leaf.pth"', + "--only rf-detr", + ), + ), +) +def test_ensure_weight_honors_no_auto_fetch( + monkeypatch, tmp_path, name, command, forbidden +): weights = _fresh_weights( monkeypatch, tmp_path, MATS_WEIGHTS_DIR=str(tmp_path), MATS_NO_AUTO_FETCH="1" ) - with pytest.raises(FileNotFoundError, match="auto-fetch is disabled"): - weights.ensure_weight("birefnet") + with pytest.raises(FileNotFoundError, match="auto-fetch is disabled") as exc_info: + weights.ensure_weight(name) + message = str(exc_info.value) + assert f"mats fetch-weights --only {name}" in message + assert "From a Git checkout with Git LFS" in message + assert command in message + assert forbidden not in message def test_ensure_weight_pulls_checkout_pointer_via_lfs(monkeypatch, tmp_path): @@ -199,30 +247,69 @@ def test_status_missing_for_excluded_checkout_pointer(monkeypatch, tmp_path): status = weights.get_weight_status("birefnet") assert status.state == "missing" assert "excluded" in status.detail.lower() + assert "--only birefnet --source lfs" in status.detail assert {s.id for s in status.sources} == {"hf", "lfs"} -def test_download_from_lfs_success(monkeypatch, tmp_path): +def test_status_rf_detr_pointer_recommends_rf_detr_repair(monkeypatch, tmp_path): weights = _fresh_weights(monkeypatch, tmp_path, MATS_WEIGHTS_DIR=str(tmp_path / "cache")) checkout = tmp_path / "checkout" (checkout / "weights").mkdir(parents=True) (checkout / ".git").mkdir() + (checkout / "weights" / "rf_detr_marker.pth").write_bytes(LFS_POINTER) monkeypatch.setattr(weights, "_REPO_ROOT", checkout) + + status = weights.get_weight_status("rf-detr") + + assert status.state == "missing" + assert "git lfs install" in status.detail + assert "--only rf-detr --source lfs" in status.detail + assert "--only birefnet" not in status.detail + + +@pytest.mark.parametrize( + ("name", "filename", "expected_args"), + ( + ( + "rf-detr", + "rf_detr_marker.pth", + ["git", "lfs", "pull", "--exclude", "weights/birefnet_leaf.pth"], + ), + ( + "birefnet", + "birefnet_leaf.pth", + ["git", "lfs", "pull", "-X", "", "-I", "weights/birefnet_leaf.pth"], + ), + ), +) +def test_download_from_lfs_success(monkeypatch, tmp_path, name, filename, expected_args): + weights = _fresh_weights(monkeypatch, tmp_path, MATS_WEIGHTS_DIR=str(tmp_path / "cache")) + checkout = tmp_path / "checkout" + (checkout / "weights").mkdir(parents=True) + (checkout / ".git").mkdir() + monkeypatch.setattr(weights, "_REPO_ROOT", checkout) + monkeypatch.setattr(weights, "_git_lfs_installed", lambda: True) monkeypatch.setattr(weights, "free_bytes", lambda path: 10 ** 12) real_bytes = b"\x80\x02" + b"x" * 4094 digest = hashlib.sha256(real_bytes).hexdigest() - monkeypatch.setitem(weights._MANIFEST["birefnet"], "size_bytes", len(real_bytes)) - monkeypatch.setitem(weights._MANIFEST["birefnet"], "sha256", digest) + monkeypatch.setitem(weights._MANIFEST[name], "size_bytes", len(real_bytes)) + monkeypatch.setitem(weights._MANIFEST[name], "sha256", digest) + calls = [] - def fake_popen(*args, **kwargs): - (checkout / "weights" / "birefnet_leaf.pth").write_bytes(real_bytes) + def fake_popen(args, **kwargs): + calls.append((args, kwargs)) + (checkout / "weights" / filename).write_bytes(real_bytes) return _FakeProcess(0) monkeypatch.setattr(weights.subprocess, "Popen", fake_popen) - assert weights._download_from_lfs("birefnet") is True - assert (checkout / "weights" / "birefnet_leaf.pth").read_bytes() == real_bytes + assert weights._download_from_lfs(name) is True + assert (checkout / "weights" / filename).read_bytes() == real_bytes + assert len(calls) == 1 + args, kwargs = calls[0] + assert args == expected_args + assert kwargs["cwd"] == checkout def test_download_from_lfs_failure(monkeypatch, tmp_path): @@ -231,6 +318,7 @@ def test_download_from_lfs_failure(monkeypatch, tmp_path): (checkout / "weights").mkdir(parents=True) (checkout / ".git").mkdir() monkeypatch.setattr(weights, "_REPO_ROOT", checkout) + monkeypatch.setattr(weights, "_git_lfs_installed", lambda: True) monkeypatch.setattr(weights, "free_bytes", lambda path: 10 ** 12) monkeypatch.setattr(weights.subprocess, "Popen", lambda *a, **k: _FakeProcess(1))