micro_sam.sam_annotator.util

  1import os
  2import pickle
  3import warnings
  4import argparse
  5from glob import glob
  6from pathlib import Path
  7from typing import List, Optional, Tuple
  8
  9import h5py
 10import napari
 11import numpy as np
 12from skimage import draw
 13from scipy.ndimage import shift
 14
 15from .. import prompt_based_segmentation, util
 16from .. import _model_settings as model_settings
 17from ..multi_dimensional_segmentation import _validate_projection
 18
 19# Green and Red
 20LABEL_COLOR_CYCLE = ["#00FF00", "#FF0000"]
 21"""@private"""
 22
 23
 24#
 25# Misc helper functions
 26#
 27
 28
 29def toggle_label(prompts):
 30    """@private"""
 31    # get the currently selected label
 32    current_properties = prompts.current_properties
 33    current_label = current_properties["label"][0]
 34    new_label = "negative" if current_label == "positive" else "positive"
 35    current_properties["label"] = np.array([new_label])
 36    prompts.current_properties = current_properties
 37    prompts.refresh()
 38    prompts.refresh_colors()
 39
 40
 41def _initialize_parser(description, with_segmentation_result=True, with_instance_segmentation=True):
 42
 43    available_models = list(util.get_model_names())
 44    available_models = ", ".join(available_models)
 45
 46    parser = argparse.ArgumentParser(description=description)
 47
 48    parser.add_argument(
 49        "-i", "--input", required=True,
 50        help="The filepath to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "
 51        "or elf.io.open_file (e.g. hdf5, zarr, mrc). For the latter you also need to pass the 'key' parameter."
 52    )
 53    parser.add_argument(
 54        "-k", "--key",
 55        help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
 56        "for a image series it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
 57    )
 58    parser.add_argument(
 59        "-e", "--embedding_path",
 60        help="The filepath for saving/loading the pre-computed image embeddings. "
 61        "It is recommended to pass this argument and store the embeddings if you want to open the annotator "
 62        "multiple times for this image. Otherwise the embeddings will be recomputed every time."
 63    )
 64
 65    if with_segmentation_result:
 66        parser.add_argument(
 67            "-s", "--segmentation_result",
 68            help="Optional filepath to a precomputed segmentation. If passed this will be used to initialize the "
 69            "'committed_objects' layer. This can be useful if you want to correct an existing segmentation or if you "
 70            "have saved intermediate results from the annotator and want to continue with your annotations. "
 71            "Supports the same file formats as 'input'."
 72        )
 73        parser.add_argument(
 74            "-sk", "--segmentation_key",
 75            help="The key for opening the segmentation data. Same rules as for 'key' apply."
 76        )
 77
 78    parser.add_argument(
 79        "-m", "--model_type", default=util._DEFAULT_MODEL,
 80        help=f"The segment anything model that will be used, one of {available_models}."
 81    )
 82    parser.add_argument(
 83        "-c", "--checkpoint", default=None,
 84        help="Checkpoint from which the SAM model will be loaded."
 85    )
 86    parser.add_argument(
 87        "--decoder_path", default=None,
 88        help="Optional checkpoint path to decoder-only weights to enable decoder-based instance segmentation."
 89    )
 90    parser.add_argument(
 91        "-d", "--device", default=None,
 92        help="The device to use for the predictor. Can be one of 'cuda', 'cpu' or 'mps' (only MAC)."
 93        "By default the most performant available device will be selected."
 94    )
 95    parser.add_argument(
 96        "--tile_shape", nargs="+", type=int, help="The tile shape for using tiled prediction", default=None
 97    )
 98    parser.add_argument(
 99        "--halo", nargs="+", type=int, help="The halo for using tiled prediction", default=None
100    )
101
102    if with_instance_segmentation:
103        parser.add_argument(
104            "--precompute_amg_state", action="store_true",
105            help="Whether to precompute the state for automatic instance segmentation. "
106            "This will lead to a longer start-up time, but the automatic instance segmentation can "
107            "be run directly once the tool has started."
108        )
109        parser.add_argument(
110            "--prefer_decoder", action="store_false",
111            help="Whether to use decoder based instance segmentation if the model "
112            "being used has an additional decoder for that purpose."
113        )
114
115    return parser
116
117
118def clear_annotations(viewer: napari.Viewer, clear_segmentations=True) -> None:
119    """@private"""
120    point_layer = viewer.layers["point_prompts"]
121    point_layer.selected_data = set(range(len(point_layer.data)))
122    point_layer.remove_selected()
123    point_layer.refresh()
124    if "prompts" in viewer.layers:
125        # Select all prompts and then remove them.
126        # This is how it worked before napari 0.5.
127        # viewer.layers["prompts"].data = []
128        viewer.layers["prompts"].selected_data = set(range(len(viewer.layers["prompts"].data)))
129        viewer.layers["prompts"].remove_selected()
130        viewer.layers["prompts"].refresh()
131    if not clear_segmentations:
132        return
133    viewer.layers["current_object"].data = np.zeros(viewer.layers["current_object"].data.shape, dtype="uint32")
134    viewer.layers["current_object"].refresh()
135
136
137def clear_annotations_slice(viewer: napari.Viewer, i: int, clear_segmentations=True) -> None:
138    """@private"""
139    point_layer = viewer.layers["point_prompts"]
140    point_layer.selected_data = set(np.flatnonzero(point_layer.data[:, 0] == i))
141    point_layer.remove_selected()
142    point_layer.refresh()
143    if "prompts" in viewer.layers:
144        prompts = viewer.layers["prompts"].data
145        prompts = [prompt for prompt in prompts if not (prompt[:, 0] == i).all()]
146        viewer.layers["prompts"].data = prompts
147        viewer.layers["prompts"].refresh()
148    if not clear_segmentations:
149        return
150    viewer.layers["current_object"].data[i] = 0
151    viewer.layers["current_object"].refresh()
152
153
154#
155# Helper functions to extract prompts from napari layers.
156#
157
158
159def point_layer_to_prompts(
160    layer: napari.layers.Points, i=None, track_id=None, with_stop_annotation=True,
161) -> Optional[Tuple[np.ndarray, np.ndarray]]:
162    """Extract point prompts for SAM from a napari point layer.
163
164    Args:
165        layer: The point layer from which to extract the prompts.
166        i: Index for the data (required for 3d or timeseries data).
167        track_id: Id of the current track (required for tracking data).
168        with_stop_annotation: Whether a single negative point will be interpreted
169            as stop annotation or just returned as normal prompt.
170
171    Returns:
172        The point coordinates for the prompts.
173        The labels (positive or negative / 1 or 0) for the prompts.
174    """
175
176    points = layer.data
177    labels = layer.properties["label"]
178    assert len(points) == len(labels)
179
180    if i is None:
181        assert points.shape[1] == 2, f"{points.shape}"
182        this_points, this_labels = points, labels
183    else:
184        assert points.shape[1] == 3, f"{points.shape}"
185        mask = np.round(points[:, 0]) == i
186        this_points = points[mask][:, 1:]
187        this_labels = labels[mask]
188    assert len(this_points) == len(this_labels)
189
190    if track_id is not None:
191        assert i is not None
192        track_ids = np.array(list(map(int, layer.properties["track_id"])))[mask]
193        track_id_mask = track_ids == track_id
194        this_labels, this_points = this_labels[track_id_mask], this_points[track_id_mask]
195    assert len(this_points) == len(this_labels)
196
197    this_labels = np.array([1 if label == "positive" else 0 for label in this_labels])
198    # a single point with a negative label is interpreted as 'stop' signal
199    # in this case we return None
200    if with_stop_annotation and (len(this_points) == 1 and this_labels[0] == 0):
201        return None
202
203    return this_points, this_labels
204
205
206def shape_layer_to_prompts(
207    layer: napari.layers.Shapes, shape: Tuple[int, int], i=None, track_id=None
208) -> Tuple[List[np.ndarray], List[Optional[np.ndarray]]]:
209    """Extract prompts for SAM from a napari shape layer.
210
211    Extracts the bounding box for 'rectangle' shapes and the bounding box and corresponding mask
212    for 'ellipse' and 'polygon' shapes.
213
214    Args:
215        prompt_layer: The napari shape layer.
216        shape: The image shape.
217        i: Index for the data (required for 3d or timeseries data).
218        track_id: Id of the current track (required for tracking data).
219
220    Returns:
221        The box prompts.
222        The mask prompts.
223    """
224
225    def _to_prompts(shape_data, shape_types):
226        boxes, masks = [], []
227
228        for data, type_ in zip(shape_data, shape_types):
229
230            if type_ == "rectangle":
231                boxes.append(data)
232                masks.append(None)
233
234            elif type_ == "ellipse":
235                boxes.append(data)
236                center = np.mean(data, axis=0)
237                radius_r = ((data[2] - data[1]) / 2)[0]
238                radius_c = ((data[1] - data[0]) / 2)[1]
239                rr, cc = draw.ellipse(center[0], center[1], radius_r, radius_c, shape=shape)
240                mask = np.zeros(shape, dtype=bool)
241                mask[rr, cc] = 1
242                masks.append(mask)
243
244            elif type_ == "polygon":
245                boxes.append(data)
246                rr, cc = draw.polygon(data[:, 0], data[:, 1], shape=shape)
247                mask = np.zeros(shape, dtype=bool)
248                mask[rr, cc] = 1
249                masks.append(mask)
250
251            else:
252                warnings.warn(f"Shape type {type_} is not supported and will be ignored.")
253
254        # map to correct box format
255        boxes = [
256            np.array([box[:, 0].min(), box[:, 1].min(), box[:, 0].max(), box[:, 1].max()]) for box in boxes
257        ]
258        return boxes, masks
259
260    shape_data, shape_types = layer.data, layer.shape_type
261    assert len(shape_data) == len(shape_types)
262    if len(shape_data) == 0:
263        return [], []
264
265    if i is not None:
266        if track_id is None:
267            prompt_selection = [j for j, data in enumerate(shape_data) if (data[:, 0] == i).all()]
268        else:
269            track_ids = np.array(list(map(int, layer.properties["track_id"])))
270            prompt_selection = [
271                j for j, (data, this_track_id) in enumerate(zip(shape_data, track_ids))
272                if ((data[:, 0] == i).all() and this_track_id == track_id)
273            ]
274
275        shape_data = [shape_data[j][:, 1:] for j in prompt_selection]
276        shape_types = [shape_types[j] for j in prompt_selection]
277
278    boxes, masks = _to_prompts(shape_data, shape_types)
279    return boxes, masks
280
281
282def prompt_layer_to_state(prompt_layer: napari.layers.Points, i: int) -> str:
283    """Get the state of the track from a point layer for a given timeframe.
284
285    Only relevant for annotator_tracking.
286
287    Args:
288        prompt_layer: The napari layer.
289        i: Timeframe of the data.
290
291    Returns:
292        The state of this frame (either "division" or "track").
293    """
294    state = prompt_layer.properties["state"]
295
296    points = prompt_layer.data
297    assert points.shape[1] == 3, f"{points.shape}"
298    mask = points[:, 0] == i
299    this_points = points[mask][:, 1:]
300    this_state = state[mask]
301    assert len(this_points) == len(this_state)
302
303    # we set the state to 'division' if at least one point in this frame has a division label
304    if any(st == "division" for st in this_state):
305        return "division"
306    else:
307        return "track"
308
309
310def prompt_layers_to_state(point_layer: napari.layers.Points, box_layer: napari.layers.Shapes, i: int) -> str:
311    """Get the state of the track from a point layer and shape layer for a given timeframe.
312
313    Only relevant for annotator_tracking.
314
315    Args:
316        point_layer: The napari point layer.
317        box_layer: The napari box layer.
318        i: Timeframe of the data.
319
320    Returns:
321        The state of this frame (either "division" or "track").
322    """
323    state = point_layer.properties["state"]
324
325    points = point_layer.data
326    assert points.shape[1] == 3, f"{points.shape}"
327    mask = points[:, 0] == i
328    if mask.sum() > 0:
329        this_state = state[mask].tolist()
330    else:
331        this_state = []
332
333    box_states = box_layer.properties["state"]
334    this_box_states = [
335        state for box, state in zip(box_layer.data, box_states)
336        if (box[:, 0] == i).all()
337    ]
338    this_state.extend(this_box_states)
339
340    # we set the state to 'division' if at least one point in this frame has a division label
341    if any(st == "division" for st in this_state):
342        return "division"
343    else:
344        return "track"
345
346
347#
348# Helper functions to run (multi-dimensional) segmentation on napari layers.
349#
350
351
352def segment_slices_with_prompts(
353    predictor, point_prompts, box_prompts, image_embeddings, shape, track_id=None, update_progress=None,
354):
355    """@private"""
356    assert len(shape) == 3
357    image_shape = shape[1:]
358    seg = np.zeros(shape, dtype="uint32")
359
360    z_values = np.round(point_prompts.data[:, 0])
361    z_values_boxes = np.concatenate([box[:1, 0] for box in box_prompts.data]) if box_prompts.data else\
362        np.zeros(0, dtype="int")
363
364    if track_id is not None:
365        track_ids_points = np.array(list(map(int, point_prompts.properties["track_id"])))
366        assert len(track_ids_points) == len(z_values)
367        z_values = z_values[track_ids_points == track_id]
368
369        if len(z_values_boxes) > 0:
370            track_ids_boxes = np.array(list(map(int, box_prompts.properties["track_id"])))
371            assert len(track_ids_boxes) == len(z_values_boxes), f"{len(track_ids_boxes)}, {len(z_values_boxes)}"
372            z_values_boxes = z_values_boxes[track_ids_boxes == track_id]
373
374    slices = np.unique(np.concatenate([z_values, z_values_boxes])).astype("int")
375    stop_lower, stop_upper = False, False
376
377    if update_progress is None:
378        def update_progress(*args):
379            pass
380
381    for i in slices:
382        points_i = point_layer_to_prompts(point_prompts, i, track_id)
383
384        # do we end the segmentation at the outer slices?
385        if points_i is None:
386
387            if i == slices[0]:  # The bottom slice is a stop slice.
388                stop_lower = True
389                seg[i] = 0
390            elif i == slices[-1]:  # The top sloce is a stop slice.
391                stop_upper = True
392                seg[i] = 0
393            else:  # We have a stop annotation somewhere in the middle. Ignore this.
394                # Remove this slice from the annotated slices, so that it is segmented via
395                # projection in the next step.
396                slices = np.setdiff1d(slices, i)
397                print(f"You have provided a stop annotation (single red point) in slice {i},")
398                print("but you have annotated slices above or below it. This stop annotation will")
399                print(f"be ignored and the slice {i} will be segmented normally.")
400
401            update_progress(1)
402            continue
403
404        boxes, masks = shape_layer_to_prompts(box_prompts, image_shape, i=i, track_id=track_id)
405        points, labels = points_i
406
407        seg_i = prompt_segmentation(
408            predictor, points, labels, boxes, masks, image_shape, multiple_box_prompts=False,
409            image_embeddings=image_embeddings, i=i
410        )
411        if seg_i is None:
412            print(f"The prompts at slice or frame {i} are invalid and the segmentation was skipped.")
413            print("This will lead to a wrong segmentation across slices or frames.")
414            print(f"Please correct the prompts in {i} and rerun the segmentation.")
415            continue
416
417        seg[i] = seg_i
418        update_progress(1)
419
420    return seg, slices, stop_lower, stop_upper
421
422
423# For advanced batching: match prompts to already segmented objects and continue segmentation.
424def _match_prompts(previous_segmentation, points, boxes, seg_ids):
425    # Create a mapping between ids and prompts.
426    batched_prompts = {}
427    # seg_boundaries = find_boundaries(previous_segmentation, mode="inner")
428    # indices = distance_transform_edt(seg_boundaries, return_distance=False, return_index=True)
429    return batched_prompts
430
431
432def _batched_interactive_segmentation(predictor, points, labels, boxes, image_embeddings, i, previous_segmentation):
433    prev_seg = previous_segmentation if i is None else previous_segmentation[i]
434    seg = np.zeros(prev_seg.shape, dtype="uint32")
435
436    # seg_ids = np.unique(previous_segmentation)
437    # assert seg_ids[0] == 0
438
439    batched_points, batched_labels = [], []
440    negative_points, negative_labels = [], []
441    for j in range(len(points)):
442        if labels[j] == 1:  # positive point
443            batched_points.append(points[j:j+1])
444            batched_labels.append(labels[j:j+1])
445        else:  # negative points
446            negative_points.append(points[j:j+1])
447            negative_labels.append(labels[j:j+1])
448
449    batched_prompts = [(None, point, label) for point, label in zip(batched_points, batched_labels)]
450    batched_prompts.extend([(box, None, None) for box in boxes])
451    batched_prompts = {i: prompt for i, prompt in enumerate(batched_prompts, 1)}
452
453    # For advanced batching: match prompts to already segmented objects and continue segmentation.
454    # (This is left here as a reference for how this can be implemented.
455    #  I have not decided yet if this is actually a good idea or not.)
456    # # If we have no objects: this is the first call for a batched segmentation.
457    # # We treat each positive point or box as a separate object.
458    # if len(seg_ids) == 1:
459    #     # Create a list of all prompts.
460    #     batched_prompts = [(None, point, label) for point, label in zip(batched_points, batched_labels)]
461    #     batched_prompts.extend([(box, None, None) for box in boxes])
462    #     batched_prompts = {i: prompt for i, prompt in enumerate(batched_prompts, 1)}
463
464    # # Otherwise we match the prompts to existing objects.
465    # else:
466    #     batched_prompts = _match_prompts(prev_seg, batched_points, boxes, seg_ids)
467
468    for seg_id, prompt in batched_prompts.items():
469        box, point, label = prompt
470        if len(negative_points) > 0:
471            if point is None:
472                point, label = negative_points, negative_labels
473            else:
474                point = np.concatenate([point] + negative_points)
475                label = np.concatenate([label] + negative_labels)
476
477        if (box is not None) and (point is not None):
478            prediction = prompt_based_segmentation.segment_from_box_and_points(
479                predictor, box, point, label, image_embeddings=image_embeddings, i=i
480            ).squeeze()
481        elif (box is not None) and (point is None):
482            prediction = prompt_based_segmentation.segment_from_box(
483                predictor, box, image_embeddings=image_embeddings, i=i
484            ).squeeze()
485        else:
486            prediction = prompt_based_segmentation.segment_from_points(
487                predictor, point, label, image_embeddings=image_embeddings, i=i
488            ).squeeze()
489
490        seg[prediction] = seg_id
491
492    return seg
493
494
495def prompt_segmentation(
496    predictor, points, labels, boxes, masks, shape, multiple_box_prompts,
497    image_embeddings=None, i=None, box_extension=0, batched=None, previous_segmentation=None,
498):
499    """@private"""
500    assert len(points) == len(labels)
501    have_points = len(points) > 0
502    have_boxes = len(boxes) > 0
503
504    # No prompts were given, return None.
505    if not have_points and not have_boxes:
506        return
507
508    # Batched interactive segmentation.
509    elif batched:
510        assert previous_segmentation is not None
511        seg = _batched_interactive_segmentation(
512            predictor, points, labels, boxes, image_embeddings, i, previous_segmentation
513        )
514
515    # Box and point prompts were given.
516    elif have_points and have_boxes:
517        if len(boxes) > 1:
518            print("You have provided point prompts and more than one box prompt.")
519            print("This setting is currently not supported.")
520            print("When providing both points and prompts you can only segment one object at a time.")
521            return
522        mask = masks[0]
523        if mask is None:
524            seg = prompt_based_segmentation.segment_from_box_and_points(
525                predictor, boxes[0], points, labels, image_embeddings=image_embeddings, i=i
526            ).squeeze()
527        else:
528            seg = prompt_based_segmentation.segment_from_mask(
529                predictor, mask, box=boxes[0], points=points, labels=labels, image_embeddings=image_embeddings, i=i
530            ).squeeze()
531
532    # Only point prompts were given.
533    elif have_points and not have_boxes:
534        seg = prompt_based_segmentation.segment_from_points(
535            predictor, points, labels, image_embeddings=image_embeddings, i=i
536        ).squeeze()
537
538    # Only box prompts were given.
539    elif not have_points and have_boxes:
540        seg = np.zeros(shape, dtype="uint32")
541
542        if len(boxes) > 1 and not multiple_box_prompts:
543            print("You have provided more than one box annotation. This is not yet supported in the 3d annotator.")
544            print("You can only segment one object at a time in 3d.")
545            return
546
547        # Batch this?
548        for seg_id, (box, mask) in enumerate(zip(boxes, masks), 1):
549            if mask is None:
550                prediction = prompt_based_segmentation.segment_from_box(
551                    predictor, box, image_embeddings=image_embeddings, i=i
552                ).squeeze()
553            else:
554                prediction = prompt_based_segmentation.segment_from_mask(
555                    predictor, mask, box=box, image_embeddings=image_embeddings, i=i,
556                    box_extension=box_extension,
557                ).squeeze()
558            seg[prediction] = seg_id
559
560    return seg
561
562
563def _compute_movement(seg, t0, t1):
564
565    def compute_center(t):
566        # computation with center of mass
567        center = np.where(seg[t] == 1)
568        center = np.array([np.mean(center[0]), np.mean(center[1])])
569        return center
570
571    center0 = compute_center(t0)
572    center1 = compute_center(t1)
573
574    move = center0 - center1
575    return move.astype("float64")
576
577
578def _shift_object(mask, motion_model):
579    mask_shifted = np.zeros_like(mask)
580    shift(mask, motion_model, output=mask_shifted, order=0, prefilter=False)
581    return mask_shifted
582
583
584def track_from_prompts(
585    point_prompts, box_prompts, seg, predictor, slices, image_embeddings,
586    stop_upper, threshold, projection, motion_smoothing=0.5, box_extension=0, update_progress=None,
587):
588    """@private
589    """
590    use_box, use_mask, use_points, use_single_point = _validate_projection(projection)
591
592    if update_progress is None:
593        def update_progress(*args):
594            pass
595
596    # shift the segmentation based on the motion model and update the motion model
597    def _update_motion_model(seg, t, t0, motion_model):
598        if t in (t0, t0 + 1):  # this is the first or second frame, we don't have a motion yet
599            pass
600        elif t == t0 + 2:  # this the third frame, we initialize the motion model
601            current_move = _compute_movement(seg, t - 1, t - 2)
602            motion_model = current_move
603        else:  # we already have a motion model and update it
604            current_move = _compute_movement(seg, t - 1, t - 2)
605            alpha = motion_smoothing
606            motion_model = alpha * motion_model + (1 - alpha) * current_move
607
608        return motion_model
609
610    has_division = False
611    motion_model = None
612    verbose = False
613
614    t0 = int(slices.min())
615    t = t0 + 1
616    while True:
617
618        # update the motion model
619        motion_model = _update_motion_model(seg, t, t0, motion_model)
620
621        # use the segmentation from prompts if we are in a slice with prompts
622        if t in slices:
623            seg_prev = None
624            seg_t = seg[t]
625            # currently using the box layer doesn't work for keeping track of the track state
626            # track_state = prompt_layers_to_state(point_prompts, box_prompts, t)
627            track_state = prompt_layer_to_state(point_prompts, t)
628
629        # otherwise project the mask (under the motion model) and segment the next slice from the mask
630        else:
631            if verbose:
632                print(f"Tracking object in frame {t} with movement {motion_model}")
633
634            seg_prev = seg[t - 1]
635            # shift the segmentation according to the motion model
636            if motion_model is not None:
637                seg_prev = _shift_object(seg_prev, motion_model)
638
639            seg_t = prompt_based_segmentation.segment_from_mask(
640                predictor, seg_prev, image_embeddings=image_embeddings, i=t,
641                use_mask=use_mask, use_box=use_box, use_points=use_points,
642                box_extension=box_extension, use_single_point=use_single_point,
643            )
644            track_state = "track"
645
646            # are we beyond the last slice with prompt?
647            # if no: we continue tracking because we know we need to connect to a future frame
648            # if yes: we only continue tracking if overlaps are above the threshold
649            if t < slices[-1]:
650                seg_prev = None
651
652            update_progress(1)
653
654        if (threshold is not None) and (seg_prev is not None):
655            iou = util.compute_iou(seg_prev, seg_t)
656            if iou < threshold:
657                msg = f"Segmentation stopped at frame {t} due to IOU {iou} < {threshold}."
658                print(msg)
659                break
660
661        # stop if we have a division
662        if track_state == "division":
663            has_division = True
664            break
665
666        seg[t] = seg_t
667        t += 1
668
669        # stop tracking if we have stop upper set (i.e. single negative point was set to indicate stop track)
670        if t == slices[-1] and stop_upper:
671            break
672
673        # stop if we are at the last slce
674        if t == seg.shape[0]:
675            break
676
677    return seg, has_division
678
679
680def _sync_embedding_widget(widget, model_type, save_path, checkpoint_path, device, tile_shape, halo):
681
682    # Update the index for model family, eg. 'Natural Images (SAM)', 'Light Microscopy', etc.
683    supported_dropdown_maps = {
684        "lm": "Light Microscopy",
685        "em_organelles": "Electron Microscopy",
686        "medical_imaging": "Medical Imaging",
687        "histopathology": "Histopathology",
688    }
689
690    model_family = "Natural Images (SAM)"  # If no suffix patterns match, stick to 'Natural Images (SAM)' family.
691    for k, v in supported_dropdown_maps.items():
692        if model_type.endswith(k):
693            model_family = v
694            break
695
696    index = widget.model_family_dropdown.findText(model_family)
697    if index > 0:
698        widget.model_family_dropdown.setCurrentIndex(index)
699
700    # Update the index for model size, eg. 'base', 'tiny', etc.
701    size_map = {"t": "tiny", "b": "base", "l": "large", "h": "huge"}
702    model_size = size_map[model_type[4]]
703
704    index = widget.model_size_dropdown.findText(model_size)
705    if index > 0:
706        widget.model_size_dropdown.setCurrentIndex(index)
707
708    if save_path is not None and isinstance(save_path, str):
709        widget.embeddings_save_path_param.setText(str(save_path))
710
711    if checkpoint_path is not None:
712        widget.custom_weights_param.setText(str(checkpoint_path))
713
714    if device is not None:
715        widget.device = device
716        index = widget.device_dropdown.findText(device)
717        widget.device_dropdown.setCurrentIndex(index)
718
719    if tile_shape is not None:
720        widget.tile_x_param.setValue(tile_shape[0])
721        widget.tile_y_param.setValue(tile_shape[1])
722
723    if halo is not None:
724        widget.halo_x_param.setValue(halo[0])
725        widget.halo_y_param.setValue(halo[1])
726
727
728# Read parameters from checkpoint path if it is given instead.
729def _sync_autosegment_widget(widget, model_type, checkpoint_path, update_decoder=None):
730    if update_decoder is not None:
731        widget._reset_segmentation_mode(update_decoder)
732
733    if widget.with_decoder:
734        settings = model_settings.AIS_SETTINGS.get(model_type, {})
735        params = ("center_distance_thresh", "boundary_distance_thresh")
736        for param in params:
737            if param in settings:
738                getattr(widget, f"{param}_param").setValue(settings[param])
739    else:
740        settings = model_settings.AMG_SETTINGS.get(model_type, {})
741        params = ("pred_iou_thresh", "stability_score_thresh", "min_object_size")
742        for param in params:
743            if param in settings:
744                getattr(widget, f"{param}_param").setValue(settings[param])
745
746
747# Read parameters from checkpoint path if it is given instead.
748def _sync_ndsegment_widget(widget, model_type, checkpoint_path):
749    settings = model_settings.ND_SEGMENT_SETTINGS.get(model_type, {})
750
751    if "projection_mode" in settings:
752        projection_mode = settings["projection_mode"]
753        widget.projection = projection_mode
754        index = widget.projection_dropdown.findText(projection_mode)
755        if index > 0:
756            widget.projection_dropdown.setCurrentIndex(index)
757
758    params = ("iou_threshold", "box_extension")
759    for param in params:
760        if param in settings:
761            getattr(widget, f"{param}_param").setValue(settings[param])
762
763
764def _load_amg_state(embedding_path):
765    if embedding_path is None or not os.path.exists(embedding_path):
766        return {"cache_folder": None}
767
768    cache_folder = os.path.join(embedding_path, "amg_state")
769    os.makedirs(cache_folder, exist_ok=True)
770    amg_state = {"cache_folder": cache_folder}
771
772    state_paths = glob(os.path.join(cache_folder, "*.pkl"))
773    for path in state_paths:
774        with open(path, "rb") as f:
775            state = pickle.load(f)
776        i = int(Path(path).stem.split("-")[-1])
777        amg_state[i] = state
778    return amg_state
779
780
781def _load_is_state(embedding_path):
782    if embedding_path is None or not os.path.exists(embedding_path):
783        return {"cache_path": None}
784
785    cache_path = os.path.join(embedding_path, "is_state.h5")
786    is_state = {"cache_path": cache_path}
787
788    with h5py.File(cache_path, "a") as f:
789        for name, g in f.items():
790            i = int(name.split("-")[-1])
791            state = {
792                "foreground": g["foreground"][:],
793                "boundary_distances": g["boundary_distances"][:],
794                "center_distances": g["center_distances"][:],
795            }
796            is_state[i] = state
797
798    return is_state