Write the paint logic once: Python that both the server and the browser run

cellmap_flow/ai_annotate/backends/vertex_gemini.py

@@ -33,7 +33,7 @@
 from cellmap_flow.ai_annotate.backends.base import SegmentRequest, SegmentResult
 from cellmap_flow.ai_annotate.errors import INSTALL_HINT, AIAnnotateError
 from cellmap_flow.ai_annotate.images import decode_model_image
-from cellmap_flow.ai_annotate.mask_extraction import extract_mask
+from cellmap_flow.portable.ai_annotate import model_mask
 
 logger = logging.getLogger(__name__)
 
@@ -128,11 +128,11 @@ def segment(self, request: SegmentRequest, model: str) -> SegmentResult:
         model_image = decode_model_image(image_bytes)
         if model_image.size != sent.size:
             _warn_if_aspect_differs(model_image.size, sent.size)
-            model_image = model_image.resize(sent.size, Image.BILINEAR)
-        # Drop the padding _square added: what is left lines up with the plane.
-        model_image = model_image.crop((0, 0, *request.image.size))
-        mask = extract_mask(model_image, request.target_rgb)
-        return SegmentResult(mask=mask, model_image=model_image, model=model, usage=_usage(response))
+        # Back to the sent size, the padding _square added dropped, and the
+        # target colour's pixels taken: what is left lines up with the plane
+        # (model_mask, which the browser's provider calls too).
+        mask, aligned = model_mask(np.asarray(model_image), sent.size[::-1], request.image.size[::-1], request.target_rgb)
+        return SegmentResult(mask=mask, model_image=Image.fromarray(aligned), model=model, usage=_usage(response))
 
     def _generate(self, client, model, contents, config):
         """``generate_content``, retried on 429/503 only, at most ``MAX_ATTEMPTS`` times."""

cellmap_flow/ai_annotate/mask_extraction.py

@@ -3,29 +3,20 @@
 Adapted from ask-to-mask (``_normalize_to_uint8``/``_slice_to_rgb`` and
 ``extract_mask``'s colour-difference mode): the model is asked to paint the
 structure one colour and everything else black, and the mask is the pixels
-that came back that colour.
+that came back that colour. The arithmetic is
+``cellmap_flow.portable.ai_annotate``'s (``normalize_to_uint8``,
+``extract_mask``; ``model_mask`` lines a reply up with the plane first),
+which the browser runs too; here are the Pillow images the server's
+backends and staging pass around.
 """
 
 from __future__ import annotations
 
 import numpy as np
 from PIL import Image
 
-
-def normalize_to_uint8(data: np.ndarray) -> np.ndarray:
-    """A 2D array as uint8 grey levels.
-
-    uint8 passes through unchanged. Anything else is clipped to its 1st-99th
-    percentile over the plane and stretched to 0-255, as raw EM is not always
-    8-bit and a few saturated pixels should not flatten the contrast.
-    """
-    if data.dtype == np.uint8:
-        return data
-    p_lo, p_hi = np.percentile(data, (1, 99))
-    if p_hi <= p_lo:
-        p_hi = p_lo + 1
-    clipped = np.clip(data.astype(np.float32), p_lo, p_hi)
-    return ((clipped - p_lo) / (p_hi - p_lo) * 255).astype(np.uint8)
+from cellmap_flow.portable import ai_annotate as portable
+from cellmap_flow.portable.ai_annotate import normalize_to_uint8  # noqa: F401  (re-exported)
 
 
 def slice_to_rgb(data_2d: np.ndarray) -> Image.Image:
@@ -35,37 +26,12 @@ def slice_to_rgb(data_2d: np.ndarray) -> Image.Image:
 
 
 def extract_mask(
-    output_image: Image.Image,
+    output_image,
     target_rgb: tuple[int, int, int],
     threshold: float = 200.0,
 ) -> np.ndarray:
-    """The pixels of the model's reply painted ``target_rgb``, as a bool (H, W) mask.
-
-    Each pixel scores how strongly it shows the target colour: its weakest
-    channel among those the target has at full strength, minus its strongest
-    deviation elsewhere -- a channel the target has off (0) being on, or a
-    channel the target has at part strength being away from that strength.
-    For red (255, 0, 0) that is ``R - max(G, B)``: red that dominates. A
-    pixel counts when its score reaches ``threshold``.
-
-    The part-strength term is what lets orange (255, 128, 0) or purple
-    (128, 0, 255) be found at all: counting their half channel as "on" would
-    cap the score at 128, under the threshold, and counting it as "off"
-    would reject the colour itself.
-    """
+    """The pixels of the model's reply (a Pillow image or an RGB array)
+    painted ``target_rgb``, as a bool (H, W) mask
+    (``portable.ai_annotate.extract_mask`` says how a pixel is scored)."""
     out = np.asarray(output_image.convert("RGB") if isinstance(output_image, Image.Image) else output_image)
-    out = out.astype(np.float32)
-    target = np.asarray(target_rgb, dtype=np.float32)
-
-    full = target >= 255
-    off = target <= 0
-    part = ~full & ~off
-
-    strength = np.min(out[:, :, full], axis=-1) if full.any() else np.full(out.shape[:2], 255.0, np.float32)
-    deviations = []
-    if off.any():
-        deviations.append(np.max(out[:, :, off], axis=-1))
-    if part.any():
-        deviations.append(np.max(np.abs(out[:, :, part] - target[part]), axis=-1))
-    score = strength - np.max(deviations, axis=0) if deviations else strength
-    return score >= threshold
+    return portable.extract_mask(out, target_rgb, threshold)

cellmap_flow/ai_annotate/pipeline.py

@@ -4,31 +4,24 @@
 grid it planned and laid out as on screen; ``build_request`` makes it the
 request a backend sends; ``mask_to_write_shape`` cuts the mask the backend
 returns into its connected objects and puts them onto the write box's
-voxels, and ``labels_for_box`` turns those into the labels ``fill.paint_box``
-writes: one id per object, background elsewhere, as a seed from a model's
-prediction is labelled.
+voxels, and ``labels_for_box`` (``accept_labels`` for Accept) turns those
+into the labels ``fill.paint_box`` writes: one id per object, background
+elsewhere, as a seed from a model's prediction is labelled.
+
+The maths is ``cellmap_flow.portable.ai_annotate``'s, which the browser runs
+too (its resampling is Pillow's, in both modes); what is the server's alone
+is reading the raw through ``ImageDataInterface``.
 """
 
 import numpy as np
-from PIL import Image
-
-from cellmap_flow.ai_annotate.geometry import PlanePlan, raw_box, screen_shape
-
 
-def _resize(plane: np.ndarray, shape) -> np.ndarray:
-    """``plane`` resampled to ``shape`` (rows, cols); itself when it already is that.
-
-    Shrinking averages each pixel's footprint (PIL's box filter), so fine
-    membranes fade rather than alias; growing (the coarse axis of an
-    anisotropic plane) interpolates linearly.
-    """
-    rows, cols = (int(s) for s in shape)
-    if plane.shape == (rows, cols):
-        return plane
-    shrinking = plane.shape[0] >= rows and plane.shape[1] >= cols
-    image = Image.fromarray(np.ascontiguousarray(plane, dtype=np.float32))
-    resized = image.resize((cols, rows), Image.BOX if shrinking else Image.BILINEAR)
-    return np.asarray(resized, dtype=np.float32)
+from cellmap_flow.ai_annotate.geometry import PlanePlan, raw_box
+from cellmap_flow.portable.ai_annotate import (  # noqa: F401  (re-exported for the routes)
+    accept_labels,
+    labels_for_box,
+    mask_to_write_shape,
+    plane_image,
+)
 
 
 def _raw(volume: dict, voxel_size):
@@ -53,7 +46,7 @@ def raw_extent_nm(volume: dict):
 
 def read_plane(volume: dict, plan: PlanePlan) -> np.ndarray:
     """The raw plane ``plan`` reads, as a 2D array of ``plan.image_shape``,
-    laid out as on screen (``PlanePlan.to_screen``).
+    laid out as on screen (``portable.ai_annotate.plane_image``).
 
     Read as stored (no input normalization: the image is scaled to 8 bits
     for the model on its own) at the plan's input voxel size, from the level
@@ -64,9 +57,7 @@ def read_plane(volume: dict, plan: PlanePlan) -> np.ndarray:
     """
     idi = _raw(volume, plan.input_voxel_size)
     box = raw_box(plan, idi.voxel_size, idi.offset)
-    data = np.asarray(idi.read_box(box))
-    plane = np.take(data, 0, axis=plan.depth_axis)
-    return _resize(plan.to_screen(plane), plan.image_shape)
+    return plane_image(np.asarray(idi.read_box(box)), plan)
 
 
 def build_request(plane: np.ndarray, plan: PlanePlan, profile, prompt_override=None):
@@ -89,81 +80,3 @@ def build_request(plane: np.ndarray, plan: PlanePlan, profile, prompt_override=N
         target_rgb=tuple(int(c) for c in profile.rgb),
         click_px=plan.click_px,
     )
-
-
-def mask_to_write_shape(mask: np.ndarray, plan: PlanePlan) -> np.ndarray:
-    """The objects of ``mask`` (the image's shape, as on screen) on the write
-    box's in-plane voxels, in z, y, x order (``PlanePlan.write_shape``):
-    int32, 0 for background and 1..n for the objects, numbered in the box's
-    raster order as ``scipy.ndimage.label`` numbers a mask.
-
-    Objects are the mask's 4-connected pieces (scipy's default, so pieces
-    touching only at a corner stay apart, as the model was asked to leave
-    them), found at the image's own resolution, where the model drew the
-    lines between them: on voxels coarser than the pixels a line can fall
-    between the pixels the voxels take, and the objects either side of it
-    still keep ids of their own. Then nearest neighbour, pixel centre to
-    voxel centre: each annotation voxel takes the object of the pixel at
-    its centre, so the mask keeps hard edges and lands where it was drawn
-    whatever the two sizes are. An object no voxel centre falls in is lost.
-    """
-    from scipy.ndimage import label
-
-    mask = np.asarray(mask, dtype=bool)
-    if mask.ndim != 2:
-        raise ValueError(f"The mask must be 2D, got shape {mask.shape}")
-    objects, _ = label(mask)
-    rows, cols = screen_shape(plan.write_shape, plan.depth_axis, plan.display_dimensions)
-    if objects.shape != (rows, cols):
-        objects = np.asarray(Image.fromarray(objects.astype(np.int32)).resize((cols, rows), Image.NEAREST))
-    objects = plan.from_screen(objects)
-    # Renumbered by first voxel, so fresh ids go to the objects in the order
-    # labelling the box itself would give them, whatever the screen layout.
-    ids, first = np.unique(objects, return_index=True)
-    order = ids[ids > 0][np.argsort(first[ids > 0])]
-    number = np.zeros(int(objects.max(initial=0)) + 1, dtype=np.int32)
-    number[order] = np.arange(1, order.size + 1)
-    return number[objects]
-
-
-def labels_for_box(objects: np.ndarray, existing: np.ndarray, depth_axis: int, count_up: bool,
-                   beside=None) -> np.ndarray:
-    """The labels for the write box: an id per object of ``objects``, 1 elsewhere.
-
-    ``objects`` is the in-plane objects (``mask_to_write_shape``), 0 for
-    background, ``existing`` the one-voxel-thick box as the volume holds
-    it. The objects are labelled by ``fill.seed_labels``: an object the user
-    has already painted part of keeps that id, the rest get fresh ones,
-    counting up for an instance target's ``count_up``.
-
-    ``beside`` is the planes on either side of the box along the depth axis,
-    as the volume holds them (each shaped like ``existing``, or None at the
-    volume's edge). Planes are annotated one at a time, so an object usually
-    continues one in the plane next to it: where this plane is unpainted,
-    the ids beside it count as painted here, and an object takes the id it
-    overlaps most, rather than a fresh one that an instance target would
-    read as a boundary between the two planes. Fresh ids skip those planes'
-    ids too, so a new object does not take the id of a different object
-    right next to it, which an instance target would read as one object.
-    """
-    from cellmap_flow.finetune.session.fill import seed_labels
-
-    objects = np.expand_dims(np.asarray(objects), depth_axis)
-    if objects.shape != existing.shape:
-        raise ValueError(f"The mask makes a box of {objects.shape}; the box holds {existing.shape}")
-    planes = [np.asarray(p, dtype=existing.dtype) for p in (beside or ()) if p is not None]
-    if any(p.shape != existing.shape for p in planes):
-        raise ValueError(f"The planes beside the box must be shaped {existing.shape}")
-    if not planes:
-        return seed_labels(objects, existing, count_up=count_up)
-    seen = existing.copy()
-    for plane in planes:
-        carry = (seen == 0) & (plane >= 2)
-        seen[carry] = plane[carry]
-    # The planes beside are stacked after the box with no objects in them:
-    # seed_labels labels only the box's objects, but its fresh ids see every
-    # id in the stack.
-    stack_existing = np.concatenate([seen, *planes], axis=depth_axis)
-    stack_objects = np.concatenate([objects, *[np.zeros_like(objects)] * len(planes)], axis=depth_axis)
-    labels = seed_labels(stack_objects, stack_existing, count_up=count_up)
-    return np.take(labels, [0], axis=depth_axis)

cellmap_flow/ai_annotate/staging.py

@@ -26,6 +26,7 @@
 from PIL import Image
 
 from cellmap_flow.ai_annotate.errors import AIAnnotateError
+from cellmap_flow.portable import ai_annotate as portable
 
 ANNOTATE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
 STAGING_DIRNAME = ".ai_annotate"
@@ -36,7 +37,7 @@
 META_JSON = "meta.json"
 
 # How strongly the overlay tints the mask with the target colour.
-OVERLAY_ALPHA = 0.5
+OVERLAY_ALPHA = portable.OVERLAY_ALPHA
 
 
 def new_annotate_id() -> str:
@@ -142,14 +143,10 @@ def overlay(input_image: Image.Image, mask_write: np.ndarray, target_rgb) -> Ima
 
     The mask is the write box's, laid out as the image is (``to_screen``)
     and scaled to it by nearest neighbour, so the overlay shows what
-    accepting writes, not the model's own image.
+    accepting writes, not the model's own image
+    (``portable.ai_annotate.overlay``, which the browser's review uses too).
     """
-    rgb = np.asarray(input_image.convert("RGB"), dtype=np.float32)
-    mask = Image.fromarray(np.asarray(mask_write, dtype=np.uint8) * 255).resize(input_image.size, Image.NEAREST)
-    on = np.asarray(mask) > 127
-    tint = np.asarray(target_rgb, dtype=np.float32)
-    rgb[on] = (1 - OVERLAY_ALPHA) * rgb[on] + OVERLAY_ALPHA * tint
-    return Image.fromarray(np.round(rgb).astype(np.uint8))
+    return Image.fromarray(portable.overlay(np.asarray(input_image.convert("RGB")), mask_write, target_rgb, OVERLAY_ALPHA))
 
 
 def preview(corrections_dir, annotate_id) -> dict:

cellmap_flow/dashboard/routes/finetune/ai_annotate.py

@@ -687,7 +687,7 @@ def _resend_job(ai, job, config, segment_request):
 def _planes_beside(state, volume_id, lo, hi, depth_axis):
     """The planes on either side of the box ``[lo, hi)`` along ``depth_axis``, as MinIO serves them.
 
-    For ``pipeline.labels_for_box``; None for a side past the volume's edge.
+    For ``pipeline.accept_labels``; None for a side past the volume's edge.
     Read just before the write, so a stroke painted in them since the plane
     was staged counts.
     """
@@ -742,20 +742,14 @@ def ai_annotate_accept():
         beside = _planes_beside(state, volume_id, lo, hi, plan.depth_axis)
 
         def labels_for(existing):
-            # As a seed: an instance target on a uint16/uint32 volume counts
+            # As a seed (an instance target on a uint16/uint32 volume counts
             # its ids up, a uint8 volume reuses free ones; an object
-            # continuing one in the plane beside keeps its id.
-            labels = pipeline.labels_for_box(
-                mask_write, existing, plan.depth_axis,
-                count_up=instance_target and existing.dtype.itemsize > 1, beside=beside,
-            )
-            if body.overwrite:
-                # Overwrite lets the model's objects replace what is under
-                # them, but its background only fills unlabelled voxels: an
-                # object already painted that the model missed is kept, not
-                # erased to background.
-                labels[(labels == 1) & (existing != 0)] = 0
-            return labels
+            # continuing one in the plane beside keeps its id), and with
+            # overwrite the model's background only filling unlabelled
+            # voxels: portable/ai_annotate.py's accept_labels, which the
+            # browser's Accept runs too.
+            return pipeline.accept_labels(mask_write, existing, plan.depth_axis, instance_target, body.overwrite,
+                                          beside=beside)
 
         n_foreground, n_background, n_overwritten = paint_box(
             state, volume_id, lo, hi, labels_for, volume.get("zarr_path"),
@@ -769,7 +763,7 @@ def labels_for(existing):
             logger.warning(f"Could not accept AI annotation {job['annotate_id']}: {e}")
             return _error("The staged annotation or the served volume is missing; reject it and run again.", 409)
         if isinstance(e, ValueError):
-            # Our own messages (labels_for_box, fill's ids running out), which
+            # Our own messages (accept_labels, fill's ids running out), which
             # say what to do: the seed route answers them the same way.
             logger.warning(f"Could not accept AI annotation {job['annotate_id']}: {e}")
             return _error(str(e), 409)

cellmap_flow/dashboard/routes/finetune/common.py

@@ -12,7 +12,11 @@
 from cellmap_flow.finetune.session.minio import MINIO_PROXY_URL_ENV, proxied_url
 from cellmap_flow.finetune.session.store import SessionStore
 from cellmap_flow.models.geometry import channel_names_of
+# INSTANCE_TARGETS: the targets trained on instances, whose annotation
+# volumes hold an id per object; imported from here by the routes.
+from cellmap_flow.portable.volume import INSTANCE_TARGETS, annotation_dtype  # noqa: F401
 from cellmap_flow.serving.client import fetch_model_info, running_job_host
+from cellmap_flow.portable.label_it import affinity_offsets
 
 logger = logging.getLogger(__name__)
 
@@ -191,25 +195,18 @@ def declared_output_type(model_config):
     return None
 
 
-# The targets trained on instances, whose annotation volumes hold an id per
-# object rather than foreground and background alone.
-INSTANCE_TARGETS = ("affinities", "flows")
-
-
 def annotation_dtype_for(model_config):
-    """The labels' dtype of a new annotation volume for ``model_config``.
-
-    uint8 (ids up to 255) for a model trained on foreground and background;
-    uint16 (up to 65535) for one trained on instances (affinities, Cellpose's
-    flows), whose seeds number each object: 282 of a Cellpose model's masks
-    in one box did not fit a uint8 volume.
-    """
+    """The labels' dtype of a new annotation volume for ``model_config``:
+    ``cellmap_flow.portable.volume.annotation_dtype`` (the browser's rule too)
+    of what it trains on (``autodetect_output_type``), uint16 for the
+    instance targets in INSTANCE_TARGETS, else uint8; uint8 when that cannot
+    be told."""
     try:
         output_type, _ = autodetect_output_type(model_config, None, None)
     except Exception as e:
         logger.info(f"Could not tell what {getattr(model_config, 'name', 'the model')} trains on ({e}): uint8 labels")
         return "uint8"
-    return "uint16" if output_type in INSTANCE_TARGETS else "uint8"
+    return annotation_dtype(output_type)
 
 
 def autodetect_output_type(model_config, output_type, offsets):
@@ -265,13 +262,13 @@ def autodetect_output_type(model_config, output_type, offsets):
                 except Exception as e:
                     logger.debug(f"Could not resolve channels for autodetect: {e}")
 
-            if channels and any("_aff" in channel for channel in channels):
+            # The one rule for a model known by its channel names alone
+            # (portable/label_it.affinity_offsets): the browser's Python
+            # applies it to its manifest's names, for training and seeding,
+            # so both modes treat such a model the same way.
+            default_offsets = affinity_offsets(channels) if channels else None
+            if default_offsets is not None:
                 resolved_output_type = "affinities"
-                n_aff = sum(1 for channel in channels if "_aff" in channel)
-                default_offsets = [
-                    [1 if axis == index else 0 for axis in range(3)]
-                    for index in range(min(n_aff, 3))
-                ]
                 resolved_offsets = json.dumps(default_offsets)
                 logger.info(
                     f"Auto-detected output_type='affinities' from "