support more formats

_utils/image_io.py

@@ -32,6 +32,31 @@ _CHANNEL = ("C", "S")      # C = separate channel planes, S = interleaved RGB sa
 _UNKNOWN = ("Q", "I")      # file named no axis; only these may be *guessed* as channels
 
 
+TIFF_EXTENSIONS = (
+    sorted("." + e for e in tifffile.TIFF.FILE_EXTENSIONS if "." not in e)
+    if tifffile is not None else [".tif", ".tiff"]
+)
+
+
+def _read_tiff(path):
+    """Read a TIFF-family file into (array, axes). Raises if not readable."""
+    with tifffile.TiffFile(path) as tif:
+        series = tif.series[0]
+        page = tif.pages[0]
+        arr = np.asarray(series.asarray())
+        axes = str(series.axes)
+
+        is_palette = (getattr(page, "photometric", None) == tifffile.PHOTOMETRIC.PALETTE
+                      and not any(a in _CHANNEL for a in axes))
+        if is_palette and page.colormap is not None:
+            colormap = np.asarray(page.colormap)  # (3, 2**bits), uint16
+            rgb = np.moveaxis(colormap[:, arr], 0, -1)  # (..., H, W, 3)
+            # TIFF colormaps are 16-bit; scale down to 8-bit (65535/255=257).
+            arr = np.round(rgb / 257.0).astype(np.uint8)
+            axes = axes + "S"  # the LUT added an RGB sample axis
+        return arr, axes
+
+
 def _read_array(path):
     """Load an image file into (array, axes) where axes names each dimension.
 
@@ -40,23 +65,12 @@ def _read_array(path):
     "8-bit Color", palette PNG/GIF) are expanded through their color lookup
     table so we return true RGB, not bare indices.
     """
-    ext = os.path.splitext(path)[1].lower()
-    if ext in (".tif", ".tiff") and tifffile is not None:
-        with tifffile.TiffFile(path) as tif:
-            series = tif.series[0]
-            page = tif.pages[0]
-            arr = np.asarray(series.asarray())
-            axes = str(series.axes)
-            is_palette = getattr(page, "photometric", None) == tifffile.PHOTOMETRIC.PALETTE
-            if is_palette and page.colormap is not None:
-                colormap = np.asarray(page.colormap)  # (3, 2**bits), uint16
-                rgb = np.moveaxis(colormap[:, arr], 0, -1)  # (..., H, W, 3)
-                # TIFF colormaps are 16-bit; scale down to 8-bit (65535/255=257).
-                arr = np.round(rgb / 257.0).astype(np.uint8)
-                axes = axes + "S"  # the LUT added an RGB sample axis
-            return arr, axes
-    # Everything else (and if tifffile is unavailable): PIL. Expand palette
-    # images to RGB so the LUT is applied instead of returning bare indices.
+    if tifffile is not None:
+        try:
+            return _read_tiff(path)
+        except Exception:  # noqa: BLE001 - not a TIFF, or tifffile cannot parse it
+            pass
+
     img = Image.open(path)
     if img.mode in ("P", "PA"):
         img = img.convert("RGB")
@@ -66,8 +80,7 @@ def _read_array(path):
 
 def _shape_axes(path):
     """Return (shape, axes) from metadata only - no pixel decode."""
-    ext = os.path.splitext(path)[1].lower()
-    if ext in (".tif", ".tiff") and tifffile is not None:
+    if tifffile is not None:
         try:
             with tifffile.TiffFile(path) as tif:
                 series = tif.series[0]
@@ -114,12 +127,17 @@ def _plan_axes(shape, axes):
     return shape, axes, guessed
 
 
-def _reduce_to_hwc(arr, axes, frame=0):
+def _reduce_to_hwc(arr, axes, frame=0, sub_frame=None):
     """Collapse a labelled array to 2D (H, W) or 3D (H, W, C).
 
     The channel axis (named by the file, or inferred by ``_plan_axes`` only when
-    the file names none) is moved last and kept whole. The remaining stack axes
-    (Z / time / page) are indexed: the first by ``frame``, any deeper ones by 0.
+    the file names none) is moved last and kept whole. Stack axes (Z / time /
+    page) are handled in order:
+      * the first is indexed by ``frame`` (the frame slider);
+      * the second is indexed by ``sub_frame`` if given, else collapsed by a
+        maximum-intensity projection (the natural default for a focal stack -
+        keeps the brightest signal across planes rather than an arbitrary one);
+      * any deeper ones are projected too.
     """
     axes = list(axes)
 
@@ -141,27 +159,45 @@ def _reduce_to_hwc(arr, axes, frame=0):
         axes.append(axes.pop(ci))
 
     target = 3 if ci is not None else 2
-    first = True
+    stack = 0
     while len(axes) > target:
-        idx = int(frame) if first else 0
-        idx = max(0, min(idx, arr.shape[0] - 1))
-        arr = arr[idx]
+        if stack == 0:
+            idx = max(0, min(int(frame), arr.shape[0] - 1))
+            arr = arr[idx]
+        elif stack == 1 and sub_frame is not None:
+            idx = max(0, min(int(sub_frame), arr.shape[0] - 1))
+            arr = arr[idx]
+        else:
+            arr = arr.max(axis=0)  # maximum-intensity projection over this axis
         axes.pop(0)
-        first = False
+        stack += 1
 
     return arr
 
 
+# tifffile axis letters -> words a microscopist uses, for messages and slider
+# labels. Anything unmapped (an unnamed 'Q'/'I' axis, a raw page axis) is just a
+# "frame".
+_AXIS_LABELS = {"T": "timepoint", "Z": "z-plane"}
+
+
+def _axis_label(a):
+    return _AXIS_LABELS.get(a, "frame")
+
+
 class ImageInfo(NamedTuple):
     """How a file's dimensions were interpreted (read from metadata only).
 
-    frames   - length of the first stack (Z/T) axis, 1 if not a stack
-    channels - length of the channel axis, 1 if single-channel
-    axes     - the file's own axes string, e.g. 'CZYX' ('Q' = unnamed)
-    guessed  - True if the channel axis was inferred rather than read
-    shape    - the file's raw shape as stored, e.g. (1, 4, 1, 1024, 1024)
-    width    - pixels along the X axis (0 if unknown)
-    height   - pixels along the Y axis (0 if unknown)
+    frames    - length of the first stack (Z/T) axis, 1 if not a stack
+    channels  - length of the channel axis, 1 if single-channel
+    axes      - the file's own axes string, e.g. 'CZYX' ('Q' = unnamed)
+    guessed   - True if the channel axis was inferred rather than read
+    shape     - the file's raw shape as stored, e.g. (1, 4, 1, 1024, 1024)
+    width     - pixels along the X axis (0 if unknown)
+    height    - pixels along the Y axis (0 if unknown)
+    sub_frames- length of a *second* stack axis (e.g. Z in a time+Z file), 1 if
+                there is only one stack axis
+    frame_label / sub_label - the words for those two axes ('timepoint', ...)
     """
     frames: int
     channels: int
@@ -170,6 +206,9 @@ class ImageInfo(NamedTuple):
     shape: tuple
     width: int
     height: int
+    sub_frames: int = 1
+    frame_label: str = "frame"
+    sub_label: str = "frame"
 
 
 def inspect_image(path):
@@ -179,8 +218,13 @@ def inspect_image(path):
         planned_shape, planned_axes, guessed = _plan_axes(shape, axes)
         has_c = bool(planned_axes) and planned_axes[-1] in _CHANNEL
         channels = int(planned_shape[-1]) if has_c else 1
-        target = 3 if has_c else 2
-        frames = int(planned_shape[0]) if len(planned_axes) > target else 1
+
+        stack = [(a, d) for a, d in zip(planned_axes, planned_shape)
+                 if a not in ("Y", "X") and a not in _CHANNEL]
+        frames = int(stack[0][1]) if stack else 1
+        frame_label = _axis_label(stack[0][0]) if stack else "frame"
+        sub_frames = int(stack[1][1]) if len(stack) > 1 else 1
+        sub_label = _axis_label(stack[1][0]) if len(stack) > 1 else "frame"
 
         # Spatial size comes from the named Y/X axes, so it stays correct
         # whatever order the other axes are in.
@@ -193,7 +237,8 @@ def inspect_image(path):
         if not (width and height):  # no Y/X named; fall back to the header probe
             width, height = image_size(path)
 
-        return ImageInfo(frames, channels, str(axes), guessed, tuple(shape), width, height)
+        return ImageInfo(frames, channels, str(axes), guessed, tuple(shape),
+                         width, height, sub_frames, frame_label, sub_label)
     except Exception:  # noqa: BLE001
         return ImageInfo(1, 1, "", False, (), 0, 0)
 
@@ -209,8 +254,7 @@ def array_nbytes(path):
     Computed from shape + dtype in the header - no pixel decode - so it is safe
     to call on a file that is too large to open. Returns 0 if unknown.
     """
-    ext = os.path.splitext(path)[1].lower()
-    if ext in (".tif", ".tiff") and tifffile is not None:
+    if tifffile is not None:
         try:
             with tifffile.TiffFile(path) as tif:
                 series = tif.series[0]
@@ -232,8 +276,7 @@ def image_size(path):
     that would be too large to load. Returns (0, 0) if the size cannot be
     determined, so callers treat it as "unknown" and proceed.
     """
-    ext = os.path.splitext(path)[1].lower()
-    if ext in (".tif", ".tiff") and tifffile is not None:
+    if tifffile is not None:
         try:
             with tifffile.TiffFile(path) as tif:
                 page = tif.pages[0]
@@ -261,10 +304,10 @@ def _to_rgb(arr):
     return arr[:, :, :3]
 
 
-def _load_rgb(path, frame=0):
+def _load_rgb(path, frame=0, sub_frame=None):
     """Read a file and reduce it to an (H, W, 3) array plus its source dtype."""
     raw, axes = _read_array(path)
-    arr = _reduce_to_hwc(raw, axes, frame=frame)
+    arr = _reduce_to_hwc(raw, axes, frame=frame, sub_frame=sub_frame)
     arr = _to_rgb(arr)
     return arr, raw.dtype
 
@@ -306,7 +349,7 @@ def _save_png(arr, out_path=None, out_dir=None, base=None, suffix="_std.png"):
     return out_path
 
 
-def standardize_image(path, frame=0, out_dir=None, out_path=None):
+def standardize_image(path, frame=0, sub_frame=None, out_dir=None, out_path=None):
     """Standardize any image/TIFF to a canonical 8-bit RGB PNG.
 
     Used for both the preview and the model: one frame, <=3 channels, with a
@@ -315,7 +358,11 @@ def standardize_image(path, frame=0, out_dir=None, out_path=None):
 
     Args:
         path: path to the input image (TIFF, PNG, JPG, ...).
-        frame: which frame of a stack to use (0-based; ignored for non-stacks).
+        frame: which frame of the first stack axis to use (0-based; ignored for
+            non-stacks).
+        sub_frame: which plane of a second stack axis (e.g. Z in a time+Z file)
+            to use (0-based). None means combine those planes by a maximum-
+            intensity projection.
         out_dir: optional directory for the output PNG.
         out_path: optional explicit output file path (overrides out_dir).
 
@@ -324,7 +371,7 @@ def standardize_image(path, frame=0, out_dir=None, out_path=None):
         returned unchanged so callers degrade gracefully.
     """
     try:
-        arr, dtype = _load_rgb(path, frame=frame)
+        arr, dtype = _load_rgb(path, frame=frame, sub_frame=sub_frame)
         arr = _to_uint8(arr, stretch=(dtype != np.uint8))
         base = os.path.splitext(os.path.basename(path))[0]
         return _save_png(arr, out_path=out_path, out_dir=out_dir, base=base, suffix="_std.png")

app.py

@@ -23,7 +23,7 @@ from inference_count import load_model as load_count_model, run as run_count
 from inference_track import load_model as load_track_model, run as run_track
 from _utils.image_io import (
     standardize_image, inspect_image, image_size, array_nbytes,
-    RECOMMENDED_SIZE, WARN_SIZE, MAX_SIZE, MAX_READ_BYTES,
+    RECOMMENDED_SIZE, WARN_SIZE, MAX_SIZE, MAX_READ_BYTES, TIFF_EXTENSIONS,
 )
 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -408,6 +408,8 @@ def _describe_read(info):
 
     if info.frames > 1:
         lines.append(f"Showing first frame by default (you can use the stack frame slider to pick another).")
+        if info.sub_frames > 1:
+            lines.append(f"Showing the first of {info.sub_frames} {info.sub_label}s by default (use the {info.sub_label} slider to pick another).")
 
     lines.append(
         "If wrong, please convert your image to a 8-bit RGB/Grayscale image file before uploading (there are available tools such as ImageJ/Fiji)."
@@ -417,26 +419,30 @@ def _describe_read(info):
 
 def prepare_uploaded_image(annot_value):
     """On annotator upload: standardize frame 0 for preview and configure the
-    stack frame slider.
+    stack sliders.
 
     Browsers cannot render 16-bit / float / multi-page TIFFs, so we standardize
     to an 8-bit RGB PNG. If the file is a stack, notify the user and reveal a
-    frame slider (default frame 1); otherwise keep it hidden.
+    frame slider (default frame 1); a file with a second stack axis (e.g. Z in a
+    time+Z series) also gets a z-plane slider (0 = max-intensity projection).
+    Sliders that don't apply stay hidden.
 
-    Returns (annotator_value, raw_path, frame_slider_update).
+    Returns (annotator_value, raw_path, frame_slider_update, zplane_slider_update).
     """
+    hidden = gr.update(visible=False, value=1)
     img_path = _annot_path(annot_value)
     if not img_path:
-        return annot_value, None, gr.update(visible=False, value=1)
+        return annot_value, None, hidden, hidden
 
     # Clear the annotator on refusal, so a rejected file cannot reach inference.
     rejected = check_image(img_path)
     if rejected:
         gr.Warning(rejected, duration=None, title="❌ Cannot use this file")
-        return None, None, gr.update(visible=False, value=1)
+        return None, None, hidden, hidden
 
     info = inspect_image(img_path)
-    display = standardize_image(img_path, frame=0)
+    display = standardize_image(img_path, frame=0,
+                                sub_frame=0 if info.sub_frames > 1 else None)
 
     # Stay quiet unless something non-obvious happened; a plain RGB or grayscale
     # image needs no explanation.
@@ -444,15 +450,19 @@ def prepare_uploaded_image(annot_value):
         gr.Info(_describe_read(info), duration=None, title="📚 Image Loading Info")
 
     slider = (gr.update(visible=True, maximum=info.frames, value=1) if info.frames > 1
-              else gr.update(visible=False, value=1))
-    return display, img_path, slider
+              else hidden)
+    zslider = (gr.update(visible=True, maximum=info.sub_frames, value=1,
+                         label=f"🔬 {info.sub_label.capitalize()} (choose plane to use)")
+               if info.sub_frames > 1 else hidden)
+    return display, img_path, slider, zslider
 
 
-def select_frame(raw_path, frame_num):
-    """Re-render the annotator preview for the chosen stack frame (1-based)."""
+def select_frame(raw_path, frame_num, zplane=1):
+    """Re-render the annotator preview for the chosen frame and z-plane (both
+    1-based; z-plane is ignored for files with only one stack axis)."""
     if not raw_path:
         return gr.update()
-    return standardize_image(raw_path, frame=int(frame_num) - 1)
+    return standardize_image(raw_path, frame=int(frame_num) - 1, sub_frame=int(zplane) - 1)
 
 
 @spaces.GPU
@@ -1470,6 +1480,11 @@ with gr.Blocks(
                         label="📚 Stack frame (choose frame to use)",
                         visible=False,
                     )
+                    seg_zplane_slider = gr.Slider(
+                        minimum=1, maximum=1, step=1, value=1,
+                        label="🔬 Z-plane (choose plane to use)",
+                        visible=False,
+                    )
 
                     # Example Images Gallery
                     gr.Markdown("📁 Example Image Gallery", elem_classes="frame-heading")
@@ -1497,7 +1512,7 @@ with gr.Blocks(
                     gr.Markdown("➕ Upload New Example Image to Gallery", elem_classes="frame-heading")
                     image_uploader = gr.File(
                         show_label=False,
-                        file_types=["image", ".tif", ".tiff"],
+                        file_types=["image"] + list(TIFF_EXTENSIONS),
                         type="filepath"
                     )
                     add_to_gallery_btn = gr.Button("➕ Add to Gallery", variant="secondary", size="sm")
@@ -1555,13 +1570,14 @@ with gr.Blocks(
             annotator.upload(
                 fn=prepare_uploaded_image,
                 inputs=annotator,
-                outputs=[annotator, seg_raw_state, seg_frame_slider]
-            )
-            seg_frame_slider.release(
-                fn=select_frame,
-                inputs=[seg_raw_state, seg_frame_slider],
-                outputs=annotator
+                outputs=[annotator, seg_raw_state, seg_frame_slider, seg_zplane_slider]
             )
+            for _slider in (seg_frame_slider, seg_zplane_slider):
+                _slider.release(
+                    fn=select_frame,
+                    inputs=[seg_raw_state, seg_frame_slider, seg_zplane_slider],
+                    outputs=annotator
+                )
 
             # click event for segmentation
             run_seg_btn.click(
@@ -1577,9 +1593,10 @@ with gr.Blocks(
 
             # click event for clear button
             clear_btn.click(
-                fn=lambda: (None, {}, None, gr.update(visible=False, value=1)),
+                fn=lambda: (None, {}, None, gr.update(visible=False, value=1),
+                            gr.update(visible=False, value=1)),
                 inputs=None,
-                outputs=[annotator, seg_vis_state, seg_raw_state, seg_frame_slider]
+                outputs=[annotator, seg_vis_state, seg_raw_state, seg_frame_slider, seg_zplane_slider]
             )
             
             # init Gallery with example images
@@ -1621,12 +1638,12 @@ with gr.Blocks(
                 if evt.index is not None and evt.index < len(all_imgs):
                     item = all_imgs[evt.index]
                     return prepare_uploaded_image(_GALLERY_RAW_SOURCE.get(item, item))
-                return None, None, gr.update(visible=False, value=1)
+                return None, None, gr.update(visible=False, value=1), gr.update(visible=False, value=1)
 
             example_gallery.select(
                 fn=load_from_gallery,
                 inputs=user_uploaded_examples,
-                outputs=[annotator, seg_raw_state, seg_frame_slider]
+                outputs=[annotator, seg_raw_state, seg_frame_slider, seg_zplane_slider]
             )
             
             # click event for submitting feedback
@@ -1704,6 +1721,11 @@ with gr.Blocks(
                         label="📚 Stack frame (choose frame to use)",
                         visible=False,
                     )
+                    count_zplane_slider = gr.Slider(
+                        minimum=1, maximum=1, step=1, value=1,
+                        label="🔬 Z-plane (choose plane to use)",
+                        visible=False,
+                    )
 
                     # Example gallery with "add" functionality
                     gr.Markdown("📁 Example Image Gallery", elem_classes="frame-heading")
@@ -1734,7 +1756,7 @@ with gr.Blocks(
                     gr.Markdown("➕ Add Example Image to Gallery", elem_classes="frame-heading")
                     count_image_uploader = gr.File(
                         show_label=False,
-                        file_types=["image"],
+                        file_types=["image"] + list(TIFF_EXTENSIONS),
                         type="filepath"
                     )
                     add_to_count_gallery_btn = gr.Button("➕ Add to Gallery", variant="secondary", size="sm")
@@ -1828,13 +1850,14 @@ with gr.Blocks(
             count_annotator.upload(
                 fn=prepare_uploaded_image,
                 inputs=count_annotator,
-                outputs=[count_annotator, count_raw_state, count_frame_slider]
-            )
-            count_frame_slider.release(
-                fn=select_frame,
-                inputs=[count_raw_state, count_frame_slider],
-                outputs=count_annotator
+                outputs=[count_annotator, count_raw_state, count_frame_slider, count_zplane_slider]
             )
+            for _slider in (count_frame_slider, count_zplane_slider):
+                _slider.release(
+                    fn=select_frame,
+                    inputs=[count_raw_state, count_frame_slider, count_zplane_slider],
+                    outputs=count_annotator
+                )
 
             # When user selects from gallery, load into annotator
             def load_from_count_gallery(evt: gr.SelectData, all_imgs):
@@ -1843,12 +1866,12 @@ with gr.Blocks(
                     selected_img = all_imgs[evt.index]
                     print(f"📸 Loading image from gallery: {selected_img}")
                     return prepare_uploaded_image(_GALLERY_RAW_SOURCE.get(selected_img, selected_img))
-                return None, None, gr.update(visible=False, value=1)
+                return None, None, gr.update(visible=False, value=1), gr.update(visible=False, value=1)
 
             count_example_gallery.select(
                 fn=load_from_count_gallery,
                 inputs=count_user_examples,
-                outputs=[count_annotator, count_raw_state, count_frame_slider]
+                outputs=[count_annotator, count_raw_state, count_frame_slider, count_zplane_slider]
             )
 
             # Run counting
@@ -1865,9 +1888,10 @@ with gr.Blocks(
 
             # Clear selection
             clear_btn.click(
-                fn=lambda: (None, {}, None, gr.update(visible=False, value=1)),
+                fn=lambda: (None, {}, None, gr.update(visible=False, value=1),
+                            gr.update(visible=False, value=1)),
                 inputs=None,
-                outputs=[count_annotator, count_vis_state, count_raw_state, count_frame_slider]
+                outputs=[count_annotator, count_vis_state, count_raw_state, count_frame_slider, count_zplane_slider]
             )
 
             # Submit feedback