Upload 78 files

.gitattributes

@@ -35,3 +35,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 003_img.png filter=lfs diff=lfs merge=lfs -text
 1977_Well_F-5_Field_1.png filter=lfs diff=lfs merge=lfs -text
+example_imgs/1977_Well_F-5_Field_1_seg.png filter=lfs diff=lfs merge=lfs -text
+example_imgs/1977_Well_F-5_Field_1.png filter=lfs diff=lfs merge=lfs -text

_utils/attn_utils.py

@@ -0,0 +1,592 @@
+import abc
+
+import cv2
+import numpy as np
+import torch
+from IPython.display import display
+from PIL import Image
+from typing import Union, Tuple, List
+from einops import rearrange, repeat
+import math
+from torch import nn, einsum
+from inspect import isfunction
+from diffusers.utils import logging
+try:
+    from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+except:
+    from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
+
+try:
+    from diffusers.models.cross_attention import CrossAttention
+except:
+    from diffusers.models.attention_processor import Attention as CrossAttention
+
+MAX_NUM_WORDS = 77
+LOW_RESOURCE = False 
+
+class CountingCrossAttnProcessor1:
+
+    def __init__(self, attnstore, place_in_unet):
+        super().__init__()
+        self.attnstore = attnstore
+        self.place_in_unet = place_in_unet
+
+    def __call__(self, attn_layer: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, dim = hidden_states.shape
+        h = attn_layer.heads
+        q = attn_layer.to_q(hidden_states)
+        is_cross = encoder_hidden_states is not None
+        context = encoder_hidden_states if is_cross else hidden_states
+        k = attn_layer.to_k(context)
+        v = attn_layer.to_v(context)
+        # q = attn_layer.reshape_heads_to_batch_dim(q)
+        # k = attn_layer.reshape_heads_to_batch_dim(k)
+        # v = attn_layer.reshape_heads_to_batch_dim(v)
+        # q = attn_layer.head_to_batch_dim(q)
+        # k = attn_layer.head_to_batch_dim(k)
+        # v = attn_layer.head_to_batch_dim(v)
+        q = self.head_to_batch_dim(q, h)
+        k = self.head_to_batch_dim(k, h)
+        v = self.head_to_batch_dim(v, h)
+
+        sim = torch.einsum("b i d, b j d -> b i j", q, k) * attn_layer.scale
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.reshape(batch_size, -1)
+            max_neg_value = -torch.finfo(sim.dtype).max
+            attention_mask = attention_mask[:, None, :].repeat(h, 1, 1)
+            sim.masked_fill_(~attention_mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        attn_ = sim.softmax(dim=-1).clone()
+        # softmax = nn.Softmax(dim=-1)
+        # attn_ = softmax(sim)
+        self.attnstore(attn_, is_cross, self.place_in_unet)
+        out = torch.einsum("b i j, b j d -> b i d", attn_, v)
+        # out = attn_layer.batch_to_head_dim(out)
+        out = self.batch_to_head_dim(out, h)
+
+        if type(attn_layer.to_out) is torch.nn.modules.container.ModuleList:
+            to_out = attn_layer.to_out[0]
+        else:
+            to_out = attn_layer.to_out
+
+        out = to_out(out)
+        return out
+    
+    def batch_to_head_dim(self, tensor, head_size):
+        # head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+        return tensor
+
+    def head_to_batch_dim(self, tensor, head_size, out_dim=3):
+        # head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3)
+
+        if out_dim == 3:
+            tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)
+
+        return tensor
+
+
+def register_attention_control(model, controller):
+
+    attn_procs = {}
+    cross_att_count = 0
+    for name in model.unet.attn_processors.keys():
+        cross_attention_dim = None if name.endswith("attn1.processor") else model.unet.config.cross_attention_dim
+        if name.startswith("mid_block"):
+            hidden_size = model.unet.config.block_out_channels[-1]
+            place_in_unet = "mid"
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(model.unet.config.block_out_channels))[block_id]
+            place_in_unet = "up"
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = model.unet.config.block_out_channels[block_id]
+            place_in_unet = "down"
+        else:
+            continue
+
+        cross_att_count += 1
+        # attn_procs[name] = AttendExciteCrossAttnProcessor(
+        #     attnstore=controller, place_in_unet=place_in_unet
+        # )
+        attn_procs[name] = CountingCrossAttnProcessor1(
+            attnstore=controller, place_in_unet=place_in_unet
+        )
+
+    model.unet.set_attn_processor(attn_procs)
+    controller.num_att_layers = cross_att_count
+
+def register_hier_output(model):
+    self = model.unet
+    from ldm.modules.diffusionmodules.util import checkpoint, timestep_embedding
+    logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+    def forward(sample, timestep=None, encoder_hidden_states=None, class_labels=None, timestep_cond=None, 
+                attention_mask=None, cross_attention_kwargs=None, added_cond_kwargs=None, down_block_additional_residuals=None,
+                mid_block_additional_residual=None, encoder_attention_mask=None, return_dict=True):
+
+        out_list = []
+
+        
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+                # `Timesteps` does not contain any weights and will always return f32 tensors
+                # there might be better ways to encapsulate this.
+                class_labels = class_labels.to(dtype=sample.dtype)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
+
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        if self.config.addition_embed_type == "text":
+            aug_emb = self.add_embedding(encoder_hidden_states)
+        elif self.config.addition_embed_type == "text_image":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            image_embs = added_cond_kwargs.get("image_embeds")
+            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
+            aug_emb = self.add_embedding(text_embs, image_embs)
+        elif self.config.addition_embed_type == "text_time":
+            # SDXL - style
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            aug_emb = self.add_embedding(image_embs)
+        elif self.config.addition_embed_type == "image_hint":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            hint = added_cond_kwargs.get("hint")
+            aug_emb, hint = self.add_embedding(image_embs, hint)
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
+            # Kadinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
+        # 2. pre-process
+        sample = self.conv_in(sample) # 1, 320, 64, 64
+
+        # 2.5 GLIGEN position net
+        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
+
+        # 3. down
+        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None
+
+        down_block_res_samples = (sample,)
+
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)
+
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    sample += down_block_additional_residuals.pop(0)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                cross_attention_kwargs=cross_attention_kwargs,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+            # To support T2I-Adapter-XL
+            if (
+                is_adapter
+                and len(down_block_additional_residuals) > 0
+                and sample.shape == down_block_additional_residuals[0].shape
+            ):
+                sample += down_block_additional_residuals.pop(0)
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    scale=lora_scale,
+                )
+
+            # if i in [1, 4, 7]:
+            out_list.append(sample)
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample), out_list
+    
+    self.forward = forward
+
+
+class AttentionControl(abc.ABC):
+
+    def step_callback(self, x_t):
+        return x_t
+
+    def between_steps(self):
+        return
+
+    @property
+    def num_uncond_att_layers(self):
+        return 0
+
+    @abc.abstractmethod
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        raise NotImplementedError
+
+    def __call__(self, attn, is_cross: bool, place_in_unet: str):
+        if self.cur_att_layer >= self.num_uncond_att_layers:
+            # self.forward(attn, is_cross, place_in_unet)
+            if LOW_RESOURCE:
+                attn = self.forward(attn, is_cross, place_in_unet)
+            else:
+                h = attn.shape[0]
+                attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
+        self.cur_att_layer += 1
+        if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
+            self.cur_att_layer = 0
+            self.cur_step += 1
+            self.between_steps()
+        return attn
+
+    def reset(self):
+        self.cur_step = 0
+        self.cur_att_layer = 0
+
+    def __init__(self):
+        self.cur_step = 0
+        self.num_att_layers = -1
+        self.cur_att_layer = 0
+
+
+class EmptyControl(AttentionControl):
+
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        return attn
+
+
+class AttentionStore(AttentionControl):
+
+    @staticmethod
+    def get_empty_store():
+        return {"down_cross": [], "mid_cross": [], "up_cross": [],
+                "down_self": [], "mid_self": [], "up_self": []}
+
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
+        if attn.shape[1] <= self.max_size ** 2:  # avoid memory overhead
+            self.step_store[key].append(attn)
+        return attn
+
+    def between_steps(self):
+        self.attention_store = self.step_store
+        if self.save_global_store:
+            with torch.no_grad():
+                if len(self.global_store) == 0:
+                    self.global_store = self.step_store
+                else:
+                    for key in self.global_store:
+                        for i in range(len(self.global_store[key])):
+                            self.global_store[key][i] += self.step_store[key][i].detach()
+        self.step_store = self.get_empty_store()
+        self.step_store = self.get_empty_store()
+
+    def get_average_attention(self):
+        average_attention = self.attention_store
+        return average_attention
+
+    def get_average_global_attention(self):
+        average_attention = {key: [item / self.cur_step for item in self.global_store[key]] for key in
+                             self.attention_store}
+        return average_attention
+
+    def reset(self):
+        super(AttentionStore, self).reset()
+        self.step_store = self.get_empty_store()
+        self.attention_store = {}
+        self.global_store = {}
+
+    def __init__(self, max_size=32, save_global_store=False):
+        '''
+        Initialize an empty AttentionStore
+        :param step_index: used to visualize only a specific step in the diffusion process
+        '''
+        super(AttentionStore, self).__init__()
+        self.save_global_store = save_global_store
+        self.max_size = max_size
+        self.step_store = self.get_empty_store()
+        self.attention_store = {}
+        self.global_store = {}
+        self.curr_step_index = 0
+
+def aggregate_attention(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
+    out = []
+    attention_maps = attention_store.get_average_attention()
+    num_pixels = res ** 2
+    for location in from_where:
+        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
+            if item.shape[1] == num_pixels:
+                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
+                out.append(cross_maps)
+    out = torch.cat(out, dim=0)
+    out = out.sum(0) / out.shape[0]
+    return out
+
+
+def show_cross_attention(tokenizer, prompts, attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):
+    tokens = tokenizer.encode(prompts[select])
+    decoder = tokenizer.decode
+    attention_maps = aggregate_attention(attention_store, res, from_where, True, select)
+    images = []
+    for i in range(len(tokens)):
+        image = attention_maps[:, :, i]
+        image = 255 * image / image.max()
+        image = image.unsqueeze(-1).expand(*image.shape, 3)
+        image = image.numpy().astype(np.uint8)
+        image = np.array(Image.fromarray(image).resize((256, 256)))
+        image = text_under_image(image, decoder(int(tokens[i])))
+        images.append(image)
+    view_images(np.stack(images, axis=0))
+    
+
+def show_self_attention_comp(attention_store: AttentionStore, res: int, from_where: List[str],
+                        max_com=10, select: int = 0):
+    attention_maps = aggregate_attention(attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))
+    u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
+    images = []
+    for i in range(max_com):
+        image = vh[i].reshape(res, res)
+        image = image - image.min()
+        image = 255 * image / image.max()
+        image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
+        image = Image.fromarray(image).resize((256, 256))
+        image = np.array(image)
+        images.append(image)
+    view_images(np.concatenate(images, axis=1))
+
+def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
+    h, w, c = image.shape
+    offset = int(h * .2)
+    img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
+    img[:h] = image
+    textsize = cv2.getTextSize(text, font, 1, 2)[0]
+    text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
+    cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
+    return img
+
+
+def view_images(images, num_rows=1, offset_ratio=0.02):
+    if type(images) is list:
+        num_empty = len(images) % num_rows
+    elif images.ndim == 4:
+        num_empty = images.shape[0] % num_rows
+    else:
+        images = [images]
+        num_empty = 0
+
+    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
+    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
+    num_items = len(images)
+
+    h, w, c = images[0].shape
+    offset = int(h * offset_ratio)
+    num_cols = num_items // num_rows
+    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
+                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
+    for i in range(num_rows):
+        for j in range(num_cols):
+            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
+                i * num_cols + j]
+
+    pil_img = Image.fromarray(image_)
+    display(pil_img)
+
+def self_cross_attn(self_attn, cross_attn):
+    res = self_attn.shape[0]
+    assert res == cross_attn.shape[0]
+    # cross attn [res, res] -> [res*res]
+    cross_attn_ = cross_attn.reshape([res*res])
+    # self_attn [res, res, res*res]
+    self_cross_attn = cross_attn_ * self_attn
+    self_cross_attn = self_cross_attn.mean(-1).unsqueeze(0).unsqueeze(0)
+    return self_cross_attn

_utils/attn_utils_new.py

@@ -0,0 +1,610 @@
+import abc
+
+import cv2
+import numpy as np
+import torch
+from IPython.display import display
+from PIL import Image
+from typing import Union, Tuple, List
+from einops import rearrange, repeat
+import math
+from torch import nn, einsum
+from inspect import isfunction
+from diffusers.utils import logging
+try:
+    from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+except:
+    from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
+try:
+    from diffusers.models.cross_attention import CrossAttention
+except:
+    from diffusers.models.attention_processor import Attention as CrossAttention
+from typing import Any, Dict, List, Optional, Tuple, Union
+MAX_NUM_WORDS = 77
+LOW_RESOURCE = False 
+
+class CountingCrossAttnProcessor1:
+
+    def __init__(self, attnstore, place_in_unet):
+        super().__init__()
+        self.attnstore = attnstore
+        self.place_in_unet = place_in_unet
+
+    def __call__(self, attn_layer: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, dim = hidden_states.shape
+        h = attn_layer.heads
+        q = attn_layer.to_q(hidden_states)
+        is_cross = encoder_hidden_states is not None
+        context = encoder_hidden_states if is_cross else hidden_states
+        k = attn_layer.to_k(context)
+        v = attn_layer.to_v(context)
+        # q = attn_layer.reshape_heads_to_batch_dim(q)
+        # k = attn_layer.reshape_heads_to_batch_dim(k)
+        # v = attn_layer.reshape_heads_to_batch_dim(v)
+        # q = attn_layer.head_to_batch_dim(q)
+        # k = attn_layer.head_to_batch_dim(k)
+        # v = attn_layer.head_to_batch_dim(v)
+        q = self.head_to_batch_dim(q, h)
+        k = self.head_to_batch_dim(k, h)
+        v = self.head_to_batch_dim(v, h)
+
+        sim = torch.einsum("b i d, b j d -> b i j", q, k) * attn_layer.scale
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.reshape(batch_size, -1)
+            max_neg_value = -torch.finfo(sim.dtype).max
+            attention_mask = attention_mask[:, None, :].repeat(h, 1, 1)
+            sim.masked_fill_(~attention_mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        attn_ = sim.softmax(dim=-1).clone()
+        # softmax = nn.Softmax(dim=-1)
+        # attn_ = softmax(sim)
+        self.attnstore(attn_, is_cross, self.place_in_unet)
+        out = torch.einsum("b i j, b j d -> b i d", attn_, v)
+        # out = attn_layer.batch_to_head_dim(out)
+        out = self.batch_to_head_dim(out, h)
+
+        if type(attn_layer.to_out) is torch.nn.modules.container.ModuleList:
+            to_out = attn_layer.to_out[0]
+        else:
+            to_out = attn_layer.to_out
+
+        out = to_out(out)
+        return out
+    
+    def batch_to_head_dim(self, tensor, head_size):
+        # head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+        return tensor
+
+    def head_to_batch_dim(self, tensor, head_size, out_dim=3):
+        # head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3)
+
+        if out_dim == 3:
+            tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)
+
+        return tensor
+
+
+def register_attention_control(model, controller):
+
+    attn_procs = {}
+    cross_att_count = 0
+    for name in model.unet.attn_processors.keys():
+        cross_attention_dim = None if name.endswith("attn1.processor") else model.unet.config.cross_attention_dim
+        if name.startswith("mid_block"):
+            hidden_size = model.unet.config.block_out_channels[-1]
+            place_in_unet = "mid"
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(model.unet.config.block_out_channels))[block_id]
+            place_in_unet = "up"
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = model.unet.config.block_out_channels[block_id]
+            place_in_unet = "down"
+        else:
+            continue
+
+        cross_att_count += 1
+        # attn_procs[name] = AttendExciteCrossAttnProcessor(
+        #     attnstore=controller, place_in_unet=place_in_unet
+        # )
+        attn_procs[name] = CountingCrossAttnProcessor1(
+            attnstore=controller, place_in_unet=place_in_unet
+        )
+
+    model.unet.set_attn_processor(attn_procs)
+    controller.num_att_layers = cross_att_count
+
+def register_hier_output(model):
+    self = model.unet
+    logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+    def forward(sample, timestep=None, encoder_hidden_states=None, class_labels=None, timestep_cond=None, 
+                attention_mask=None, cross_attention_kwargs=None, added_cond_kwargs=None, down_block_additional_residuals=None,
+                mid_block_additional_residual=None, encoder_attention_mask=None, return_dict=True):
+
+        out_list = []
+
+        
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+                # `Timesteps` does not contain any weights and will always return f32 tensors
+                # there might be better ways to encapsulate this.
+                class_labels = class_labels.to(dtype=sample.dtype)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
+
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        if self.config.addition_embed_type == "text":
+            aug_emb = self.add_embedding(encoder_hidden_states)
+        elif self.config.addition_embed_type == "text_image":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            image_embs = added_cond_kwargs.get("image_embeds")
+            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
+            aug_emb = self.add_embedding(text_embs, image_embs)
+        elif self.config.addition_embed_type == "text_time":
+            # SDXL - style
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            aug_emb = self.add_embedding(image_embs)
+        elif self.config.addition_embed_type == "image_hint":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            hint = added_cond_kwargs.get("hint")
+            aug_emb, hint = self.add_embedding(image_embs, hint)
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
+            # Kadinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
+        # 2. pre-process
+        sample = self.conv_in(sample) # 1, 320, 64, 64
+
+        # 2.5 GLIGEN position net
+        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
+
+        # 3. down
+        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None
+
+        down_block_res_samples = (sample,)
+
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)
+
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    sample += down_block_additional_residuals.pop(0)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                cross_attention_kwargs=cross_attention_kwargs,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+            # To support T2I-Adapter-XL
+            if (
+                is_adapter
+                and len(down_block_additional_residuals) > 0
+                and sample.shape == down_block_additional_residuals[0].shape
+            ):
+                sample += down_block_additional_residuals.pop(0)
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    scale=lora_scale,
+                )
+
+            out_list.append(sample)
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample), out_list
+    
+    self.forward = forward
+
+
+
+
+
+
+
+class AttentionControl(abc.ABC):
+
+    def step_callback(self, x_t):
+        return x_t
+
+    def between_steps(self):
+        return
+
+    @property
+    def num_uncond_att_layers(self):
+        return 0
+
+    @abc.abstractmethod
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        raise NotImplementedError
+
+    def __call__(self, attn, is_cross: bool, place_in_unet: str):
+        if self.cur_att_layer >= self.num_uncond_att_layers:
+            # self.forward(attn, is_cross, place_in_unet)
+            if LOW_RESOURCE:
+                attn = self.forward(attn, is_cross, place_in_unet)
+            else:
+                h = attn.shape[0]
+                attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
+        self.cur_att_layer += 1
+        if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
+            self.cur_att_layer = 0
+            self.cur_step += 1
+            self.between_steps()
+        return attn
+
+    def reset(self):
+        self.cur_step = 0
+        self.cur_att_layer = 0
+
+    def __init__(self):
+        self.cur_step = 0
+        self.num_att_layers = -1
+        self.cur_att_layer = 0
+
+
+class EmptyControl(AttentionControl):
+
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        return attn
+
+
+class AttentionStore(AttentionControl):
+
+    @staticmethod
+    def get_empty_store():
+        return {"down_cross": [], "mid_cross": [], "up_cross": [],
+                "down_self": [], "mid_self": [], "up_self": []}
+
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
+        if attn.shape[1] <= self.max_size ** 2:  # avoid memory overhead
+            self.step_store[key].append(attn)
+        return attn
+
+    def between_steps(self):
+        self.attention_store = self.step_store
+        if self.save_global_store:
+            with torch.no_grad():
+                if len(self.global_store) == 0:
+                    self.global_store = self.step_store
+                else:
+                    for key in self.global_store:
+                        for i in range(len(self.global_store[key])):
+                            self.global_store[key][i] += self.step_store[key][i].detach()
+        self.step_store = self.get_empty_store()
+        self.step_store = self.get_empty_store()
+
+    def get_average_attention(self):
+        average_attention = self.attention_store
+        return average_attention
+
+    def get_average_global_attention(self):
+        average_attention = {key: [item / self.cur_step for item in self.global_store[key]] for key in
+                             self.attention_store}
+        return average_attention
+
+    def reset(self):
+        super(AttentionStore, self).reset()
+        self.step_store = self.get_empty_store()
+        self.attention_store = {}
+        self.global_store = {}
+
+    def __init__(self, max_size=32, save_global_store=False):
+        '''
+        Initialize an empty AttentionStore
+        :param step_index: used to visualize only a specific step in the diffusion process
+        '''
+        super(AttentionStore, self).__init__()
+        self.save_global_store = save_global_store
+        self.max_size = max_size
+        self.step_store = self.get_empty_store()
+        self.attention_store = {}
+        self.global_store = {}
+        self.curr_step_index = 0
+
+def aggregate_attention(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
+    out = []
+    attention_maps = attention_store.get_average_attention()
+    num_pixels = res ** 2
+    for location in from_where:
+        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
+            if item.shape[1] == num_pixels:
+                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
+                out.append(cross_maps)
+    out = torch.cat(out, dim=0)
+    out = out.sum(0) / out.shape[0]
+    return out
+
+def aggregate_attention1(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
+    out = []
+    attention_maps = attention_store.get_average_attention()
+    num_pixels = res ** 2
+    for location in from_where:
+        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
+            if item.shape[1] == num_pixels:
+                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
+                out.append(cross_maps)
+    # out = torch.cat(out, dim=0)
+    # out = out.sum(0) / out.shape[0]
+    out = out[1]
+    out = out.sum(0) / out.shape[0]
+    return out
+
+
+def show_cross_attention(tokenizer, prompts, attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):
+    tokens = tokenizer.encode(prompts[select])
+    decoder = tokenizer.decode
+    attention_maps = aggregate_attention(attention_store, res, from_where, True, select)
+    images = []
+    for i in range(len(tokens)):
+        image = attention_maps[:, :, i]
+        image = 255 * image / image.max()
+        image = image.unsqueeze(-1).expand(*image.shape, 3)
+        image = image.numpy().astype(np.uint8)
+        image = np.array(Image.fromarray(image).resize((256, 256)))
+        image = text_under_image(image, decoder(int(tokens[i])))
+        images.append(image)
+    view_images(np.stack(images, axis=0))
+    
+
+def show_self_attention_comp(attention_store: AttentionStore, res: int, from_where: List[str],
+                        max_com=10, select: int = 0):
+    attention_maps = aggregate_attention(attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))
+    u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
+    images = []
+    for i in range(max_com):
+        image = vh[i].reshape(res, res)
+        image = image - image.min()
+        image = 255 * image / image.max()
+        image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
+        image = Image.fromarray(image).resize((256, 256))
+        image = np.array(image)
+        images.append(image)
+    view_images(np.concatenate(images, axis=1))
+
+def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
+    h, w, c = image.shape
+    offset = int(h * .2)
+    img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
+    img[:h] = image
+    textsize = cv2.getTextSize(text, font, 1, 2)[0]
+    text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
+    cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
+    return img
+
+
+def view_images(images, num_rows=1, offset_ratio=0.02):
+    if type(images) is list:
+        num_empty = len(images) % num_rows
+    elif images.ndim == 4:
+        num_empty = images.shape[0] % num_rows
+    else:
+        images = [images]
+        num_empty = 0
+
+    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
+    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
+    num_items = len(images)
+
+    h, w, c = images[0].shape
+    offset = int(h * offset_ratio)
+    num_cols = num_items // num_rows
+    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
+                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
+    for i in range(num_rows):
+        for j in range(num_cols):
+            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
+                i * num_cols + j]
+
+    pil_img = Image.fromarray(image_)
+    display(pil_img)
+
+def self_cross_attn(self_attn, cross_attn):
+    cross_attn = cross_attn.squeeze()
+    res = self_attn.shape[0]
+    assert res == cross_attn.shape[-1]
+    # cross attn [res, res] -> [res*res]
+    cross_attn_ = cross_attn.reshape([res*res])
+    # self_attn [res, res, res*res]
+    self_cross_attn = cross_attn_ * self_attn
+    self_cross_attn = self_cross_attn.mean(-1).unsqueeze(0).unsqueeze(0)
+    return self_cross_attn

_utils/config.yaml

@@ -0,0 +1,15 @@
+attn_dist_mode: v0
+attn_positional_bias: rope
+attn_positional_bias_n_spatial: 16
+causal_norm: quiet_softmax
+coord_dim: 2
+d_model: 320
+dropout: 0.0
+feat_dim: 7
+feat_embed_per_dim: 8
+nhead: 4
+num_decoder_layers: 6
+num_encoder_layers: 6
+pos_embed_per_dim: 32
+spatial_pos_cutoff: 256
+window: 4

_utils/example_config.yaml

@@ -0,0 +1,20 @@
+batch_size: 1
+crop_size:
+- 256
+- 256
+detection_folders:
+- TRA
+dropout: 0.01
+example_images: False  # Slow
+input_train:
+- data/ctc/Fluo-N2DL-HeLa/01
+input_val:
+- data/ctc/Fluo-N2DL-HeLa/02
+max_tokens: 2048
+name: example
+ndim: 2
+num_decoder_layers: 5
+num_encoder_layers: 5
+outdir: runs
+distributed: False
+window: 4

_utils/load_models.py

@@ -0,0 +1,16 @@
+from config import RunConfig
+import torch
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
+import torch.nn as nn
+
+def load_stable_diffusion_model(config: RunConfig):
+    device = torch.device('cpu')
+
+    if config.sd_2_1:
+        stable_diffusion_version = "stabilityai/stable-diffusion-2-1-base"
+    else:
+        stable_diffusion_version = "CompVis/stable-diffusion-v1-4"
+    # stable = StableCountingPipeline.from_pretrained(stable_diffusion_version).to(device)
+    stable = StableDiffusionPipeline.from_pretrained(stable_diffusion_version).to(device)
+    return stable
+

_utils/load_track_data.py

@@ -0,0 +1,104 @@
+import os
+from glob import glob
+from pathlib import Path
+from natsort import natsorted
+from PIL import Image
+import numpy as np
+import tifffile
+import skimage.io as io
+import torchvision.transforms as T
+import cv2
+from tqdm import tqdm
+from models.tra_post_model.trackastra.utils import normalize_01, normalize
+IMG_SIZE = 512
+
+def _load_tiffs(folder: Path, dtype=None):
+    """Load a sequence of tiff files from a folder into a 3D numpy array."""
+    images = glob(str(folder / "*.tif"))
+    test_data = tifffile.imread(images[0])
+    if len(test_data.shape) == 3:
+        turn_gray = True
+    else:
+        turn_gray = False
+    end_frame = len(images)
+    if not turn_gray:
+        x = np.stack([
+            tifffile.imread(f).astype(dtype)
+            for f in tqdm(
+                sorted(folder.glob("*.tif"))[0 : end_frame : 1],
+                leave=False,
+                desc=f"Loading [0:{end_frame}]",
+            )
+        ])
+    else:
+        x = []
+        for f in tqdm(
+            sorted(folder.glob("*.tif"))[0 : end_frame : 1],
+            leave=False,
+            desc=f"Loading [0:{end_frame}]",
+        ):
+            img = tifffile.imread(f).astype(dtype)
+            if img.ndim == 3:
+                if img.shape[-1] > 3:
+                    img = img[..., :3]
+                img = (0.299 * img[..., 0] + 0.587 * img[..., 1] + 0.114 * img[..., 2])
+            x.append(img)
+        x = np.stack(x)
+    return x
+
+
+def load_track_images(file_dir):
+    
+    # suffix_ = [".png", ".tif", ".tiff", ".jpg"]
+    assert len(glob(file_dir + "/*.tif")) > 0, f"No tif images found in {file_dir}"
+    images = natsorted(glob(file_dir + "/*.tif"))
+    imgs = []
+    imgs_raw = []
+    images_stable = []
+    # load images for seg and track
+    for img_path in tqdm(images, desc="Loading images"):
+        img = tifffile.imread(img_path)
+        img_raw = io.imread(img_path)
+    
+        if img.dtype == 'uint16':
+            img = ((img - img.min()) / (img.max() - img.min() + 1e-6) * 255).astype(np.uint8)
+            img = np.stack([img] * 3, axis=-1)
+            w, h = img.shape[1], img.shape[0]
+        else:
+            img = Image.open(img_path).convert("RGB")
+            w, h = img.size
+
+        img = T.Compose([
+            T.ToTensor(),
+            T.Resize((IMG_SIZE, IMG_SIZE)),
+        ])(img)
+
+        image_stable = img - 0.5
+        img = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)
+
+
+        imgs.append(img)
+        imgs_raw.append(img_raw)
+        images_stable.append(image_stable)
+
+    height = h
+    width = w
+    imgs = np.stack(imgs, axis=0)
+    imgs_raw = np.stack(imgs_raw, axis=0)
+    images_stable = np.stack(images_stable, axis=0)
+
+    # track data
+    imgs_ = _load_tiffs(Path(file_dir), dtype=np.float32)
+    imgs_01 = np.stack([
+                normalize_01(_x) for _x in tqdm(imgs_, desc="Normalizing", leave=False)
+            ])
+    imgs_ = np.stack([
+                normalize(_x) for _x in tqdm(imgs_, desc="Normalizing", leave=False)
+            ])
+
+    return imgs, imgs_raw, images_stable, imgs_, imgs_01, height, width
+
+if __name__ == "__main__":
+    file_dir = "data/2D+Time/DIC-C2DH-HeLa/train/DIC-C2DH-HeLa/02"
+    imgs, imgs_raw, images_stable, imgs_, imgs_01, height, width = load_track_images(file_dir)
+    print(imgs.shape, imgs_raw.shape, images_stable.shape, imgs_.shape, imgs_01.shape, height, width)

_utils/misc_helper.py

@@ -0,0 +1,37 @@
+import logging
+import os
+import random
+import shutil
+from collections.abc import Mapping
+from datetime import datetime
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+
+def basicConfig(*args, **kwargs):
+    return
+
+
+# To prevent duplicate logs, we mask this baseConfig setting
+logging.basicConfig = basicConfig
+
+
+def create_logger(name, log_file, level=logging.INFO):
+    log = logging.getLogger(name)
+    formatter = logging.Formatter(
+        "[%(asctime)s][%(filename)15s][line:%(lineno)4d][%(levelname)8s] %(message)s"
+    )
+    fh = logging.FileHandler(log_file)
+    fh.setFormatter(formatter)
+    sh = logging.StreamHandler()
+    sh.setFormatter(formatter)
+    log.setLevel(level)
+    log.addHandler(fh)
+    log.addHandler(sh)
+    return log
+
+def get_current_time():
+    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+    return current_time

_utils/seg_eval.py

@@ -0,0 +1,61 @@
+import torch
+
+
+def iou_torch(inst1, inst2):
+    inter = torch.logical_and(inst1, inst2).sum().float()
+    union = torch.logical_or(inst1, inst2).sum().float()
+    if union == 0:
+        return torch.tensor(float('nan'))
+    return inter / union
+
+def get_instances_torch(mask):
+    # 返回所有非背景的 instance mask（布尔型）
+    ids = torch.unique(mask)
+    return [(mask == i) for i in ids if i != 0]
+
+def compute_instance_miou(pred_mask, gt_mask):
+    # pred_mask 和 gt_mask 都是 torch.Tensor, shape [H, W], 整数类型
+    pred_instances = get_instances_torch(pred_mask)
+    gt_instances = get_instances_torch(gt_mask)
+
+    ious = []
+    for gt in gt_instances:
+        best_iou = torch.tensor(0.0).to(pred_mask.device)
+        for pred in pred_instances:
+            i = iou_torch(pred, gt)
+            if i > best_iou:
+                best_iou = i
+        ious.append(best_iou)
+    
+    # 处理空情况
+    if len(ious) == 0:
+        return torch.tensor(float('nan'))
+    return torch.nanmean(torch.stack(ious))
+
+from torch import Tensor
+
+
+def dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon: float = 1e-6):
+    # Average of Dice coefficient for all batches, or for a single mask
+    assert input.size() == target.size()
+    assert input.dim() == 3 or not reduce_batch_first
+
+    sum_dim = (-1, -2) if input.dim() == 2 or not reduce_batch_first else (-1, -2, -3)
+
+    inter = 2 * (input * target).sum(dim=sum_dim)
+    sets_sum = input.sum(dim=sum_dim) + target.sum(dim=sum_dim)
+    sets_sum = torch.where(sets_sum == 0, inter, sets_sum)
+
+    dice = (inter + epsilon) / (sets_sum + epsilon)
+    return dice.mean()
+
+
+def multiclass_dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon: float = 1e-6):
+    # Average of Dice coefficient for all classes
+    return dice_coeff(input.flatten(0, 1), target.flatten(0, 1), reduce_batch_first, epsilon)
+
+
+def dice_loss(input: Tensor, target: Tensor, multiclass: bool = False):
+    # Dice loss (objective to minimize) between 0 and 1
+    fn = multiclass_dice_coeff if multiclass else dice_coeff
+    return 1 - fn(input, target, reduce_batch_first=True)

_utils/track_args.py

@@ -0,0 +1,157 @@
+import configargparse
+
+
+def parse_train_args():
+    parser = configargparse.ArgumentParser(
+        formatter_class=configargparse.ArgumentDefaultsHelpFormatter,
+        config_file_parser_class=configargparse.YAMLConfigFileParser,
+        allow_abbrev=False,
+    )
+    parser.add_argument(
+        "-c",
+        "--config",
+        default="_utils/example_config.yaml",
+        is_config_file=True,
+        help="config file path",
+    )
+    parser.add_argument("--device", type=str, choices=["cuda", "cpu"], default="cuda")
+    parser.add_argument("-o", "--outdir", type=str, default="runs")
+    parser.add_argument("--name", type=str, help="Name to append to timestamp")
+    parser.add_argument("--timestamp", type=bool, default=True)
+    parser.add_argument(
+        "-m",
+        "--model",
+        type=str,
+        default="",
+        help="load this model at start (e.g. to continue training)",
+    )
+    parser.add_argument(
+        "--ndim", type=int, default=2, help="number of spatial dimensions"
+    )
+    parser.add_argument("-d", "--d_model", type=int, default=256)
+    parser.add_argument("-w", "--window", type=int, default=10)
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--warmup_epochs", type=int, default=10)
+    parser.add_argument(
+        "--detection_folders",
+        type=str,
+        nargs="+",
+        default=["TRA"],
+        help=(
+            "Subfolders to search for detections. Defaults to `TRA`, which corresponds"
+            " to using only the GT."
+        ),
+    )
+    parser.add_argument("--downscale_temporal", type=int, default=1)
+    parser.add_argument("--downscale_spatial", type=int, default=1)
+    parser.add_argument("--spatial_pos_cutoff", type=int, default=256)
+    parser.add_argument("--from_subfolder", action="store_true")
+    # parser.add_argument("--train_samples", type=int, default=50000)
+    parser.add_argument("--num_encoder_layers", type=int, default=6)
+    parser.add_argument("--num_decoder_layers", type=int, default=6)
+    parser.add_argument("--pos_embed_per_dim", type=int, default=32)
+    parser.add_argument("--feat_embed_per_dim", type=int, default=8)
+    parser.add_argument("--dropout", type=float, default=0.00)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--max_tokens", type=int, default=None)
+    parser.add_argument("--delta_cutoff", type=int, default=2)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument(
+        "--attn_positional_bias",
+        type=str,
+        choices=["rope", "bias", "none"],
+        default="rope",
+    )
+    parser.add_argument("--attn_positional_bias_n_spatial", type=int, default=16)
+    parser.add_argument("--attn_dist_mode", default="v0")
+    parser.add_argument("--mixedp", type=bool, default=True)
+    parser.add_argument("--dry", action="store_true")
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument(
+        "--features",
+        type=str,
+        choices=[
+            "none",
+            "regionprops",
+            "regionprops2",
+            "patch",
+            "patch_regionprops",
+            "wrfeat",
+        ],
+        default="wrfeat",
+    )
+    parser.add_argument(
+        "--causal_norm",
+        type=str,
+        choices=["none", "linear", "softmax", "quiet_softmax"],
+        default="quiet_softmax",
+    )
+    parser.add_argument("--div_upweight", type=float, default=2)
+
+    parser.add_argument("--augment", type=int, default=3)
+    parser.add_argument("--tracking_frequency", type=int, default=-1)
+
+    parser.add_argument("--sanity_dist", action="store_true")
+    parser.add_argument("--preallocate", type=bool, default=False)
+    parser.add_argument("--only_prechecks", action="store_true")
+    parser.add_argument(
+        "--compress", type=bool, default=True, help="compress dataset"
+    )
+
+
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument(
+        "--logger",
+        type=str,
+        default="tensorboard",
+        choices=["tensorboard", "wandb", "none"],
+    )
+    parser.add_argument("--wandb_project", type=str, default="trackastra")
+    parser.add_argument(
+        "--crop_size",
+        type=int,
+        # required=True,
+        nargs="+",
+        default=None,
+        help="random crop size for augmentation",
+    )
+    parser.add_argument(
+        "--weight_by_ndivs",
+        type=bool,
+        default=True,
+        help="Oversample windows that contain divisions",
+    )
+    parser.add_argument(
+        "--weight_by_dataset",
+        type=bool,
+        default=False,
+        help=(
+            "Inversely weight datasets by number of samples (to counter dataset size"
+            " imbalance)"
+        ),
+    )
+
+    args, unknown_args = parser.parse_known_args()
+
+    # # Hack to allow for --input_test
+    # allowed_unknown = ["input_test"]
+    # if not set(a.split("=")[0].strip("-") for a in unknown_args).issubset(
+    #     set(allowed_unknown)
+    # ):
+    #     raise ValueError(f"Unknown args: {unknown_args}")
+
+    # pprint(vars(args))
+
+    # for backward compatibility
+    # if args.attn_positional_bias == "True":
+    #     args.attn_positional_bias = "bias"
+    # elif args.attn_positional_bias == "False":
+    #     args.attn_positional_bias = False
+
+    # if args.train_samples == 0:
+    #     raise NotImplementedError(
+    #         "--train_samples must be > 0, full dataset pass not supported."
+    #     )
+
+    return args

config.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List
+
+
+@dataclass
+class RunConfig:
+    # Guiding text prompt
+    prompt: str = "<task-prompt>"
+    # Whether to use Stable Diffusion v2.1
+    sd_2_1: bool = False
+    # Which token indices to alter with attend-and-excite
+    token_indices: List[int] = field(default_factory=lambda: [2,5])
+    # Which random seeds to use when generating
+    seeds: List[int] = field(default_factory=lambda: [42])
+    # Path to save all outputs to
+    output_path: Path = Path('./outputs')
+    # Number of denoising steps
+    n_inference_steps: int = 50
+    # Text guidance scale
+    guidance_scale: float = 7.5
+    # Number of denoising steps to apply attend-and-excite
+    max_iter_to_alter: int = 25
+    # Resolution of UNet to compute attention maps over
+    attention_res: int = 16
+    # Whether to run standard SD or attend-and-excite
+    run_standard_sd: bool = False
+    # Dictionary defining the iterations and desired thresholds to apply iterative latent refinement in
+    thresholds: Dict[int, float] = field(default_factory=lambda: {0: 0.05, 10: 0.5, 20: 0.8})
+    # Scale factor for updating the denoised latent z_t
+    scale_factor: int = 20
+    # Start and end values used for scaling the scale factor - decays linearly with the denoising timestep
+    scale_range: tuple = field(default_factory=lambda: (1.0, 0.5))
+    # Whether to apply the Gaussian smoothing before computing the maximum attention value for each subject token
+    smooth_attentions: bool = True
+    # Standard deviation for the Gaussian smoothing
+    sigma: float = 0.5
+    # Kernel size for the Gaussian smoothing
+    kernel_size: int = 3
+    # Whether to save cross attention maps for the final results
+    save_cross_attention_maps: bool = False
+
+    def __post_init__(self):
+        self.output_path.mkdir(exist_ok=True, parents=True)

counting.py

@@ -0,0 +1,337 @@
+# stable diffusion x loca
+import os
+# os.system("source /etc/network_turbo")
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+import pprint
+from typing import Any, List, Optional
+import argparse
+import pyrallis
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+import torch
+import os
+from PIL import Image
+import numpy as np
+from config import RunConfig
+from _utils import attn_utils_new as attn_utils
+from _utils.attn_utils import AttentionStore
+from _utils.misc_helper import *
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import cv2
+import warnings
+from pytorch_lightning.callbacks import ModelCheckpoint
+warnings.filterwarnings("ignore", category=UserWarning)
+import pytorch_lightning as pl
+from _utils.load_models import load_stable_diffusion_model
+from models.model import Counting_with_SD_features_loca as Counting
+from pytorch_lightning.loggers import WandbLogger
+from models.enc_model.loca_args import get_argparser as loca_get_argparser
+from models.enc_model.loca import build_model as build_loca_model
+import time
+import torchvision.transforms as T
+import skimage.io as io
+from _utils.dummy_box_gen import gen_dummy_boxes
+
+SCALE = 1
+
+
+class CountingModule(pl.LightningModule):
+    def __init__(self, use_box=True):
+        super().__init__()
+        self.use_box = use_box
+        self.config = RunConfig()   # config for stable diffusion
+        self.initialize_model()
+        
+
+    def initialize_model(self):
+        
+        # load loca model
+        loca_args = loca_get_argparser().parse_args()
+        self.loca_model = build_loca_model(loca_args)
+        # weights = torch.load("ckpt/loca_few_shot.pt")["model"]
+        # weights = {k.replace("module","") : v for k, v in weights.items()}
+        # self.loca_model.load_state_dict(weights, strict=False)
+        # del weights
+
+        self.counting_adapter = Counting(scale_factor=SCALE)
+        # if os.path.isfile(self.args.adapter_weight):
+        #     adapter_weight = torch.load(self.args.adapter_weight,map_location=torch.device('cpu'))
+        #     self.counting_adapter.load_state_dict(adapter_weight, strict=False)
+            
+        ### load stable diffusion and its controller
+        self.stable = load_stable_diffusion_model(config=self.config)
+        self.noise_scheduler = self.stable.scheduler
+        self.controller = AttentionStore(max_size=64)
+        attn_utils.register_attention_control(self.stable, self.controller)
+        attn_utils.register_hier_output(self.stable)
+
+        ##### initialize token_emb #####
+        placeholder_token = "<task-prompt>"
+        self.task_token = "repetitive objects"
+        # Add the placeholder token in tokenizer
+        num_added_tokens = self.stable.tokenizer.add_tokens(placeholder_token)
+        if num_added_tokens == 0:
+            raise ValueError(
+                f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
+                " `placeholder_token` that is not already in the tokenizer."
+            )
+        if os.path.isfile("pretrained/task_embed.pth"):
+            task_embed_from_pretrain = torch.load("pretrained/task_embed.pth")
+            placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)
+            self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))
+
+            token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data
+            token_embeds[placeholder_token_id] = task_embed_from_pretrain
+        else:
+            initializer_token = "count"
+            token_ids = self.stable.tokenizer.encode(initializer_token, add_special_tokens=False)
+            # Check if initializer_token is a single token or a sequence of tokens
+            if len(token_ids) > 1:
+                raise ValueError("The initializer token must be a single token.")
+
+            initializer_token_id = token_ids[0]
+            placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)
+
+            self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))
+
+            token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data
+            token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]
+
+        # others
+        self.placeholder_token = placeholder_token
+        self.placeholder_token_id = placeholder_token_id
+    
+
+    def move_to_device(self, device):
+        self.stable.to(device)
+        if self.loca_model is not None and self.counting_adapter is not None:
+            self.loca_model.to(device)
+            self.counting_adapter.to(device)
+        self.to(device)
+
+    def forward(self, data_path, box=None):
+        filename = data_path.split("/")[-1]
+        img = Image.open(data_path).convert("RGB")
+        width, height = img.size
+        input_image = T.Compose([T.ToTensor(), T.Resize((512, 512))])(img)
+        input_image_stable = input_image - 0.5
+        input_image = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(input_image)
+        if box is not None:
+            boxes = torch.tensor(box) / torch.tensor([width, height, width, height]) * 512  # xyxy, normalized
+            assert self.use_box == True
+        else:
+            boxes = torch.tensor([[100,100,130,130], [200,200,250,250]], dtype=torch.float32)  # dummy box
+            assert self.use_box == False
+
+        # move to device
+        input_image = input_image.unsqueeze(0).to(self.device)
+        boxes = boxes.unsqueeze(0).to(self.device)
+        input_image_stable = input_image_stable.unsqueeze(0).to(self.device)
+        
+
+
+        latents = self.stable.vae.encode(input_image_stable).latent_dist.sample().detach()
+        latents = latents * 0.18215
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(latents)
+        timesteps = torch.tensor([20], device=latents.device).long()
+        noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
+        input_ids_ = self.stable.tokenizer(
+            self.placeholder_token + " repetitive objects",
+            # "object",
+            padding="max_length",
+            truncation=True,
+            max_length=self.stable.tokenizer.model_max_length,
+            return_tensors="pt",
+        )
+        input_ids = input_ids_["input_ids"].to(self.device)
+        attention_mask = input_ids_["attention_mask"].to(self.device)
+        encoder_hidden_states = self.stable.text_encoder(input_ids, attention_mask)[0]
+
+        input_image = input_image.to(self.device)
+        boxes = boxes.to(self.device)
+        
+
+        task_loc_idx = torch.nonzero(input_ids == self.placeholder_token_id)
+        if self.use_box:
+            loca_out = self.loca_model.forward_before_reg(input_image, boxes)
+            loca_feature_bf_regression =  loca_out["feature_bf_regression"]
+            adapted_emb = self.counting_adapter.adapter(loca_feature_bf_regression, boxes)      # shape [1, 768]
+            if task_loc_idx.shape[0] == 0:
+                encoder_hidden_states[0,2,:] = adapted_emb.squeeze()  # 放在task prompt下一位
+            else:
+                encoder_hidden_states[0,task_loc_idx[0, 1]+1,:] = adapted_emb.squeeze()  # 放在task prompt下一位
+
+        # Predict the noise residual
+        noise_pred, feature_list = self.stable.unet(noisy_latents, timesteps, encoder_hidden_states)
+        noise_pred = noise_pred.sample
+        attention_store = self.controller.attention_store
+
+
+        attention_maps = []
+        exemplar_attention_maps = []
+        exemplar_attention_maps1 = []
+        exemplar_attention_maps2 = []
+        exemplar_attention_maps3 = []
+
+        cross_self_task_attn_maps = []
+        cross_self_exe_attn_maps = []
+
+        # only use 64x64 self-attention
+        self_attn_aggregate = attn_utils.aggregate_attention( # [res, res, 4096]
+                prompts=[self.config.prompt],        # 这里要改么
+                attention_store=self.controller,     
+                res=64,
+                from_where=("up", "down"),
+                is_cross=False,
+                select=0
+            )
+        self_attn_aggregate32 = attn_utils.aggregate_attention( # [res, res, 4096]
+                prompts=[self.config.prompt],        # 这里要改么
+                attention_store=self.controller,     
+                res=32,
+                from_where=("up", "down"),
+                is_cross=False,
+                select=0
+            )
+        self_attn_aggregate16 = attn_utils.aggregate_attention( # [res, res, 4096]
+                prompts=[self.config.prompt],        # 这里要改么
+                attention_store=self.controller,     
+                res=16,
+                from_where=("up", "down"),
+                is_cross=False,
+                select=0
+            )
+
+        # cross attention
+        for res in [32, 16]:
+            attn_aggregate = attn_utils.aggregate_attention( # [res, res, 77]
+                prompts=[self.config.prompt],        # 这里要改么
+                attention_store=self.controller,     
+                res=res,
+                from_where=("up", "down"),
+                is_cross=True,
+                select=0
+            )
+
+            task_attn_ = attn_aggregate[:, :, 1].unsqueeze(0).unsqueeze(0) # [1, 1, res, res]
+            attention_maps.append(task_attn_)
+            if self.use_box:
+                exemplar_attns = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0) # 取exemplar的attn
+                exemplar_attention_maps.append(exemplar_attns)
+            else:
+                exemplar_attns1 = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0)
+                exemplar_attns2 = attn_aggregate[:, :, 3].unsqueeze(0).unsqueeze(0)
+                exemplar_attns3 = attn_aggregate[:, :, 4].unsqueeze(0).unsqueeze(0)
+                exemplar_attention_maps1.append(exemplar_attns1)
+                exemplar_attention_maps2.append(exemplar_attns2)
+                exemplar_attention_maps3.append(exemplar_attns3)
+
+
+        scale_factors = [(64 // attention_maps[i].shape[-1]) for i in range(len(attention_maps))]
+        attns = torch.cat([F.interpolate(attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(attention_maps))])
+        task_attn_64 = torch.mean(attns, dim=0, keepdim=True)
+        cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)
+        cross_self_task_attn_maps.append(cross_self_task_attn)
+
+        if self.use_box:
+            scale_factors = [(64 // exemplar_attention_maps[i].shape[-1]) for i in range(len(exemplar_attention_maps))]
+            attns = torch.cat([F.interpolate(exemplar_attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps))])
+            exemplar_attn_64 = torch.mean(attns, dim=0, keepdim=True)
+
+            cross_self_exe_attn = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64)
+            cross_self_exe_attn_maps.append(cross_self_exe_attn)
+        else:
+            scale_factors = [(64 // exemplar_attention_maps1[i].shape[-1]) for i in range(len(exemplar_attention_maps1))]
+            attns = torch.cat([F.interpolate(exemplar_attention_maps1[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps1))])
+            exemplar_attn_64_1 = torch.mean(attns, dim=0, keepdim=True)
+
+            scale_factors = [(64 // exemplar_attention_maps2[i].shape[-1]) for i in range(len(exemplar_attention_maps2))]
+            attns = torch.cat([F.interpolate(exemplar_attention_maps2[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps2))])
+            exemplar_attn_64_2 = torch.mean(attns, dim=0, keepdim=True)
+
+            scale_factors = [(64 // exemplar_attention_maps3[i].shape[-1]) for i in range(len(exemplar_attention_maps3))]
+            attns = torch.cat([F.interpolate(exemplar_attention_maps3[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps3))])
+            exemplar_attn_64_3 = torch.mean(attns, dim=0, keepdim=True)
+
+
+            cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)
+            cross_self_task_attn_maps.append(cross_self_task_attn)
+
+            # if self.args.merge_exemplar == "average":
+            cross_self_exe_attn1 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_1)
+            cross_self_exe_attn2 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_2)
+            cross_self_exe_attn3 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_3)
+            exemplar_attn_64 = (exemplar_attn_64_1 + exemplar_attn_64_2 + exemplar_attn_64_3) / 3
+            cross_self_exe_attn = (cross_self_exe_attn1 + cross_self_exe_attn2 + cross_self_exe_attn3) / 3
+
+            exemplar_attn_64 = (exemplar_attn_64 - exemplar_attn_64.min()) / (exemplar_attn_64.max() - exemplar_attn_64.min() + 1e-6)
+
+        attn_stack = [exemplar_attn_64 / 2, cross_self_exe_attn / 2, exemplar_attn_64, cross_self_exe_attn]
+        attn_stack = torch.cat(attn_stack, dim=1)
+        
+        if not self.use_box:
+            
+            # cross_self_exe_attn_np = cross_self_exe_attn.detach().squeeze().cpu().numpy()
+            # boxes = gen_dummy_boxes(cross_self_exe_attn_np, max_boxes=1)
+            # boxes = boxes.to(self.device)
+
+            loca_out = self.loca_model.forward_before_reg(input_image, boxes)
+            loca_feature_bf_regression =  loca_out["feature_bf_regression"]
+        attn_out = self.loca_model.forward_reg(loca_out, attn_stack, feature_list[-1])
+        pred_density = attn_out["pred"].squeeze().cpu().numpy()
+        pred_cnt = pred_density.sum().item()
+
+        # resize pred_density to original image size
+        pred_density_rsz = cv2.resize(pred_density, (width, height), interpolation=cv2.INTER_CUBIC)
+        pred_density_rsz = pred_density_rsz / pred_density_rsz.sum() * pred_cnt
+
+        return pred_density_rsz, pred_cnt
+
+    
+def inference(data_path, box=None, save_path="./example_imgs", visualize=False):
+    if box is not None:
+        use_box = True
+    else:
+        use_box = False
+    model = CountingModule(use_box=use_box)
+    load_msg = model.load_state_dict(torch.load("pretrained/microscopy_matching_cnt.pth"), strict=True)
+    model.eval()
+    with torch.no_grad():
+        density_map, cnt = model(data_path, box)
+    
+    if visualize:
+        img = io.imread(data_path)
+        if len(img.shape) == 3 and img.shape[2] > 3:
+            img = img[:,:,:3]
+        if len(img.shape) == 2:
+            img = np.stack([img]*3, axis=-1)
+        img_show = img.squeeze()
+        density_map_show = density_map.squeeze()
+        os.makedirs(save_path, exist_ok=True)
+        filename = data_path.split("/")[-1]
+        img_show = (img_show - np.min(img_show)) / (np.max(img_show) - np.min(img_show))
+        fig, ax = plt.subplots(1,2, figsize=(12,6))
+        ax[0].imshow(img_show)
+        ax[0].axis('off')
+        ax[0].set_title(f"Input image")
+        ax[1].imshow(img_show)
+        ax[1].imshow(density_map_show, cmap='jet', alpha=0.5)  # Overlay density map with some transparency
+        ax[1].axis('off')
+        ax[1].set_title(f"Predicted density map, count: {cnt:.1f}")
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_path, filename.split(".")[0]+"_cnt.png"), dpi=300)
+        plt.close()
+    return density_map
+
+def main():
+    
+    inference(
+        data_path = "example_imgs/1977_Well_F-5_Field_1.png",
+        # box=[[150, 60, 183, 87]],
+        save_path = "./example_imgs",
+        visualize = True
+    )
+
+if __name__ == "__main__":
+    main()

models/enc_model/backbone.py

@@ -0,0 +1,64 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torchvision import models
+from torchvision.ops.misc import FrozenBatchNorm2d
+
+
+class Backbone(nn.Module):
+
+    def __init__(
+        self,
+        name: str,
+        pretrained: bool,
+        dilation: bool,
+        reduction: int,
+        swav: bool,
+        requires_grad: bool
+    ):
+
+        super(Backbone, self).__init__()
+
+        resnet = getattr(models, name)(
+            replace_stride_with_dilation=[False, False, dilation],
+            pretrained=pretrained, norm_layer=FrozenBatchNorm2d
+        )
+
+        self.backbone = resnet
+        self.reduction = reduction
+
+        if name == 'resnet50' and swav:
+            checkpoint = torch.hub.load_state_dict_from_url(
+                'https://dl.fbaipublicfiles.com/deepcluster/swav_800ep_pretrain.pth.tar',
+                map_location="cpu"
+            )
+            state_dict = {k.replace("module.", ""): v for k, v in checkpoint.items()}
+            self.backbone.load_state_dict(state_dict, strict=False)
+
+        # concatenation of layers 2, 3 and 4
+        self.num_channels = 896 if name in ['resnet18', 'resnet34'] else 3584
+
+        for n, param in self.backbone.named_parameters():
+            if 'layer2' not in n and 'layer3' not in n and 'layer4' not in n:
+                param.requires_grad_(False)
+            else:
+                param.requires_grad_(requires_grad)
+
+    def forward(self, x):
+        size = x.size(-2) // self.reduction, x.size(-1) // self.reduction
+        x = self.backbone.conv1(x)
+        x = self.backbone.bn1(x)
+        x = self.backbone.relu(x)
+        x = self.backbone.maxpool(x)
+
+        x = self.backbone.layer1(x)
+        x = layer2 = self.backbone.layer2(x)
+        x = layer3 = self.backbone.layer3(x)
+        x = layer4 = self.backbone.layer4(x)
+
+        x = torch.cat([
+            F.interpolate(f, size=size, mode='bilinear', align_corners=True)
+            for f in [layer2, layer3, layer4]
+        ], dim=1)
+
+        return x

models/enc_model/loca.py

@@ -0,0 +1,232 @@
+from .backbone import Backbone
+from .transformer import TransformerEncoder
+from .ope import OPEModule
+from .positional_encoding import PositionalEncodingsFixed
+from .regression_head import DensityMapRegressor
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class LOCA(nn.Module):
+
+    def __init__(
+        self,
+        image_size: int,
+        num_encoder_layers: int,
+        num_ope_iterative_steps: int,
+        num_objects: int,
+        emb_dim: int,
+        num_heads: int,
+        kernel_dim: int,
+        backbone_name: str,
+        swav_backbone: bool,
+        train_backbone: bool,
+        reduction: int,
+        dropout: float,
+        layer_norm_eps: float,
+        mlp_factor: int,
+        norm_first: bool,
+        activation: nn.Module,
+        norm: bool,
+        zero_shot: bool,
+    ):
+
+        super(LOCA, self).__init__()
+
+        self.emb_dim = emb_dim
+        self.num_objects = num_objects
+        self.reduction = reduction
+        self.kernel_dim = kernel_dim
+        self.image_size = image_size
+        self.zero_shot = zero_shot
+        self.num_heads = num_heads
+        self.num_encoder_layers = num_encoder_layers
+
+        self.backbone = Backbone(
+            backbone_name, pretrained=True, dilation=False, reduction=reduction,
+            swav=swav_backbone, requires_grad=train_backbone
+        )
+        self.input_proj = nn.Conv2d(
+            self.backbone.num_channels, emb_dim, kernel_size=1
+        )
+
+        if num_encoder_layers > 0:
+            self.encoder = TransformerEncoder(
+                num_encoder_layers, emb_dim, num_heads, dropout, layer_norm_eps,
+                mlp_factor, norm_first, activation, norm
+            )
+
+        self.ope = OPEModule(
+            num_ope_iterative_steps, emb_dim, kernel_dim, num_objects, num_heads,
+            reduction, layer_norm_eps, mlp_factor, norm_first, activation, norm, zero_shot
+        )
+
+        self.regression_head = DensityMapRegressor(emb_dim, reduction)
+        self.aux_heads = nn.ModuleList([
+            DensityMapRegressor(emb_dim, reduction)
+            for _ in range(num_ope_iterative_steps - 1)
+        ])
+
+        self.pos_emb = PositionalEncodingsFixed(emb_dim)
+
+        self.attn_norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.fuse = nn.Sequential(
+            nn.Conv2d(324, 256, kernel_size=1, stride=1),
+            nn.LeakyReLU(),
+            nn.LayerNorm((64, 64))
+        )
+
+        # self.fuse1 = nn.Sequential(
+        #     nn.Conv2d(322, 256, kernel_size=1, stride=1),
+        #     nn.LeakyReLU(),
+        #     nn.LayerNorm((64, 64))
+        # )
+
+    def forward_before_reg(self, x, bboxes):
+        num_objects = bboxes.size(1) if not self.zero_shot else self.num_objects
+        # backbone
+        backbone_features = self.backbone(x)
+        # prepare the encoder input
+        src = self.input_proj(backbone_features)
+        bs, c, h, w = src.size()
+        pos_emb = self.pos_emb(bs, h, w, src.device).flatten(2).permute(2, 0, 1)
+        src = src.flatten(2).permute(2, 0, 1)
+
+        # push through the encoder
+        if self.num_encoder_layers > 0:
+            image_features = self.encoder(src, pos_emb, src_key_padding_mask=None, src_mask=None)
+        else:
+            image_features = src
+
+        # prepare OPE input
+        f_e = image_features.permute(1, 2, 0).reshape(-1, self.emb_dim, h, w)
+
+        all_prototypes = self.ope(f_e, pos_emb, bboxes) # [3, 27, 1, 256]
+
+        outputs = list()
+        response_maps_list = []
+        for i in range(all_prototypes.size(0)):
+            prototypes = all_prototypes[i, ...].permute(1, 0, 2).reshape(
+                bs, num_objects, self.kernel_dim, self.kernel_dim, -1
+            ).permute(0, 1, 4, 2, 3).flatten(0, 2)[:, None, ...] # [768, 1, 3, 3]
+
+            response_maps = F.conv2d(
+                torch.cat([f_e for _ in range(num_objects)], dim=1).flatten(0, 1).unsqueeze(0),
+                prototypes,
+                bias=None,
+                padding=self.kernel_dim // 2,
+                groups=prototypes.size(0)
+            ).view(
+                bs, num_objects, self.emb_dim, h, w
+            ).max(dim=1)[0]
+
+            # # send through regression heads
+            # if i == all_prototypes.size(0) - 1:
+            #     predicted_dmaps = self.regression_head(response_maps)
+            # else:
+            #     predicted_dmaps = self.aux_heads[i](response_maps)
+            # outputs.append(predicted_dmaps)
+            response_maps_list.append(response_maps)
+
+        out = {
+            # "pred": outputs[-1],
+            "feature_bf_regression": response_maps_list[-1],
+            # "aux_pred": outputs[:-1],
+            "aux_feature_bf_regression": response_maps_list[:-1]
+        }
+
+        return out
+
+    def forward_reg(self, response_maps, attn_stack, unet_feature):
+        attn_stack = self.attn_norm(attn_stack)
+        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
+        unet_feature = unet_feature * attn_stack_mean
+        if unet_feature.shape[1] == 322:
+            unet_feature = self.fuse1(unet_feature)
+        else:
+            unet_feature = self.fuse(unet_feature)
+
+        response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
+
+        outputs = []
+        for i in range(len(response_maps)):
+            response_map = response_maps[i] + unet_feature
+            if i == len(response_maps) - 1:
+                predicted_dmaps = self.regression_head(response_map)
+            else:
+                predicted_dmaps = self.aux_heads[i](response_map)
+            outputs.append(predicted_dmaps)
+        
+        return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
+    
+    def forward_reg1(self, response_maps, self_attn):
+        # attn_stack = self.attn_norm(attn_stack)
+        # attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        # unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
+        # unet_feature = unet_feature * attn_stack_mean
+        # if unet_feature.shape[1] == 322:
+        #     unet_feature = self.fuse1(unet_feature)
+        # else:
+        #     unet_feature = self.fuse(unet_feature)
+
+        
+
+        response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
+
+        outputs = []
+        for i in range(len(response_maps)):
+            response_map = response_maps[i] + self_attn
+            if i == len(response_maps) - 1:
+                predicted_dmaps = self.regression_head(response_map)
+            else:
+                predicted_dmaps = self.aux_heads[i](response_map)
+            outputs.append(predicted_dmaps)
+        
+        return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
+    
+    def forward_reg_without_unet(self, response_maps, attn_stack):
+        # attn_stack = self.attn_norm(attn_stack)
+        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+
+        response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
+
+        outputs = []
+        for i in range(len(response_maps)):
+            response_map = response_maps[i] * attn_stack_mean * 0.5 + response_maps[i]
+            if i == len(response_maps) - 1:
+                predicted_dmaps = self.regression_head(response_map)
+            else:
+                predicted_dmaps = self.aux_heads[i](response_map)
+            outputs.append(predicted_dmaps)
+        
+        return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
+
+
+def build_model(args):
+
+    assert args.backbone in ['resnet18', 'resnet50', 'resnet101']
+    assert args.reduction in [4, 8, 16]
+
+    return LOCA(
+        image_size=args.image_size,
+        num_encoder_layers=args.num_enc_layers,
+        num_ope_iterative_steps=args.num_ope_iterative_steps,
+        num_objects=args.num_objects,
+        zero_shot=args.zero_shot,
+        emb_dim=args.emb_dim,
+        num_heads=args.num_heads,
+        kernel_dim=args.kernel_dim,
+        backbone_name=args.backbone,
+        swav_backbone=args.swav_backbone,
+        train_backbone=args.backbone_lr > 0,
+        reduction=args.reduction,
+        dropout=args.dropout,
+        layer_norm_eps=1e-5,
+        mlp_factor=8,
+        norm_first=args.pre_norm,
+        activation=nn.GELU,
+        norm=True,
+    )

models/enc_model/loca_args.py

@@ -0,0 +1,44 @@
+import argparse
+
+
+def get_argparser():
+
+    parser = argparse.ArgumentParser("LOCA parser", add_help=False)
+
+    parser.add_argument('--model_name', default='loca_few_shot', type=str)
+    parser.add_argument(
+        '--data_path',
+        default='./data/FSC147_384_V2',
+        type=str
+    )
+    parser.add_argument(
+        '--model_path',
+        default='ckpt',
+        type=str
+    )
+    parser.add_argument('--backbone', default='resnet50', type=str)
+    parser.add_argument('--swav_backbone', action='store_true', default=True)
+    parser.add_argument('--reduction', default=8, type=int)
+    parser.add_argument('--image_size', default=512, type=int)
+    parser.add_argument('--num_enc_layers', default=3, type=int)
+    parser.add_argument('--num_ope_iterative_steps', default=3, type=int)
+    parser.add_argument('--emb_dim', default=256, type=int)
+    parser.add_argument('--num_heads', default=8, type=int)
+    parser.add_argument('--kernel_dim', default=3, type=int)
+    parser.add_argument('--num_objects', default=3, type=int)
+    parser.add_argument('--epochs', default=200, type=int)
+    parser.add_argument('--resume_training', action='store_true')
+    parser.add_argument('--lr', default=1e-4, type=float)
+    parser.add_argument('--backbone_lr', default=0, type=float)
+    parser.add_argument('--lr_drop', default=200, type=int)
+    parser.add_argument('--weight_decay', default=1e-4, type=float)
+    parser.add_argument('--batch_size', default=1, type=int)
+    parser.add_argument('--dropout', default=0.1, type=float)
+    parser.add_argument('--num_workers', default=8, type=int)
+    parser.add_argument('--max_grad_norm', default=0.1, type=float)
+    parser.add_argument('--aux_weight', default=0.3, type=float)
+    parser.add_argument('--tiling_p', default=0.5, type=float)
+    parser.add_argument('--zero_shot', action='store_true')
+    parser.add_argument('--pre_norm', action='store_true', default=True)
+
+    return parser

models/enc_model/mlp.py

@@ -0,0 +1,23 @@
+from torch import nn
+
+
+class MLP(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        dropout: float,
+        activation: nn.Module
+    ):
+        super(MLP, self).__init__()
+
+        self.linear1 = nn.Linear(input_dim, hidden_dim)
+        self.linear2 = nn.Linear(hidden_dim, input_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = activation()
+
+    def forward(self, x):
+        return (
+            self.linear2(self.dropout(self.activation(self.linear1(x))))
+        )

models/enc_model/ope.py

@@ -0,0 +1,245 @@
+from .mlp import MLP
+from .positional_encoding import PositionalEncodingsFixed
+
+import torch
+from torch import nn
+
+from torchvision.ops import roi_align
+
+
+class OPEModule(nn.Module):
+
+    def __init__(
+        self,
+        num_iterative_steps: int,
+        emb_dim: int,
+        kernel_dim: int,
+        num_objects: int,
+        num_heads: int,
+        reduction: int,
+        layer_norm_eps: float,
+        mlp_factor: int,
+        norm_first: bool,
+        activation: nn.Module,
+        norm: bool,
+        zero_shot: bool,
+    ):
+
+        super(OPEModule, self).__init__()
+
+        self.num_iterative_steps = num_iterative_steps
+        self.zero_shot = zero_shot
+        self.kernel_dim = kernel_dim
+        self.num_objects = num_objects
+        self.emb_dim = emb_dim
+        self.reduction = reduction
+
+        if num_iterative_steps > 0:
+            self.iterative_adaptation = IterativeAdaptationModule(
+                num_layers=num_iterative_steps, emb_dim=emb_dim, num_heads=num_heads,
+                dropout=0, layer_norm_eps=layer_norm_eps,
+                mlp_factor=mlp_factor, norm_first=norm_first,
+                activation=activation, norm=norm,
+                zero_shot=zero_shot
+            )
+
+        if not self.zero_shot:
+            self.shape_or_objectness = nn.Sequential(
+                nn.Linear(2, 64),
+                nn.ReLU(),
+                nn.Linear(64, emb_dim),
+                nn.ReLU(),
+                nn.Linear(emb_dim, self.kernel_dim**2 * emb_dim)
+            )
+        else:
+            self.shape_or_objectness = nn.Parameter(
+                torch.empty((self.num_objects, self.kernel_dim**2, emb_dim))
+            )
+            nn.init.normal_(self.shape_or_objectness)
+
+        self.pos_emb = PositionalEncodingsFixed(emb_dim)
+
+    def forward(self, f_e, pos_emb, bboxes):
+        bs, _, h, w = f_e.size()
+        # extract the shape features or objectness
+        if not self.zero_shot:
+            box_hw = torch.zeros(bboxes.size(0), bboxes.size(1), 2).to(bboxes.device)
+            box_hw[:, :, 0] = bboxes[:, :, 2] - bboxes[:, :, 0]
+            box_hw[:, :, 1] = bboxes[:, :, 3] - bboxes[:, :, 1]
+            shape_or_objectness = self.shape_or_objectness(box_hw).reshape(
+                bs, -1, self.kernel_dim ** 2, self.emb_dim
+            ).flatten(1, 2).transpose(0, 1)
+        else:
+            shape_or_objectness = self.shape_or_objectness.expand(
+                bs, -1, -1, -1
+            ).flatten(1, 2).transpose(0, 1)
+
+        # if not zero shot add appearance
+        if not self.zero_shot:
+            # reshape bboxes into the format suitable for roi_align
+            num_of_boxes = bboxes.size(1)
+            bboxes = torch.cat([
+                torch.arange(
+                    bs, requires_grad=False
+                ).to(bboxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
+                bboxes.flatten(0, 1),
+            ], dim=1)
+            appearance = roi_align(
+                f_e,
+                boxes=bboxes, output_size=self.kernel_dim,
+                spatial_scale=1.0 / self.reduction, aligned=True
+            ).permute(0, 2, 3, 1).reshape(
+                bs, num_of_boxes * self.kernel_dim ** 2, -1
+            ).transpose(0, 1)
+        else:
+            num_of_boxes = self.num_objects
+            appearance = None
+
+        query_pos_emb = self.pos_emb(
+            bs, self.kernel_dim, self.kernel_dim, f_e.device
+        ).flatten(2).permute(2, 0, 1).repeat(num_of_boxes, 1, 1)
+
+        if self.num_iterative_steps > 0:
+            memory = f_e.flatten(2).permute(2, 0, 1)
+            all_prototypes = self.iterative_adaptation(
+                shape_or_objectness, appearance, memory, pos_emb, query_pos_emb
+            )
+        else:
+            if shape_or_objectness is not None and appearance is not None:
+                all_prototypes = (shape_or_objectness + appearance).unsqueeze(0)
+            else:
+                all_prototypes = (
+                    shape_or_objectness if shape_or_objectness is not None else appearance
+                ).unsqueeze(0)
+
+        return all_prototypes
+
+
+
+class IterativeAdaptationModule(nn.Module):
+
+    def __init__(
+        self,
+        num_layers: int,
+        emb_dim: int,
+        num_heads: int,
+        dropout: float,
+        layer_norm_eps: float,
+        mlp_factor: int,
+        norm_first: bool,
+        activation: nn.Module,
+        norm: bool,
+        zero_shot: bool
+    ):
+
+        super(IterativeAdaptationModule, self).__init__()
+
+        self.layers = nn.ModuleList([
+            IterativeAdaptationLayer(
+                emb_dim, num_heads, dropout, layer_norm_eps,
+                mlp_factor, norm_first, activation, zero_shot
+            ) for i in range(num_layers)
+        ])
+
+        self.norm = nn.LayerNorm(emb_dim, layer_norm_eps) if norm else nn.Identity()
+
+    def forward(
+        self, tgt, appearance, memory, pos_emb, query_pos_emb, tgt_mask=None, memory_mask=None,
+        tgt_key_padding_mask=None, memory_key_padding_mask=None
+    ):
+
+        output = tgt
+        outputs = list()
+        for i, layer in enumerate(self.layers):
+            output = layer(
+                output, appearance, memory, pos_emb, query_pos_emb, tgt_mask, memory_mask,
+                tgt_key_padding_mask, memory_key_padding_mask
+            )
+            outputs.append(self.norm(output))
+
+        return torch.stack(outputs)
+
+
+class IterativeAdaptationLayer(nn.Module):
+
+    def __init__(
+        self,
+        emb_dim: int,
+        num_heads: int,
+        dropout: float,
+        layer_norm_eps: float,
+        mlp_factor: int,
+        norm_first: bool,
+        activation: nn.Module,
+        zero_shot: bool
+    ):
+        super(IterativeAdaptationLayer, self).__init__()
+
+        self.norm_first = norm_first
+        self.zero_shot = zero_shot
+
+        if not self.zero_shot:
+            self.norm1 = nn.LayerNorm(emb_dim, layer_norm_eps)
+        self.norm2 = nn.LayerNorm(emb_dim, layer_norm_eps)
+        self.norm3 = nn.LayerNorm(emb_dim, layer_norm_eps)
+        if not self.zero_shot:
+            self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        if not self.zero_shot:
+            self.self_attn = nn.MultiheadAttention(emb_dim, num_heads, dropout)
+        self.enc_dec_attn = nn.MultiheadAttention(emb_dim, num_heads, dropout)
+
+        self.mlp = MLP(emb_dim, mlp_factor * emb_dim, dropout, activation)
+
+    def with_emb(self, x, emb):
+        return x if emb is None else x + emb
+
+    def forward(
+        self, tgt, appearance, memory, pos_emb, query_pos_emb, tgt_mask, memory_mask,
+        tgt_key_padding_mask, memory_key_padding_mask
+    ):
+        if self.norm_first:
+            if not self.zero_shot:
+                tgt_norm = self.norm1(tgt)
+                tgt = tgt + self.dropout1(self.self_attn(
+                    query=self.with_emb(tgt_norm, query_pos_emb),
+                    key=self.with_emb(appearance, query_pos_emb),
+                    value=appearance,
+                    attn_mask=tgt_mask,
+                    key_padding_mask=tgt_key_padding_mask
+                )[0])
+
+            tgt_norm = self.norm2(tgt)
+            tgt = tgt + self.dropout2(self.enc_dec_attn(
+                query=self.with_emb(tgt_norm, query_pos_emb),
+                key=memory+pos_emb,
+                value=memory,
+                attn_mask=memory_mask,
+                key_padding_mask=memory_key_padding_mask
+            )[0])
+            tgt_norm = self.norm3(tgt)
+            tgt = tgt + self.dropout3(self.mlp(tgt_norm))
+
+        else:
+            if not self.zero_shot:
+                tgt = self.norm1(tgt + self.dropout1(self.self_attn(
+                    query=self.with_emb(tgt, query_pos_emb),
+                    key=self.with_emb(appearance),
+                    value=appearance,
+                    attn_mask=tgt_mask,
+                    key_padding_mask=tgt_key_padding_mask
+                )[0]))
+
+            tgt = self.norm2(tgt + self.dropout2(self.enc_dec_attn(
+                query=self.with_emb(tgt, query_pos_emb),
+                key=memory+pos_emb,
+                value=memory,
+                attn_mask=memory_mask,
+                key_padding_mask=memory_key_padding_mask
+            )[0]))
+
+            tgt = self.norm3(tgt + self.dropout3(self.mlp(tgt)))
+
+        return tgt

models/enc_model/positional_encoding.py

@@ -0,0 +1,30 @@
+import torch
+from torch import nn
+
+
+class PositionalEncodingsFixed(nn.Module):
+
+    def __init__(self, emb_dim, temperature=10000):
+
+        super(PositionalEncodingsFixed, self).__init__()
+
+        self.emb_dim = emb_dim
+        self.temperature = temperature
+
+    def _1d_pos_enc(self, mask, dim):
+        temp = torch.arange(self.emb_dim // 2).float().to(mask.device)
+        temp = self.temperature ** (2 * (temp.div(2, rounding_mode='floor')) / self.emb_dim)
+
+        enc = (~mask).cumsum(dim).float().unsqueeze(-1) / temp
+        enc = torch.stack([
+            enc[..., 0::2].sin(), enc[..., 1::2].cos()
+        ], dim=-1).flatten(-2)
+
+        return enc
+
+    def forward(self, bs, h, w, device):
+        mask = torch.zeros(bs, h, w, dtype=torch.bool, requires_grad=False, device=device)
+        x = self._1d_pos_enc(mask, dim=2)
+        y = self._1d_pos_enc(mask, dim=1)
+
+        return torch.cat([y, x], dim=3).permute(0, 3, 1, 2)

models/enc_model/regression_head.py

@@ -0,0 +1,92 @@
+from torch import nn
+import torch
+
+
+class UpsamplingLayer(nn.Module):
+
+    def __init__(self, in_channels, out_channels, leaky=True):
+
+        super(UpsamplingLayer, self).__init__()
+
+        self.layer = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.LeakyReLU() if leaky else nn.ReLU(),
+            nn.UpsamplingBilinear2d(scale_factor=2)
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class DensityMapRegressor(nn.Module):
+
+    def __init__(self, in_channels, reduction):
+
+        super(DensityMapRegressor, self).__init__()
+
+        if reduction == 8:
+            self.regressor = nn.Sequential(
+                UpsamplingLayer(in_channels, 128),
+                UpsamplingLayer(128, 64),
+                UpsamplingLayer(64, 32),
+                nn.Conv2d(32, 1, kernel_size=1),
+                nn.LeakyReLU()
+            )
+        elif reduction == 16:
+            self.regressor = nn.Sequential(
+                UpsamplingLayer(in_channels, 128),
+                UpsamplingLayer(128, 64),
+                UpsamplingLayer(64, 32),
+                UpsamplingLayer(32, 16),
+                nn.Conv2d(16, 1, kernel_size=1),
+                nn.LeakyReLU()
+            )
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        return self.regressor(x)
+
+    def reset_parameters(self):
+        for module in self.modules():
+            if isinstance(module, nn.Conv2d):
+                nn.init.normal_(module.weight, std=0.01)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+
+class DensityMapRegressor_(nn.Module):
+
+    def __init__(self, in_channels, reduction):
+
+        super(DensityMapRegressor, self).__init__()
+
+        if reduction == 8:
+            self.regressor = nn.Sequential(
+                UpsamplingLayer(in_channels, 128),
+                UpsamplingLayer(128, 64),
+                UpsamplingLayer(64, 32),
+                nn.Conv2d(32, 1, kernel_size=1),
+                nn.LeakyReLU()
+            )
+        elif reduction == 16:
+            self.regressor = nn.Sequential(
+                UpsamplingLayer(in_channels, 128),
+                UpsamplingLayer(128, 64),
+                UpsamplingLayer(64, 32),
+                UpsamplingLayer(32, 16),
+                nn.Conv2d(16, 1, kernel_size=1),
+                nn.LeakyReLU()
+            )
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        return self.regressor(x)
+
+    def reset_parameters(self):
+        for module in self.modules():
+            if isinstance(module, nn.Conv2d):
+                nn.init.normal_(module.weight, std=0.01)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)

models/enc_model/transformer.py

@@ -0,0 +1,94 @@
+from .mlp import MLP
+
+from torch import nn
+
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(
+        self,
+        num_layers: int,
+        emb_dim: int,
+        num_heads: int,
+        dropout: float,
+        layer_norm_eps: float,
+        mlp_factor: int,
+        norm_first: bool,
+        activation: nn.Module,
+        norm: bool,
+    ):
+
+        super(TransformerEncoder, self).__init__()
+
+        self.layers = nn.ModuleList([
+            TransformerEncoderLayer(
+                emb_dim, num_heads, dropout, layer_norm_eps,
+                mlp_factor, norm_first, activation
+            ) for _ in range(num_layers)
+        ])
+
+        self.norm = nn.LayerNorm(emb_dim, layer_norm_eps) if norm else nn.Identity()
+
+    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
+        output = src
+        for layer in self.layers:
+            output = layer(output, pos_emb, src_mask, src_key_padding_mask)
+        return self.norm(output)
+
+
+class TransformerEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        emb_dim: int,
+        num_heads: int,
+        dropout: float,
+        layer_norm_eps: float,
+        mlp_factor: int,
+        norm_first: bool,
+        activation: nn.Module,
+    ):
+        super(TransformerEncoderLayer, self).__init__()
+
+        self.norm_first = norm_first
+
+        self.norm1 = nn.LayerNorm(emb_dim, layer_norm_eps)
+        self.norm2 = nn.LayerNorm(emb_dim, layer_norm_eps)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.self_attn = nn.MultiheadAttention(
+            emb_dim, num_heads, dropout
+        )
+        self.mlp = MLP(emb_dim, mlp_factor * emb_dim, dropout, activation)
+
+    def with_emb(self, x, emb):
+        return x if emb is None else x + emb
+
+    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
+        if self.norm_first:
+            src_norm = self.norm1(src)
+            q = k = src_norm + pos_emb
+            src = src + self.dropout1(self.self_attn(
+                query=q,
+                key=k,
+                value=src_norm,
+                attn_mask=src_mask,
+                key_padding_mask=src_key_padding_mask
+            )[0])
+
+            src_norm = self.norm2(src)
+            src = src + self.dropout2(self.mlp(src_norm))
+        else:
+            q = k = src + pos_emb
+            src = self.norm1(src + self.dropout1(self.self_attn(
+                query=q,
+                key=k,
+                value=src,
+                attn_mask=src_mask,
+                key_padding_mask=src_key_padding_mask
+            )[0]))
+
+            src = self.norm2(src + self.dropout2(self.mlp(src)))
+
+        return src

models/enc_model/unet_parts.py

@@ -0,0 +1,77 @@
+""" Parts of the U-Net model """
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)

models/model.py

@@ -0,0 +1,991 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import os
+import clip
+import sys
+from models.seg_post_model.cellpose.models import CellposeModel
+
+from torchvision.ops import roi_align
+def crop_roi_feat(feat, boxes):
+    """
+    feat: 1 x c x h x w
+    boxes: m x 4, 4: [y_tl, x_tl, y_br, x_br]
+    """
+    _, _, h, w = feat.shape
+    out_stride = 512 / h
+    boxes_scaled = boxes / out_stride
+    boxes_scaled[:, :2] = torch.floor(boxes_scaled[:, :2])  # y_tl, x_tl: floor
+    boxes_scaled[:, 2:] = torch.ceil(boxes_scaled[:, 2:])  # y_br, x_br: ceil
+    boxes_scaled[:, :2] = torch.clamp_min(boxes_scaled[:, :2], 0)
+    boxes_scaled[:, 2] = torch.clamp_max(boxes_scaled[:, 2], h)
+    boxes_scaled[:, 3] = torch.clamp_max(boxes_scaled[:, 3], w)
+    feat_boxes = []
+    for idx_box in range(0, boxes.shape[0]):
+        y_tl, x_tl, y_br, x_br = boxes_scaled[idx_box]
+        y_tl, x_tl, y_br, x_br = int(y_tl), int(x_tl), int(y_br), int(x_br)
+        feat_box = feat[:, :, y_tl : (y_br + 1), x_tl : (x_br + 1)]
+        feat_boxes.append(feat_box)
+    return feat_boxes
+
+class Counting_with_SD_features(nn.Module):
+    def __init__(self, scale_factor):
+        super(Counting_with_SD_features, self).__init__()
+        self.adapter = adapter_roi()
+        # self.regressor = regressor_with_SD_features()
+
+class Counting_with_SD_features_loca(nn.Module):
+    def __init__(self, scale_factor):
+        super(Counting_with_SD_features_loca, self).__init__()
+        self.adapter = adapter_roi_loca()
+        self.regressor = regressor_with_SD_features()
+
+
+class Counting_with_SD_features_dino_vit_c3(nn.Module):
+    def __init__(self, scale_factor, vit=None):
+        super(Counting_with_SD_features_dino_vit_c3, self).__init__()
+        self.adapter = adapter_roi_loca()
+        self.regressor = regressor_with_SD_features_seg_vit_c3()
+
+class Counting_with_SD_features_track(nn.Module):
+    def __init__(self, scale_factor, vit=None):
+        super(Counting_with_SD_features_track, self).__init__()
+        self.adapter = adapter_roi_loca()
+        self.regressor = regressor_with_SD_features_tra()
+
+class Counting_with_SD_features_loca_rand(nn.Module):
+    def __init__(self, scale_factor, num_of_roi = 3):
+        super(Counting_with_SD_features_loca_rand, self).__init__()
+        self.adapter = adapter_roi_loca_rand(num_of_roi=num_of_roi)
+        self.regressor = regressor_with_SD_features()
+
+class Counting_with_SD_features_loca_carpk(nn.Module):
+    def __init__(self, scale_factor, num_of_roi = 3):
+        super(Counting_with_SD_features_loca_carpk, self).__init__()
+        self.adapter = adapter_roi_loca_carpk(num_of_roi=num_of_roi)
+        self.regressor = regressor_with_SD_features()
+
+class Counting_with_SD_features_clip_carpk(nn.Module):
+    def __init__(self, scale_factor, num_of_roi = 3):
+        super(Counting_with_SD_features_clip_carpk, self).__init__()
+        self.adapter = adapter_roi_clip_carpk(num_of_roi=num_of_roi)
+        # self.regressor = regressor_with_SD_features()
+
+class Counting_with_SD_features_zero(nn.Module):
+    def __init__(self, scale_factor):
+        super(Counting_with_SD_features_zero, self).__init__()
+        self.adapter = adapter_roi_zero()
+        self.regressor = regressor_with_SD_features()
+
+class Counting_with_SD_features_zero_loca(nn.Module):
+    def __init__(self, scale_factor):
+        super(Counting_with_SD_features_zero_loca, self).__init__()
+        self.adapter = adapter_roi_zero_loca()
+        self.regressor = regressor_with_SD_features()
+
+class Counting_with_SD_features_zero_loca_self(nn.Module):
+    def __init__(self, scale_factor):
+        super(Counting_with_SD_features_zero_loca_self, self).__init__()
+        self.adapter = adapter_roi_zero_loca()
+        # self.regressor = regressor_with_SD_features_self()
+        self.regressor = regressor_with_SD_features_latent()
+
+class Counting_with_SD_features_loca_v2(nn.Module):
+    def __init__(self, scale_factor):
+        super(Counting_with_SD_features_loca_v2, self).__init__()
+        self.adapter = adapter_roi_loca_v2()
+        # self.regressor = regressor_with_SD_features()
+
+class adapter1(nn.Module):
+    def __init__(self):
+        super(adapter1, self).__init__()
+        self.conv1 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(128 * 64 * 64, 768)
+        self.initialize_weights()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.pool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class adapter(nn.Module):
+    def __init__(self, pool_size=[3, 3]):
+        super(adapter, self).__init__()
+        self.pool_size = pool_size
+        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(256 * 3 * 3, 768)
+        self.initialize_weights()
+
+    def forward(self, xs):
+        x_list = []
+        for x in xs:
+            x = F.adaptive_max_pool2d(x, self.pool_size, return_indices=False) # [1, 256, 3, 3]
+            x_list.append(x)
+        x_list = torch.cat(x_list, dim=0)
+        x_list = torch.mean(x_list, dim=0, keepdim=True) # [1, 256, 3, 3]
+        x = self.conv1(x_list)
+        # x = self.pool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class adapter_roi(nn.Module):
+    def __init__(self, pool_size=[3, 3]):
+        super(adapter_roi, self).__init__()
+        self.pool_size = pool_size
+        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        # self.relu = nn.ReLU()
+        # self.conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(256 * 3 * 3, 768)
+        # **new
+        self.fc1 = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(768, 768 // 4, bias=False),
+            nn.ReLU()
+            )
+        self.fc2 = nn.Sequential(
+            nn.Linear(768 // 4, 768, bias=False),
+            # nn.ReLU()
+            )
+        self.initialize_weights()
+
+    def forward(self, x, boxes):
+            num_of_boxes = boxes.shape[1]
+            rois = []
+            bs, _, h, w = x.shape
+            boxes = torch.cat([
+                torch.arange(
+                    bs, requires_grad=False
+                ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
+                boxes.flatten(0, 1),
+            ], dim=1)
+            rois = roi_align(
+                x,
+                boxes=boxes, output_size=3,
+                spatial_scale=1.0 / 8, aligned=True
+            )
+            rois = torch.mean(rois, dim=0, keepdim=True)
+            x = self.conv1(rois)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+
+            x = self.fc1(x)
+            x = self.fc2(x)
+            return x
+ 
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class adapter_roi_loca(nn.Module):
+    def __init__(self, pool_size=[3, 3]):
+        super(adapter_roi_loca, self).__init__()
+        self.pool_size = pool_size
+        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(256 * 3 * 3, 768)
+        self.initialize_weights()
+    def forward(self, x, boxes):
+            num_of_boxes = boxes.shape[1]
+            rois = []
+            bs, _, h, w = x.shape
+            if h != 512 or w != 512:
+                x = F.interpolate(x, size=(512, 512), mode='bilinear', align_corners=False)
+            if bs == 1:
+                boxes = torch.cat([
+                    torch.arange(
+                        bs, requires_grad=False
+                    ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
+                    boxes.flatten(0, 1),
+                ], dim=1)
+                rois = roi_align(
+                    x,
+                    boxes=boxes, output_size=3,
+                    spatial_scale=1.0 / 8, aligned=True
+                )
+                rois = torch.mean(rois, dim=0, keepdim=True)
+            else:
+                boxes = torch.cat([
+                    boxes.flatten(0, 1),
+                ], dim=1).split(num_of_boxes, dim=0)
+                rois = roi_align(
+                    x,
+                    boxes=boxes, output_size=3,
+                    spatial_scale=1.0 / 8, aligned=True
+                )
+                rois = rois.split(num_of_boxes, dim=0)
+                rois = torch.stack(rois, dim=0)
+                rois = torch.mean(rois, dim=1, keepdim=False)
+            x = self.conv1(rois)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+            return x
+    
+    def forward_boxes(self, x, boxes):
+            num_of_boxes = boxes.shape[1]
+            rois = []
+            bs, _, h, w = x.shape
+            if h != 512 or w != 512:
+                x = F.interpolate(x, size=(512, 512), mode='bilinear', align_corners=False)
+            if bs == 1:
+                boxes = torch.cat([
+                    torch.arange(
+                        bs, requires_grad=False
+                    ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
+                    boxes.flatten(0, 1),
+                ], dim=1)
+                rois = roi_align(
+                    x,
+                    boxes=boxes, output_size=3,
+                    spatial_scale=1.0 / 8, aligned=True
+                )
+                # rois = torch.mean(rois, dim=0, keepdim=True)
+            else:
+                raise NotImplementedError
+            x = self.conv1(rois)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+            return x
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class adapter_roi_dino(nn.Module):
+    def __init__(self, pool_size=[3, 3]):
+        super(adapter_roi_dino, self).__init__()
+        self.pool_size = pool_size
+        # self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        # self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(1024, 768)
+        self.initialize_weights()
+    def forward(self, crops, dino_model):
+            num_of_boxes = len(crops)
+            feats = []
+            for i in range(num_of_boxes):
+                with torch.no_grad():
+                    feat = dino_model(crops[i])
+
+                feats.append(feat)
+            feats = torch.cat(feats, dim=0)
+            feats = torch.mean(feats, dim=0) 
+            x = self.fc(feats)
+            return x
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+
+class adapter_roi_loca_v2(nn.Module):
+    def __init__(self, pool_size=[3, 3]):
+        super(adapter_roi_loca_v2, self).__init__()
+        self.pool_size = pool_size
+        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(256 * 3 * 3, 1024)
+        self.initialize_weights()
+    def forward(self, x, boxes):
+            rois = []
+            bs, _, h, w = x.shape
+            boxes = torch.cat([
+                torch.arange(
+                    bs, requires_grad=False
+                ).to(boxes.device).repeat_interleave(3).reshape(-1, 1),
+                boxes.flatten(0, 1),
+            ], dim=1)
+            rois = roi_align(
+                x,
+                boxes=boxes, output_size=3,
+                spatial_scale=1.0 / 8, aligned=True
+            )
+            rois = torch.mean(rois, dim=0, keepdim=True)
+            x = self.conv1(rois)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+            return x
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class adapter_roi_zero(nn.Module):
+    def __init__(self, reduction=4):
+        super(adapter_roi_zero, self).__init__()
+        self.fc1 = nn.Sequential(
+            nn.Linear(768, 768 // reduction, bias=False),
+            nn.ReLU()
+            )
+        self.fc2 = nn.Sequential(
+            nn.Linear(768 // reduction, 768, bias=False),
+            nn.ReLU()
+            )
+        self.initialize_weights()
+    def forward(self, x):
+            x1 = self.fc1(x)
+            x1 = self.fc2(x1)
+            return x + x1
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class adapter_roi_zero_loca(nn.Module):
+    def __init__(self, reduction=4):
+        super(adapter_roi_zero_loca, self).__init__()
+        self.fc1 = nn.Sequential(
+            nn.Linear(768, 768 // reduction, bias=False),
+            nn.ReLU()
+            )
+        self.fc2 = nn.Sequential(
+            nn.Linear(768 // reduction, 768, bias=False),
+            nn.ReLU()
+            )
+        
+        self.pool_size = (3, 3)
+        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(256 * 3 * 3, 768)
+
+        self.initialize_weights()
+    def forward(self, feature, boxes, class_emb):
+            x1 = self.fc1(class_emb)
+            x1 = self.fc2(x1)
+            class_emb = class_emb + x1
+
+            rois = []
+            bs, _, h, w = feature.shape
+            n_box = boxes.shape[1]
+            boxes = torch.cat([
+                torch.arange(
+                    bs, requires_grad=False
+                ).to(boxes.device).repeat_interleave(n_box).reshape(-1, 1),
+                boxes.flatten(0, 1),
+            ], dim=1)
+            rois = roi_align(
+                feature,
+                boxes=boxes, output_size=3,
+                spatial_scale=1.0 / 8, aligned=True
+            )
+            # rois = torch.mean(rois, dim=0, keepdim=True)
+            x = self.conv1(rois)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+
+            if len(class_emb.shape) == 3:
+                class_emb = class_emb.squeeze(1)
+            dist = torch.cosine_similarity(class_emb, x) # [n_box]
+            _, topk = torch.sort(dist[:10])
+            x_topk = x[topk[:3], :]
+            x_topk = torch.mean(x_topk, dim=0, keepdim=True)
+            return x_topk + class_emb
+    
+    def vis(self, feature, boxes, class_emb):
+        x1 = self.fc1(class_emb)
+        x1 = self.fc2(x1)
+        class_emb = class_emb + x1
+
+        rois = []
+        bs, _, h, w = feature.shape
+        n_box = boxes.shape[1]
+        boxes = torch.cat([
+            torch.arange(
+                bs, requires_grad=False
+            ).to(boxes.device).repeat_interleave(n_box).reshape(-1, 1),
+            boxes.flatten(0, 1),
+        ], dim=1)
+        rois = roi_align(
+            feature,
+            boxes=boxes, output_size=3,
+            spatial_scale=1.0 / 8, aligned=True
+        )
+        # rois = torch.mean(rois, dim=0, keepdim=True)
+        x = self.conv1(rois)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        if len(class_emb.shape) == 3:
+            class_emb = class_emb.squeeze(1)
+        dist = torch.cosine_similarity(class_emb, x) # [n_box]
+        _, topk = torch.sort(dist[:10])
+        x_topk = x[topk[:3], :]
+        x_topk = torch.mean(x_topk, dim=0, keepdim=True)
+        return x_topk
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class adapter_roi_loca_rand(nn.Module):
+    def __init__(self, pool_size=[3, 3],num_of_roi = 3):
+        super(adapter_roi_loca_rand, self).__init__()
+        self.pool_size = pool_size
+        self.num_of_roi = num_of_roi
+        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.fc = nn.Linear(256 * 3 * 3, 768)
+
+        # # **new
+        # self.fc1 = nn.Sequential(
+        #     nn.Linear(768, 768 // 4, bias=False),
+        #     nn.ReLU()
+        #     )
+        # self.fc2 = nn.Sequential(
+        #     nn.Linear(768 // 4, 768, bias=False),
+        #     nn.ReLU()
+        #     )
+        # # 
+        self.initialize_weights()
+    def forward(self, x, boxes, rand_boxes):
+            num_of_boxes = boxes.shape[1]
+            bs, _, h, w = x.shape
+            boxes = torch.cat([
+                torch.arange(
+                    bs, requires_grad=False
+                ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
+                boxes.flatten(0, 1),
+            ], dim=1)
+            rois = roi_align(
+                x,
+                boxes=boxes, output_size=3,
+                spatial_scale=1.0 / 8, aligned=True
+            )
+
+            # new
+            num_of_boxes = rand_boxes.shape[1]
+            bs, _, h, w = x.shape
+            rand_boxes = torch.cat([
+                torch.arange(
+                    bs, requires_grad=False
+                ).to(rand_boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
+                rand_boxes.flatten(0, 1),
+            ], dim=1)
+            rand_rois = roi_align(
+                x,
+                boxes=rand_boxes, output_size=3,
+                spatial_scale=1.0 / 8, aligned=True
+            )
+
+            rois = torch.mean(rois, dim=0, keepdim=True)
+
+            # new
+            cos = torch.nn.CosineSimilarity(dim=1)
+            dist = cos(rois.view(1, -1), rand_rois.view(num_of_boxes, -1)) # [n_box]
+            _, topk = torch.sort(-dist)
+            x_topk = rand_rois[topk[:3], ...]
+            x_topk = torch.mean(x_topk, dim=0, keepdim=True)
+
+            rois += x_topk
+
+            x = self.conv1(rois)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+            # new
+            # x = self.fc1(x)
+            # x = self.fc2(x)
+            return x
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class regressor1(nn.Module):
+    def __init__(self):
+        super(regressor1, self).__init__()
+        self.conv1 = nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(4, 1, kernel_size=3, stride=1, padding=1)
+        self.upsampler = nn.UpsamplingBilinear2d(scale_factor=2)
+        self.leaky_relu = nn.LeakyReLU()
+        self.relu = nn.ReLU()
+        self.initialize_weights()
+
+    
+
+    def forward(self, x):
+        x_ = self.conv1(x)
+        x_ = self.leaky_relu(x_)
+        x_ = self.upsampler(x_)
+        x_ = self.conv2(x_)
+        x_ = self.leaky_relu(x_)
+        x_ = self.upsampler(x_)
+        x_ = self.conv3(x_)
+        x_ = self.relu(x_)
+        out = x_
+        return out
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class regressor1(nn.Module):
+    def __init__(self):
+        super(regressor1, self).__init__()
+        self.conv1 = nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(4, 1, kernel_size=3, stride=1, padding=1)
+        self.upsampler = nn.UpsamplingBilinear2d(scale_factor=2)
+        self.leaky_relu = nn.LeakyReLU()
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x_ = self.conv1(x)
+        x_ = self.leaky_relu(x_)
+        x_ = self.upsampler(x_)
+        x_ = self.conv2(x_)
+        x_ = self.leaky_relu(x_)
+        x_ = self.upsampler(x_)
+        x_ = self.conv3(x_)
+        x_ = self.relu(x_)
+        out = x_
+        return out
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class regressor_with_SD_features(nn.Module):
+    def __init__(self):
+        super(regressor_with_SD_features, self).__init__()
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(324, 256, kernel_size=1, stride=1),
+            nn.LeakyReLU(),
+            nn.LayerNorm((64, 64))
+        )
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(128, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(64, 32, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=4, stride=2, padding=1),
+        )
+        self.conv = nn.Sequential(
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.ReLU()
+        )
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.initialize_weights()
+
+    def forward(self, attn_stack, feature_list):
+        attn_stack = self.norm(attn_stack)
+        unet_feature = feature_list[-1]
+        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        unet_feature = unet_feature * attn_stack_mean
+        unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
+        x = self.layer1(unet_feature)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        out = self.conv(x)
+        return out / 100
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class regressor_with_SD_features_seg(nn.Module):
+    def __init__(self):
+        super(regressor_with_SD_features_seg, self).__init__()
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(324, 256, kernel_size=1, stride=1),
+            nn.LeakyReLU(),
+            nn.LayerNorm((64, 64))
+        )
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(128, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(64, 32, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=4, stride=2, padding=1),
+        )
+        self.conv = nn.Sequential(
+            nn.Conv2d(32, 2, kernel_size=1),
+            # nn.ReLU()
+        )
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.initialize_weights()
+
+    def forward(self, attn_stack, feature_list):
+        attn_stack = self.norm(attn_stack)
+        unet_feature = feature_list[-1]
+        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        unet_feature = unet_feature * attn_stack_mean
+        unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
+        x = self.layer1(unet_feature)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        out = self.conv(x)
+        return out
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+from models.enc_model.unet_parts import *
+
+
+class regressor_with_SD_features_seg_vit_c3(nn.Module):
+    def __init__(self, n_channels=3, n_classes=2, bilinear=False):
+        super(regressor_with_SD_features_seg_vit_c3, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.inc_0 = nn.Conv2d(n_channels, 3, kernel_size=3, padding=1)
+        self.vit_model = CellposeModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)
+        self.vit = self.vit_model.net
+
+    def forward(self, img, attn_stack, feature_list):
+        attn_stack = attn_stack[:, [1,3], ...]
+        attn_stack = self.norm(attn_stack)
+        unet_feature = feature_list[-1]
+        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)
+
+        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]
+
+        if x.shape[-1] != 512:
+            x = F.interpolate(x, size=(512, 512), mode="bilinear")
+        x = self.inc_0(x)
+
+
+        
+        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())[0]
+        out = torch.from_numpy(out).unsqueeze(0).to(x.device)
+        return out
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class regressor_with_SD_features_tra(nn.Module):
+    def __init__(self, n_channels=2, n_classes=2, bilinear=False):
+        super(regressor_with_SD_features_tra, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+
+        # segmentation
+        self.inc_0 = nn.Conv2d(3, 3, kernel_size=3, padding=1)
+        self.vit_model = CellposeModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)
+        self.vit = self.vit_model.net
+
+        self.inc_1 = nn.Conv2d(n_channels, 1, kernel_size=3, padding=1)
+        self.mlp = nn.Linear(64 * 64, 320)
+        # self.vit = self.vit_model.net.float()
+
+    def forward_seg(self, img, attn_stack, feature_list, mask, training=False):
+        attn_stack = attn_stack[:, [1,3], ...]
+        attn_stack = self.norm(attn_stack)
+        unet_feature = feature_list[-1]
+        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)
+        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]
+
+        if x.shape[-1] != 512:
+            x = F.interpolate(x, size=(512, 512), mode="bilinear")
+        x = self.inc_0(x)
+        feat = x
+        
+        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())[0]
+        out = torch.from_numpy(out).unsqueeze(0).to(x.device)
+        return out, 0., feat
+
+    def forward(self, attn_prev, feature_list_prev, attn_after, feature_list_after):
+        assert attn_prev.shape == attn_after.shape, "attn_prev and attn_after must have the same shape"
+        n_instances = attn_prev.shape[0] 
+        attn_prev = self.norm(attn_prev) # [n_instances, 1, 64, 64]
+        attn_after = self.norm(attn_after)
+        
+        x = torch.cat([attn_prev, attn_after], dim=1)  # n_instances, 2, 64, 64
+        
+        x = self.inc_1(x)
+        x = x.view(1, n_instances, -1)  # Flatten the tensor to [n_instances, 64*64*4]
+        x = self.mlp(x)  # Apply the MLP to get the output
+        
+        return x  # Output shape will be [n_instances, 4]
+
+        
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+
+class regressor_with_SD_features_inst_seg_unet(nn.Module):
+    def __init__(self, n_channels=8, n_classes=3, bilinear=False):
+        super(regressor_with_SD_features_inst_seg_unet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.inc_0 = (DoubleConv(n_channels, 3))
+        self.inc = (DoubleConv(3, 64))
+        self.down1 = (Down(64, 128))
+        self.down2 = (Down(128, 256))
+        self.down3 = (Down(256, 512))
+        factor = 2 if bilinear else 1
+        self.down4 = (Down(512, 1024 // factor))
+        self.up1 = (Up(1024, 512 // factor, bilinear))
+        self.up2 = (Up(512, 256 // factor, bilinear))
+        self.up3 = (Up(256, 128 // factor, bilinear))
+        self.up4 = (Up(128, 64, bilinear))
+        self.outc = (OutConv(64, n_classes))
+
+    def forward(self, img, attn_stack, feature_list):
+        attn_stack = self.norm(attn_stack)
+        unet_feature = feature_list[-1]
+        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)
+        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        unet_feature_mean = unet_feature_mean * attn_stack_mean
+        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]
+        if x.shape[-1] != 512:
+            x = F.interpolate(x, size=(512, 512), mode="bilinear")
+        x = torch.cat([img, x], dim=1) # [1, 8, 512, 512]
+        x = self.inc_0(x)
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        out = self.outc(x)
+        return out
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class regressor_with_SD_features_self(nn.Module):
+    def __init__(self):
+        super(regressor_with_SD_features_self, self).__init__()
+        self.layer = nn.Sequential(
+            nn.Conv2d(4096, 1024, kernel_size=1, stride=1),
+            nn.LeakyReLU(),
+            nn.LayerNorm((64, 64)),
+            nn.Conv2d(1024, 256, kernel_size=1, stride=1),
+            nn.LeakyReLU(),
+            nn.LayerNorm((64, 64)),
+        )
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(128, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(64, 32, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=4, stride=2, padding=1),
+        )
+        self.conv = nn.Sequential(
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.ReLU()
+        )
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.initialize_weights()
+
+    def forward(self, self_attn):
+        self_attn = self_attn.permute(2, 0, 1)
+        self_attn = self.layer(self_attn)
+        return self_attn
+        # attn_stack = self.norm(attn_stack)
+        # unet_feature = feature_list[-1]
+        # attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        # unet_feature = unet_feature * attn_stack_mean
+        # unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
+        # x = self.layer(unet_feature)
+        # x = self.layer2(x)
+        # x = self.layer3(x)
+        # x = self.layer4(x)
+        # out = self.conv(x)
+        # return out / 100
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class regressor_with_SD_features_latent(nn.Module):
+    def __init__(self):
+        super(regressor_with_SD_features_latent, self).__init__()
+        self.layer = nn.Sequential(
+            nn.Conv2d(4, 256, kernel_size=1, stride=1),
+            nn.LeakyReLU(),
+            nn.LayerNorm((64, 64))
+        )
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(128, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=4, stride=2, padding=1),
+        )
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(64, 32, kernel_size=3, padding=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=4, stride=2, padding=1),
+        )
+        self.conv = nn.Sequential(
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.ReLU()
+        )
+        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
+        self.initialize_weights()
+
+    def forward(self, self_attn):
+        # self_attn = self_attn.permute(2, 0, 1)
+        self_attn = self.layer(self_attn)
+        return self_attn
+        # attn_stack = self.norm(attn_stack)
+        # unet_feature = feature_list[-1]
+        # attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+        # unet_feature = unet_feature * attn_stack_mean
+        # unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
+        # x = self.layer(unet_feature)
+        # x = self.layer2(x)
+        # x = self.layer3(x)
+        # x = self.layer4(x)
+        # out = self.conv(x)
+        # return out / 100
+    
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+
+
+
+class regressor_with_deconv(nn.Module):
+    def __init__(self):
+        super(regressor_with_deconv, self).__init__()
+        self.conv1 = nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(4, 1, kernel_size=3, stride=1, padding=1)
+        self.deconv1 = nn.ConvTranspose2d(4, 4, kernel_size=4, stride=2, padding=1)
+        self.deconv2 = nn.ConvTranspose2d(4, 4, kernel_size=4, stride=2, padding=1)
+        self.leaky_relu = nn.LeakyReLU()
+        self.relu = nn.ReLU()
+        self.initialize_weights()
+
+    def forward(self, x):
+        x_ = self.conv1(x)
+        x_ = self.leaky_relu(x_)
+        x_ = self.deconv1(x_)
+        x_ = self.conv2(x_)
+        x_ = self.leaky_relu(x_)
+        x_ = self.deconv2(x_)
+        x_ = self.conv3(x_)
+        x_ = self.relu(x_)
+        out = x_
+        return out
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+

models/seg_post_model/cellpose/__init__.py

@@ -0,0 +1 @@
+from .version import version, version_str

models/seg_post_model/cellpose/__main__.py

@@ -0,0 +1,272 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os, time
+import numpy as np
+from tqdm import tqdm
+from cellpose import utils, models, io, train
+from .version import version_str
+from cellpose.cli import get_arg_parser
+
+try:
+    from cellpose.gui import gui3d, gui
+    GUI_ENABLED = True
+except ImportError as err:
+    GUI_ERROR = err
+    GUI_ENABLED = False
+    GUI_IMPORT = True
+except Exception as err:
+    GUI_ENABLED = False
+    GUI_ERROR = err
+    GUI_IMPORT = False
+    raise
+
+import logging
+
+
+def main():
+    """ Run cellpose from command line
+    """
+
+    args = get_arg_parser().parse_args()  # this has to be in a separate file for autodoc to work
+
+    if args.version:
+        print(version_str)
+        return
+
+    ######## if no image arguments are provided, run GUI or add model and exit ########
+    if len(args.dir) == 0 and len(args.image_path) == 0:
+        if args.add_model:
+            io.add_model(args.add_model)
+            return
+        else:
+            if not GUI_ENABLED:
+                print("GUI ERROR: %s" % GUI_ERROR)
+                if GUI_IMPORT:
+                    print(
+                        "GUI FAILED: GUI dependencies may not be installed, to install, run"
+                    )
+                    print("     pip install 'cellpose[gui]'")
+            else:
+                if args.Zstack:
+                    gui3d.run()
+                else:
+                    gui.run()
+            return
+
+    ############################## run cellpose on images ##############################
+    if args.verbose:
+        from .io import logger_setup
+        logger, log_file = logger_setup()
+    else:
+        print(
+            ">>>> !LOGGING OFF BY DEFAULT! To see cellpose progress, set --verbose")
+        print("No --verbose => no progress or info printed")
+        logger = logging.getLogger(__name__)
+
+
+    # find images
+    if len(args.img_filter) > 0:
+        image_filter = args.img_filter
+    else:
+        image_filter = None
+
+    device, gpu = models.assign_device(use_torch=True, gpu=args.use_gpu,
+                                        device=args.gpu_device)
+
+    if args.pretrained_model is None or args.pretrained_model == "None" or args.pretrained_model == "False" or args.pretrained_model == "0":
+        pretrained_model = "cpsam"
+        logger.warning("training from scratch is disabled, using 'cpsam' model")
+    else:
+        pretrained_model = args.pretrained_model
+
+    # Warn users about old arguments from CP3:
+    if args.pretrained_model_ortho:
+        logger.warning(
+            "the '--pretrained_model_ortho' flag is deprecated in v4.0.1+ and no longer used")
+    if args.train_size:
+        logger.warning("the '--train_size' flag is deprecated in v4.0.1+ and no longer used")
+    if args.chan or args.chan2:
+        logger.warning('--chan and --chan2 are deprecated, all channels are used by default')
+    if args.all_channels:
+        logger.warning("the '--all_channels' flag is deprecated in v4.0.1+ and no longer used")
+    if args.restore_type:
+        logger.warning("the '--restore_type' flag is deprecated in v4.0.1+ and no longer used")
+    if args.transformer:
+        logger.warning("the '--tranformer' flag is deprecated in v4.0.1+ and no longer used")
+    if args.invert:
+        logger.warning("the '--invert' flag is deprecated in v4.0.1+ and no longer used")
+    if args.chan2_restore:
+        logger.warning("the '--chan2_restore' flag is deprecated in v4.0.1+ and no longer used")
+    if args.diam_mean:
+        logger.warning("the '--diam_mean' flag is deprecated in v4.0.1+ and no longer used")
+    if args.train_size:
+        logger.warning("the '--train_size' flag is deprecated in v4.0.1+ and no longer used")
+    
+    if args.norm_percentile is not None:
+        value1, value2 = args.norm_percentile
+        normalize = {'percentile': (float(value1), float(value2))} 
+    else:
+        normalize = (not args.no_norm)
+
+    if args.save_each:
+        if not args.save_every:
+            raise ValueError("ERROR: --save_each requires --save_every")
+
+    if len(args.image_path) > 0 and args.train:
+        raise ValueError("ERROR: cannot train model with single image input")
+
+    ## Run evaluation on images
+    if not args.train:
+        _evaluate_cellposemodel_cli(args, logger, image_filter, device, pretrained_model, normalize)
+    
+    ## Train a model ##
+    else:
+        _train_cellposemodel_cli(args, logger, image_filter, device, pretrained_model, normalize)
+
+
+def _train_cellposemodel_cli(args, logger, image_filter, device, pretrained_model, normalize):
+    test_dir = None if len(args.test_dir) == 0 else args.test_dir
+    images, labels, image_names, train_probs = None, None, None, None
+    test_images, test_labels, image_names_test, test_probs = None, None, None, None
+    compute_flows = False
+    if len(args.file_list) > 0:
+        if os.path.exists(args.file_list):
+            dat = np.load(args.file_list, allow_pickle=True).item()
+            image_names = dat["train_files"]
+            image_names_test = dat.get("test_files", None)
+            train_probs = dat.get("train_probs", None)
+            test_probs = dat.get("test_probs", None)
+            compute_flows = dat.get("compute_flows", False)
+            load_files = False
+        else:
+            logger.critical(f"ERROR: {args.file_list} does not exist")
+    else:
+        output = io.load_train_test_data(args.dir, test_dir, image_filter,
+                                                args.mask_filter,
+                                                args.look_one_level_down)
+        images, labels, image_names, test_images, test_labels, image_names_test = output
+        load_files = True
+
+    # initialize model
+    model = models.CellposeModel(device=device, pretrained_model=pretrained_model)
+
+    # train segmentation model
+    cpmodel_path = train.train_seg(
+            model.net, images, labels, train_files=image_names,
+            test_data=test_images, test_labels=test_labels,
+            test_files=image_names_test, train_probs=train_probs,
+            test_probs=test_probs, compute_flows=compute_flows,
+            load_files=load_files, normalize=normalize,
+            channel_axis=args.channel_axis, 
+            learning_rate=args.learning_rate, weight_decay=args.weight_decay,
+            SGD=args.SGD, n_epochs=args.n_epochs, batch_size=args.train_batch_size,
+            min_train_masks=args.min_train_masks,
+            nimg_per_epoch=args.nimg_per_epoch,
+            nimg_test_per_epoch=args.nimg_test_per_epoch,
+            save_path=os.path.realpath(args.dir), 
+            save_every=args.save_every,
+            save_each=args.save_each,
+            model_name=args.model_name_out)[0]
+    model.pretrained_model = cpmodel_path
+    logger.info(">>>> model trained and saved to %s" % cpmodel_path)
+    return model
+
+
+def _evaluate_cellposemodel_cli(args, logger, imf, device, pretrained_model, normalize):
+    # Check with user if they REALLY mean to run without saving anything
+    if not args.train:
+        saving_something = args.save_png or args.save_tif or args.save_flows or args.save_txt
+
+    tic = time.time()
+    if len(args.dir) > 0:
+        image_names = io.get_image_files(
+                args.dir, args.mask_filter, imf=imf,
+                look_one_level_down=args.look_one_level_down)
+    else:
+        if os.path.exists(args.image_path):
+            image_names = [args.image_path]
+        else:
+            raise ValueError(f"ERROR: no file found at {args.image_path}")
+    nimg = len(image_names)
+
+    if args.savedir:
+        if not os.path.exists(args.savedir):
+            raise FileExistsError(f"--savedir {args.savedir} does not exist")
+        
+    logger.info(
+            ">>>> running cellpose on %d images using all channels" % nimg)
+
+    # handle built-in model exceptions
+    model = models.CellposeModel(device=device, pretrained_model=pretrained_model,)
+
+    tqdm_out = utils.TqdmToLogger(logger, level=logging.INFO)
+
+    channel_axis = args.channel_axis
+    z_axis = args.z_axis
+
+    for image_name in tqdm(image_names, file=tqdm_out):
+        if args.do_3D or args.stitch_threshold > 0.:
+            logger.info('loading image as 3D zstack')
+            image = io.imread_3D(image_name)
+            if channel_axis is None:
+                channel_axis = 3
+            if z_axis is None:
+                z_axis = 0
+                
+        else:
+            image = io.imread_2D(image_name)
+        out = model.eval(
+                image, 
+                diameter=args.diameter, 
+                do_3D=args.do_3D,
+                augment=args.augment, 
+                flow_threshold=args.flow_threshold,
+                cellprob_threshold=args.cellprob_threshold,
+                stitch_threshold=args.stitch_threshold, 
+                min_size=args.min_size,
+                batch_size=args.batch_size,
+                bsize=args.bsize,
+                resample=not args.no_resample,
+                normalize=normalize,
+                channel_axis=channel_axis, 
+                z_axis=z_axis,
+                anisotropy=args.anisotropy, 
+                niter=args.niter,
+                flow3D_smooth=args.flow3D_smooth)
+        masks, flows = out[:2]
+
+        if args.exclude_on_edges:
+            masks = utils.remove_edge_masks(masks)
+        if not args.no_npy:
+            io.masks_flows_to_seg(image, masks, flows, image_name,
+                                        imgs_restore=None, 
+                                        restore_type=None,
+                                        ratio=1.)
+        if saving_something:
+            suffix = "_cp_masks"
+            if args.output_name is not None: 
+                    # (1) If `savedir` is not defined, then must have a non-zero `suffix`
+                if args.savedir is None and len(args.output_name) > 0:
+                    suffix = args.output_name
+                elif args.savedir is not None and not os.path.samefile(args.savedir, args.dir):
+                        # (2) If `savedir` is defined, and different from `dir` then                              
+                        # takes the value passed as a param. (which can be empty string)
+                    suffix = args.output_name
+
+            io.save_masks(image, masks, flows, image_name,
+                                suffix=suffix, png=args.save_png,
+                                tif=args.save_tif, save_flows=args.save_flows,
+                                save_outlines=args.save_outlines,
+                                dir_above=args.dir_above, savedir=args.savedir,
+                                save_txt=args.save_txt, in_folders=args.in_folders,
+                                save_mpl=args.save_mpl)
+        if args.save_rois:
+            io.save_rois(masks, image_name)
+    logger.info(">>>> completed in %0.3f sec" % (time.time() - tic))
+
+    return model
+
+
+if __name__ == "__main__":
+    main()

models/seg_post_model/cellpose/cli.py

@@ -0,0 +1,240 @@
+"""
+Copyright © 2023 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu and Michael Rariden.
+"""
+
+import argparse
+
+
+def get_arg_parser():
+    """ Parses command line arguments for cellpose main function
+
+    Note: this function has to be in a separate file to allow autodoc to work for CLI.
+    The autodoc_mock_imports in conf.py does not work for sphinx-argparse sometimes,
+    see https://github.com/ashb/sphinx-argparse/issues/9#issue-1097057823
+    """
+
+    parser = argparse.ArgumentParser(description="Cellpose Command Line Parameters")
+
+    # misc settings
+    parser.add_argument("--version", action="store_true",
+                        help="show cellpose version info")
+    parser.add_argument(
+        "--verbose", action="store_true",
+        help="show information about running and settings and save to log")
+    parser.add_argument("--Zstack", action="store_true", help="run GUI in 3D mode")
+
+    # settings for CPU vs GPU
+    hardware_args = parser.add_argument_group("Hardware Arguments")
+    hardware_args.add_argument("--use_gpu", action="store_true",
+                               help="use gpu if torch with cuda installed")
+    hardware_args.add_argument(
+        "--gpu_device", required=False, default="0", type=str,
+        help="which gpu device to use, use an integer for torch, or mps for M1")
+    
+    # settings for locating and formatting images
+    input_img_args = parser.add_argument_group("Input Image Arguments")
+    input_img_args.add_argument("--dir", default=[], type=str,
+                                help="folder containing data to run or train on.")
+    input_img_args.add_argument(
+        "--image_path", default=[], type=str, help=
+        "if given and --dir not given, run on single image instead of folder (cannot train with this option)"
+    )
+    input_img_args.add_argument(
+        "--look_one_level_down", action="store_true",
+        help="run processing on all subdirectories of current folder")
+    input_img_args.add_argument("--img_filter", default=[], type=str,
+                                help="end string for images to run on")
+    input_img_args.add_argument(
+        "--channel_axis", default=None, type=int,
+        help="axis of image which corresponds to image channels")
+    input_img_args.add_argument("--z_axis", default=None, type=int,
+                                help="axis of image which corresponds to Z dimension")
+    
+    # TODO: remove deprecated in future version
+    input_img_args.add_argument(
+        "--chan", default=0, type=int, help=
+        "Deprecated in v4.0.1+, not used. ")
+    input_img_args.add_argument(
+        "--chan2", default=0, type=int, help=
+        'Deprecated in v4.0.1+, not used. ')
+    input_img_args.add_argument("--invert", action="store_true", help=
+        'Deprecated in v4.0.1+, not used. ')
+    input_img_args.add_argument(
+        "--all_channels", action="store_true", help=
+        'Deprecated in v4.0.1+, not used. ')
+
+    # model settings
+    model_args = parser.add_argument_group("Model Arguments")
+    model_args.add_argument("--pretrained_model", required=False, default="cpsam",
+                            type=str,
+                            help="model to use for running or starting training")
+    model_args.add_argument(
+        "--add_model", required=False, default=None, type=str,
+        help="model path to copy model to hidden .cellpose folder for using in GUI/CLI")
+    model_args.add_argument("--pretrained_model_ortho", required=False, default=None,
+                            type=str,
+                            help="Deprecated in v4.0.1+, not used. ")
+    
+    # TODO: remove deprecated in future version
+    model_args.add_argument("--restore_type", required=False, default=None, type=str, help=
+        'Deprecated in v4.0.1+, not used. ')
+    model_args.add_argument("--chan2_restore", action="store_true", help=
+        'Deprecated in v4.0.1+, not used. ')
+    model_args.add_argument(
+        "--transformer", action="store_true", help=
+        "use transformer backbone (pretrained_model from Cellpose3 is transformer_cp3)")
+    
+    # algorithm settings
+    algorithm_args = parser.add_argument_group("Algorithm Arguments")
+    algorithm_args.add_argument("--no_norm", action="store_true",
+                                help="do not normalize images (normalize=False)")
+    algorithm_args.add_argument(
+        '--norm_percentile',
+        nargs=2,  # Require exactly two values
+        metavar=('VALUE1', 'VALUE2'),
+        help="Provide two float values to set norm_percentile (e.g., --norm_percentile 1 99)"
+    )
+    algorithm_args.add_argument(
+        "--do_3D", action="store_true",
+        help="process images as 3D stacks of images (nplanes x nchan x Ly x Lx")
+    algorithm_args.add_argument(
+        "--diameter", required=False, default=None, type=float, help=
+        "use to resize cells to the training diameter (30 pixels)"
+    )
+    algorithm_args.add_argument(
+        "--stitch_threshold", required=False, default=0.0, type=float,
+        help="compute masks in 2D then stitch together masks with IoU>0.9 across planes"
+    )
+    algorithm_args.add_argument(
+        "--min_size", required=False, default=15, type=int,
+        help="minimum number of pixels per mask, can turn off with -1")
+    algorithm_args.add_argument(
+        "--flow3D_smooth", required=False, default=0, type=float,
+        help="stddev of gaussian for smoothing of dP for dynamics in 3D, default of 0 means no smoothing")
+    algorithm_args.add_argument(
+        "--flow_threshold", default=0.4, type=float, help=
+        "flow error threshold, 0 turns off this optional QC step. Default: %(default)s")
+    algorithm_args.add_argument(
+        "--cellprob_threshold", default=0, type=float,
+        help="cellprob threshold, default is 0, decrease to find more and larger masks")
+    algorithm_args.add_argument(
+        "--niter", default=0, type=int, help=
+        "niter, number of iterations for dynamics for mask creation, default of 0 means it is proportional to diameter, set to a larger number like 2000 for very long ROIs"
+    )
+    algorithm_args.add_argument("--anisotropy", required=False, default=1.0, type=float,
+                                help="anisotropy of volume in 3D")
+    algorithm_args.add_argument("--exclude_on_edges", action="store_true",
+                                help="discard masks which touch edges of image")
+    algorithm_args.add_argument(
+        "--augment", action="store_true",
+        help="tiles image with overlapping tiles and flips overlapped regions to augment"
+    )
+    algorithm_args.add_argument("--batch_size", default=8, type=int,
+                               help="inference batch size. Default: %(default)s")
+
+    # TODO: remove deprecated in future version
+    algorithm_args.add_argument(
+        "--no_resample", action="store_true", 
+        help="disables flows/cellprob resampling to original image size before computing masks. Using this flag will make more masks more jagged with larger diameter settings.")
+    algorithm_args.add_argument(
+        "--no_interp", action="store_true",
+        help="do not interpolate when running dynamics (was default)")
+
+    # output settings
+    output_args = parser.add_argument_group("Output Arguments")
+    output_args.add_argument(
+        "--save_png", action="store_true",
+        help="save masks as png")
+    output_args.add_argument(
+        "--save_tif", action="store_true",
+        help="save masks as tif")
+    output_args.add_argument(
+        "--output_name", default=None, type=str,
+        help="suffix for saved masks, default is _cp_masks, can be empty if `savedir` used and different of `dir`")
+    output_args.add_argument("--no_npy", action="store_true",
+                             help="suppress saving of npy")
+    output_args.add_argument(
+        "--savedir", default=None, type=str, help=
+        "folder to which segmentation results will be saved (defaults to input image directory)"
+    )
+    output_args.add_argument(
+        "--dir_above", action="store_true", help=
+        "save output folders adjacent to image folder instead of inside it (off by default)"
+    )
+    output_args.add_argument("--in_folders", action="store_true",
+                             help="flag to save output in folders (off by default)")
+    output_args.add_argument(
+        "--save_flows", action="store_true", help=
+        "whether or not to save RGB images of flows when masks are saved (disabled by default)"
+    )
+    output_args.add_argument(
+        "--save_outlines", action="store_true", help=
+        "whether or not to save RGB outline images when masks are saved (disabled by default)"
+    )
+    output_args.add_argument(
+        "--save_rois", action="store_true",
+        help="whether or not to save ImageJ compatible ROI archive (disabled by default)"
+    )
+    output_args.add_argument(
+        "--save_txt", action="store_true",
+        help="flag to enable txt outlines for ImageJ (disabled by default)")
+    output_args.add_argument(
+        "--save_mpl", action="store_true",
+        help="save a figure of image/mask/flows using matplotlib (disabled by default). "
+        "This is slow, especially with large images.")
+
+    # training settings
+    training_args = parser.add_argument_group("Training Arguments")
+    training_args.add_argument("--train", action="store_true",
+                               help="train network using images in dir")
+    training_args.add_argument("--test_dir", default=[], type=str,
+                               help="folder containing test data (optional)")
+    training_args.add_argument(
+        "--file_list", default=[], type=str, help=
+        "path to list of files for training and testing and probabilities for each image (optional)"
+    )
+    training_args.add_argument(
+        "--mask_filter", default="_masks", type=str, help=
+        "end string for masks to run on. use '_seg.npy' for manual annotations from the GUI. Default: %(default)s"
+    )
+    training_args.add_argument("--learning_rate", default=1e-5, type=float,
+                               help="learning rate. Default: %(default)s")
+    training_args.add_argument("--weight_decay", default=0.1, type=float,
+                               help="weight decay. Default: %(default)s")
+    training_args.add_argument("--n_epochs", default=100, type=int,
+                               help="number of epochs. Default: %(default)s")
+    training_args.add_argument("--train_batch_size", default=1, type=int,
+                               help="training batch size. Default: %(default)s")
+    training_args.add_argument("--bsize", default=256, type=int,
+                               help="block size for tiles. Default: %(default)s")
+    training_args.add_argument(
+        "--nimg_per_epoch", default=None, type=int,
+        help="number of train images per epoch. Default is to use all train images.")
+    training_args.add_argument(
+        "--nimg_test_per_epoch", default=None, type=int,
+        help="number of test images per epoch. Default is to use all test images.")
+    training_args.add_argument(
+        "--min_train_masks", default=5, type=int, help=
+        "minimum number of masks a training image must have to be used. Default: %(default)s"
+    )
+    training_args.add_argument("--SGD", default=0, type=int, 
+                               help="Deprecated in v4.0.1+, not used - AdamW used instead. ")
+    training_args.add_argument(
+        "--save_every", default=100, type=int,
+        help="number of epochs to skip between saves. Default: %(default)s")
+    training_args.add_argument(
+        "--save_each", action="store_true",
+        help="wether or not to save each epoch. Must also use --save_every. (default: False)")
+    training_args.add_argument(
+        "--model_name_out", default=None, type=str,
+        help="Name of model to save as, defaults to name describing model architecture. "
+        "Model is saved in the folder specified by --dir in models subfolder.")
+    
+    # TODO: remove deprecated in future version
+    training_args.add_argument(
+        "--diam_mean", default=30., type=float, help=
+        'Deprecated in v4.0.1+, not used. ')
+    training_args.add_argument("--train_size", action="store_true", help=
+        'Deprecated in v4.0.1+, not used. ')
+
+    return parser

models/seg_post_model/cellpose/core.py

@@ -0,0 +1,322 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import logging
+import numpy as np
+from tqdm import trange
+from . import transforms, utils
+
+import torch
+
+TORCH_ENABLED = True
+
+core_logger = logging.getLogger(__name__)
+tqdm_out = utils.TqdmToLogger(core_logger, level=logging.INFO)
+
+
+def use_gpu(gpu_number=0, use_torch=True):
+    """ 
+    Check if GPU is available for use.
+
+    Args:
+        gpu_number (int): The index of the GPU to be used. Default is 0.
+        use_torch (bool): Whether to use PyTorch for GPU check. Default is True.
+
+    Returns:
+        bool: True if GPU is available, False otherwise.
+
+    Raises:
+        ValueError: If use_torch is False, as cellpose only runs with PyTorch now.
+    """
+    if use_torch:
+        return _use_gpu_torch(gpu_number)
+    else:
+        raise ValueError("cellpose only runs with PyTorch now")
+
+
+def _use_gpu_torch(gpu_number=0):
+    """
+    Checks if CUDA or MPS is available and working with PyTorch.
+
+    Args:
+        gpu_number (int): The GPU device number to use (default is 0).
+
+    Returns:
+        bool: True if CUDA or MPS is available and working, False otherwise.
+    """
+    try:
+        device = torch.device("cuda:" + str(gpu_number))
+        _ = torch.zeros((1,1)).to(device)
+        core_logger.info("** TORCH CUDA version installed and working. **")
+        return True
+    except:
+        pass
+    try:
+        device = torch.device('mps:' + str(gpu_number))
+        _ = torch.zeros((1,1)).to(device)
+        core_logger.info('** TORCH MPS version installed and working. **')
+        return True
+    except:
+        core_logger.info('Neither TORCH CUDA nor MPS version not installed/working.')
+        return False
+
+
+def assign_device(use_torch=True, gpu=False, device=0):
+    """
+    Assigns the device (CPU or GPU or mps) to be used for computation.
+
+    Args:
+        use_torch (bool, optional): Whether to use torch for GPU detection. Defaults to True.
+        gpu (bool, optional): Whether to use GPU for computation. Defaults to False.
+        device (int or str, optional): The device index or name to be used. Defaults to 0.
+
+    Returns:
+        torch.device, bool (True if GPU is used, False otherwise)
+    """
+
+    if isinstance(device, str):
+        if device != "mps" or not(gpu and torch.backends.mps.is_available()):
+            device = int(device)
+    if gpu and use_gpu(use_torch=True):
+        try:
+            if torch.cuda.is_available():
+                device = torch.device(f'cuda:{device}')
+                core_logger.info(">>>> using GPU (CUDA)")
+                gpu = True
+                cpu = False
+        except:
+            gpu = False
+            cpu = True
+        try:
+            if torch.backends.mps.is_available():
+                device = torch.device('mps')
+                core_logger.info(">>>> using GPU (MPS)")
+                gpu = True
+                cpu = False
+        except:
+            gpu = False
+            cpu = True
+    else:
+        device = torch.device('cpu')
+        core_logger.info('>>>> using CPU')
+        gpu = False
+        cpu = True
+    
+    if cpu:
+        device = torch.device("cpu")
+        core_logger.info(">>>> using CPU")
+        gpu = False
+    return device, gpu
+
+
+def _to_device(x, device, dtype=torch.float32):
+    """
+    Converts the input tensor or numpy array to the specified device.
+
+    Args:
+        x (torch.Tensor or numpy.ndarray): The input tensor or numpy array.
+        device (torch.device): The target device.
+
+    Returns:
+        torch.Tensor: The converted tensor on the specified device.
+    """
+    if not isinstance(x, torch.Tensor):
+        X = torch.from_numpy(x).to(device, dtype=dtype)
+        return X
+    else:
+        return x
+
+
+def _from_device(X):
+    """
+    Converts a PyTorch tensor from the device to a NumPy array on the CPU.
+
+    Args:
+        X (torch.Tensor): The input PyTorch tensor.
+
+    Returns:
+        numpy.ndarray: The converted NumPy array.
+    """
+    # The cast is so numpy conversion always works
+    x = X.detach().cpu().to(torch.float32).numpy()
+    return x
+
+
+def _forward(net, x, feat=None):
+    """Converts images to torch tensors, runs the network model, and returns numpy arrays.
+
+    Args:
+        net (torch.nn.Module): The network model.
+        x (numpy.ndarray): The input images.
+
+    Returns:
+        Tuple[numpy.ndarray, numpy.ndarray]: The output predictions (flows and cellprob) and style features.
+    """
+    X = _to_device(x, device=net.device, dtype=net.dtype)
+    if feat is not None:
+        feat = _to_device(feat, device=net.device, dtype=net.dtype)
+    net.eval()
+    with torch.no_grad():
+        y, style = net(X, feat=feat)[:2]
+    del X
+    y = _from_device(y)
+    style = _from_device(style)
+    return y, style
+
+
+def run_net(net, imgi, feat=None, batch_size=8, augment=False, tile_overlap=0.1, bsize=224,
+            rsz=None):
+    """ 
+    Run network on stack of images.
+    
+    (faster if augment is False)
+
+    Args:
+        net (class): cellpose network (model.net)
+        imgi (np.ndarray): The input image or stack of images of size [Lz x Ly x Lx x nchan].
+        batch_size (int, optional): Number of tiles to run in a batch. Defaults to 8.
+        rsz (float, optional): Resize coefficient(s) for image. Defaults to 1.0.
+        augment (bool, optional): Tiles image with overlapping tiles and flips overlapped regions to augment. Defaults to False.
+        tile_overlap (float, optional): Fraction of overlap of tiles when computing flows. Defaults to 0.1.
+        bsize (int, optional): Size of tiles to use in pixels [bsize x bsize]. Defaults to 224.
+
+    Returns:
+        Tuple[numpy.ndarray, numpy.ndarray]: outputs of network y and style. If tiled `y` is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
+            y[...,0] is Y flow; y[...,1] is X flow; y[...,2] is cell probability. 
+            style is a 1D array of size 256 summarizing the style of the image, if tiled `style` is averaged over tiles.
+    """
+    # run network
+    Lz, Ly0, Lx0, nchan = imgi.shape 
+    if rsz is not None:
+        if not isinstance(rsz, list) and not isinstance(rsz, np.ndarray):
+            rsz = [rsz, rsz]
+        Lyr, Lxr = int(Ly0 * rsz[0]), int(Lx0 * rsz[1])
+    else:
+        Lyr, Lxr = Ly0, Lx0 # 512, 512
+    
+    ly, lx = bsize, bsize # 256, 256
+    ypad1, ypad2, xpad1, xpad2 = transforms.get_pad_yx(Lyr, Lxr, min_size=(bsize, bsize)) # 8
+    Ly, Lx = Lyr + ypad1 + ypad2, Lxr + xpad1 + xpad2 # 528, 528
+    pads = np.array([[0, 0], [ypad1, ypad2], [xpad1, xpad2]])
+    
+    if augment:
+        ny = max(2, int(np.ceil(2. * Ly / bsize)))
+        nx = max(2, int(np.ceil(2. * Lx / bsize)))
+    else:
+        ny = 1 if Ly <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Ly / bsize)) # 3
+        nx = 1 if Lx <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Lx / bsize)) # 3
+    
+    
+    # run multiple slices at the same time
+    ntiles = ny * nx
+    nimgs = max(1, batch_size // ntiles) # number of imgs to run in the same batch, 1
+    niter = int(np.ceil(Lz / nimgs)) # 1
+    ziterator = (trange(niter, file=tqdm_out, mininterval=30) 
+                    if niter > 10 or Lz > 1 else range(niter))
+    for k in ziterator:
+        inds = np.arange(k * nimgs, min(Lz, (k + 1) * nimgs))
+        IMGa = np.zeros((ntiles * len(inds), nchan, ly, lx), "float32") # 9, 3, 256, 256
+        if feat is not None:
+            FEATa = np.zeros((ntiles * len(inds), nchan, ly, lx), "float32") # 9, 256
+        else: 
+            FEATa = None
+        for i, b in enumerate(inds):
+            # pad image for net so Ly and Lx are divisible by 4
+            imgb = transforms.resize_image(imgi[b], rsz=rsz) if rsz is not None else imgi[b].copy()
+            imgb = np.pad(imgb.transpose(2,0,1), pads, mode="constant") # 3, 528, 528
+
+            IMG, ysub, xsub, Lyt, Lxt = transforms.make_tiles(
+                imgb, bsize=bsize, augment=augment,
+                tile_overlap=tile_overlap) # IMG: 3, 3, 3, 256, 256
+            IMGa[i * ntiles : (i+1) * ntiles] = np.reshape(IMG, 
+                                            (ny * nx, nchan, ly, lx))
+            if feat is not None:
+                featb = transforms.resize_image(feat[b], rsz=rsz) if rsz is not None else feat[b].copy()
+                featb = np.pad(featb.transpose(2,0,1), pads, mode="constant")
+                FEAT, ysub, xsub, Lyt, Lxt = transforms.make_tiles(
+                    featb, bsize=bsize, augment=augment,
+                    tile_overlap=tile_overlap)          
+                FEATa[i * ntiles : (i+1) * ntiles] = np.reshape(FEAT,
+                                            (ny * nx, nchan, ly, lx))
+                
+        # run network
+        for j in range(0, IMGa.shape[0], batch_size):
+            bslc = slice(j, min(j + batch_size, IMGa.shape[0]))
+            ya0, stylea0 = _forward(net, IMGa[bslc], feat=FEATa[bslc] if FEATa is not None else None)
+            if j == 0:
+                nout = ya0.shape[1]
+                ya = np.zeros((IMGa.shape[0], nout, ly, lx), "float32")
+                stylea = np.zeros((IMGa.shape[0], 256), "float32")
+            ya[bslc] = ya0
+            stylea[bslc] = stylea0
+
+        # average tiles
+        for i, b in enumerate(inds):
+            if i==0 and k==0:
+                yf = np.zeros((Lz, nout, Ly, Lx), "float32")
+                styles = np.zeros((Lz, 256), "float32")
+            y = ya[i * ntiles : (i + 1) * ntiles]
+            if augment:
+                y = np.reshape(y, (ny, nx, 3, ly, lx))
+                y = transforms.unaugment_tiles(y)
+                y = np.reshape(y, (-1, 3, ly, lx))
+            yfi = transforms.average_tiles(y, ysub, xsub, Lyt, Lxt)
+            yf[b] = yfi[:, :imgb.shape[-2], :imgb.shape[-1]]
+            stylei = stylea[i * ntiles:(i + 1) * ntiles].sum(axis=0)
+            stylei /= (stylei**2).sum()**0.5
+            styles[b] = stylei
+    # slices from padding
+    yf = yf[:, :, ypad1 : Ly-ypad2, xpad1 : Lx-xpad2]
+    yf = yf.transpose(0,2,3,1)   
+    return yf, np.array(styles)
+
+
+def run_3D(net, imgs, batch_size=8, augment=False,
+           tile_overlap=0.1, bsize=224, net_ortho=None,
+           progress=None):
+    """ 
+    Run network on image z-stack.
+    
+    (faster if augment is False)
+
+    Args:
+        imgs (np.ndarray): The input image stack of size [Lz x Ly x Lx x nchan].
+        batch_size (int, optional): Number of tiles to run in a batch. Defaults to 8.
+        rsz (float, optional): Resize coefficient(s) for image. Defaults to 1.0.
+        anisotropy (float, optional): for 3D segmentation, optional rescaling factor (e.g. set to 2.0 if Z is sampled half as dense as X or Y). Defaults to None.
+        augment (bool, optional): Tiles image with overlapping tiles and flips overlapped regions to augment. Defaults to False.
+        tile_overlap (float, optional): Fraction of overlap of tiles when computing flows. Defaults to 0.1.
+        bsize (int, optional): Size of tiles to use in pixels [bsize x bsize]. Defaults to 224.
+        net_ortho (class, optional): cellpose network for orthogonal ZY and ZX planes. Defaults to None.
+        progress (QProgressBar, optional): pyqt progress bar. Defaults to None.
+
+    Returns:
+        Tuple[numpy.ndarray, numpy.ndarray]: outputs of network y and style. If tiled `y` is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
+            y[...,0] is Z flow; y[...,1] is Y flow; y[...,2] is X flow; y[...,3] is cell probability. 
+            style is a 1D array of size 256 summarizing the style of the image, if tiled `style` is averaged over tiles.
+    """
+    sstr = ["YX", "ZY", "ZX"]
+    pm = [(0, 1, 2, 3), (1, 0, 2, 3), (2, 0, 1, 3)]
+    ipm = [(0, 1, 2), (1, 0, 2), (1, 2, 0)]
+    cp = [(1, 2), (0, 2), (0, 1)]
+    cpy = [(0, 1), (0, 1), (0, 1)]
+    shape = imgs.shape[:-1]
+    yf = np.zeros((*shape, 4), "float32")
+    for p in range(3):
+        xsl = imgs.transpose(pm[p])
+        # per image
+        core_logger.info("running %s: %d planes of size (%d, %d)" %
+                         (sstr[p], shape[pm[p][0]], shape[pm[p][1]], shape[pm[p][2]]))
+        y, style = run_net(net,
+                           xsl, batch_size=batch_size, augment=augment, 
+                           bsize=bsize, tile_overlap=tile_overlap, 
+                           rsz=None)
+        yf[..., -1] += y[..., -1].transpose(ipm[p])
+        for j in range(2):
+            yf[..., cp[p][j]] += y[..., cpy[p][j]].transpose(ipm[p])
+        y = None; del y
+    
+        if progress is not None:
+            progress.setValue(25 + 15 * p)
+    
+    return yf, style

models/seg_post_model/cellpose/denoise.py

@@ -0,0 +1,1474 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os, time, datetime
+import numpy as np
+from scipy.stats import mode
+import cv2
+import torch
+from torch import nn
+from torch.nn.functional import conv2d, interpolate
+from tqdm import trange
+from pathlib import Path
+
+import logging
+
+denoise_logger = logging.getLogger(__name__)
+
+from cellpose import transforms, utils, io
+from cellpose.core import run_net
+from cellpose.models import CellposeModel, model_path, normalize_default, assign_device
+
+MODEL_NAMES = []
+for ctype in ["cyto3", "cyto2", "nuclei"]:
+    for ntype in ["denoise", "deblur", "upsample", "oneclick"]:
+        MODEL_NAMES.append(f"{ntype}_{ctype}")
+        if ctype != "cyto3":
+            for ltype in ["per", "seg", "rec"]:
+                MODEL_NAMES.append(f"{ntype}_{ltype}_{ctype}")
+    if ctype != "cyto3":
+        MODEL_NAMES.append(f"aniso_{ctype}")
+
+criterion = nn.MSELoss(reduction="mean")
+criterion2 = nn.BCEWithLogitsLoss(reduction="mean")
+
+
+def deterministic(seed=0):
+    """ set random seeds to create test data """
+    import random
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
+    np.random.seed(seed)  # Numpy module.
+    random.seed(seed)  # Python random module.
+    torch.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def loss_fn_rec(lbl, y):
+    """ loss function between true labels lbl and prediction y """
+    loss = 80. * criterion(y, lbl)
+    return loss
+
+
+def loss_fn_seg(lbl, y):
+    """ loss function between true labels lbl and prediction y """
+    veci = 5. * lbl[:, 1:]
+    lbl = (lbl[:, 0] > .5).float()
+    loss = criterion(y[:, :2], veci)
+    loss /= 2.
+    loss2 = criterion2(y[:, 2], lbl)
+    loss = loss + loss2
+    return loss
+
+
+def get_sigma(Tdown):
+    """ Calculates the correlation matrices across channels for the perceptual loss.
+
+    Args:
+        Tdown (list): List of tensors output by each downsampling block of network.
+
+    Returns:
+        list: List of correlations for each input tensor.
+    """
+    Tnorm = [x - x.mean((-2, -1), keepdim=True) for x in Tdown]
+    Tnorm = [x / x.std((-2, -1), keepdim=True) for x in Tnorm]
+    Sigma = [
+        torch.einsum("bnxy, bmxy -> bnm", x, x) / (x.shape[-2] * x.shape[-1])
+        for x in Tnorm
+    ]
+    return Sigma
+
+
+def imstats(X, net1):
+    """
+    Calculates the image correlation matrices for the perceptual loss.
+
+    Args:
+        X (torch.Tensor): Input image tensor.
+        net1: Cellpose net.
+
+    Returns:
+        list: A list of tensors of correlation matrices.
+    """
+    _, _, Tdown = net1(X)
+    Sigma = get_sigma(Tdown)
+    Sigma = [x.detach() for x in Sigma]
+    return Sigma
+
+
+def loss_fn_per(img, net1, yl):
+    """
+    Calculates the perceptual loss function for image restoration.
+
+    Args:
+        img (torch.Tensor): Input image tensor (noisy/blurry/downsampled).
+        net1 (torch.nn.Module): Perceptual loss net (Cellpose segmentation net).
+        yl (torch.Tensor): Clean image tensor.
+
+    Returns:
+        torch.Tensor: Mean perceptual loss.
+    """
+    Sigma = imstats(img, net1)
+    sd = [x.std((1, 2)) + 1e-6 for x in Sigma]
+    Sigma_test = get_sigma(yl)
+    losses = torch.zeros(len(Sigma[0]), device=img.device)
+    for k in range(len(Sigma)):
+        losses = losses + (((Sigma_test[k] - Sigma[k])**2).mean((1, 2)) / sd[k]**2)
+    return losses.mean()
+
+
+def test_loss(net0, X, net1=None, img=None, lbl=None, lam=[1., 1.5, 0.]):
+    """
+    Calculates the test loss for image restoration tasks.
+
+    Args:
+        net0 (torch.nn.Module): The image restoration network.
+        X (torch.Tensor): The input image tensor.
+        net1 (torch.nn.Module, optional): The segmentation network for segmentation or perceptual loss. Defaults to None.
+        img (torch.Tensor, optional): Clean image tensor for perceptual or reconstruction loss. Defaults to None.
+        lbl (torch.Tensor, optional): The ground truth flows/cellprob tensor for segmentation loss. Defaults to None.
+        lam (list, optional): The weights for different loss components (perceptual, segmentation, reconstruction). Defaults to [1., 1.5, 0.].
+
+    Returns:
+        tuple: A tuple containing the total loss and the perceptual loss.
+    """
+    net0.eval()
+    if net1 is not None:
+        net1.eval()
+    loss, loss_per = torch.zeros(1, device=X.device), torch.zeros(1, device=X.device)
+
+    with torch.no_grad():
+        img_dn = net0(X)[0]
+        if lam[2] > 0.:
+            loss += lam[2] * loss_fn_rec(img, img_dn)
+        if lam[1] > 0. or lam[0] > 0.:
+            y, _, ydown = net1(img_dn)
+        if lam[1] > 0.:
+            loss += lam[1] * loss_fn_seg(lbl, y)
+        if lam[0] > 0.:
+            loss_per = loss_fn_per(img, net1, ydown)
+            loss += lam[0] * loss_per
+    return loss, loss_per
+
+
+def train_loss(net0, X, net1=None, img=None, lbl=None, lam=[1., 1.5, 0.]):
+    """
+    Calculates the train loss for image restoration tasks.
+
+    Args:
+        net0 (torch.nn.Module): The image restoration network.
+        X (torch.Tensor): The input image tensor.
+        net1 (torch.nn.Module, optional): The segmentation network for segmentation or perceptual loss. Defaults to None.
+        img (torch.Tensor, optional): Clean image tensor for perceptual or reconstruction loss. Defaults to None.
+        lbl (torch.Tensor, optional): The ground truth flows/cellprob tensor for segmentation loss. Defaults to None.
+        lam (list, optional): The weights for different loss components (perceptual, segmentation, reconstruction). Defaults to [1., 1.5, 0.].
+
+    Returns:
+        tuple: A tuple containing the total loss and the perceptual loss.
+    """
+    net0.train()
+    if net1 is not None:
+        net1.eval()
+    loss, loss_per = torch.zeros(1, device=X.device), torch.zeros(1, device=X.device)
+
+    img_dn = net0(X)[0]
+    if lam[2] > 0.:
+        loss += lam[2] * loss_fn_rec(img, img_dn)
+    if lam[1] > 0. or lam[0] > 0.:
+        y, _, ydown = net1(img_dn)
+    if lam[1] > 0.:
+        loss += lam[1] * loss_fn_seg(lbl, y)
+    if lam[0] > 0.:
+        loss_per = loss_fn_per(img, net1, ydown)
+        loss += lam[0] * loss_per
+    return loss, loss_per
+
+
+def img_norm(imgi):
+    """
+    Normalizes the input image by subtracting the 1st percentile and dividing by the difference between the 99th and 1st percentiles.
+
+    Args:
+        imgi (torch.Tensor): Input image tensor.
+
+    Returns:
+        torch.Tensor: Normalized image tensor.
+    """
+    shape = imgi.shape
+    imgi = imgi.reshape(imgi.shape[0], imgi.shape[1], -1)
+    perc = torch.quantile(imgi, torch.tensor([0.01, 0.99], device=imgi.device), dim=-1,
+                          keepdim=True)
+    for k in range(imgi.shape[1]):
+        hask = (perc[1, :, k, 0] - perc[0, :, k, 0]) > 1e-3
+        imgi[hask, k] -= perc[0, hask, k]
+        imgi[hask, k] /= (perc[1, hask, k] - perc[0, hask, k])
+    imgi = imgi.reshape(shape)
+    return imgi
+
+
+def add_noise(lbl, alpha=4, beta=0.7, poisson=0.7, blur=0.7, gblur=1.0, downsample=0.7,
+              ds_max=7, diams=None, pscale=None, iso=True, sigma0=None, sigma1=None,
+              ds=None, uniform_blur=False, partial_blur=False):
+    """Adds noise to the input image.
+
+    Args:
+        lbl (torch.Tensor): The input image tensor of shape (nimg, nchan, Ly, Lx).
+        alpha (float, optional): The shape parameter of the gamma distribution used for generating poisson noise. Defaults to 4.
+        beta (float, optional): The rate parameter of the gamma distribution used for generating poisson noise. Defaults to 0.7.
+        poisson (float, optional): The probability of adding poisson noise to the image. Defaults to 0.7.
+        blur (float, optional): The probability of adding gaussian blur to the image. Defaults to 0.7.
+        gblur (float, optional): The scale factor for the gaussian blur. Defaults to 1.0.
+        downsample (float, optional): The probability of downsampling the image. Defaults to 0.7.
+        ds_max (int, optional): The maximum downsampling factor. Defaults to 7.
+        diams (torch.Tensor, optional): The diameter of the objects in the image. Defaults to None.
+        pscale (torch.Tensor, optional): The scale factor for the poisson noise, instead of sampling. Defaults to None.
+        iso (bool, optional): Whether to use isotropic gaussian blur. Defaults to True.
+        sigma0 (torch.Tensor, optional): The standard deviation of the gaussian filter for the Y axis, instead of sampling. Defaults to None.
+        sigma1 (torch.Tensor, optional): The standard deviation of the gaussian filter for the X axis, instead of sampling. Defaults to None.
+        ds (torch.Tensor, optional): The downsampling factor for each image, instead of sampling. Defaults to None.
+
+    Returns:
+        torch.Tensor: The noisy image tensor of the same shape as the input image.
+    """
+    device = lbl.device
+    imgi = torch.zeros_like(lbl)
+    Ly, Lx = lbl.shape[-2:]
+
+    diams = diams if diams is not None else 30. * torch.ones(len(lbl), device=device)
+    #ds0 = 1 if ds is None else ds.item()
+    ds = ds * torch.ones(
+        (len(lbl),), device=device, dtype=torch.long) if ds is not None else ds
+
+    # downsample
+    ii = []
+    idownsample = np.random.rand(len(lbl)) < downsample
+    if (ds is None and idownsample.sum() > 0.) or not iso:
+        ds = torch.ones(len(lbl), dtype=torch.long, device=device)
+        ds[idownsample] = torch.randint(2, ds_max + 1, size=(idownsample.sum(),),
+                                        device=device)
+        ii = torch.nonzero(ds > 1).flatten()
+    elif ds is not None and (ds > 1).sum():
+        ii = torch.nonzero(ds > 1).flatten()
+
+    # add gaussian blur
+    iblur = torch.rand(len(lbl), device=device) < blur
+    iblur[ii] = True
+    if iblur.sum() > 0:
+        if sigma0 is None:
+            if uniform_blur and iso:
+                xr = torch.rand(len(lbl), device=device)
+                if len(ii) > 0:
+                    xr[ii] = ds[ii].float() / 2. / gblur
+                sigma0 = diams[iblur] / 30. * gblur * (1 / gblur + (1 - 1 / gblur) * xr[iblur])
+                sigma1 = sigma0.clone()
+            elif not iso:
+                xr = torch.rand(len(lbl), device=device)
+                if len(ii) > 0:
+                    xr[ii] = (ds[ii].float()) / gblur
+                    xr[ii] = xr[ii] + torch.rand(len(ii), device=device) * 0.7 - 0.35
+                    xr[ii] = torch.clip(xr[ii], 0.05, 1.5)
+                sigma0 = diams[iblur] / 30. * gblur * xr[iblur]
+                sigma1 = sigma0.clone() / 10.
+            else:
+                xrand = np.random.exponential(1, size=iblur.sum())
+                xrand = np.clip(xrand * 0.5, 0.1, 1.0)
+                xrand *= gblur
+                sigma0 = diams[iblur] / 30. * 5. * torch.from_numpy(xrand).float().to(
+                    device)
+                sigma1 = sigma0.clone()
+        else:
+            sigma0 = sigma0 * torch.ones((iblur.sum(),), device=device)
+            sigma1 = sigma1 * torch.ones((iblur.sum(),), device=device)
+
+        # create gaussian filter
+        xr = max(8, sigma0.max().long() * 2)
+        gfilt0 = torch.exp(-torch.arange(-xr + 1, xr, device=device)**2 /
+                           (2 * sigma0.unsqueeze(-1)**2))
+        gfilt0 /= gfilt0.sum(axis=-1, keepdims=True)
+        gfilt1 = torch.zeros_like(gfilt0)
+        gfilt1[sigma1 == sigma0] = gfilt0[sigma1 == sigma0]
+        gfilt1[sigma1 != sigma0] = torch.exp(
+            -torch.arange(-xr + 1, xr, device=device)**2 /
+            (2 * sigma1[sigma1 != sigma0].unsqueeze(-1)**2))
+        gfilt1[sigma1 == 0] = 0.
+        gfilt1[sigma1 == 0, xr] = 1.
+        gfilt1 /= gfilt1.sum(axis=-1, keepdims=True)
+        gfilt = torch.einsum("ck,cl->ckl", gfilt0, gfilt1)
+        gfilt /= gfilt.sum(axis=(1, 2), keepdims=True)
+
+        lbl_blur = conv2d(lbl[iblur].transpose(1, 0), gfilt.unsqueeze(1),
+                             padding=gfilt.shape[-1] // 2,
+                             groups=gfilt.shape[0]).transpose(1, 0)
+        if partial_blur:
+            #yc, xc = np.random.randint(100, Ly-100), np.random.randint(100, Lx-100)
+            imgi[iblur] = lbl[iblur].clone()
+            Lxc = int(Lx * 0.85)
+            ym, xm = torch.meshgrid(torch.zeros(Ly, dtype=torch.float32), 
+                                    torch.arange(0, Lxc, dtype=torch.float32), 
+                        indexing="ij")
+            mask = torch.exp(-(ym**2 + xm**2) / 2*(0.001**2))
+            mask -= mask.min()
+            mask /= mask.max()
+            lbl_blur_crop = lbl_blur[:, :, :, :Lxc]
+            imgi[iblur, :, :, :Lxc] = (lbl_blur_crop * mask + 
+                                (1-mask) * imgi[iblur, :, :, :Lxc])
+        else:
+            imgi[iblur] = lbl_blur
+
+    imgi[~iblur] = lbl[~iblur]
+
+    # apply downsample
+    for k in ii:
+        i0 = imgi[k:k + 1, :, ::ds[k], ::ds[k]] if iso else imgi[k:k + 1, :, ::ds[k]]
+        imgi[k] = interpolate(i0, size=lbl[k].shape[-2:], mode="bilinear")
+
+    # add poisson noise
+    ipoisson = np.random.rand(len(lbl)) < poisson
+    if ipoisson.sum() > 0:
+        if pscale is None:
+            pscale = torch.zeros(len(lbl))
+            m = torch.distributions.gamma.Gamma(alpha, beta)
+            pscale = torch.clamp(m.rsample(sample_shape=(ipoisson.sum(),)), 1.)
+            #pscale = torch.clamp(20 * (torch.rand(size=(len(lbl),), device=lbl.device)), 1.5)
+            pscale = pscale.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).to(device)
+        else:
+            pscale = pscale * torch.ones((ipoisson.sum(), 1, 1, 1), device=device)
+        imgi[ipoisson] = torch.poisson(pscale * imgi[ipoisson])
+    imgi[~ipoisson] = imgi[~ipoisson]
+
+    # renormalize
+    imgi = img_norm(imgi)
+
+    return imgi
+
+
+def random_rotate_and_resize_noise(data, labels=None, diams=None, poisson=0.7, blur=0.7,
+                                   downsample=0.0, beta=0.7, gblur=1.0, diam_mean=30,
+                                   ds_max=7, uniform_blur=False, iso=True, rotate=True,
+                                   device=torch.device("cuda"), xy=(224, 224),
+                                   nchan_noise=1, keep_raw=True):
+    """
+    Applies random rotation, resizing, and noise to the input data.
+
+    Args:
+        data (numpy.ndarray): The input data.
+        labels (numpy.ndarray, optional): The flow and cellprob labels associated with the data. Defaults to None.
+        diams (float, optional): The diameter of the objects. Defaults to None.
+        poisson (float, optional): The Poisson noise probability. Defaults to 0.7.
+        blur (float, optional): The blur probability. Defaults to 0.7.
+        downsample (float, optional): The downsample probability. Defaults to 0.0.
+        beta (float, optional): The beta value for the poisson noise distribution. Defaults to 0.7.
+        gblur (float, optional): The Gaussian blur level. Defaults to 1.0.
+        diam_mean (float, optional): The mean diameter. Defaults to 30.
+        ds_max (int, optional): The maximum downsample value. Defaults to 7.
+        iso (bool, optional): Whether to apply isotropic augmentation. Defaults to True.
+        rotate (bool, optional): Whether to apply rotation augmentation. Defaults to True.
+        device (torch.device, optional): The device to use. Defaults to torch.device("cuda").
+        xy (tuple, optional): The size of the output image. Defaults to (224, 224).
+        nchan_noise (int, optional): The number of channels to add noise to. Defaults to 1.
+        keep_raw (bool, optional): Whether to keep the raw image. Defaults to True.
+
+    Returns:
+        torch.Tensor: The augmented image and augmented noisy/blurry/downsampled version of image.
+        torch.Tensor: The augmented labels.
+        float: The scale factor applied to the image.
+    """
+    if device == None:
+        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('mps') if torch.backends.mps.is_available() else None
+    
+    diams = 30 if diams is None else diams
+    random_diam = diam_mean * (2**(2 * np.random.rand(len(data)) - 1))
+    random_rsc = diams / random_diam  #/ random_diam
+    #rsc /= random_scale
+    xy0 = (340, 340)
+    nchan = data[0].shape[0]
+    data_new = np.zeros((len(data), (1 + keep_raw) * nchan, xy0[0], xy0[1]), "float32")
+    labels_new = np.zeros((len(data), 3, xy0[0], xy0[1]), "float32")
+    for i in range(
+            len(data)):  #, (sc, img, lbl) in enumerate(zip(random_rsc, data, labels)):
+        sc = random_rsc[i]
+        img = data[i]
+        lbl = labels[i] if labels is not None else None
+        # create affine transform to resize
+        Ly, Lx = img.shape[-2:]
+        dxy = np.maximum(0, np.array([Lx / sc - xy0[1], Ly / sc - xy0[0]]))
+        dxy = (np.random.rand(2,) - .5) * dxy
+        cc = np.array([Lx / 2, Ly / 2])
+        cc1 = cc - np.array([Lx - xy0[1], Ly - xy0[0]]) / 2 + dxy
+        pts1 = np.float32([cc, cc + np.array([1, 0]), cc + np.array([0, 1])])
+        pts2 = np.float32(
+            [cc1, cc1 + np.array([1, 0]) / sc, cc1 + np.array([0, 1]) / sc])
+        M = cv2.getAffineTransform(pts1, pts2)
+
+        # apply to image
+        for c in range(nchan):
+            img_rsz = cv2.warpAffine(img[c], M, xy0, flags=cv2.INTER_LINEAR)
+            #img_noise = add_noise(torch.from_numpy(img_rsz).to(device).unsqueeze(0)).cpu().numpy().squeeze(0)
+            data_new[i, c] = img_rsz
+            if keep_raw:
+                data_new[i, c + nchan] = img_rsz
+
+        if lbl is not None:
+            # apply to labels
+            labels_new[i, 0] = cv2.warpAffine(lbl[0], M, xy0, flags=cv2.INTER_NEAREST)
+            labels_new[i, 1] = cv2.warpAffine(lbl[1], M, xy0, flags=cv2.INTER_LINEAR)
+            labels_new[i, 2] = cv2.warpAffine(lbl[2], M, xy0, flags=cv2.INTER_LINEAR)
+
+    rsc = random_diam / diam_mean
+
+    # add noise before augmentations
+    img = torch.from_numpy(data_new).to(device)
+    img = torch.clamp(img, 0.)
+    # just add noise to cyto if nchan_noise=1
+    img[:, :nchan_noise] = add_noise(
+        img[:, :nchan_noise], poisson=poisson, blur=blur, ds_max=ds_max, iso=iso,
+        downsample=downsample, beta=beta, gblur=gblur,
+        diams=torch.from_numpy(random_diam).to(device).float())
+    # img -= img.mean(dim=(-2,-1), keepdim=True)
+    # img /= img.std(dim=(-2,-1), keepdim=True) + 1e-3
+    img = img.cpu().numpy()
+
+    # augmentations
+    img, lbl, scale = transforms.random_rotate_and_resize(
+        img,
+        Y=labels_new,
+        xy=xy,
+        rotate=False if not iso else rotate,
+        #(iso and downsample==0),
+        rescale=rsc,
+        scale_range=0.5)
+    img = torch.from_numpy(img).to(device)
+    lbl = torch.from_numpy(lbl).to(device)
+
+    return img, lbl, scale
+
+
+def one_chan_cellpose(device, model_type="cyto2", pretrained_model=None):
+    """
+    Creates a Cellpose network with a single input channel.
+
+    Args:
+        device (str): The device to run the network on.
+        model_type (str, optional): The type of Cellpose model to use. Defaults to "cyto2".
+        pretrained_model (str, optional): The path to a pretrained model file. Defaults to None.
+
+    Returns:
+        torch.nn.Module: The Cellpose network with a single input channel.
+    """
+    if pretrained_model is not None and not os.path.exists(pretrained_model):
+        model_type = pretrained_model
+        pretrained_model = None
+    nbase = [32, 64, 128, 256]
+    nchan = 1
+    net1 = resnet_torch.CPnet([nchan, *nbase], nout=3, sz=3).to(device)
+    filename = model_path(model_type,
+                          0) if pretrained_model is None else pretrained_model
+    weights = torch.load(filename, weights_only=True)
+    zp = 0
+    print(filename)
+    for name in net1.state_dict():
+        if ("res_down_0.conv.conv_0" not in name and
+                #"output" not in name and
+                "res_down_0.proj" not in name and name != "diam_mean" and
+                name != "diam_labels"):
+            net1.state_dict()[name].copy_(weights[name])
+        elif "res_down_0" in name:
+            if len(weights[name].shape) > 0:
+                new_weight = torch.zeros_like(net1.state_dict()[name])
+                if weights[name].shape[0] == 2:
+                    new_weight[:] = weights[name][0]
+                elif len(weights[name].shape) > 1 and weights[name].shape[1] == 2:
+                    new_weight[:, zp] = weights[name][:, 0]
+                else:
+                    new_weight = weights[name]
+            else:
+                new_weight = weights[name]
+            net1.state_dict()[name].copy_(new_weight)
+    return net1
+
+
+class CellposeDenoiseModel():
+    """ model to run Cellpose and Image restoration """
+
+    def __init__(self, gpu=False, pretrained_model=False, model_type=None,
+                 restore_type="denoise_cyto3", nchan=2,
+                 chan2_restore=False, device=None):
+
+        self.dn = DenoiseModel(gpu=gpu, model_type=restore_type, chan2=chan2_restore,
+                               device=device)
+        self.cp = CellposeModel(gpu=gpu, model_type=model_type, nchan=nchan,
+                                pretrained_model=pretrained_model, device=device)
+
+    def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
+             normalize=True, rescale=None, diameter=None, tile_overlap=0.1,
+             augment=False, resample=True, invert=False, flow_threshold=0.4,
+             cellprob_threshold=0.0, do_3D=False, anisotropy=None, stitch_threshold=0.0,
+             min_size=15, niter=None, interp=True, bsize=224, flow3D_smooth=0):
+        """
+        Restore array or list of images using the image restoration model, and then segment.
+
+        Args:
+            x (list, np.ndarry): can be list of 2D/3D/4D images, or array of 2D/3D/4D images
+            batch_size (int, optional): number of 224x224 patches to run simultaneously on the GPU
+                (can make smaller or bigger depending on GPU memory usage). Defaults to 8.
+            channels (list, optional): list of channels, either of length 2 or of length number of images by 2.
+                First element of list is the channel to segment (0=grayscale, 1=red, 2=green, 3=blue).
+                Second element of list is the optional nuclear channel (0=none, 1=red, 2=green, 3=blue).
+                For instance, to segment grayscale images, input [0,0]. To segment images with cells
+                in green and nuclei in blue, input [2,3]. To segment one grayscale image and one
+                image with cells in green and nuclei in blue, input [[0,0], [2,3]].
+                Defaults to None.
+            channel_axis (int, optional): channel axis in element of list x, or of np.ndarray x. 
+                if None, channels dimension is attempted to be automatically determined. Defaults to None.
+            z_axis  (int, optional): z axis in element of list x, or of np.ndarray x. 
+                if None, z dimension is attempted to be automatically determined. Defaults to None.
+            normalize (bool, optional): if True, normalize data so 0.0=1st percentile and 1.0=99th percentile of image intensities in each channel; 
+                can also pass dictionary of parameters (all keys are optional, default values shown): 
+                    - "lowhigh"=None : pass in normalization values for 0.0 and 1.0 as list [low, high] (if not None, all following parameters ignored)
+                    - "sharpen"=0 ; sharpen image with high pass filter, recommended to be 1/4-1/8 diameter of cells in pixels
+                    - "normalize"=True ; run normalization (if False, all following parameters ignored)
+                    - "percentile"=None : pass in percentiles to use as list [perc_low, perc_high]
+                    - "tile_norm"=0 ; compute normalization in tiles across image to brighten dark areas, to turn on set to window size in pixels (e.g. 100)
+                    - "norm3D"=False ; compute normalization across entire z-stack rather than plane-by-plane in stitching mode.
+                Defaults to True.
+            rescale (float, optional): resize factor for each image, if None, set to 1.0;
+                (only used if diameter is None). Defaults to None.
+            diameter (float, optional):  diameter for each image, 
+                if diameter is None, set to diam_mean or diam_train if available. Defaults to None.
+            tile_overlap (float, optional): fraction of overlap of tiles when computing flows. Defaults to 0.1.
+            augment (bool, optional): augment tiles by flipping and averaging for segmentation. Defaults to False.
+            resample (bool, optional): run dynamics at original image size (will be slower but create more accurate boundaries). Defaults to True.
+            invert (bool, optional): invert image pixel intensity before running network. Defaults to False.
+            flow_threshold (float, optional): flow error threshold (all cells with errors below threshold are kept) (not used for 3D). Defaults to 0.4.
+            cellprob_threshold (float, optional): all pixels with value above threshold kept for masks, decrease to find more and larger masks. Defaults to 0.0.
+            do_3D (bool, optional): set to True to run 3D segmentation on 3D/4D image input. Defaults to False.
+            anisotropy (float, optional): for 3D segmentation, optional rescaling factor (e.g. set to 2.0 if Z is sampled half as dense as X or Y). Defaults to None.
+            stitch_threshold (float, optional): if stitch_threshold>0.0 and not do_3D, masks are stitched in 3D to return volume segmentation. Defaults to 0.0.
+            min_size (int, optional): all ROIs below this size, in pixels, will be discarded. Defaults to 15.
+            flow3D_smooth (int, optional): if do_3D and flow3D_smooth>0, smooth flows with gaussian filter of this stddev. Defaults to 0.
+            niter (int, optional): number of iterations for dynamics computation. if None, it is set proportional to the diameter. Defaults to None.
+            interp (bool, optional): interpolate during 2D dynamics (not available in 3D) . Defaults to True.
+            
+        Returns:
+            A tuple containing (masks, flows, styles, imgs); masks: labelled image(s), where 0=no masks; 1,2,...=mask labels; 
+            flows: list of lists: flows[k][0] = XY flow in HSV 0-255; flows[k][1] = XY(Z) flows at each pixel; flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics); flows[k][3] = final pixel locations after Euler integration;
+            styles: style vector summarizing each image of size 256;
+            imgs: Restored images.
+        """
+        
+        if isinstance(normalize, dict):
+            normalize_params = {**normalize_default, **normalize}
+        elif not isinstance(normalize, bool):
+            raise ValueError("normalize parameter must be a bool or a dict")
+        else:
+            normalize_params = normalize_default
+            normalize_params["normalize"] = normalize
+        normalize_params["invert"] = invert
+
+        img_restore = self.dn.eval(x, batch_size=batch_size, channels=channels,
+                                   channel_axis=channel_axis, z_axis=z_axis,
+                                   do_3D=do_3D, 
+                                   normalize=normalize_params, rescale=rescale,
+                                   diameter=diameter,
+                                   tile_overlap=tile_overlap, bsize=bsize)
+
+        # turn off special normalization for segmentation
+        normalize_params = normalize_default
+
+        # change channels for segmentation
+        if channels is not None:
+            channels_new = [0, 0] if channels[0] == 0 else [1, 2]
+        else:
+            channels_new = None
+        # change diameter if self.ratio > 1 (upsampled to self.dn.diam_mean)
+        diameter = self.dn.diam_mean if self.dn.ratio > 1 else diameter
+        masks, flows, styles = self.cp.eval(
+            img_restore, batch_size=batch_size, channels=channels_new, channel_axis=-1,
+            z_axis=0 if not isinstance(img_restore, list) and img_restore.ndim > 3 and img_restore.shape[0] > 0 else None,
+            normalize=normalize_params, rescale=rescale, diameter=diameter,
+            tile_overlap=tile_overlap, augment=augment, resample=resample,
+            invert=invert, flow_threshold=flow_threshold,
+            cellprob_threshold=cellprob_threshold, do_3D=do_3D, anisotropy=anisotropy,
+            stitch_threshold=stitch_threshold, min_size=min_size, niter=niter,
+            interp=interp, bsize=bsize)
+
+        return masks, flows, styles, img_restore
+
+
+class DenoiseModel():
+    """
+    DenoiseModel class for denoising images using Cellpose denoising model.
+
+    Args:
+        gpu (bool, optional): Whether to use GPU for computation. Defaults to False.
+        pretrained_model (bool or str or Path, optional): Pretrained model to use for denoising.
+            Can be a string or path. Defaults to False.
+        nchan (int, optional): Number of channels in the input images, all Cellpose 3 models were trained with nchan=1. Defaults to 1.
+        model_type (str, optional): Type of pretrained model to use ("denoise_cyto3", "deblur_cyto3", "upsample_cyto3", ...). Defaults to None.
+        chan2 (bool, optional): Whether to use a separate model for the second channel. Defaults to False.
+        diam_mean (float, optional): Mean diameter of the objects in the images. Defaults to 30.0.
+        device (torch.device, optional): Device to use for computation. Defaults to None.
+
+    Attributes:
+        nchan (int): Number of channels in the input images.
+        diam_mean (float): Mean diameter of the objects in the images.
+        net (CPnet): Cellpose network for denoising.
+        pretrained_model (bool or str or Path): Pretrained model path to use for denoising.
+        net_chan2 (CPnet or None): Cellpose network for the second channel, if applicable.
+        net_type (str): Type of the denoising network.
+
+    Methods:
+        eval(x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
+                normalize=True, rescale=None, diameter=None, tile=True, tile_overlap=0.1)
+            Denoise array or list of images using the denoising model.
+
+        _eval(net, x, normalize=True, rescale=None, diameter=None, tile=True,
+                tile_overlap=0.1)
+            Run denoising model on a single channel.
+    """
+
+    def __init__(self, gpu=False, pretrained_model=False, nchan=1, model_type=None,
+                 chan2=False, diam_mean=30., device=None):
+        self.nchan = nchan
+        if pretrained_model and (not isinstance(pretrained_model, str) and
+                                 not isinstance(pretrained_model, Path)):
+            raise ValueError("pretrained_model must be a string or path")
+
+        self.diam_mean = diam_mean
+        builtin = True
+        if model_type is not None or (pretrained_model and
+                                      not os.path.exists(pretrained_model)):
+            pretrained_model_string = model_type if model_type is not None else "denoise_cyto3"
+            if ~np.any([pretrained_model_string == s for s in MODEL_NAMES]):
+                pretrained_model_string = "denoise_cyto3"
+            pretrained_model = model_path(pretrained_model_string)
+            if (pretrained_model and not os.path.exists(pretrained_model)):
+                denoise_logger.warning("pretrained model has incorrect path")
+            denoise_logger.info(f">> {pretrained_model_string} << model set to be used")
+            self.diam_mean = 17. if "nuclei" in pretrained_model_string else 30.
+        else:
+            if pretrained_model:
+                builtin = False
+                pretrained_model_string = pretrained_model
+                denoise_logger.info(f">>>> loading model {pretrained_model_string}")
+
+        # assign network device
+        if device is None:
+            sdevice, gpu = assign_device(use_torch=True, gpu=gpu)
+        self.device = device if device is not None else sdevice
+        if device is not None:
+            device_gpu = self.device.type == "cuda"
+        self.gpu = gpu if device is None else device_gpu
+        
+        # create network
+        self.nchan = nchan
+        self.nclasses = 1
+        nbase = [32, 64, 128, 256]
+        self.nchan = nchan
+        self.nbase = [nchan, *nbase]
+
+        self.net = CPnet(self.nbase, self.nclasses, sz=3, 
+                         max_pool=True, diam_mean=diam_mean).to(self.device)
+
+        self.pretrained_model = pretrained_model
+        self.net_chan2 = None
+        if self.pretrained_model:
+            self.net.load_model(self.pretrained_model, device=self.device)
+            denoise_logger.info(
+                f">>>> model diam_mean = {self.diam_mean: .3f} (ROIs rescaled to this size during training)"
+            )
+            if chan2 and builtin:
+                chan2_path = model_path(
+                    os.path.split(self.pretrained_model)[-1].split("_")[0] + "_nuclei")
+                print(f"loading model for chan2: {os.path.split(str(chan2_path))[-1]}")
+                self.net_chan2 = CPnet(self.nbase, self.nclasses, sz=3,
+                                        max_pool=True,
+                                       diam_mean=17.).to(self.device)
+                self.net_chan2.load_model(chan2_path, device=self.device)
+        self.net_type = "cellpose_denoise"
+
+    def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
+             normalize=True, rescale=None, diameter=None, tile=True, do_3D=False,
+             tile_overlap=0.1, bsize=224):
+        """
+        Restore array or list of images using the image restoration model.
+
+        Args:
+            x (list, np.ndarry): can be list of 2D/3D/4D images, or array of 2D/3D/4D images
+            batch_size (int, optional): number of 224x224 patches to run simultaneously on the GPU
+                (can make smaller or bigger depending on GPU memory usage). Defaults to 8.
+            channels (list, optional): list of channels, either of length 2 or of length number of images by 2.
+                First element of list is the channel to segment (0=grayscale, 1=red, 2=green, 3=blue).
+                Second element of list is the optional nuclear channel (0=none, 1=red, 2=green, 3=blue).
+                For instance, to segment grayscale images, input [0,0]. To segment images with cells
+                in green and nuclei in blue, input [2,3]. To segment one grayscale image and one
+                image with cells in green and nuclei in blue, input [[0,0], [2,3]].
+                Defaults to None.
+            channel_axis (int, optional): channel axis in element of list x, or of np.ndarray x. 
+                if None, channels dimension is attempted to be automatically determined. Defaults to None.
+            z_axis  (int, optional): z axis in element of list x, or of np.ndarray x. 
+                if None, z dimension is attempted to be automatically determined. Defaults to None.
+            normalize (bool, optional): if True, normalize data so 0.0=1st percentile and 1.0=99th percentile of image intensities in each channel; 
+                can also pass dictionary of parameters (all keys are optional, default values shown): 
+                    - "lowhigh"=None : pass in normalization values for 0.0 and 1.0 as list [low, high] (if not None, all following parameters ignored)
+                    - "sharpen"=0 ; sharpen image with high pass filter, recommended to be 1/4-1/8 diameter of cells in pixels
+                    - "normalize"=True ; run normalization (if False, all following parameters ignored)
+                    - "percentile"=None : pass in percentiles to use as list [perc_low, perc_high]
+                    - "tile_norm"=0 ; compute normalization in tiles across image to brighten dark areas, to turn on set to window size in pixels (e.g. 100)
+                    - "norm3D"=False ; compute normalization across entire z-stack rather than plane-by-plane in stitching mode.
+                Defaults to True.
+            rescale (float, optional): resize factor for each image, if None, set to 1.0;
+                (only used if diameter is None). Defaults to None.
+            diameter (float, optional):  diameter for each image, 
+                if diameter is None, set to diam_mean or diam_train if available. Defaults to None.
+            tile_overlap (float, optional): fraction of overlap of tiles when computing flows. Defaults to 0.1.
+              
+        Returns:
+            list: A list of 2D/3D arrays of restored images
+
+        """
+        if isinstance(x, list) or x.squeeze().ndim == 5:
+            tqdm_out = utils.TqdmToLogger(denoise_logger, level=logging.INFO)
+            nimg = len(x)
+            iterator = trange(nimg, file=tqdm_out,
+                              mininterval=30) if nimg > 1 else range(nimg)
+            imgs = []
+            for i in iterator:
+                imgi = self.eval(
+                    x[i], batch_size=batch_size,
+                    channels=channels[i] if channels is not None and
+                    ((len(channels) == len(x) and
+                      (isinstance(channels[i], list) or
+                       isinstance(channels[i], np.ndarray)) and len(channels[i]) == 2))
+                    else channels, channel_axis=channel_axis, z_axis=z_axis,
+                    normalize=normalize,
+                    do_3D=do_3D,
+                    rescale=rescale[i] if isinstance(rescale, list) or
+                    isinstance(rescale, np.ndarray) else rescale,
+                    diameter=diameter[i] if isinstance(diameter, list) or
+                    isinstance(diameter, np.ndarray) else diameter, 
+                    tile_overlap=tile_overlap, bsize=bsize)
+                imgs.append(imgi)
+            if isinstance(x, np.ndarray):
+                imgs = np.array(imgs)
+            return imgs
+
+        else:
+            # reshape image
+            x = transforms.convert_image(x, channels, channel_axis=channel_axis,
+                                         z_axis=z_axis, do_3D=do_3D, nchan=None)
+            if x.ndim < 4:
+                squeeze = True
+                x = x[np.newaxis, ...]
+            else:
+                squeeze = False
+
+            # may need to interpolate image before running upsampling
+            self.ratio = 1.
+            if "upsample" in self.pretrained_model:
+                Ly, Lx = x.shape[-3:-1]
+                if diameter is not None and 3 <= diameter < self.diam_mean:
+                    self.ratio = self.diam_mean / diameter
+                    denoise_logger.info(
+                        f"upsampling image to {self.diam_mean} pixel diameter ({self.ratio:0.2f} times)"
+                    )
+                    Lyr, Lxr = int(Ly * self.ratio), int(Lx * self.ratio)
+                    x = transforms.resize_image(x, Ly=Lyr, Lx=Lxr)
+                else:
+                    denoise_logger.warning(
+                        f"not interpolating image before upsampling because diameter is set >= {self.diam_mean}"
+                    )
+                    #raise ValueError(f"diameter is set to {diameter}, needs to be >=3 and < {self.dn.diam_mean}")
+
+            self.batch_size = batch_size
+
+            if diameter is not None and diameter > 0:
+                rescale = self.diam_mean / diameter
+            elif rescale is None:
+                rescale = 1.0
+
+            if np.ptp(x[..., -1]) < 1e-3 or (channels is not None and channels[-1] == 0):
+                x = x[..., :1]
+
+            for c in range(x.shape[-1]):
+                rescale0 = rescale * 30. / 17. if c == 1 else rescale
+                if c == 0 or self.net_chan2 is None:
+                    x[...,
+                      c] = self._eval(self.net, x[..., c:c + 1], batch_size=batch_size,
+                                      normalize=normalize, rescale=rescale0, 
+                                      tile_overlap=tile_overlap, bsize=bsize)[...,0]
+                else:
+                    x[...,
+                      c] = self._eval(self.net_chan2, x[...,
+                                                        c:c + 1], batch_size=batch_size,
+                                      normalize=normalize, rescale=rescale0, 
+                                      tile_overlap=tile_overlap, bsize=bsize)[...,0]
+            x = x[0] if squeeze else x
+        return x
+
+    def _eval(self, net, x, batch_size=8, normalize=True, rescale=None,
+              tile_overlap=0.1, bsize=224):
+        """
+        Run image restoration model on a single channel.
+
+        Args:
+            x (list, np.ndarry): can be list of 2D/3D/4D images, or array of 2D/3D/4D images
+            batch_size (int, optional): number of 224x224 patches to run simultaneously on the GPU
+                (can make smaller or bigger depending on GPU memory usage). Defaults to 8.
+            normalize (bool, optional): if True, normalize data so 0.0=1st percentile and 1.0=99th percentile of image intensities in each channel; 
+                can also pass dictionary of parameters (all keys are optional, default values shown): 
+                    - "lowhigh"=None : pass in normalization values for 0.0 and 1.0 as list [low, high] (if not None, all following parameters ignored)
+                    - "sharpen"=0 ; sharpen image with high pass filter, recommended to be 1/4-1/8 diameter of cells in pixels
+                    - "normalize"=True ; run normalization (if False, all following parameters ignored)
+                    - "percentile"=None : pass in percentiles to use as list [perc_low, perc_high]
+                    - "tile_norm"=0 ; compute normalization in tiles across image to brighten dark areas, to turn on set to window size in pixels (e.g. 100)
+                    - "norm3D"=False ; compute normalization across entire z-stack rather than plane-by-plane in stitching mode.
+                Defaults to True.
+            rescale (float, optional): resize factor for each image, if None, set to 1.0;
+                (only used if diameter is None). Defaults to None.
+            tile_overlap (float, optional): fraction of overlap of tiles when computing flows. Defaults to 0.1.
+            
+        Returns:
+            list: A list of 2D/3D arrays of restored images
+
+        """
+        if isinstance(normalize, dict):
+            normalize_params = {**normalize_default, **normalize}
+        elif not isinstance(normalize, bool):
+            raise ValueError("normalize parameter must be a bool or a dict")
+        else:
+            normalize_params = normalize_default
+            normalize_params["normalize"] = normalize
+
+        tic = time.time()
+        shape = x.shape
+        nimg = shape[0]
+
+        do_normalization = True if normalize_params["normalize"] else False
+
+        img = np.asarray(x)
+        if do_normalization:
+            img = transforms.normalize_img(img, **normalize_params)
+        if rescale != 1.0:
+            img = transforms.resize_image(img, rsz=rescale)
+        yf, style = run_net(self.net, img, bsize=bsize,
+                            tile_overlap=tile_overlap)
+        yf = transforms.resize_image(yf, shape[1], shape[2])
+        imgs = yf
+        del yf, style
+
+        # imgs = np.zeros((*x.shape[:-1], 1), np.float32)
+        # for i in iterator:
+        #     img = np.asarray(x[i])
+        #     if do_normalization:
+        #         img = transforms.normalize_img(img, **normalize_params)
+        #     if rescale != 1.0:
+        #         img = transforms.resize_image(img, rsz=[rescale, rescale])
+        #         if img.ndim == 2:
+        #             img = img[:, :, np.newaxis]
+        #     yf, style = run_net(net, img, batch_size=batch_size, augment=False,
+        #                         tile=tile, tile_overlap=tile_overlap, bsize=bsize)
+        #     img = transforms.resize_image(yf, Ly=x.shape[-3], Lx=x.shape[-2])
+
+        #     if img.ndim == 2:
+        #         img = img[:, :, np.newaxis]
+        #     imgs[i] = img
+        #     del yf, style
+        net_time = time.time() - tic
+        if nimg > 1:
+            denoise_logger.info("imgs denoised in %2.2fs" % (net_time))
+
+        return imgs
+
+
+def train(net, train_data=None, train_labels=None, train_files=None, test_data=None,
+          test_labels=None, test_files=None, train_probs=None, test_probs=None,
+          lam=[1., 1.5, 0.], scale_range=0.5, seg_model_type="cyto2", save_path=None,
+          save_every=100, save_each=False, poisson=0.7, beta=0.7, blur=0.7, gblur=1.0,
+          iso=True, uniform_blur=False, downsample=0., ds_max=7,
+          learning_rate=0.005, n_epochs=500,
+          weight_decay=0.00001, batch_size=8, nimg_per_epoch=None,
+          nimg_test_per_epoch=None, model_name=None):
+
+    # net properties
+    device = net.device
+    nchan = net.nchan
+    diam_mean = net.diam_mean.item()
+
+    args = np.array([poisson, beta, blur, gblur, downsample])
+    if args.ndim == 1:
+        args = args[:, np.newaxis]
+    poisson, beta, blur, gblur, downsample = args
+    nnoise = len(poisson)
+
+    d = datetime.datetime.now()
+    if save_path is not None:
+        if model_name is None:  
+            filename = ""
+            lstrs = ["per", "seg", "rec"]
+            for k, (l, s) in enumerate(zip(lam, lstrs)):
+                filename += f"{s}_{l:.2f}_"
+            if not iso:
+                filename += "aniso_"
+            if poisson.sum() > 0:
+                filename += "poisson_"
+            if blur.sum() > 0:
+                filename += "blur_"
+            if downsample.sum() > 0:
+                filename += "downsample_"
+            filename += d.strftime("%Y_%m_%d_%H_%M_%S.%f")
+            filename = os.path.join(save_path, filename)
+        else:
+            filename = os.path.join(save_path, model_name)
+        print(filename)
+    for i in range(len(poisson)):
+        denoise_logger.info(
+            f"poisson: {poisson[i]: 0.2f}, beta: {beta[i]: 0.2f}, blur: {blur[i]: 0.2f}, gblur: {gblur[i]: 0.2f}, downsample: {downsample[i]: 0.2f}"
+        )
+    net1 = one_chan_cellpose(device=device, pretrained_model=seg_model_type)
+
+    learning_rate_const = learning_rate
+    LR = np.linspace(0, learning_rate_const, 10)
+    LR = np.append(LR, learning_rate_const * np.ones(n_epochs - 100))
+    for i in range(10):
+        LR = np.append(LR, LR[-1] / 2 * np.ones(10))
+    learning_rate = LR
+
+    batch_size = 8
+    optimizer = torch.optim.AdamW(net.parameters(), lr=learning_rate[0],
+                                  weight_decay=weight_decay)
+    if train_data is not None:
+        nimg = len(train_data)
+        diam_train = np.array(
+            [utils.diameters(train_labels[k])[0] for k in trange(len(train_labels))])
+        diam_train[diam_train < 5] = 5.
+        if test_data is not None:
+            diam_test = np.array(
+                [utils.diameters(test_labels[k])[0] for k in trange(len(test_labels))])
+            diam_test[diam_test < 5] = 5.
+            nimg_test = len(test_data)
+    else:
+        nimg = len(train_files)
+        denoise_logger.info(">>> using files instead of loading dataset")
+        train_labels_files = [str(tf)[:-4] + f"_flows.tif" for tf in train_files]
+        denoise_logger.info(">>> computing diameters")
+        diam_train = np.array([
+            utils.diameters(io.imread(train_labels_files[k])[0])[0]
+            for k in trange(len(train_labels_files))
+        ])
+        diam_train[diam_train < 5] = 5.
+        if test_files is not None:
+            nimg_test = len(test_files)
+            test_labels_files = [str(tf)[:-4] + f"_flows.tif" for tf in test_files]
+            diam_test = np.array([
+                utils.diameters(io.imread(test_labels_files[k])[0])[0]
+                for k in trange(len(test_labels_files))
+            ])
+            diam_test[diam_test < 5] = 5.
+    train_probs = 1. / nimg * np.ones(nimg,
+                                      "float64") if train_probs is None else train_probs
+    if test_files is not None or test_data is not None:
+        test_probs = 1. / nimg_test * np.ones(
+            nimg_test, "float64") if test_probs is None else test_probs
+
+    tic = time.time()
+
+    nimg_per_epoch = nimg if nimg_per_epoch is None else nimg_per_epoch
+    if test_files is not None or test_data is not None:
+        nimg_test_per_epoch = nimg_test if nimg_test_per_epoch is None else nimg_test_per_epoch
+
+    nbatch = 0
+    train_losses, test_losses = [], []
+    for iepoch in range(n_epochs):
+        np.random.seed(iepoch)
+        rperm = np.random.choice(np.arange(0, nimg), size=(nimg_per_epoch,),
+                                 p=train_probs)
+        torch.manual_seed(iepoch)
+        np.random.seed(iepoch)
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = learning_rate[iepoch]
+        lavg, lavg_per, nsum = 0, 0, 0
+        for ibatch in range(0, nimg_per_epoch, batch_size * nnoise):
+            inds = rperm[ibatch : ibatch + batch_size * nnoise]
+            if train_data is None:
+                imgs = [np.maximum(0, io.imread(train_files[i])[:nchan]) for i in inds]
+                lbls = [io.imread(train_labels_files[i])[1:] for i in inds]
+            else:
+                imgs = [train_data[i][:nchan] for i in inds]
+                lbls = [train_labels[i][1:] for i in inds]
+            #inoise = nbatch % nnoise
+            rnoise = np.random.permutation(nnoise)
+            for i, inoise in enumerate(rnoise):
+                if i * batch_size < len(imgs):
+                    imgi, lbli, scale = random_rotate_and_resize_noise(
+                        imgs[i * batch_size : (i + 1) * batch_size], 
+                        lbls[i * batch_size : (i + 1) * batch_size],
+                        diam_train[inds][i * batch_size : (i + 1) * batch_size].copy(), 
+                        poisson=poisson[inoise],
+                        beta=beta[inoise], gblur=gblur[inoise], blur=blur[inoise], iso=iso,
+                        downsample=downsample[inoise], uniform_blur=uniform_blur,
+                        diam_mean=diam_mean, ds_max=ds_max,
+                        device=device)
+                    if i == 0:
+                        img = imgi 
+                        lbl = lbli 
+                    else:
+                        img = torch.cat((img, imgi), axis=0)
+                        lbl = torch.cat((lbl, lbli), axis=0)
+
+            if nnoise > 0:
+                iperm = np.random.permutation(img.shape[0])
+                img, lbl = img[iperm], lbl[iperm]
+            
+            for i in range(nnoise):
+                optimizer.zero_grad()
+                imgi = img[i * batch_size: (i + 1) * batch_size]
+                lbli = lbl[i * batch_size: (i + 1) * batch_size]
+                if imgi.shape[0] > 0:
+                    loss, loss_per = train_loss(net, imgi[:, :nchan], net1=net1,
+                                            img=imgi[:, nchan:], lbl=lbli, lam=lam)
+                    loss.backward()
+                    optimizer.step()
+                    lavg += loss.item() * imgi.shape[0]
+                    lavg_per += loss_per.item() * imgi.shape[0]
+
+            nsum += len(img)
+            nbatch += 1
+
+        if iepoch % 5 == 0 or iepoch < 10:
+            lavg = lavg / nsum
+            lavg_per = lavg_per / nsum
+            if test_data is not None or test_files is not None:
+                lavgt, nsum = 0., 0
+                np.random.seed(42)
+                rperm = np.random.choice(np.arange(0, nimg_test),
+                                         size=(nimg_test_per_epoch,), p=test_probs)
+                inoise = iepoch % nnoise
+                torch.manual_seed(inoise)
+                for ibatch in range(0, nimg_test_per_epoch, batch_size):
+                    inds = rperm[ibatch:ibatch + batch_size]
+                    if test_data is None:
+                        imgs = [
+                            np.maximum(0,
+                                       io.imread(test_files[i])[:nchan]) for i in inds
+                        ]
+                        lbls = [io.imread(test_labels_files[i])[1:] for i in inds]
+                    else:
+                        imgs = [test_data[i][:nchan] for i in inds]
+                        lbls = [test_labels[i][1:] for i in inds]
+                    img, lbl, scale = random_rotate_and_resize_noise(
+                        imgs, lbls, diam_test[inds].copy(), poisson=poisson[inoise],
+                        beta=beta[inoise], blur=blur[inoise], gblur=gblur[inoise],
+                        iso=iso, downsample=downsample[inoise], uniform_blur=uniform_blur,
+                        diam_mean=diam_mean, ds_max=ds_max, device=device)
+                    loss, loss_per = test_loss(net, img[:, :nchan], net1=net1,
+                                               img=img[:, nchan:], lbl=lbl, lam=lam)
+
+                    lavgt += loss.item() * img.shape[0]
+                    nsum += len(img)
+                lavgt = lavgt / nsum
+                denoise_logger.info(
+                    "Epoch %d, Time %4.1fs, Loss %0.3f, loss_per %0.3f, Loss Test %0.3f, LR %2.4f"
+                    % (iepoch, time.time() - tic, lavg, lavg_per, lavgt,
+                       learning_rate[iepoch]))
+                test_losses.append(lavgt)
+            else:
+                denoise_logger.info(
+                    "Epoch %d, Time %4.1fs, Loss %0.3f, loss_per %0.3f, LR %2.4f" %
+                    (iepoch, time.time() - tic, lavg, lavg_per, learning_rate[iepoch]))
+            train_losses.append(lavg)
+            
+        if save_path is not None:
+            if iepoch == n_epochs - 1 or (iepoch % save_every == 0 and iepoch != 0):
+                if save_each:  #separate files as model progresses
+                    filename0 = str(filename) + f"_epoch_{iepoch:%04d}"
+                else:
+                    filename0 = filename
+                denoise_logger.info(f"saving network parameters to {filename0}")
+                net.save_model(filename0)
+        else:
+            filename = save_path
+
+    return filename, train_losses, test_losses
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="cellpose parameters")
+
+    input_img_args = parser.add_argument_group("input image arguments")
+    input_img_args.add_argument("--dir", default=[], type=str,
+                                help="folder containing data to run or train on.")
+    input_img_args.add_argument("--img_filter", default=[], type=str,
+                                help="end string for images to run on")
+
+    model_args = parser.add_argument_group("model arguments")
+    model_args.add_argument("--pretrained_model", default=[], type=str,
+                            help="pretrained denoising model")
+
+    training_args = parser.add_argument_group("training arguments")
+    training_args.add_argument("--test_dir", default=[], type=str,
+                               help="folder containing test data (optional)")
+    training_args.add_argument("--file_list", default=[], type=str,
+                               help="npy file containing list of train and test files")
+    training_args.add_argument("--seg_model_type", default="cyto2", type=str,
+                               help="model to use for seg training loss")
+    training_args.add_argument(
+        "--noise_type", default=[], type=str,
+        help="noise type to use (if input, then other noise params are ignored)")
+    training_args.add_argument("--poisson", default=0.8, type=float,
+                               help="fraction of images to add poisson noise to")
+    training_args.add_argument("--beta", default=0.7, type=float,
+                               help="scale of poisson noise")
+    training_args.add_argument("--blur", default=0., type=float,
+                               help="fraction of images to blur")
+    training_args.add_argument("--gblur", default=1.0, type=float,
+                               help="scale of gaussian blurring stddev")
+    training_args.add_argument("--downsample", default=0., type=float,
+                               help="fraction of images to downsample")
+    training_args.add_argument("--ds_max", default=7, type=int,
+                               help="max downsampling factor")
+    training_args.add_argument("--lam_per", default=1.0, type=float,
+                               help="weighting of perceptual loss")
+    training_args.add_argument("--lam_seg", default=1.5, type=float,
+                               help="weighting of segmentation loss")
+    training_args.add_argument("--lam_rec", default=0., type=float,
+                               help="weighting of reconstruction loss")
+    training_args.add_argument(
+        "--diam_mean", default=30., type=float, help=
+        "mean diameter to resize cells to during training -- if starting from pretrained models it cannot be changed from 30.0"
+    )
+    training_args.add_argument("--learning_rate", default=0.001, type=float,
+                               help="learning rate. Default: %(default)s")
+    training_args.add_argument("--n_epochs", default=2000, type=int,
+                               help="number of epochs. Default: %(default)s")
+    training_args.add_argument(
+        "--save_each", default=False, action="store_true",
+        help="save each epoch as separate model")
+    training_args.add_argument(
+        "--nimg_per_epoch", default=0, type=int,
+        help="number of images per epoch. Default is length of training images")
+    training_args.add_argument(
+        "--nimg_test_per_epoch", default=0, type=int,
+        help="number of test images per epoch. Default is length of testing images")
+
+    io.logger_setup()
+
+    args = parser.parse_args()
+    lams = [args.lam_per, args.lam_seg, args.lam_rec]
+    print("lam", lams)
+
+    if len(args.noise_type) > 0:
+        noise_type = args.noise_type
+        uniform_blur = False
+        iso = True
+        if noise_type == "poisson":
+            poisson = 0.8
+            blur = 0.
+            downsample = 0.
+            beta = 0.7
+            gblur = 1.0
+        elif noise_type == "blur_expr":
+            poisson = 0.8
+            blur = 0.8
+            downsample = 0.
+            beta = 0.1
+            gblur = 0.5
+        elif noise_type == "blur":
+            poisson = 0.8
+            blur = 0.8
+            downsample = 0.
+            beta = 0.1
+            gblur = 10.0
+            uniform_blur = True
+        elif noise_type == "downsample_expr":
+            poisson = 0.8
+            blur = 0.8
+            downsample = 0.8
+            beta = 0.03
+            gblur = 1.0
+        elif noise_type == "downsample":
+            poisson = 0.8
+            blur = 0.8
+            downsample = 0.8
+            beta = 0.03
+            gblur = 5.0
+            uniform_blur = True
+        elif noise_type == "all":
+            poisson = [0.8, 0.8, 0.8]
+            blur = [0., 0.8, 0.8]
+            downsample = [0., 0., 0.8]
+            beta = [0.7, 0.1, 0.03]
+            gblur = [0., 10.0, 5.0]
+            uniform_blur = True
+        elif noise_type == "aniso":
+            poisson = 0.8
+            blur = 0.8
+            downsample = 0.8
+            beta = 0.1
+            gblur = args.ds_max * 1.5
+            iso = False
+        else:
+            raise ValueError(f"{noise_type} noise_type is not supported")
+    else:
+        poisson, beta = args.poisson, args.beta
+        blur, gblur = args.blur, args.gblur
+        downsample = args.downsample
+
+    pretrained_model = None if len(
+        args.pretrained_model) == 0 else args.pretrained_model
+    model = DenoiseModel(gpu=True, nchan=1, diam_mean=args.diam_mean,
+                         pretrained_model=pretrained_model)
+
+    train_data, labels, train_files, train_probs = None, None, None, None
+    test_data, test_labels, test_files, test_probs = None, None, None, None
+    if len(args.file_list) == 0:
+        output = io.load_train_test_data(args.dir, args.test_dir, "_img", "_masks", 0)
+        images, labels, image_names, test_images, test_labels, image_names_test = output
+        train_data = []
+        for i in range(len(images)):
+            img = images[i].astype("float32")
+            if img.ndim > 2:
+                img = img[0]
+            train_data.append(
+                np.maximum(transforms.normalize99(img), 0)[np.newaxis, :, :])
+        if len(args.test_dir) > 0:
+            test_data = []
+            for i in range(len(test_images)):
+                img = test_images[i].astype("float32")
+                if img.ndim > 2:
+                    img = img[0]
+                test_data.append(
+                    np.maximum(transforms.normalize99(img), 0)[np.newaxis, :, :])
+        save_path = os.path.join(args.dir, "../models/")
+    else:
+        root = args.dir
+        denoise_logger.info(
+            ">>> using file_list (assumes images are normalized and have flows!)")
+        dat = np.load(args.file_list, allow_pickle=True).item()
+        train_files = dat["train_files"]
+        test_files = dat["test_files"]
+        train_probs = dat["train_probs"] if "train_probs" in dat else None
+        test_probs = dat["test_probs"] if "test_probs" in dat else None
+        if str(train_files[0])[:len(str(root))] != str(root):
+            for i in range(len(train_files)):
+                new_path = root / Path(*train_files[i].parts[-3:])
+                if i == 0:
+                    print(f"changing path from {train_files[i]} to {new_path}")
+                train_files[i] = new_path
+
+            for i in range(len(test_files)):
+                new_path = root / Path(*test_files[i].parts[-3:])
+                test_files[i] = new_path
+        save_path = os.path.join(args.dir, "models/")
+
+    os.makedirs(save_path, exist_ok=True)
+
+    nimg_per_epoch = None if args.nimg_per_epoch == 0 else args.nimg_per_epoch
+    nimg_test_per_epoch = None if args.nimg_test_per_epoch == 0 else args.nimg_test_per_epoch
+
+    model_path = train(
+        model.net, train_data=train_data, train_labels=labels, train_files=train_files,
+        test_data=test_data, test_labels=test_labels, test_files=test_files,
+        train_probs=train_probs, test_probs=test_probs, poisson=poisson, beta=beta,
+        blur=blur, gblur=gblur, downsample=downsample, ds_max=args.ds_max,
+        iso=iso, uniform_blur=uniform_blur, n_epochs=args.n_epochs,
+        learning_rate=args.learning_rate,
+        lam=lams, 
+        seg_model_type=args.seg_model_type, nimg_per_epoch=nimg_per_epoch,
+        nimg_test_per_epoch=nimg_test_per_epoch, save_path=save_path)
+
+
+def seg_train_noisy(model, train_data, train_labels, test_data=None, test_labels=None,
+                    poisson=0.8, blur=0.0, downsample=0.0, save_path=None,
+                    save_every=100, save_each=False, learning_rate=0.2, n_epochs=500,
+                    momentum=0.9, weight_decay=0.00001, SGD=True, batch_size=8,
+                    nimg_per_epoch=None, diameter=None, rescale=True, z_masking=False,
+                    model_name=None):
+    """ train function uses loss function model.loss_fn in models.py
+    
+    (data should already be normalized)
+
+    """
+
+    d = datetime.datetime.now()
+
+    model.n_epochs = n_epochs
+    if isinstance(learning_rate, (list, np.ndarray)):
+        if isinstance(learning_rate, np.ndarray) and learning_rate.ndim > 1:
+            raise ValueError("learning_rate.ndim must equal 1")
+        elif len(learning_rate) != n_epochs:
+            raise ValueError(
+                "if learning_rate given as list or np.ndarray it must have length n_epochs"
+            )
+        model.learning_rate = learning_rate
+        model.learning_rate_const = mode(learning_rate)[0][0]
+    else:
+        model.learning_rate_const = learning_rate
+        # set learning rate schedule
+        if SGD:
+            LR = np.linspace(0, model.learning_rate_const, 10)
+            if model.n_epochs > 250:
+                LR = np.append(
+                    LR, model.learning_rate_const * np.ones(model.n_epochs - 100))
+                for i in range(10):
+                    LR = np.append(LR, LR[-1] / 2 * np.ones(10))
+            else:
+                LR = np.append(
+                    LR,
+                    model.learning_rate_const * np.ones(max(0, model.n_epochs - 10)))
+        else:
+            LR = model.learning_rate_const * np.ones(model.n_epochs)
+        model.learning_rate = LR
+
+    model.batch_size = batch_size
+    model._set_optimizer(model.learning_rate[0], momentum, weight_decay, SGD)
+    model._set_criterion()
+
+    nimg = len(train_data)
+
+    # compute average cell diameter
+    if diameter is None:
+        diam_train = np.array(
+            [utils.diameters(train_labels[k][0])[0] for k in range(len(train_labels))])
+        diam_train_mean = diam_train[diam_train > 0].mean()
+        model.diam_labels = diam_train_mean
+        if rescale:
+            diam_train[diam_train < 5] = 5.
+            if test_data is not None:
+                diam_test = np.array([
+                    utils.diameters(test_labels[k][0])[0]
+                    for k in range(len(test_labels))
+                ])
+                diam_test[diam_test < 5] = 5.
+            denoise_logger.info(">>>> median diameter set to = %d" % model.diam_mean)
+    elif rescale:
+        diam_train_mean = diameter
+        model.diam_labels = diameter
+        denoise_logger.info(">>>> median diameter set to = %d" % model.diam_mean)
+        diam_train = diameter * np.ones(len(train_labels), "float32")
+        if test_data is not None:
+            diam_test = diameter * np.ones(len(test_labels), "float32")
+
+    denoise_logger.info(
+        f">>>> mean of training label mask diameters (saved to model) {diam_train_mean:.3f}"
+    )
+    model.net.diam_labels.data = torch.ones(1, device=model.device) * diam_train_mean
+
+    nchan = train_data[0].shape[0]
+    denoise_logger.info(">>>> training network with %d channel input <<<<" % nchan)
+    denoise_logger.info(">>>> LR: %0.5f, batch_size: %d, weight_decay: %0.5f" %
+                        (model.learning_rate_const, model.batch_size, weight_decay))
+
+    if test_data is not None:
+        denoise_logger.info(f">>>> ntrain = {nimg}, ntest = {len(test_data)}")
+    else:
+        denoise_logger.info(f">>>> ntrain = {nimg}")
+
+    tic = time.time()
+
+    lavg, nsum = 0, 0
+
+    if save_path is not None:
+        _, file_label = os.path.split(save_path)
+        file_path = os.path.join(save_path, "models/")
+
+        if not os.path.exists(file_path):
+            os.makedirs(file_path)
+    else:
+        denoise_logger.warning("WARNING: no save_path given, model not saving")
+
+    ksave = 0
+
+    # get indices for each epoch for training
+    np.random.seed(0)
+    inds_all = np.zeros((0,), "int32")
+    if nimg_per_epoch is None or nimg > nimg_per_epoch:
+        nimg_per_epoch = nimg
+    denoise_logger.info(f">>>> nimg_per_epoch = {nimg_per_epoch}")
+    while len(inds_all) < n_epochs * nimg_per_epoch:
+        rperm = np.random.permutation(nimg)
+        inds_all = np.hstack((inds_all, rperm))
+
+    for iepoch in range(model.n_epochs):
+        if SGD:
+            model._set_learning_rate(model.learning_rate[iepoch])
+        np.random.seed(iepoch)
+        rperm = inds_all[iepoch * nimg_per_epoch:(iepoch + 1) * nimg_per_epoch]
+        for ibatch in range(0, nimg_per_epoch, batch_size):
+            inds = rperm[ibatch:ibatch + batch_size]
+            imgi, lbl, scale = random_rotate_and_resize_noise(
+                [train_data[i] for i in inds], [train_labels[i][1:] for i in inds],
+                poisson=poisson, blur=blur, downsample=downsample,
+                diams=diam_train[inds], diam_mean=model.diam_mean)
+            imgi = imgi[:, :1]  # keep noisy only
+            if z_masking:
+                nc = imgi.shape[1]
+                nb = imgi.shape[0]
+                ncmin = (np.random.rand(nb) > 0.25) * (np.random.randint(
+                    nc // 2 - 1, size=nb))
+                ncmax = nc - (np.random.rand(nb) > 0.25) * (np.random.randint(
+                    nc // 2 - 1, size=nb))
+                for b in range(nb):
+                    imgi[b, :ncmin[b]] = 0
+                    imgi[b, ncmax[b]:] = 0
+
+            train_loss = model._train_step(imgi, lbl)
+            lavg += train_loss
+            nsum += len(imgi)
+
+        if iepoch % 10 == 0 or iepoch == 5:
+            lavg = lavg / nsum
+            if test_data is not None:
+                lavgt, nsum = 0., 0
+                np.random.seed(42)
+                rperm = np.arange(0, len(test_data), 1, int)
+                for ibatch in range(0, len(test_data), batch_size):
+                    inds = rperm[ibatch:ibatch + batch_size]
+                    imgi, lbl, scale = random_rotate_and_resize_noise(
+                        [test_data[i] for i in inds],
+                        [test_labels[i][1:] for i in inds], poisson=poisson, blur=blur,
+                        downsample=downsample, diams=diam_test[inds],
+                        diam_mean=model.diam_mean)
+                    imgi = imgi[:, :1]  # keep noisy only
+                    test_loss = model._test_eval(imgi, lbl)
+                    lavgt += test_loss
+                    nsum += len(imgi)
+
+                denoise_logger.info(
+                    "Epoch %d, Time %4.1fs, Loss %2.4f, Loss Test %2.4f, LR %2.4f" %
+                    (iepoch, time.time() - tic, lavg, lavgt / nsum,
+                     model.learning_rate[iepoch]))
+            else:
+                denoise_logger.info(
+                    "Epoch %d, Time %4.1fs, Loss %2.4f, LR %2.4f" %
+                    (iepoch, time.time() - tic, lavg, model.learning_rate[iepoch]))
+
+            lavg, nsum = 0, 0
+
+        if save_path is not None:
+            if iepoch == model.n_epochs - 1 or iepoch % save_every == 1:
+                # save model at the end
+                if save_each:  #separate files as model progresses
+                    if model_name is None:
+                        filename = "{}_{}_{}_{}".format(
+                            model.net_type, file_label,
+                            d.strftime("%Y_%m_%d_%H_%M_%S.%f"), "epoch_" + str(iepoch))
+                    else:
+                        filename = "{}_{}".format(model_name, "epoch_" + str(iepoch))
+                else:
+                    if model_name is None:
+                        filename = "{}_{}_{}".format(model.net_type, file_label,
+                                                     d.strftime("%Y_%m_%d_%H_%M_%S.%f"))
+                    else:
+                        filename = model_name
+                filename = os.path.join(file_path, filename)
+                ksave += 1
+                denoise_logger.info(f"saving network parameters to {filename}")
+                model.net.save_model(filename)
+        else:
+            filename = save_path
+
+    return filename

models/seg_post_model/cellpose/dynamics.py

@@ -0,0 +1,691 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os
+from scipy.ndimage import find_objects, center_of_mass, mean
+import torch
+import numpy as np
+import tifffile
+from tqdm import trange
+import fastremap
+
+import logging
+
+dynamics_logger = logging.getLogger(__name__)
+
+from . import utils
+
+import torch
+import torch.nn.functional as F
+
+def _extend_centers_gpu(neighbors, meds, isneighbor, shape, n_iter=200, 
+                        device=torch.device("cpu")):
+    """Runs diffusion on GPU to generate flows for training images or quality control.
+
+    Args:
+        neighbors (torch.Tensor): 9 x pixels in masks.
+        meds (torch.Tensor): Mask centers.
+        isneighbor (torch.Tensor): Valid neighbor boolean 9 x pixels.
+        shape (tuple): Shape of the tensor.
+        n_iter (int, optional): Number of iterations. Defaults to 200.
+        device (torch.device, optional): Device to run the computation on. Defaults to torch.device("cpu").
+
+    Returns:
+        torch.Tensor: Generated flows.
+
+    """
+    if torch.prod(torch.tensor(shape)) > 4e7 or device.type == "mps":
+        T = torch.zeros(shape, dtype=torch.float, device=device)
+    else:
+        T = torch.zeros(shape, dtype=torch.double, device=device)
+
+    for i in range(n_iter):
+        T[tuple(meds.T)] += 1
+        Tneigh = T[tuple(neighbors)]
+        Tneigh *= isneighbor
+        T[tuple(neighbors[:, 0])] = Tneigh.mean(axis=0)
+    del meds, isneighbor, Tneigh
+
+    if T.ndim == 2:
+        grads = T[neighbors[0, [2, 1, 4, 3]], neighbors[1, [2, 1, 4, 3]]]
+        del neighbors
+        dy = grads[0] - grads[1]
+        dx = grads[2] - grads[3]
+        del grads
+        mu_torch = np.stack((dy.cpu().squeeze(0), dx.cpu().squeeze(0)), axis=-2)
+    else:
+        grads = T[tuple(neighbors[:, 1:])]
+        del neighbors
+        dz = grads[0] - grads[1]
+        dy = grads[2] - grads[3]
+        dx = grads[4] - grads[5]
+        del grads
+        mu_torch = np.stack(
+            (dz.cpu().squeeze(0), dy.cpu().squeeze(0), dx.cpu().squeeze(0)), axis=-2)
+    return mu_torch
+
+def center_of_mass(mask):
+    yi, xi = np.nonzero(mask)
+    ymean = int(np.round(yi.sum() / len(yi)))
+    xmean = int(np.round(xi.sum() / len(xi)))
+    if not ((yi==ymean) * (xi==xmean)).sum():
+        # center is closest point to (ymean, xmean) within mask
+        imin = ((xi - xmean)**2 + (yi - ymean)**2).argmin()
+        ymean = yi[imin]
+        xmean = xi[imin]
+    
+    return ymean, xmean
+
+def get_centers(masks, slices):
+    centers = [center_of_mass(masks[slices[i]]==(i+1)) for i in range(len(slices))]
+    centers = np.array([np.array([centers[i][0] + slices[i][0].start, centers[i][1] + slices[i][1].start]) 
+                    for i in range(len(slices))])
+    exts = np.array([(slc[0].stop - slc[0].start) + (slc[1].stop - slc[1].start) + 2 for slc in slices])
+    return centers, exts
+
+
+def masks_to_flows_gpu(masks, device=torch.device("cpu"), niter=None):
+    """Convert masks to flows using diffusion from center pixel.
+
+    Center of masks where diffusion starts is defined by pixel closest to median within the mask.
+
+    Args:
+        masks (int, 2D or 3D array): Labelled masks. 0=NO masks; 1,2,...=mask labels.
+        device (torch.device, optional): The device to run the computation on. Defaults to torch.device("cpu").
+        niter (int, optional): Number of iterations for the diffusion process. Defaults to None.
+
+    Returns:
+        np.ndarray: A 4D array representing the flows for each pixel in Z, X, and Y.
+       
+
+    Returns:
+        A tuple containing (mu, meds_p). mu is float 3D or 4D array of flows in (Z)XY. 
+        meds_p are cell centers.
+    """
+    if device is None:
+        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('mps') if torch.backends.mps.is_available() else None
+
+    if masks.max() > 0:
+        Ly0, Lx0 = masks.shape
+        Ly, Lx = Ly0 + 2, Lx0 + 2
+        
+        masks_padded = torch.from_numpy(masks.astype("int64")).to(device)
+        masks_padded = F.pad(masks_padded, (1, 1, 1, 1))
+        shape = masks_padded.shape
+        
+        ### get mask pixel neighbors
+        y, x = torch.nonzero(masks_padded, as_tuple=True)
+        y = y.int()
+        x = x.int()
+        neighbors = torch.zeros((2, 9, y.shape[0]), dtype=torch.int, device=device)
+        yxi = [[0, -1, 1, 0, 0, -1, -1, 1, 1], [0, 0, 0, -1, 1, -1, 1, -1, 1]]
+        for i in range(9):
+            neighbors[0, i] = y + yxi[0][i]
+            neighbors[1, i] = x + yxi[1][i]
+        isneighbor = torch.ones((9, y.shape[0]), dtype=torch.bool, device=device)
+        m0 = masks_padded[neighbors[0, 0], neighbors[1, 0]]
+        for i in range(1, 9):
+            isneighbor[i] = masks_padded[neighbors[0, i], neighbors[1, i]] == m0
+        del m0, masks_padded
+        
+        ### get center-of-mass within cell
+        slices = find_objects(masks)
+        centers, ext = get_centers(masks, slices)
+        meds_p = torch.from_numpy(centers).to(device).long()
+        meds_p += 1  # for padding
+
+        ### run diffusion
+        n_iter = 2 * ext.max() if niter is None else niter
+        mu = _extend_centers_gpu(neighbors, meds_p, isneighbor, shape, n_iter=n_iter,
+                                device=device)
+        mu = mu.astype("float64")
+
+        # new normalization
+        mu /= (1e-60 + (mu**2).sum(axis=0)**0.5)
+
+        # put into original image
+        mu0 = np.zeros((2, Ly0, Lx0))
+        mu0[:, y.cpu().numpy() - 1, x.cpu().numpy() - 1] = mu
+    else:
+        # no masks, return empty flows
+        mu0 = np.zeros((2, masks.shape[0], masks.shape[1]))
+    return mu0
+
+def masks_to_flows_gpu_3d(masks, device=None, niter=None):
+    """Convert masks to flows using diffusion from center pixel.
+
+    Args:
+        masks (int, 2D or 3D array): Labelled masks. 0=NO masks; 1,2,...=mask labels.
+        device (torch.device, optional): The device to run the computation on. Defaults to None.
+        niter (int, optional): Number of iterations for the diffusion process. Defaults to None.
+
+    Returns:
+        np.ndarray: A 4D array representing the flows for each pixel in Z, X, and Y.
+        
+    """
+    if device is None:
+        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('mps') if torch.backends.mps.is_available() else None
+
+    Lz0, Ly0, Lx0 = masks.shape
+    Lz, Ly, Lx = Lz0 + 2, Ly0 + 2, Lx0 + 2
+
+    masks_padded = torch.from_numpy(masks.astype("int64")).to(device)
+    masks_padded = F.pad(masks_padded, (1, 1, 1, 1, 1, 1))
+
+    # get mask pixel neighbors
+    z, y, x = torch.nonzero(masks_padded).T
+    neighborsZ = torch.stack((z, z + 1, z - 1, z, z, z, z))
+    neighborsY = torch.stack((y, y, y, y + 1, y - 1, y, y), axis=0)
+    neighborsX = torch.stack((x, x, x, x, x, x + 1, x - 1), axis=0)
+
+    neighbors = torch.stack((neighborsZ, neighborsY, neighborsX), axis=0)
+
+    # get mask centers
+    slices = find_objects(masks)
+
+    centers = np.zeros((masks.max(), 3), "int")
+    for i, si in enumerate(slices):
+        if si is not None:
+            sz, sy, sx = si
+            zi, yi, xi = np.nonzero(masks[sz, sy, sx] == (i + 1))
+            zi = zi.astype(np.int32) + 1  # add padding
+            yi = yi.astype(np.int32) + 1  # add padding
+            xi = xi.astype(np.int32) + 1  # add padding
+            zmed = np.mean(zi)
+            ymed = np.mean(yi)
+            xmed = np.mean(xi)
+            imin = np.argmin((zi - zmed)**2 + (xi - xmed)**2 + (yi - ymed)**2)
+            zmed = zi[imin]
+            ymed = yi[imin]
+            xmed = xi[imin]
+            centers[i, 0] = zmed + sz.start
+            centers[i, 1] = ymed + sy.start
+            centers[i, 2] = xmed + sx.start
+
+    # get neighbor validator (not all neighbors are in same mask)
+    neighbor_masks = masks_padded[tuple(neighbors)]
+    isneighbor = neighbor_masks == neighbor_masks[0]
+    ext = np.array(
+        [[sz.stop - sz.start + 1, sy.stop - sy.start + 1, sx.stop - sx.start + 1]
+         for sz, sy, sx in slices])
+    n_iter = 6 * (ext.sum(axis=1)).max() if niter is None else niter
+
+    # run diffusion
+    shape = masks_padded.shape
+    mu = _extend_centers_gpu(neighbors, centers, isneighbor, shape, n_iter=n_iter,
+                             device=device)
+    # normalize
+    mu /= (1e-60 + (mu**2).sum(axis=0)**0.5)
+
+    # put into original image
+    mu0 = np.zeros((3, Lz0, Ly0, Lx0))
+    mu0[:, z.cpu().numpy() - 1, y.cpu().numpy() - 1, x.cpu().numpy() - 1] = mu
+    return mu0
+
+def labels_to_flows(labels, files=None, device=None, redo_flows=False, niter=None,
+                    return_flows=True):
+    """Converts labels (list of masks or flows) to flows for training model.
+
+    Args:
+        labels (list of ND-arrays): The labels to convert. labels[k] can be 2D or 3D. If [3 x Ly x Lx], 
+            it is assumed that flows were precomputed. Otherwise, labels[k][0] or labels[k] (if 2D) 
+            is used to create flows and cell probabilities.
+        files (list of str, optional): The files to save the flows to. If provided, flows are saved to 
+            files to be reused. Defaults to None.
+        device (str, optional): The device to use for computation. Defaults to None.
+        redo_flows (bool, optional): Whether to recompute the flows. Defaults to False.
+        niter (int, optional): The number of iterations for computing flows. Defaults to None.
+
+    Returns:
+        list of [4 x Ly x Lx] arrays: The flows for training the model. flows[k][0] is labels[k], 
+        flows[k][1] is cell distance transform, flows[k][2] is Y flow, flows[k][3] is X flow, 
+        and flows[k][4] is heat distribution.
+    """
+    nimg = len(labels)
+    if labels[0].ndim < 3:
+        labels = [labels[n][np.newaxis, :, :] for n in range(nimg)]
+
+    flows = []
+    # flows need to be recomputed
+    if labels[0].shape[0] == 1 or labels[0].ndim < 3 or redo_flows:
+        dynamics_logger.info("computing flows for labels")
+
+        # compute flows; labels are fixed here to be unique, so they need to be passed back
+        # make sure labels are unique!
+        labels = [fastremap.renumber(label, in_place=True)[0] for label in labels]
+        iterator = trange if nimg > 1 else range
+        for n in iterator(nimg):
+            labels[n][0] = fastremap.renumber(labels[n][0], in_place=True)[0]
+            vecn = masks_to_flows_gpu(labels[n][0].astype(int), device=device, niter=niter)
+
+            # concatenate labels, distance transform, vector flows, heat (boundary and mask are computed in augmentations)
+            flow = np.concatenate((labels[n], labels[n] > 0.5, vecn),
+                                  axis=0).astype(np.float32)
+            if files is not None:
+                file_name = os.path.splitext(files[n])[0]
+                tifffile.imwrite(file_name + "_flows.tif", flow)
+            if return_flows:
+                flows.append(flow)
+    else:
+        dynamics_logger.info("flows precomputed")
+        if return_flows:
+            flows = [labels[n].astype(np.float32) for n in range(nimg)]
+    return flows
+
+
+def flow_error(maski, dP_net, device=None):
+    """Error in flows from predicted masks vs flows predicted by network run on image.
+
+    This function serves to benchmark the quality of masks. It works as follows:
+    1. The predicted masks are used to create a flow diagram.
+    2. The mask-flows are compared to the flows that the network predicted.
+
+    If there is a discrepancy between the flows, it suggests that the mask is incorrect.
+    Masks with flow_errors greater than 0.4 are discarded by default. This setting can be
+    changed in Cellpose.eval or CellposeModel.eval.
+
+    Args:
+        maski (np.ndarray, int): Masks produced from running dynamics on dP_net, where 0=NO masks; 1,2... are mask labels.
+        dP_net (np.ndarray, float): ND flows where dP_net.shape[1:] = maski.shape.
+
+    Returns:
+        A tuple containing (flow_errors, dP_masks): flow_errors (np.ndarray, float): Mean squared error between predicted flows and flows from masks; 
+        dP_masks (np.ndarray, float): ND flows produced from the predicted masks.
+    """
+    if dP_net.shape[1:] != maski.shape:
+        print("ERROR: net flow is not same size as predicted masks")
+        return
+
+    # flows predicted from estimated masks
+    dP_masks = masks_to_flows_gpu(maski, device=device)
+    # difference between predicted flows vs mask flows
+    flow_errors = np.zeros(maski.max())
+    for i in range(dP_masks.shape[0]):
+        flow_errors += mean((dP_masks[i] - dP_net[i] / 5.)**2, maski,
+                            index=np.arange(1,
+                                            maski.max() + 1))
+
+    return flow_errors, dP_masks
+
+
+def steps_interp(dP, inds, niter, device=torch.device("cpu")):
+    """ Run dynamics of pixels to recover masks in 2D/3D, with interpolation between pixel values.
+
+    Euler integration of dynamics dP for niter steps.
+
+    Args:
+        p (numpy.ndarray): Array of shape (n_points, 2 or 3) representing the initial pixel locations.
+        dP (numpy.ndarray): Array of shape (2, Ly, Lx) or (3, Lz, Ly, Lx) representing the flow field.
+        niter (int): Number of iterations to perform.
+        device (torch.device, optional): Device to use for computation. Defaults to None.
+
+    Returns:
+        numpy.ndarray: Array of shape (n_points, 2) or (n_points, 3) representing the final pixel locations.
+
+    Raises:
+        None
+
+    """
+    
+    shape = dP.shape[1:]
+    ndim = len(shape)
+    
+    pt = torch.zeros((*[1]*ndim, len(inds[0]), ndim), dtype=torch.float32, device=device)
+    im = torch.zeros((1, ndim, *shape), dtype=torch.float32, device=device)
+    # Y and X dimensions, flipped X-1, Y-1
+    # pt is [1 1 1 3 n_points]
+    for n in range(ndim):
+        if ndim==3:
+            pt[0, 0, 0, :, ndim - n - 1] = torch.from_numpy(inds[n]).to(device, dtype=torch.float32)
+        else:
+            pt[0, 0, :, ndim - n - 1] = torch.from_numpy(inds[n]).to(device, dtype=torch.float32)
+        im[0, ndim - n - 1] = torch.from_numpy(dP[n]).to(device, dtype=torch.float32)
+    shape = np.array(shape)[::-1].astype("float") - 1  
+    
+    # normalize pt between  0 and  1, normalize the flow
+    for k in range(ndim):
+        im[:, k] *= 2. / shape[k]
+        pt[..., k] /= shape[k]
+
+    # normalize to between -1 and 1
+    pt *= 2 
+    pt -= 1
+    
+    # dynamics
+    for t in range(niter):
+        dPt = torch.nn.functional.grid_sample(im, pt, align_corners=False)
+        for k in range(ndim):  #clamp the final pixel locations
+            pt[..., k] = torch.clamp(pt[..., k] + dPt[:, k], -1., 1.)
+
+    #undo the normalization from before, reverse order of operations
+    pt += 1 
+    pt *= 0.5
+    for k in range(ndim):
+        pt[..., k] *= shape[k]
+
+    if ndim==3:
+        pt = pt[..., [2, 1, 0]].squeeze()
+        pt = pt.unsqueeze(0) if pt.ndim==1 else pt 
+        return pt.T
+    else:
+        pt = pt[..., [1, 0]].squeeze()
+        pt = pt.unsqueeze(0) if pt.ndim==1 else pt
+        return pt.T
+
+def follow_flows(dP, inds, niter=200, device=torch.device("cpu")):
+    """ Run dynamics to recover masks in 2D or 3D.
+
+    Pixels are represented as a meshgrid. Only pixels with non-zero cell-probability
+    are used (as defined by inds).
+
+    Args:
+        dP (np.ndarray): Flows [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
+        mask (np.ndarray, optional): Pixel mask to seed masks. Useful when flows have low magnitudes.
+        niter (int, optional): Number of iterations of dynamics to run. Default is 200.
+        interp (bool, optional): Interpolate during 2D dynamics (not available in 3D). Default is True.
+        device (torch.device, optional): Device to use for computation. Default is None.
+
+    Returns:
+        A tuple containing (p, inds): p (np.ndarray): Final locations of each pixel after dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx]; 
+        inds (np.ndarray): Indices of pixels used for dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
+    """
+    shape = np.array(dP.shape[1:]).astype(np.int32)
+    ndim = len(inds)
+    
+    p = steps_interp(dP, inds, niter, device=device)
+        
+    return p
+
+
+def remove_bad_flow_masks(masks, flows, threshold=0.4, device=torch.device("cpu")):
+    """Remove masks which have inconsistent flows.
+
+    Uses metrics.flow_error to compute flows from predicted masks 
+    and compare flows to predicted flows from the network. Discards 
+    masks with flow errors greater than the threshold.
+
+    Args:
+        masks (int, 2D or 3D array): Labelled masks, 0=NO masks; 1,2,...=mask labels,
+            size [Ly x Lx] or [Lz x Ly x Lx].
+        flows (float, 3D or 4D array): Flows [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
+        threshold (float, optional): Masks with flow error greater than threshold are discarded.
+            Default is 0.4.
+
+    Returns:
+        masks (int, 2D or 3D array): Masks with inconsistent flow masks removed,
+            0=NO masks; 1,2,...=mask labels, size [Ly x Lx] or [Lz x Ly x Lx].
+    """
+    device0 = device
+    if masks.size > 10000 * 10000 and (device is not None and device.type == "cuda"):
+
+        major_version, minor_version = torch.__version__.split(".")[:2]
+        torch.cuda.empty_cache()
+        if major_version == "1" and int(minor_version) < 10:
+            # for PyTorch version lower than 1.10
+            def mem_info():
+                total_mem = torch.cuda.get_device_properties(device0.index).total_memory
+                used_mem = torch.cuda.memory_allocated(device0.index)
+                free_mem = total_mem - used_mem
+                return total_mem, free_mem
+        else:
+            # for PyTorch version 1.10 and above
+            def mem_info():
+                free_mem, total_mem = torch.cuda.mem_get_info(device0.index)
+                return total_mem, free_mem
+        total_mem, free_mem = mem_info()
+        if masks.size * 32 > free_mem:
+            dynamics_logger.warning(
+                "WARNING: image is very large, not using gpu to compute flows from masks for QC step flow_threshold"
+            )
+            dynamics_logger.info("turn off QC step with flow_threshold=0 if too slow")
+            device0 = torch.device("cpu")
+
+    merrors, _ = flow_error(masks, flows, device0)
+    badi = 1 + (merrors > threshold).nonzero()[0]
+    masks[np.isin(masks, badi)] = 0
+    return masks
+
+
+def max_pool1d(h, kernel_size=5, axis=1, out=None):
+    """ memory efficient max_pool thanks to Mark Kittisopikul 
+    
+    for stride=1, padding=kernel_size//2, requires odd kernel_size >= 3
+
+    """
+    if out is None:
+        out = h.clone()
+    else:
+        out.copy_(h)
+
+    nd = h.shape[axis]    
+    k0 = kernel_size // 2
+    for d in range(-k0, k0+1):
+        if axis==1:
+            mv = out[:, max(-d,0):min(nd-d,nd)]
+            hv = h[:, max(d,0):min(nd+d,nd)]
+        elif axis==2:
+            mv = out[:, :, max(-d,0):min(nd-d,nd)]
+            hv = h[:,  :, max(d,0):min(nd+d,nd)]
+        elif axis==3:
+            mv = out[:, :, :, max(-d,0):min(nd-d,nd)]
+            hv = h[:, :,  :, max(d,0):min(nd+d,nd)]
+        torch.maximum(mv, hv, out=mv)
+    return out
+
+def max_pool_nd(h, kernel_size=5):
+    """ memory efficient max_pool in 2d or 3d """
+    ndim = h.ndim - 1
+    hmax = max_pool1d(h, kernel_size=kernel_size, axis=1)
+    hmax2 = max_pool1d(hmax, kernel_size=kernel_size, axis=2)
+    if ndim==2:
+        del hmax
+        return hmax2
+    else:
+        hmax = max_pool1d(hmax2, kernel_size=kernel_size, axis=3, out=hmax)
+        del hmax2 
+        return hmax
+
+def get_masks_torch(pt, inds, shape0, rpad=20, max_size_fraction=0.4):
+    """Create masks using pixel convergence after running dynamics.
+
+    Makes a histogram of final pixel locations p, initializes masks 
+    at peaks of histogram and extends the masks from the peaks so that
+    they include all pixels with more than 2 final pixels p. Discards 
+    masks with flow errors greater than the threshold. 
+
+    Parameters:
+        p (float32, 3D or 4D array): Final locations of each pixel after dynamics,
+            size [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
+        iscell (bool, 2D or 3D array): If iscell is not None, set pixels that are 
+            iscell False to stay in their original location.
+        rpad (int, optional): Histogram edge padding. Default is 20.
+        max_size_fraction (float, optional): Masks larger than max_size_fraction of
+            total image size are removed. Default is 0.4.
+
+    Returns:
+        M0 (int, 2D or 3D array): Masks with inconsistent flow masks removed, 
+            0=NO masks; 1,2,...=mask labels, size [Ly x Lx] or [Lz x Ly x Lx].
+    """
+    
+    ndim = len(shape0)
+    device = pt.device
+    
+    rpad = 20
+    pt += rpad
+    pt = torch.clamp(pt, min=0)
+    for i in range(len(pt)):
+        pt[i] = torch.clamp(pt[i], max=shape0[i]+rpad-1)
+
+    # # add extra padding to make divisible by 5
+    # shape = tuple((np.ceil((shape0 + 2*rpad)/5) * 5).astype(int))
+    shape = tuple(np.array(shape0) + 2*rpad)
+
+    # sparse coo torch
+    coo = torch.sparse_coo_tensor(pt, torch.ones(pt.shape[1], device=pt.device, dtype=torch.int), 
+                                shape)
+    h1 = coo.to_dense()
+    del coo
+
+    hmax1 = max_pool_nd(h1.unsqueeze(0), kernel_size=5)
+    hmax1 = hmax1.squeeze()
+    seeds1 = torch.nonzero((h1 - hmax1 > -1e-6) * (h1 > 10))
+    del hmax1
+    if len(seeds1) == 0:
+        dynamics_logger.warning("no seeds found in get_masks_torch - no masks found.")
+        return np.zeros(shape0, dtype="uint16")
+    
+    npts = h1[tuple(seeds1.T)]
+    isort1 = npts.argsort()
+    seeds1 = seeds1[isort1]
+
+    n_seeds = len(seeds1)
+    h_slc = torch.zeros((n_seeds, *[11]*ndim), device=seeds1.device)
+    for k in range(n_seeds):
+        slc = tuple([slice(seeds1[k][j]-5, seeds1[k][j]+6) for j in range(ndim)])
+        h_slc[k] = h1[slc]
+    del h1
+    seed_masks = torch.zeros((n_seeds, *[11]*ndim), device=seeds1.device)
+    if ndim==2:
+        seed_masks[:,5,5] = 1
+    else:
+        seed_masks[:,5,5,5] = 1
+    
+    for iter in range(5):
+        # extend
+        seed_masks = max_pool_nd(seed_masks, kernel_size=3)
+        seed_masks *= h_slc > 2
+    del h_slc 
+    seeds_new = [tuple((torch.nonzero(seed_masks[k]) + seeds1[k] - 5).T) 
+            for k in range(n_seeds)]
+    del seed_masks 
+    
+    dtype = torch.int32 if n_seeds < 2**16 else torch.int64
+    M1 = torch.zeros(shape, dtype=dtype, device=device)
+    for k in range(n_seeds):
+        M1[seeds_new[k]] = 1 + k
+
+    M1 = M1[tuple(pt)]
+    M1 = M1.cpu().numpy()
+
+    dtype = "uint16" if n_seeds < 2**16 else "uint32"
+    M0 = np.zeros(shape0, dtype=dtype)
+    M0[inds] = M1
+        
+    # remove big masks
+    uniq, counts = fastremap.unique(M0, return_counts=True)
+    big = np.prod(shape0) * max_size_fraction
+    bigc = uniq[counts > big]
+    if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):
+        M0 = fastremap.mask(M0, bigc)
+    fastremap.renumber(M0, in_place=True)  #convenient to guarantee non-skipped labels
+    M0 = M0.reshape(tuple(shape0))
+    
+    #print(f"mem used: {torch.cuda.memory_allocated()/1e9:.3f} gb, max mem used: {torch.cuda.max_memory_allocated()/1e9:.3f} gb")
+    return M0
+
+
+def resize_and_compute_masks(dP, cellprob, niter=200, cellprob_threshold=0.0,
+                             flow_threshold=0.4, do_3D=False, min_size=15,
+                             max_size_fraction=0.4, resize=None, device=torch.device("cpu")):
+    """Compute masks using dynamics from dP and cellprob, and resizes masks if resize is not None.
+
+    Args:
+        dP (numpy.ndarray): The dynamics flow field array.
+        cellprob (numpy.ndarray): The cell probability array.
+        p (numpy.ndarray, optional): The pixels on which to run dynamics. Defaults to None
+        niter (int, optional): The number of iterations for mask computation. Defaults to 200.
+        cellprob_threshold (float, optional): The threshold for cell probability. Defaults to 0.0.
+        flow_threshold (float, optional): The threshold for quality control metrics. Defaults to 0.4.
+        interp (bool, optional): Whether to interpolate during dynamics computation. Defaults to True.
+        do_3D (bool, optional): Whether to perform mask computation in 3D. Defaults to False.
+        min_size (int, optional): The minimum size of the masks. Defaults to 15.
+        max_size_fraction (float, optional): Masks larger than max_size_fraction of
+            total image size are removed. Default is 0.4.
+        resize (tuple, optional): The desired size for resizing the masks. Defaults to None.
+        device (torch.device, optional): The device to use for computation. Defaults to torch.device("cpu").
+
+    Returns:
+        tuple: A tuple containing the computed masks and the final pixel locations.
+    """
+    mask = compute_masks(dP, cellprob, niter=niter,
+                            cellprob_threshold=cellprob_threshold,
+                            flow_threshold=flow_threshold, do_3D=do_3D,
+                            max_size_fraction=max_size_fraction, 
+                            device=device)
+
+    if resize is not None:
+        dynamics_logger.warning("Resizing is depricated in v4.0.1+")
+
+    mask = utils.fill_holes_and_remove_small_masks(mask, min_size=min_size)
+
+    return mask
+
+
+def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
+                  flow_threshold=0.4, do_3D=False, min_size=-1,
+                  max_size_fraction=0.4, device=torch.device("cpu")):
+    """Compute masks using dynamics from dP and cellprob.
+
+    Args:
+        dP (numpy.ndarray): The dynamics flow field array.
+        cellprob (numpy.ndarray): The cell probability array.
+        p (numpy.ndarray, optional): The pixels on which to run dynamics. Defaults to None
+        niter (int, optional): The number of iterations for mask computation. Defaults to 200.
+        cellprob_threshold (float, optional): The threshold for cell probability. Defaults to 0.0.
+        flow_threshold (float, optional): The threshold for quality control metrics. Defaults to 0.4.
+        interp (bool, optional): Whether to interpolate during dynamics computation. Defaults to True.
+        do_3D (bool, optional): Whether to perform mask computation in 3D. Defaults to False.
+        min_size (int, optional): The minimum size of the masks. Defaults to 15.
+        max_size_fraction (float, optional): Masks larger than max_size_fraction of
+            total image size are removed. Default is 0.4.
+        device (torch.device, optional): The device to use for computation. Defaults to torch.device("cpu").
+
+    Returns:
+        tuple: A tuple containing the computed masks and the final pixel locations.
+    """
+    
+    if (cellprob > cellprob_threshold).sum():  #mask at this point is a cell cluster binary map, not labels
+        inds = np.nonzero(cellprob > cellprob_threshold)
+        if len(inds[0]) == 0:
+            dynamics_logger.info("No cell pixels found.")
+            shape = cellprob.shape
+            mask = np.zeros(shape, "uint16")
+            return mask
+
+        p_final = follow_flows(dP * (cellprob > cellprob_threshold) / 5., 
+                               inds=inds, niter=niter, 
+                                device=device)
+        if not torch.is_tensor(p_final):
+            p_final = torch.from_numpy(p_final).to(device, dtype=torch.int)
+        else:
+            p_final = p_final.int()
+        # calculate masks
+        if device.type == "mps":
+            p_final = p_final.to(torch.device("cpu"))
+        mask = get_masks_torch(p_final, inds, dP.shape[1:], 
+                               max_size_fraction=max_size_fraction)
+        del p_final
+        # flow thresholding factored out of get_masks
+        if not do_3D:
+            if mask.max() > 0 and flow_threshold is not None and flow_threshold > 0:
+                # make sure labels are unique at output of get_masks
+                mask = remove_bad_flow_masks(mask, dP, threshold=flow_threshold,
+                                             device=device)
+
+        if mask.max() < 2**16 and mask.dtype != "uint16":
+            mask = mask.astype("uint16")
+
+    else:  # nothing to compute, just make it compatible
+        dynamics_logger.info("No cell pixels found.")
+        shape = cellprob.shape
+        mask = np.zeros(cellprob.shape, "uint16")
+        return mask
+    
+    if min_size > 0:
+        mask = utils.fill_holes_and_remove_small_masks(mask, min_size=min_size)
+
+    if mask.dtype == np.uint32:
+        dynamics_logger.warning(
+            "more than 65535 masks in image, masks returned as np.uint32")
+
+    return mask

models/seg_post_model/cellpose/export.py

@@ -0,0 +1,405 @@
+"""Auxiliary module for bioimageio format export
+
+Example usage:
+
+```bash
+#!/bin/bash
+
+# Define default paths and parameters
+DEFAULT_CHANNELS="1 0"
+DEFAULT_PATH_PRETRAINED_MODEL="/home/qinyu/models/cp/cellpose_residual_on_style_on_concatenation_off_1135_rest_2023_05_04_23_41_31.252995"
+DEFAULT_PATH_README="/home/qinyu/models/cp/README.md"
+DEFAULT_LIST_PATH_COVER_IMAGES="/home/qinyu/images/cp/cellpose_raw_and_segmentation.jpg /home/qinyu/images/cp/cellpose_raw_and_probability.jpg /home/qinyu/images/cp/cellpose_raw.jpg"
+DEFAULT_MODEL_ID="philosophical-panda"
+DEFAULT_MODEL_ICON="🐼"
+DEFAULT_MODEL_VERSION="0.1.0"
+DEFAULT_MODEL_NAME="My Cool Cellpose"
+DEFAULT_MODEL_DOCUMENTATION="A cool Cellpose model trained for my cool dataset."
+DEFAULT_MODEL_AUTHORS='[{"name": "Qin Yu", "affiliation": "EMBL", "github_user": "qin-yu", "orcid": "0000-0002-4652-0795"}]'
+DEFAULT_MODEL_CITE='[{"text": "For more details of the model itself, see the manuscript", "doi": "10.1242/dev.202800", "url": null}]'
+DEFAULT_MODEL_TAGS="cellpose 3d 2d"
+DEFAULT_MODEL_LICENSE="MIT"
+DEFAULT_MODEL_REPO="https://github.com/kreshuklab/go-nuclear"
+
+# Run the Python script with default parameters
+python export.py \
+    --channels $DEFAULT_CHANNELS \
+    --path_pretrained_model "$DEFAULT_PATH_PRETRAINED_MODEL" \
+    --path_readme "$DEFAULT_PATH_README" \
+    --list_path_cover_images $DEFAULT_LIST_PATH_COVER_IMAGES \
+    --model_version "$DEFAULT_MODEL_VERSION" \
+    --model_name "$DEFAULT_MODEL_NAME" \
+    --model_documentation "$DEFAULT_MODEL_DOCUMENTATION" \
+    --model_authors "$DEFAULT_MODEL_AUTHORS" \
+    --model_cite "$DEFAULT_MODEL_CITE" \
+    --model_tags $DEFAULT_MODEL_TAGS \
+    --model_license "$DEFAULT_MODEL_LICENSE" \
+    --model_repo "$DEFAULT_MODEL_REPO"
+```
+"""
+
+import os
+import sys
+import json
+import argparse
+from pathlib import Path
+from urllib.parse import urlparse
+
+import torch
+import numpy as np
+
+from cellpose.io import imread
+from cellpose.utils import download_url_to_file
+from cellpose.transforms import pad_image_ND, normalize_img, convert_image
+from cellpose.vit_sam import CPnetBioImageIO
+
+from bioimageio.spec.model.v0_5 import (
+    ArchitectureFromFileDescr,
+    Author,
+    AxisId,
+    ChannelAxis,
+    CiteEntry,
+    Doi,
+    FileDescr,
+    Identifier,
+    InputTensorDescr,
+    IntervalOrRatioDataDescr,
+    LicenseId,
+    ModelDescr,
+    ModelId,
+    OrcidId,
+    OutputTensorDescr,
+    ParameterizedSize,
+    PytorchStateDictWeightsDescr,
+    SizeReference,
+    SpaceInputAxis,
+    SpaceOutputAxis,
+    TensorId,
+    TorchscriptWeightsDescr,
+    Version,
+    WeightsDescr,
+)
+# Define ARBITRARY_SIZE if it is not available in the module
+try:
+    from bioimageio.spec.model.v0_5 import ARBITRARY_SIZE
+except ImportError:
+    ARBITRARY_SIZE = ParameterizedSize(min=1, step=1)
+
+from bioimageio.spec.common import HttpUrl
+from bioimageio.spec import save_bioimageio_package
+from bioimageio.core import test_model
+
+DEFAULT_CHANNELS = [2, 1]
+DEFAULT_NORMALIZE_PARAMS = {
+    "axis": -1,
+    "lowhigh": None,
+    "percentile": None,
+    "normalize": True,
+    "norm3D": False,
+    "sharpen_radius": 0,
+    "smooth_radius": 0,
+    "tile_norm_blocksize": 0,
+    "tile_norm_smooth3D": 1,
+    "invert": False,
+}
+IMAGE_URL = "http://www.cellpose.org/static/data/rgb_3D.tif"
+
+
+def download_and_normalize_image(path_dir_temp, channels=DEFAULT_CHANNELS):
+    """
+    Download and normalize image.
+    """
+    filename = os.path.basename(urlparse(IMAGE_URL).path)
+    path_image = path_dir_temp / filename
+    if not path_image.exists():
+        sys.stderr.write(f'Downloading: "{IMAGE_URL}" to {path_image}\n')
+        download_url_to_file(IMAGE_URL, path_image)
+    img = imread(path_image).astype(np.float32)
+    img = convert_image(img, channels, channel_axis=1, z_axis=0, do_3D=False, nchan=2)
+    img = normalize_img(img, **DEFAULT_NORMALIZE_PARAMS)
+    img = np.transpose(img, (0, 3, 1, 2))
+    img, _, _ = pad_image_ND(img)
+    return img
+
+
+def load_bioimageio_cpnet_model(path_model_weight, nchan=2):
+    cpnet_kwargs = {
+        "nout": 3,
+    }
+    cpnet_biio = CPnetBioImageIO(**cpnet_kwargs)
+    state_dict_cuda = torch.load(path_model_weight, map_location=torch.device("cpu"), weights_only=True)
+    cpnet_biio.load_state_dict(state_dict_cuda)
+    cpnet_biio.eval()  # crucial for the prediction results
+    return cpnet_biio, cpnet_kwargs
+
+
+def descr_gen_input(path_test_input, nchan=2):
+    input_axes = [
+        SpaceInputAxis(id=AxisId("z"), size=ARBITRARY_SIZE),
+        ChannelAxis(channel_names=[Identifier(f"c{i+1}") for i in range(nchan)]),
+        SpaceInputAxis(id=AxisId("y"), size=ParameterizedSize(min=16, step=16)),
+        SpaceInputAxis(id=AxisId("x"), size=ParameterizedSize(min=16, step=16)),
+    ]
+    data_descr = IntervalOrRatioDataDescr(type="float32")
+    path_test_input = Path(path_test_input)
+    descr_input = InputTensorDescr(
+        id=TensorId("raw"),
+        axes=input_axes,
+        test_tensor=FileDescr(source=path_test_input),
+        data=data_descr,
+    )
+    return descr_input
+
+
+def descr_gen_output_flow(path_test_output):
+    output_axes_output_tensor = [
+        SpaceOutputAxis(id=AxisId("z"), size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("z"))),
+        ChannelAxis(channel_names=[Identifier("flow1"), Identifier("flow2"), Identifier("flow3")]),
+        SpaceOutputAxis(id=AxisId("y"), size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("y"))),
+        SpaceOutputAxis(id=AxisId("x"), size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("x"))),
+    ]
+    path_test_output = Path(path_test_output)
+    descr_output = OutputTensorDescr(
+        id=TensorId("flow"),
+        axes=output_axes_output_tensor,
+        test_tensor=FileDescr(source=path_test_output),
+    )
+    return descr_output
+
+
+def descr_gen_output_downsampled(path_dir_temp, nbase=None):
+    if nbase is None:
+        nbase = [32, 64, 128, 256]
+
+    output_axes_downsampled_tensors = [
+        [
+            SpaceOutputAxis(id=AxisId("z"), size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("z"))),
+            ChannelAxis(channel_names=[Identifier(f"feature{i+1}") for i in range(base)]),
+            SpaceOutputAxis(
+                id=AxisId("y"),
+                size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("y")),
+                scale=2**offset,
+            ),
+            SpaceOutputAxis(
+                id=AxisId("x"),
+                size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("x")),
+                scale=2**offset,
+            ),
+        ]
+        for offset, base in enumerate(nbase)
+    ]
+    path_downsampled_tensors = [
+        Path(path_dir_temp / f"test_downsampled_{i}.npy") for i in range(len(output_axes_downsampled_tensors))
+    ]
+    descr_output_downsampled_tensors = [
+        OutputTensorDescr(
+            id=TensorId(f"downsampled_{i}"),
+            axes=axes,
+            test_tensor=FileDescr(source=path),
+        )
+        for i, (axes, path) in enumerate(zip(output_axes_downsampled_tensors, path_downsampled_tensors))
+    ]
+    return descr_output_downsampled_tensors
+
+
+def descr_gen_output_style(path_test_style, nchannel=256):
+    output_axes_style_tensor = [
+        SpaceOutputAxis(id=AxisId("z"), size=SizeReference(tensor_id=TensorId("raw"), axis_id=AxisId("z"))),
+        ChannelAxis(channel_names=[Identifier(f"feature{i+1}") for i in range(nchannel)]),
+    ]
+    path_style_tensor = Path(path_test_style)
+    descr_output_style_tensor = OutputTensorDescr(
+        id=TensorId("style"),
+        axes=output_axes_style_tensor,
+        test_tensor=FileDescr(source=path_style_tensor),
+    )
+    return descr_output_style_tensor
+
+
+def descr_gen_arch(cpnet_kwargs, path_cpnet_wrapper=None):
+    if path_cpnet_wrapper is None:
+        path_cpnet_wrapper = Path(__file__).parent / "resnet_torch.py"
+    pytorch_architecture = ArchitectureFromFileDescr(
+        callable=Identifier("CPnetBioImageIO"),
+        source=Path(path_cpnet_wrapper),
+        kwargs=cpnet_kwargs,
+    )
+    return pytorch_architecture
+
+
+def descr_gen_documentation(path_doc, markdown_text):
+    with open(path_doc, "w") as f:
+        f.write(markdown_text)
+
+
+def package_to_bioimageio(
+    path_pretrained_model,
+    path_save_trace,
+    path_readme,
+    list_path_cover_images,
+    descr_input,
+    descr_output,
+    descr_output_downsampled_tensors,
+    descr_output_style_tensor,
+    pytorch_version,
+    pytorch_architecture,
+    model_id,
+    model_icon,
+    model_version,
+    model_name,
+    model_documentation,
+    model_authors,
+    model_cite,
+    model_tags,
+    model_license,
+    model_repo,
+):
+    """Package model description to BioImage.IO format."""
+    my_model_descr = ModelDescr(
+        id=ModelId(model_id) if model_id is not None else None,
+        id_emoji=model_icon,
+        version=Version(model_version),
+        name=model_name,
+        description=model_documentation,
+        authors=[
+            Author(
+                name=author["name"],
+                affiliation=author["affiliation"],
+                github_user=author["github_user"],
+                orcid=OrcidId(author["orcid"]),
+            )
+            for author in model_authors
+        ],
+        cite=[CiteEntry(text=cite["text"], doi=Doi(cite["doi"]), url=cite["url"]) for cite in model_cite],
+        covers=[Path(img) for img in list_path_cover_images],
+        license=LicenseId(model_license),
+        tags=model_tags,
+        documentation=Path(path_readme),
+        git_repo=HttpUrl(model_repo),
+        inputs=[descr_input],
+        outputs=[descr_output, descr_output_style_tensor] + descr_output_downsampled_tensors,
+        weights=WeightsDescr(
+            pytorch_state_dict=PytorchStateDictWeightsDescr(
+                source=Path(path_pretrained_model),
+                architecture=pytorch_architecture,
+                pytorch_version=pytorch_version,
+            ),
+            torchscript=TorchscriptWeightsDescr(
+                source=Path(path_save_trace),
+                pytorch_version=pytorch_version,
+                parent="pytorch_state_dict",  # these weights were converted from the pytorch_state_dict weights.
+            ),
+        ),
+    )
+
+    return my_model_descr
+
+
+def parse_args():
+    # fmt: off
+    parser = argparse.ArgumentParser(description="BioImage.IO model packaging for Cellpose")
+    parser.add_argument("--channels", nargs=2, default=[2, 1], type=int, help="Cyto-only = [2, 0], Cyto + Nuclei = [2, 1], Nuclei-only = [1, 0]")
+    parser.add_argument("--path_pretrained_model", required=True, type=str, help="Path to pretrained model file, e.g., cellpose_residual_on_style_on_concatenation_off_1135_rest_2023_05_04_23_41_31.252995")
+    parser.add_argument("--path_readme", required=True, type=str, help="Path to README file")
+    parser.add_argument("--list_path_cover_images", nargs='+', required=True, type=str, help="List of paths to cover images")
+    parser.add_argument("--model_id", type=str, help="Model ID, provide if already exists", default=None)
+    parser.add_argument("--model_icon", type=str, help="Model icon, provide if already exists", default=None)
+    parser.add_argument("--model_version", required=True, type=str, help="Model version, new model should be 0.1.0")
+    parser.add_argument("--model_name", required=True, type=str, help="Model name, e.g., My Cool Cellpose")
+    parser.add_argument("--model_documentation", required=True, type=str, help="Model documentation, e.g., A cool Cellpose model trained for my cool dataset.")
+    parser.add_argument("--model_authors", required=True, type=str, help="Model authors in JSON format, e.g., '[{\"name\": \"Qin Yu\", \"affiliation\": \"EMBL\", \"github_user\": \"qin-yu\", \"orcid\": \"0000-0002-4652-0795\"}]'")
+    parser.add_argument("--model_cite", required=True, type=str, help="Model citation in JSON format, e.g., '[{\"text\": \"For more details of the model itself, see the manuscript\", \"doi\": \"10.1242/dev.202800\", \"url\": null}]'")
+    parser.add_argument("--model_tags", nargs='+', required=True, type=str, help="Model tags, e.g., cellpose 3d 2d")
+    parser.add_argument("--model_license", required=True, type=str, help="Model license, e.g., MIT")
+    parser.add_argument("--model_repo", required=True, type=str, help="Model repository URL")
+    return parser.parse_args()
+    # fmt: on
+
+
+def main():
+    args = parse_args()
+
+    # Parse user-provided paths and arguments
+    channels = args.channels
+    model_cite = json.loads(args.model_cite)
+    model_authors = json.loads(args.model_authors)
+
+    path_readme = Path(args.path_readme)
+    path_pretrained_model = Path(args.path_pretrained_model)
+    list_path_cover_images = [Path(path_image) for path_image in args.list_path_cover_images]
+
+    # Auto-generated paths
+    path_cpnet_wrapper = Path(__file__).resolve().parent / "resnet_torch.py"
+    path_dir_temp = Path(__file__).resolve().parent.parent / "models" / path_pretrained_model.stem
+    path_dir_temp.mkdir(parents=True, exist_ok=True)
+
+    path_save_trace = path_dir_temp / "cp_traced.pt"
+    path_test_input = path_dir_temp / "test_input.npy"
+    path_test_output = path_dir_temp / "test_output.npy"
+    path_test_style = path_dir_temp / "test_style.npy"
+    path_bioimageio_package = path_dir_temp / "cellpose_model.zip"
+
+    # Download test input image
+    img_np = download_and_normalize_image(path_dir_temp, channels=channels)
+    np.save(path_test_input, img_np)
+    img = torch.tensor(img_np).float()
+
+    # Load model
+    cpnet_biio, cpnet_kwargs = load_bioimageio_cpnet_model(path_pretrained_model)
+
+    # Test model and save output
+    tuple_output_tensor = cpnet_biio(img)
+    np.save(path_test_output, tuple_output_tensor[0].detach().numpy())
+    np.save(path_test_style, tuple_output_tensor[1].detach().numpy())
+    for i, t in enumerate(tuple_output_tensor[2:]):
+        np.save(path_dir_temp / f"test_downsampled_{i}.npy", t.detach().numpy())
+
+    # Save traced model
+    model_traced = torch.jit.trace(cpnet_biio, img)
+    model_traced.save(path_save_trace)
+
+    # Generate model description
+    descr_input = descr_gen_input(path_test_input)
+    descr_output = descr_gen_output_flow(path_test_output)
+    descr_output_downsampled_tensors = descr_gen_output_downsampled(path_dir_temp, nbase=cpnet_biio.nbase[1:])
+    descr_output_style_tensor = descr_gen_output_style(path_test_style, cpnet_biio.nbase[-1])
+    pytorch_version = Version(torch.__version__)
+    pytorch_architecture = descr_gen_arch(cpnet_kwargs, path_cpnet_wrapper)
+
+    # Package model
+    my_model_descr = package_to_bioimageio(
+        path_pretrained_model,
+        path_save_trace,
+        path_readme,
+        list_path_cover_images,
+        descr_input,
+        descr_output,
+        descr_output_downsampled_tensors,
+        descr_output_style_tensor,
+        pytorch_version,
+        pytorch_architecture,
+        args.model_id,
+        args.model_icon,
+        args.model_version,
+        args.model_name,
+        args.model_documentation,
+        model_authors,
+        model_cite,
+        args.model_tags,
+        args.model_license,
+        args.model_repo,
+    )
+
+    # Test model
+    summary = test_model(my_model_descr, weight_format="pytorch_state_dict")
+    summary.display()
+    summary = test_model(my_model_descr, weight_format="torchscript")
+    summary.display()
+
+    # Save BioImage.IO package
+    package_path = save_bioimageio_package(my_model_descr, output_path=Path(path_bioimageio_package))
+    print("package path:", package_path)
+
+
+if __name__ == "__main__":
+    main()

models/seg_post_model/cellpose/gui/gui.py

@@ -0,0 +1,2007 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer, Michael Rariden and Marius Pachitariu.
+"""
+
+import sys, os, pathlib, warnings, datetime, time, copy
+
+from qtpy import QtGui, QtCore
+from superqt import QRangeSlider, QCollapsible
+from qtpy.QtWidgets import QScrollArea, QMainWindow, QApplication, QWidget, QScrollBar, \
+    QComboBox, QGridLayout, QPushButton, QFrame, QCheckBox, QLabel, QProgressBar, \
+        QLineEdit, QMessageBox, QGroupBox, QMenu, QAction
+import pyqtgraph as pg
+
+import numpy as np
+from scipy.stats import mode
+import cv2
+
+from . import guiparts, menus, io
+from .. import models, core, dynamics, version, train
+from ..utils import download_url_to_file, masks_to_outlines, diameters
+from ..io import get_image_files, imsave, imread
+from ..transforms import resize_image, normalize99, normalize99_tile, smooth_sharpen_img
+from ..models import normalize_default
+from ..plot import disk
+
+try:
+    import matplotlib.pyplot as plt
+    MATPLOTLIB = True
+except:
+    MATPLOTLIB = False
+
+Horizontal = QtCore.Qt.Orientation.Horizontal
+
+
+class Slider(QRangeSlider):
+
+    def __init__(self, parent, name, color):
+        super().__init__(Horizontal)
+        self.setEnabled(False)
+        self.valueChanged.connect(lambda: self.levelChanged(parent))
+        self.name = name
+
+        self.setStyleSheet(""" QSlider{
+                             background-color: transparent;
+                             }
+        """)
+        self.show()
+
+    def levelChanged(self, parent):
+        parent.level_change(self.name)
+
+
+class QHLine(QFrame):
+
+    def __init__(self):
+        super(QHLine, self).__init__()
+        self.setFrameShape(QFrame.HLine)
+        self.setLineWidth(8)
+
+
+def make_bwr():
+    # make a bwr colormap
+    b = np.append(255 * np.ones(128), np.linspace(0, 255, 128)[::-1])[:, np.newaxis]
+    r = np.append(np.linspace(0, 255, 128), 255 * np.ones(128))[:, np.newaxis]
+    g = np.append(np.linspace(0, 255, 128),
+                  np.linspace(0, 255, 128)[::-1])[:, np.newaxis]
+    color = np.concatenate((r, g, b), axis=-1).astype(np.uint8)
+    bwr = pg.ColorMap(pos=np.linspace(0.0, 255, 256), color=color)
+    return bwr
+
+
+def make_spectral():
+    # make spectral colormap
+    r = np.array([
+        0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80,
+        84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 128, 128, 128, 128, 128,
+        128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 120, 112, 104, 96, 88,
+        80, 72, 64, 56, 48, 40, 32, 24, 16, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 7, 11, 15, 19, 23,
+        27, 31, 35, 39, 43, 47, 51, 55, 59, 63, 67, 71, 75, 79, 83, 87, 91, 95, 99, 103,
+        107, 111, 115, 119, 123, 127, 131, 135, 139, 143, 147, 151, 155, 159, 163, 167,
+        171, 175, 179, 183, 187, 191, 195, 199, 203, 207, 211, 215, 219, 223, 227, 231,
+        235, 239, 243, 247, 251, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255
+    ])
+    g = np.array([
+        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 9, 8, 8, 7, 7, 6, 6, 5, 5, 5, 4, 4, 3, 3,
+        2, 2, 1, 1, 0, 0, 0, 7, 15, 23, 31, 39, 47, 55, 63, 71, 79, 87, 95, 103, 111,
+        119, 127, 135, 143, 151, 159, 167, 175, 183, 191, 199, 207, 215, 223, 231, 239,
+        247, 255, 247, 239, 231, 223, 215, 207, 199, 191, 183, 175, 167, 159, 151, 143,
+        135, 128, 129, 131, 132, 134, 135, 137, 139, 140, 142, 143, 145, 147, 148, 150,
+        151, 153, 154, 156, 158, 159, 161, 162, 164, 166, 167, 169, 170, 172, 174, 175,
+        177, 178, 180, 181, 183, 185, 186, 188, 189, 191, 193, 194, 196, 197, 199, 201,
+        202, 204, 205, 207, 208, 210, 212, 213, 215, 216, 218, 220, 221, 223, 224, 226,
+        228, 229, 231, 232, 234, 235, 237, 239, 240, 242, 243, 245, 247, 248, 250, 251,
+        253, 255, 251, 247, 243, 239, 235, 231, 227, 223, 219, 215, 211, 207, 203, 199,
+        195, 191, 187, 183, 179, 175, 171, 167, 163, 159, 155, 151, 147, 143, 139, 135,
+        131, 127, 123, 119, 115, 111, 107, 103, 99, 95, 91, 87, 83, 79, 75, 71, 67, 63,
+        59, 55, 51, 47, 43, 39, 35, 31, 27, 23, 19, 15, 11, 7, 3, 0, 8, 16, 24, 32, 41,
+        49, 57, 65, 74, 82, 90, 98, 106, 115, 123, 131, 139, 148, 156, 164, 172, 180,
+        189, 197, 205, 213, 222, 230, 238, 246, 254
+    ])
+    b = np.array([
+        0, 7, 15, 23, 31, 39, 47, 55, 63, 71, 79, 87, 95, 103, 111, 119, 127, 135, 143,
+        151, 159, 167, 175, 183, 191, 199, 207, 215, 223, 231, 239, 247, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+        255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 251, 247,
+        243, 239, 235, 231, 227, 223, 219, 215, 211, 207, 203, 199, 195, 191, 187, 183,
+        179, 175, 171, 167, 163, 159, 155, 151, 147, 143, 139, 135, 131, 128, 126, 124,
+        122, 120, 118, 116, 114, 112, 110, 108, 106, 104, 102, 100, 98, 96, 94, 92, 90,
+        88, 86, 84, 82, 80, 78, 76, 74, 72, 70, 68, 66, 64, 62, 60, 58, 56, 54, 52, 50,
+        48, 46, 44, 42, 40, 38, 36, 34, 32, 30, 28, 26, 24, 22, 20, 18, 16, 14, 12, 10,
+        8, 6, 4, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 8, 16, 24, 32, 41, 49, 57, 65, 74,
+        82, 90, 98, 106, 115, 123, 131, 139, 148, 156, 164, 172, 180, 189, 197, 205,
+        213, 222, 230, 238, 246, 254
+    ])
+    color = (np.vstack((r, g, b)).T).astype(np.uint8)
+    spectral = pg.ColorMap(pos=np.linspace(0.0, 255, 256), color=color)
+    return spectral
+
+
+def make_cmap(cm=0):
+    # make a single channel colormap
+    r = np.arange(0, 256)
+    color = np.zeros((256, 3))
+    color[:, cm] = r
+    color = color.astype(np.uint8)
+    cmap = pg.ColorMap(pos=np.linspace(0.0, 255, 256), color=color)
+    return cmap
+
+
+def run(image=None):
+    from ..io import logger_setup
+    logger, log_file = logger_setup()
+    # Always start by initializing Qt (only once per application)
+    warnings.filterwarnings("ignore")
+    app = QApplication(sys.argv)
+    icon_path = pathlib.Path.home().joinpath(".cellpose", "logo.png")
+    guip_path = pathlib.Path.home().joinpath(".cellpose", "cellposeSAM_gui.png")
+    if not icon_path.is_file():
+        cp_dir = pathlib.Path.home().joinpath(".cellpose")
+        cp_dir.mkdir(exist_ok=True)
+        print("downloading logo")
+        download_url_to_file(
+            "https://www.cellpose.org/static/images/cellpose_transparent.png",
+            icon_path, progress=True)
+    if not guip_path.is_file():
+        print("downloading help window image")
+        download_url_to_file("https://www.cellpose.org/static/images/cellposeSAM_gui.png",
+                             guip_path, progress=True)
+    icon_path = str(icon_path.resolve())
+    app_icon = QtGui.QIcon()
+    app_icon.addFile(icon_path, QtCore.QSize(16, 16))
+    app_icon.addFile(icon_path, QtCore.QSize(24, 24))
+    app_icon.addFile(icon_path, QtCore.QSize(32, 32))
+    app_icon.addFile(icon_path, QtCore.QSize(48, 48))
+    app_icon.addFile(icon_path, QtCore.QSize(64, 64))
+    app_icon.addFile(icon_path, QtCore.QSize(256, 256))
+    app.setWindowIcon(app_icon)
+    app.setStyle("Fusion")
+    app.setPalette(guiparts.DarkPalette())
+    MainW(image=image, logger=logger)
+    ret = app.exec_()
+    sys.exit(ret)
+
+
+class MainW(QMainWindow):
+
+    def __init__(self, image=None, logger=None):
+        super(MainW, self).__init__()
+
+        self.logger = logger
+        pg.setConfigOptions(imageAxisOrder="row-major")
+        self.setGeometry(50, 50, 1200, 1000)
+        self.setWindowTitle(f"cellpose v{version}")
+        self.cp_path = os.path.dirname(os.path.realpath(__file__))
+        app_icon = QtGui.QIcon()
+        icon_path = pathlib.Path.home().joinpath(".cellpose", "logo.png")
+        icon_path = str(icon_path.resolve())
+        app_icon.addFile(icon_path, QtCore.QSize(16, 16))
+        app_icon.addFile(icon_path, QtCore.QSize(24, 24))
+        app_icon.addFile(icon_path, QtCore.QSize(32, 32))
+        app_icon.addFile(icon_path, QtCore.QSize(48, 48))
+        app_icon.addFile(icon_path, QtCore.QSize(64, 64))
+        app_icon.addFile(icon_path, QtCore.QSize(256, 256))
+        self.setWindowIcon(app_icon)
+        # rgb(150,255,150)
+        self.setStyleSheet(guiparts.stylesheet())
+
+        menus.mainmenu(self)
+        menus.editmenu(self)
+        menus.modelmenu(self)
+        menus.helpmenu(self)
+
+        self.stylePressed = """QPushButton {Text-align: center; 
+                             background-color: rgb(150,50,150); 
+                             border-color: white;
+                             color:white;}
+                            QToolTip { 
+                           background-color: black; 
+                           color: white; 
+                           border: black solid 1px
+                           }"""
+        self.styleUnpressed = """QPushButton {Text-align: center; 
+                               background-color: rgb(50,50,50);
+                                border-color: white;
+                               color:white;}
+                                QToolTip { 
+                           background-color: black; 
+                           color: white; 
+                           border: black solid 1px
+                           }"""
+        self.loaded = False
+
+        # ---- MAIN WIDGET LAYOUT ---- #
+        self.cwidget = QWidget(self)
+        self.lmain = QGridLayout()
+        self.cwidget.setLayout(self.lmain)
+        self.setCentralWidget(self.cwidget)
+        self.lmain.setVerticalSpacing(0)
+        self.lmain.setContentsMargins(0, 0, 0, 10)
+
+        self.imask = 0
+        self.scrollarea = QScrollArea()
+        self.scrollarea.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOn)
+        self.scrollarea.setStyleSheet("""QScrollArea { border: none }""")
+        self.scrollarea.setWidgetResizable(True)
+        self.swidget = QWidget(self)
+        self.scrollarea.setWidget(self.swidget)
+        self.l0 = QGridLayout()
+        self.swidget.setLayout(self.l0)
+        b = self.make_buttons()
+        self.lmain.addWidget(self.scrollarea, 0, 0, 39, 9)
+
+        # ---- drawing area ---- #
+        self.win = pg.GraphicsLayoutWidget()
+
+        self.lmain.addWidget(self.win, 0, 9, 40, 30)
+
+        self.win.scene().sigMouseClicked.connect(self.plot_clicked)
+        self.win.scene().sigMouseMoved.connect(self.mouse_moved)
+        self.make_viewbox()
+        self.lmain.setColumnStretch(10, 1)
+        bwrmap = make_bwr()
+        self.bwr = bwrmap.getLookupTable(start=0.0, stop=255.0, alpha=False)
+        self.cmap = []
+        # spectral colormap
+        self.cmap.append(make_spectral().getLookupTable(start=0.0, stop=255.0,
+                                                        alpha=False))
+        # single channel colormaps
+        for i in range(3):
+            self.cmap.append(
+                make_cmap(i).getLookupTable(start=0.0, stop=255.0, alpha=False))
+
+        if MATPLOTLIB:
+            self.colormap = (plt.get_cmap("gist_ncar")(np.linspace(0.0, .9, 1000000)) *
+                             255).astype(np.uint8)
+            np.random.seed(42)  # make colors stable
+            self.colormap = self.colormap[np.random.permutation(1000000)]
+        else:
+            np.random.seed(42)  # make colors stable
+            self.colormap = ((np.random.rand(1000000, 3) * 0.8 + 0.1) * 255).astype(
+                np.uint8)
+        self.NZ = 1
+        self.restore = None
+        self.ratio = 1.
+        self.reset()
+
+        # This needs to go after .reset() is called to get state fully set up:
+        self.autobtn.checkStateChanged.connect(self.compute_saturation_if_checked)
+
+        self.load_3D = False
+
+        # if called with image, load it
+        if image is not None:
+            self.filename = image
+            io._load_image(self, self.filename)
+
+        # training settings
+        d = datetime.datetime.now()
+        self.training_params = {
+            "model_index": 0,
+            "learning_rate": 1e-5,
+            "weight_decay": 0.1,
+            "n_epochs": 100,
+            "model_name": "cpsam" + d.strftime("_%Y%m%d_%H%M%S"),
+        }
+
+        self.stitch_threshold = 0.
+        self.flow3D_smooth = 0.
+        self.anisotropy = 1.
+        self.min_size = 15
+
+        self.setAcceptDrops(True)
+        self.win.show()
+        self.show()
+
+    def help_window(self):
+        HW = guiparts.HelpWindow(self)
+        HW.show()
+
+    def train_help_window(self):
+        THW = guiparts.TrainHelpWindow(self)
+        THW.show()
+
+    def gui_window(self):
+        EG = guiparts.ExampleGUI(self)
+        EG.show()
+
+    def make_buttons(self):
+        self.boldfont = QtGui.QFont("Arial", 11, QtGui.QFont.Bold)
+        self.boldmedfont = QtGui.QFont("Arial", 9, QtGui.QFont.Bold)
+        self.medfont = QtGui.QFont("Arial", 9)
+        self.smallfont = QtGui.QFont("Arial", 8)
+
+        b = 0
+        self.satBox = QGroupBox("Views")
+        self.satBox.setFont(self.boldfont)
+        self.satBoxG = QGridLayout()
+        self.satBox.setLayout(self.satBoxG)
+        self.l0.addWidget(self.satBox, b, 0, 1, 9)
+
+        widget_row = 0
+        self.view = 0  # 0=image, 1=flowsXY, 2=flowsZ, 3=cellprob
+        self.color = 0  # 0=RGB, 1=gray, 2=R, 3=G, 4=B
+        self.RGBDropDown = QComboBox()
+        self.RGBDropDown.addItems(
+            ["RGB", "red=R", "green=G", "blue=B", "gray", "spectral"])
+        self.RGBDropDown.setFont(self.medfont)
+        self.RGBDropDown.currentIndexChanged.connect(self.color_choose)
+        self.satBoxG.addWidget(self.RGBDropDown, widget_row, 0, 1, 3)
+
+        label = QLabel("<p>[&uarr; / &darr; or W/S]</p>")
+        label.setFont(self.smallfont)
+        self.satBoxG.addWidget(label, widget_row, 3, 1, 3)
+        label = QLabel("[R / G / B \n toggles color ]")
+        label.setFont(self.smallfont)
+        self.satBoxG.addWidget(label, widget_row, 6, 1, 3)
+
+        widget_row += 1
+        self.ViewDropDown = QComboBox()
+        self.ViewDropDown.addItems(["image", "gradXY", "cellprob", "restored"])
+        self.ViewDropDown.setFont(self.medfont)
+        self.ViewDropDown.model().item(3).setEnabled(False)
+        self.ViewDropDown.currentIndexChanged.connect(self.update_plot)
+        self.satBoxG.addWidget(self.ViewDropDown, widget_row, 0, 2, 3)
+
+        label = QLabel("[pageup / pagedown]")
+        label.setFont(self.smallfont)
+        self.satBoxG.addWidget(label, widget_row, 3, 1, 5)
+
+        widget_row += 2
+        label = QLabel("")
+        label.setToolTip(
+            "NOTE: manually changing the saturation bars does not affect normalization in segmentation"
+        )
+        self.satBoxG.addWidget(label, widget_row, 0, 1, 5)
+
+        self.autobtn = QCheckBox("auto-adjust saturation")
+        self.autobtn.setToolTip("sets scale-bars as normalized for segmentation")
+        self.autobtn.setFont(self.medfont)
+        self.autobtn.setChecked(True)
+        self.satBoxG.addWidget(self.autobtn, widget_row, 1, 1, 8)
+
+        widget_row += 1
+        self.sliders = []
+        colors = [[255, 0, 0], [0, 255, 0], [0, 0, 255], [100, 100, 100]]
+        colornames = ["red", "Chartreuse", "DodgerBlue"]
+        names = ["red", "green", "blue"]
+        for r in range(3):
+            widget_row += 1
+            if r == 0:
+                label = QLabel('<font color="gray">gray/</font><br>red')
+            else:
+                label = QLabel(names[r] + ":")
+            label.setStyleSheet(f"color: {colornames[r]}")
+            label.setFont(self.boldmedfont)
+            self.satBoxG.addWidget(label, widget_row, 0, 1, 2)
+            self.sliders.append(Slider(self, names[r], colors[r]))
+            self.sliders[-1].setMinimum(-.1)
+            self.sliders[-1].setMaximum(255.1)
+            self.sliders[-1].setValue([0, 255])
+            self.sliders[-1].setToolTip(
+                "NOTE: manually changing the saturation bars does not affect normalization in segmentation"
+            )
+            self.satBoxG.addWidget(self.sliders[-1], widget_row, 2, 1, 7)
+
+        b += 1
+        self.drawBox = QGroupBox("Drawing")
+        self.drawBox.setFont(self.boldfont)
+        self.drawBoxG = QGridLayout()
+        self.drawBox.setLayout(self.drawBoxG)
+        self.l0.addWidget(self.drawBox, b, 0, 1, 9)
+        self.autosave = True
+
+        widget_row = 0
+        self.brush_size = 3
+        self.BrushChoose = QComboBox()
+        self.BrushChoose.addItems(["1", "3", "5", "7", "9"])
+        self.BrushChoose.currentIndexChanged.connect(self.brush_choose)
+        self.BrushChoose.setFixedWidth(40)
+        self.BrushChoose.setFont(self.medfont)
+        self.drawBoxG.addWidget(self.BrushChoose, widget_row, 3, 1, 2)
+        label = QLabel("brush size:")
+        label.setFont(self.medfont)
+        self.drawBoxG.addWidget(label, widget_row, 0, 1, 3)
+
+        widget_row += 1
+        # turn off masks
+        self.layer_off = False
+        self.masksOn = True
+        self.MCheckBox = QCheckBox("MASKS ON [X]")
+        self.MCheckBox.setFont(self.medfont)
+        self.MCheckBox.setChecked(True)
+        self.MCheckBox.toggled.connect(self.toggle_masks)
+        self.drawBoxG.addWidget(self.MCheckBox, widget_row, 0, 1, 5)
+
+        widget_row += 1
+        # turn off outlines
+        self.outlinesOn = False  # turn off by default
+        self.OCheckBox = QCheckBox("outlines on [Z]")
+        self.OCheckBox.setFont(self.medfont)
+        self.drawBoxG.addWidget(self.OCheckBox, widget_row, 0, 1, 5)
+        self.OCheckBox.setChecked(False)
+        self.OCheckBox.toggled.connect(self.toggle_masks)
+
+        widget_row += 1
+        self.SCheckBox = QCheckBox("single stroke")
+        self.SCheckBox.setFont(self.medfont)
+        self.SCheckBox.setChecked(True)
+        self.SCheckBox.toggled.connect(self.autosave_on)
+        self.SCheckBox.setEnabled(True)
+        self.drawBoxG.addWidget(self.SCheckBox, widget_row, 0, 1, 5)
+
+        # buttons for deleting multiple cells
+        self.deleteBox = QGroupBox("delete multiple ROIs")
+        self.deleteBox.setStyleSheet("color: rgb(200, 200, 200)")
+        self.deleteBox.setFont(self.medfont)
+        self.deleteBoxG = QGridLayout()
+        self.deleteBox.setLayout(self.deleteBoxG)
+        self.drawBoxG.addWidget(self.deleteBox, 0, 5, 4, 4)
+        self.MakeDeletionRegionButton = QPushButton("region-select")
+        self.MakeDeletionRegionButton.clicked.connect(self.remove_region_cells)
+        self.deleteBoxG.addWidget(self.MakeDeletionRegionButton, 0, 0, 1, 4)
+        self.MakeDeletionRegionButton.setFont(self.smallfont)
+        self.MakeDeletionRegionButton.setFixedWidth(70)
+        self.DeleteMultipleROIButton = QPushButton("click-select")
+        self.DeleteMultipleROIButton.clicked.connect(self.delete_multiple_cells)
+        self.deleteBoxG.addWidget(self.DeleteMultipleROIButton, 1, 0, 1, 4)
+        self.DeleteMultipleROIButton.setFont(self.smallfont)
+        self.DeleteMultipleROIButton.setFixedWidth(70)
+        self.DoneDeleteMultipleROIButton = QPushButton("done")
+        self.DoneDeleteMultipleROIButton.clicked.connect(
+            self.done_remove_multiple_cells)
+        self.deleteBoxG.addWidget(self.DoneDeleteMultipleROIButton, 2, 0, 1, 2)
+        self.DoneDeleteMultipleROIButton.setFont(self.smallfont)
+        self.DoneDeleteMultipleROIButton.setFixedWidth(35)
+        self.CancelDeleteMultipleROIButton = QPushButton("cancel")
+        self.CancelDeleteMultipleROIButton.clicked.connect(self.cancel_remove_multiple)
+        self.deleteBoxG.addWidget(self.CancelDeleteMultipleROIButton, 2, 2, 1, 2)
+        self.CancelDeleteMultipleROIButton.setFont(self.smallfont)
+        self.CancelDeleteMultipleROIButton.setFixedWidth(35)
+
+        b += 1
+        widget_row = 0
+        self.segBox = QGroupBox("Segmentation")
+        self.segBoxG = QGridLayout()
+        self.segBox.setLayout(self.segBoxG)
+        self.l0.addWidget(self.segBox, b, 0, 1, 9)
+        self.segBox.setFont(self.boldfont)
+
+        widget_row += 1
+
+        # use GPU
+        self.useGPU = QCheckBox("use GPU")
+        self.useGPU.setToolTip(
+            "if you have specially installed the <i>cuda</i> version of torch, then you can activate this"
+        )
+        self.useGPU.setFont(self.medfont)
+        self.check_gpu()
+        self.segBoxG.addWidget(self.useGPU, widget_row, 0, 1, 3)
+
+        # compute segmentation with general models
+        self.net_text = ["run CPSAM"]
+        nett = ["cellpose super-generalist model"]
+
+        self.StyleButtons = []
+        jj = 4
+        for j in range(len(self.net_text)):
+            self.StyleButtons.append(
+                guiparts.ModelButton(self, self.net_text[j], self.net_text[j]))
+            w = 5
+            self.segBoxG.addWidget(self.StyleButtons[-1], widget_row, jj, 1, w)
+            jj += w
+            self.StyleButtons[-1].setToolTip(nett[j])
+
+        widget_row += 1
+        self.ncells = guiparts.ObservableVariable(0)
+        self.roi_count = QLabel()
+        self.roi_count.setFont(self.boldfont)
+        self.roi_count.setAlignment(QtCore.Qt.AlignLeft)
+        self.ncells.valueChanged.connect(
+            lambda n: self.roi_count.setText(f'{str(n)} ROIs')
+        )
+
+        self.segBoxG.addWidget(self.roi_count, widget_row, 0, 1, 4)
+
+        self.progress = QProgressBar(self)
+        self.segBoxG.addWidget(self.progress, widget_row, 4, 1, 5)
+
+        widget_row += 1
+
+        ############################### Segmentation settings ###############################
+        self.additional_seg_settings_qcollapsible = QCollapsible("additional settings")
+        self.additional_seg_settings_qcollapsible.setFont(self.medfont)
+        self.additional_seg_settings_qcollapsible._toggle_btn.setFont(self.medfont)
+        self.segmentation_settings = guiparts.SegmentationSettings(self.medfont)
+        self.additional_seg_settings_qcollapsible.setContent(self.segmentation_settings)
+        self.segBoxG.addWidget(self.additional_seg_settings_qcollapsible, widget_row, 0, 1, 9)
+
+        # connect edits to image processing steps: 
+        self.segmentation_settings.diameter_box.editingFinished.connect(self.update_scale)
+        self.segmentation_settings.flow_threshold_box.returnPressed.connect(self.compute_cprob)
+        self.segmentation_settings.cellprob_threshold_box.returnPressed.connect(self.compute_cprob)
+        self.segmentation_settings.niter_box.returnPressed.connect(self.compute_cprob)
+
+        # Needed to do this for the drop down to not be open on startup
+        self.additional_seg_settings_qcollapsible._toggle_btn.setChecked(True)
+        self.additional_seg_settings_qcollapsible._toggle_btn.setChecked(False)
+
+        b += 1
+        self.modelBox = QGroupBox("user-trained models")
+        self.modelBoxG = QGridLayout()
+        self.modelBox.setLayout(self.modelBoxG)
+        self.l0.addWidget(self.modelBox, b, 0, 1, 9)
+        self.modelBox.setFont(self.boldfont)
+        # choose models
+        self.ModelChooseC = QComboBox()
+        self.ModelChooseC.setFont(self.medfont)
+        current_index = 0
+        self.ModelChooseC.addItems(["custom models"])
+        if len(self.model_strings) > 0:
+            self.ModelChooseC.addItems(self.model_strings)
+        self.ModelChooseC.setFixedWidth(175)
+        self.ModelChooseC.setCurrentIndex(current_index)
+        tipstr = 'add or train your own models in the "Models" file menu and choose model here'
+        self.ModelChooseC.setToolTip(tipstr)
+        self.ModelChooseC.activated.connect(lambda: self.model_choose(custom=True))
+        self.modelBoxG.addWidget(self.ModelChooseC, widget_row, 0, 1, 8)
+
+        # compute segmentation w/ custom model
+        self.ModelButtonC = QPushButton(u"run")
+        self.ModelButtonC.setFont(self.medfont)
+        self.ModelButtonC.setFixedWidth(35)
+        self.ModelButtonC.clicked.connect(
+            lambda: self.compute_segmentation(custom=True))
+        self.modelBoxG.addWidget(self.ModelButtonC, widget_row, 8, 1, 1)
+        self.ModelButtonC.setEnabled(False)
+
+
+        b += 1
+        self.filterBox = QGroupBox("Image filtering")
+        self.filterBox.setFont(self.boldfont)
+        self.filterBox_grid_layout = QGridLayout()
+        self.filterBox.setLayout(self.filterBox_grid_layout)
+        self.l0.addWidget(self.filterBox, b, 0, 1, 9)
+
+        widget_row = 0
+        
+        # Filtering
+        self.FilterButtons = []
+        nett = [
+            "clear restore/filter",
+            "filter image (settings below)",
+        ]
+        self.filter_text = ["none", 
+                             "filter", 
+                             ]
+        self.restore = None
+        self.ratio = 1.
+        jj = 0
+        w = 3
+        for j in range(len(self.filter_text)):
+            self.FilterButtons.append(
+                guiparts.FilterButton(self, self.filter_text[j]))
+            self.filterBox_grid_layout.addWidget(self.FilterButtons[-1], widget_row, jj, 1, w)
+            self.FilterButtons[-1].setFixedWidth(75)
+            self.FilterButtons[-1].setToolTip(nett[j])
+            self.FilterButtons[-1].setFont(self.medfont)
+            widget_row += 1 if j%2==1 else 0
+            jj = 0 if j%2==1 else jj + w
+
+        self.save_norm = QCheckBox("save restored/filtered image")
+        self.save_norm.setFont(self.medfont)
+        self.save_norm.setToolTip("save restored/filtered image in _seg.npy file")
+        self.save_norm.setChecked(True)
+
+        widget_row += 2
+
+        self.filtBox = QCollapsible("custom filter settings")
+        self.filtBox._toggle_btn.setFont(self.medfont)
+        self.filtBoxG = QGridLayout()
+        _content = QWidget()
+        _content.setLayout(self.filtBoxG)
+        _content.setMaximumHeight(0)
+        _content.setMinimumHeight(0)
+        self.filtBox.setContent(_content)
+        self.filterBox_grid_layout.addWidget(self.filtBox, widget_row, 0, 1, 9)
+
+        self.filt_vals = [0., 0., 0., 0.]
+        self.filt_edits = []
+        labels = [
+            "sharpen\nradius", "smooth\nradius", "tile_norm\nblocksize",
+            "tile_norm\nsmooth3D"
+        ]
+        tooltips = [
+            "set size of surround-subtraction filter for sharpening image",
+            "set size of gaussian filter for smoothing image",
+            "set size of tiles to use to normalize image",
+            "set amount of smoothing of normalization values across planes"
+        ]
+
+        for p in range(4):
+            label = QLabel(f"{labels[p]}:")
+            label.setToolTip(tooltips[p])
+            label.setFont(self.medfont)
+            self.filtBoxG.addWidget(label, widget_row + p // 2, 4 * (p % 2), 1, 2)
+            self.filt_edits.append(QLineEdit())
+            self.filt_edits[p].setText(str(self.filt_vals[p]))
+            self.filt_edits[p].setFixedWidth(40)
+            self.filt_edits[p].setFont(self.medfont)
+            self.filtBoxG.addWidget(self.filt_edits[p], widget_row + p // 2, 4 * (p % 2) + 2, 1,
+                                    2)
+            self.filt_edits[p].setToolTip(tooltips[p])
+
+        widget_row += 3
+        self.norm3D_cb = QCheckBox("norm3D")
+        self.norm3D_cb.setFont(self.medfont)
+        self.norm3D_cb.setChecked(True)
+        self.norm3D_cb.setToolTip("run same normalization across planes")
+        self.filtBoxG.addWidget(self.norm3D_cb, widget_row, 0, 1, 3)
+
+
+        return b
+
+    def level_change(self, r):
+        r = ["red", "green", "blue"].index(r)
+        if self.loaded:
+            sval = self.sliders[r].value()
+            self.saturation[r][self.currentZ] = sval
+            if not self.autobtn.isChecked():
+                for r in range(3):
+                    for i in range(len(self.saturation[r])):
+                        self.saturation[r][i] = self.saturation[r][self.currentZ]
+            self.update_plot()
+
+    def keyPressEvent(self, event):
+        if self.loaded:
+            if not (event.modifiers() &
+                    (QtCore.Qt.ControlModifier | QtCore.Qt.ShiftModifier |
+                     QtCore.Qt.AltModifier) or self.in_stroke):
+                updated = False
+                if len(self.current_point_set) > 0:
+                    if event.key() == QtCore.Qt.Key_Return:
+                        self.add_set()
+                else:
+                    nviews = self.ViewDropDown.count() - 1
+                    nviews += int(
+                        self.ViewDropDown.model().item(self.ViewDropDown.count() -
+                                                       1).isEnabled())
+                    if event.key() == QtCore.Qt.Key_X:
+                        self.MCheckBox.toggle()
+                    if event.key() == QtCore.Qt.Key_Z:
+                        self.OCheckBox.toggle()
+                    if event.key() == QtCore.Qt.Key_Left or event.key(
+                    ) == QtCore.Qt.Key_A:
+                        self.get_prev_image()
+                    elif event.key() == QtCore.Qt.Key_Right or event.key(
+                    ) == QtCore.Qt.Key_D:
+                        self.get_next_image()
+                    elif event.key() == QtCore.Qt.Key_PageDown:
+                        self.view = (self.view + 1) % (nviews)
+                        self.ViewDropDown.setCurrentIndex(self.view)
+                    elif event.key() == QtCore.Qt.Key_PageUp:
+                        self.view = (self.view - 1) % (nviews)
+                        self.ViewDropDown.setCurrentIndex(self.view)
+
+                # can change background or stroke size if cell not finished
+                if event.key() == QtCore.Qt.Key_Up or event.key() == QtCore.Qt.Key_W:
+                    self.color = (self.color - 1) % (6)
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_Down or event.key(
+                ) == QtCore.Qt.Key_S:
+                    self.color = (self.color + 1) % (6)
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_R:
+                    if self.color != 1:
+                        self.color = 1
+                    else:
+                        self.color = 0
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_G:
+                    if self.color != 2:
+                        self.color = 2
+                    else:
+                        self.color = 0
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_B:
+                    if self.color != 3:
+                        self.color = 3
+                    else:
+                        self.color = 0
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif (event.key() == QtCore.Qt.Key_Comma or
+                      event.key() == QtCore.Qt.Key_Period):
+                    count = self.BrushChoose.count()
+                    gci = self.BrushChoose.currentIndex()
+                    if event.key() == QtCore.Qt.Key_Comma:
+                        gci = max(0, gci - 1)
+                    else:
+                        gci = min(count - 1, gci + 1)
+                    self.BrushChoose.setCurrentIndex(gci)
+                    self.brush_choose()
+                if not updated:
+                    self.update_plot()
+        if event.key() == QtCore.Qt.Key_Minus or event.key() == QtCore.Qt.Key_Equal:
+            self.p0.keyPressEvent(event)
+
+    def autosave_on(self):
+        if self.SCheckBox.isChecked():
+            self.autosave = True
+        else:
+            self.autosave = False
+
+    def check_gpu(self, torch=True):
+        # also decide whether or not to use torch
+        self.useGPU.setChecked(False)
+        self.useGPU.setEnabled(False)
+        if core.use_gpu(use_torch=True):
+            self.useGPU.setEnabled(True)
+            self.useGPU.setChecked(True)
+        else:
+            self.useGPU.setStyleSheet("color: rgb(80,80,80);")
+
+
+    def model_choose(self, custom=False):
+        index = self.ModelChooseC.currentIndex(
+        ) if custom else self.ModelChooseB.currentIndex()
+        if index > 0:
+            if custom:
+                model_name = self.ModelChooseC.currentText()
+            else:
+                model_name = self.net_names[index - 1]
+            print(f"GUI_INFO: selected model {model_name}, loading now")
+            self.initialize_model(model_name=model_name, custom=custom)
+
+    def toggle_scale(self):
+        if self.scale_on:
+            self.p0.removeItem(self.scale)
+            self.scale_on = False
+        else:
+            self.p0.addItem(self.scale)
+            self.scale_on = True
+
+    def enable_buttons(self):
+        if len(self.model_strings) > 0:
+            self.ModelButtonC.setEnabled(True)
+        for i in range(len(self.StyleButtons)):
+            self.StyleButtons[i].setEnabled(True)
+
+        for i in range(len(self.FilterButtons)):
+            self.FilterButtons[i].setEnabled(True)
+        if self.load_3D:
+            self.FilterButtons[-2].setEnabled(False)
+
+        self.newmodel.setEnabled(True)
+        self.loadMasks.setEnabled(True)
+
+        for n in range(self.nchan):
+            self.sliders[n].setEnabled(True)
+        for n in range(self.nchan, 3):
+            self.sliders[n].setEnabled(True)
+
+        self.toggle_mask_ops()
+
+        self.update_plot()
+        self.setWindowTitle(self.filename)
+
+    def disable_buttons_removeROIs(self):
+        if len(self.model_strings) > 0:
+            self.ModelButtonC.setEnabled(False)
+        for i in range(len(self.StyleButtons)):
+            self.StyleButtons[i].setEnabled(False)
+        self.newmodel.setEnabled(False)
+        self.loadMasks.setEnabled(False)
+        self.saveSet.setEnabled(False)
+        self.savePNG.setEnabled(False)
+        self.saveFlows.setEnabled(False)
+        self.saveOutlines.setEnabled(False)
+        self.saveROIs.setEnabled(False)
+
+        self.MakeDeletionRegionButton.setEnabled(False)
+        self.DeleteMultipleROIButton.setEnabled(False)
+        self.DoneDeleteMultipleROIButton.setEnabled(True)
+        self.CancelDeleteMultipleROIButton.setEnabled(True)
+
+    def toggle_mask_ops(self):
+        self.update_layer()
+        self.toggle_saving()
+        self.toggle_removals()
+
+    def toggle_saving(self):
+        if self.ncells > 0:
+            self.saveSet.setEnabled(True)
+            self.savePNG.setEnabled(True)
+            self.saveFlows.setEnabled(True)
+            self.saveOutlines.setEnabled(True)
+            self.saveROIs.setEnabled(True)
+        else:
+            self.saveSet.setEnabled(False)
+            self.savePNG.setEnabled(False)
+            self.saveFlows.setEnabled(False)
+            self.saveOutlines.setEnabled(False)
+            self.saveROIs.setEnabled(False)
+
+    def toggle_removals(self):
+        if self.ncells > 0:
+            self.ClearButton.setEnabled(True)
+            self.remcell.setEnabled(True)
+            self.undo.setEnabled(True)
+            self.MakeDeletionRegionButton.setEnabled(True)
+            self.DeleteMultipleROIButton.setEnabled(True)
+            self.DoneDeleteMultipleROIButton.setEnabled(False)
+            self.CancelDeleteMultipleROIButton.setEnabled(False)
+        else:
+            self.ClearButton.setEnabled(False)
+            self.remcell.setEnabled(False)
+            self.undo.setEnabled(False)
+            self.MakeDeletionRegionButton.setEnabled(False)
+            self.DeleteMultipleROIButton.setEnabled(False)
+            self.DoneDeleteMultipleROIButton.setEnabled(False)
+            self.CancelDeleteMultipleROIButton.setEnabled(False)
+
+    def remove_action(self):
+        if self.selected > 0:
+            self.remove_cell(self.selected)
+
+    def undo_action(self):
+        if (len(self.strokes) > 0 and self.strokes[-1][0][0] == self.currentZ):
+            self.remove_stroke()
+        else:
+            # remove previous cell
+            if self.ncells > 0:
+                self.remove_cell(self.ncells.get())
+
+    def undo_remove_action(self):
+        self.undo_remove_cell()
+
+    def get_files(self):
+        folder = os.path.dirname(self.filename)
+        mask_filter = "_masks"
+        images = get_image_files(folder, mask_filter)
+        fnames = [os.path.split(images[k])[-1] for k in range(len(images))]
+        f0 = os.path.split(self.filename)[-1]
+        idx = np.nonzero(np.array(fnames) == f0)[0][0]
+        return images, idx
+
+    def get_prev_image(self):
+        images, idx = self.get_files()
+        idx = (idx - 1) % len(images)
+        io._load_image(self, filename=images[idx])
+
+    def get_next_image(self, load_seg=True):
+        images, idx = self.get_files()
+        idx = (idx + 1) % len(images)
+        io._load_image(self, filename=images[idx], load_seg=load_seg)
+
+    def dragEnterEvent(self, event):
+        if event.mimeData().hasUrls():
+            event.accept()
+        else:
+            event.ignore()
+
+    def dropEvent(self, event):
+        files = [u.toLocalFile() for u in event.mimeData().urls()]
+        if os.path.splitext(files[0])[-1] == ".npy":
+            io._load_seg(self, filename=files[0], load_3D=self.load_3D)
+        else:
+            io._load_image(self, filename=files[0], load_seg=True, load_3D=self.load_3D)
+
+    def toggle_masks(self):
+        if self.MCheckBox.isChecked():
+            self.masksOn = True
+        else:
+            self.masksOn = False
+        if self.OCheckBox.isChecked():
+            self.outlinesOn = True
+        else:
+            self.outlinesOn = False
+        if not self.masksOn and not self.outlinesOn:
+            self.p0.removeItem(self.layer)
+            self.layer_off = True
+        else:
+            if self.layer_off:
+                self.p0.addItem(self.layer)
+            self.draw_layer()
+            self.update_layer()
+        if self.loaded:
+            self.update_plot()
+            self.update_layer()
+
+    def make_viewbox(self):
+        self.p0 = guiparts.ViewBoxNoRightDrag(parent=self, lockAspect=True,
+                                              name="plot1", border=[100, 100,
+                                                                    100], invertY=True)
+        self.p0.setCursor(QtCore.Qt.CrossCursor)
+        self.brush_size = 3
+        self.win.addItem(self.p0, 0, 0, rowspan=1, colspan=1)
+        self.p0.setMenuEnabled(False)
+        self.p0.setMouseEnabled(x=True, y=True)
+        self.img = pg.ImageItem(viewbox=self.p0, parent=self)
+        self.img.autoDownsample = False
+        self.layer = guiparts.ImageDraw(viewbox=self.p0, parent=self)
+        self.layer.setLevels([0, 255])
+        self.scale = pg.ImageItem(viewbox=self.p0, parent=self)
+        self.scale.setLevels([0, 255])
+        self.p0.scene().contextMenuItem = self.p0
+        self.Ly, self.Lx = 512, 512
+        self.p0.addItem(self.img)
+        self.p0.addItem(self.layer)
+        self.p0.addItem(self.scale)
+
+    def reset(self):
+        # ---- start sets of points ---- #
+        self.selected = 0
+        self.nchan = 3
+        self.loaded = False
+        self.channel = [0, 1]
+        self.current_point_set = []
+        self.in_stroke = False
+        self.strokes = []
+        self.stroke_appended = True
+        self.resize = False
+        self.ncells.reset()
+        self.zdraw = []
+        self.removed_cell = []
+        self.cellcolors = np.array([255, 255, 255])[np.newaxis, :]
+
+        # -- zero out image stack -- #
+        self.opacity = 128  # how opaque masks should be
+        self.outcolor = [200, 200, 255, 200]
+        self.NZ, self.Ly, self.Lx = 1, 256, 256
+        self.saturation = self.saturation if hasattr(self, 'saturation') else []
+
+        # only adjust the saturation if auto-adjust is on: 
+        if self.autobtn.isChecked():
+            for r in range(3):
+                self.saturation.append([[0, 255] for n in range(self.NZ)])
+                self.sliders[r].setValue([0, 255])
+                self.sliders[r].setEnabled(False)
+                self.sliders[r].show()
+        self.currentZ = 0
+        self.flows = [[], [], [], [], [[]]]
+        # masks matrix
+        # image matrix with a scale disk
+        self.stack = np.zeros((1, self.Ly, self.Lx, 3))
+        self.Lyr, self.Lxr = self.Ly, self.Lx
+        self.Ly0, self.Lx0 = self.Ly, self.Lx
+        self.radii = 0 * np.ones((self.Ly, self.Lx, 4), np.uint8)
+        self.layerz = 0 * np.ones((self.Ly, self.Lx, 4), np.uint8)
+        self.cellpix = np.zeros((1, self.Ly, self.Lx), np.uint16)
+        self.outpix = np.zeros((1, self.Ly, self.Lx), np.uint16)
+        self.ismanual = np.zeros(0, "bool")
+
+        # -- set menus to default -- #
+        self.color = 0
+        self.RGBDropDown.setCurrentIndex(self.color)
+        self.view = 0
+        self.ViewDropDown.setCurrentIndex(0)
+        self.ViewDropDown.model().item(self.ViewDropDown.count() - 1).setEnabled(False)
+        self.delete_restore()
+
+        self.clear_all()
+
+        self.filename = []
+        self.loaded = False
+        self.recompute_masks = False
+
+        self.deleting_multiple = False
+        self.removing_cells_list = []
+        self.removing_region = False
+        self.remove_roi_obj = None
+
+    def delete_restore(self):
+        """ delete restored imgs but don't reset settings """
+        if hasattr(self, "stack_filtered"):
+            del self.stack_filtered
+        if hasattr(self, "cellpix_orig"):
+            self.cellpix = self.cellpix_orig.copy()
+            self.outpix = self.outpix_orig.copy()
+            del self.outpix_orig, self.outpix_resize
+            del self.cellpix_orig, self.cellpix_resize
+
+    def clear_restore(self):
+        """ delete restored imgs and reset settings """
+        print("GUI_INFO: clearing restored image")
+        self.ViewDropDown.model().item(self.ViewDropDown.count() - 1).setEnabled(False)
+        if self.ViewDropDown.currentIndex() == self.ViewDropDown.count() - 1:
+            self.ViewDropDown.setCurrentIndex(0)
+        self.delete_restore()
+        self.restore = None
+        self.ratio = 1.
+        self.set_normalize_params(self.get_normalize_params())
+
+    def brush_choose(self):
+        self.brush_size = self.BrushChoose.currentIndex() * 2 + 1
+        if self.loaded:
+            self.layer.setDrawKernel(kernel_size=self.brush_size)
+            self.update_layer()
+
+    def clear_all(self):
+        self.prev_selected = 0
+        self.selected = 0
+        if self.restore and "upsample" in self.restore:
+            self.layerz = 0 * np.ones((self.Lyr, self.Lxr, 4), np.uint8)
+            self.cellpix = np.zeros((self.NZ, self.Lyr, self.Lxr), np.uint16)
+            self.outpix = np.zeros((self.NZ, self.Lyr, self.Lxr), np.uint16)
+            self.cellpix_resize = self.cellpix.copy()
+            self.outpix_resize = self.outpix.copy()
+            self.cellpix_orig = np.zeros((self.NZ, self.Ly0, self.Lx0), np.uint16)
+            self.outpix_orig = np.zeros((self.NZ, self.Ly0, self.Lx0), np.uint16)
+        else:
+            self.layerz = 0 * np.ones((self.Ly, self.Lx, 4), np.uint8)
+            self.cellpix = np.zeros((self.NZ, self.Ly, self.Lx), np.uint16)
+            self.outpix = np.zeros((self.NZ, self.Ly, self.Lx), np.uint16)
+
+        self.cellcolors = np.array([255, 255, 255])[np.newaxis, :]
+        self.ncells.reset()
+        self.toggle_removals()
+        self.update_scale()
+        self.update_layer()
+
+    def select_cell(self, idx):
+        self.prev_selected = self.selected
+        self.selected = idx
+        if self.selected > 0:
+            z = self.currentZ
+            self.layerz[self.cellpix[z] == idx] = np.array(
+                [255, 255, 255, self.opacity])
+            self.update_layer()
+
+    def select_cell_multi(self, idx):
+        if idx > 0:
+            z = self.currentZ
+            self.layerz[self.cellpix[z] == idx] = np.array(
+                [255, 255, 255, self.opacity])
+            self.update_layer()
+
+    def unselect_cell(self):
+        if self.selected > 0:
+            idx = self.selected
+            if idx < (self.ncells.get() + 1):
+                z = self.currentZ
+                self.layerz[self.cellpix[z] == idx] = np.append(
+                    self.cellcolors[idx], self.opacity)
+                if self.outlinesOn:
+                    self.layerz[self.outpix[z] == idx] = np.array(self.outcolor).astype(
+                        np.uint8)
+                    #[0,0,0,self.opacity])
+                self.update_layer()
+        self.selected = 0
+
+    def unselect_cell_multi(self, idx):
+        z = self.currentZ
+        self.layerz[self.cellpix[z] == idx] = np.append(self.cellcolors[idx],
+                                                        self.opacity)
+        if self.outlinesOn:
+            self.layerz[self.outpix[z] == idx] = np.array(self.outcolor).astype(
+                np.uint8)
+            # [0,0,0,self.opacity])
+        self.update_layer()
+
+    def remove_cell(self, idx):
+        if isinstance(idx, (int, np.integer)):
+            idx = [idx]
+        # because the function remove_single_cell updates the state of the cellpix and outpix arrays
+        # by reindexing cells to avoid gaps in the indices, we need to remove the cells in reverse order
+        # so that the indices are correct
+        idx.sort(reverse=True)
+        for i in idx:
+            self.remove_single_cell(i)
+        self.ncells -= len(idx)  # _save_sets uses ncells
+        self.update_layer()
+
+        if self.ncells == 0:
+            self.ClearButton.setEnabled(False)
+        if self.NZ == 1:
+            io._save_sets_with_check(self)
+
+
+    def remove_single_cell(self, idx):
+        # remove from manual array
+        self.selected = 0
+        if self.NZ > 1:
+            zextent = ((self.cellpix == idx).sum(axis=(1, 2)) > 0).nonzero()[0]
+        else:
+            zextent = [0]
+        for z in zextent:
+            cp = self.cellpix[z] == idx
+            op = self.outpix[z] == idx
+            # remove from self.cellpix and self.outpix
+            self.cellpix[z, cp] = 0
+            self.outpix[z, op] = 0
+            if z == self.currentZ:
+                # remove from mask layer
+                self.layerz[cp] = np.array([0, 0, 0, 0])
+
+        # reduce other pixels by -1
+        self.cellpix[self.cellpix > idx] -= 1
+        self.outpix[self.outpix > idx] -= 1
+
+        if self.NZ == 1:
+            self.removed_cell = [
+                self.ismanual[idx - 1], self.cellcolors[idx],
+                np.nonzero(cp),
+                np.nonzero(op)
+            ]
+            self.redo.setEnabled(True)
+            ar, ac = self.removed_cell[2]
+            d = datetime.datetime.now()
+            self.track_changes.append(
+                [d.strftime("%m/%d/%Y, %H:%M:%S"), "removed mask", [ar, ac]])
+        # remove cell from lists
+        self.ismanual = np.delete(self.ismanual, idx - 1)
+        self.cellcolors = np.delete(self.cellcolors, [idx], axis=0)
+        del self.zdraw[idx - 1]
+        print("GUI_INFO: removed cell %d" % (idx - 1))
+
+    def remove_region_cells(self):
+        if self.removing_cells_list:
+            for idx in self.removing_cells_list:
+                self.unselect_cell_multi(idx)
+            self.removing_cells_list.clear()
+        self.disable_buttons_removeROIs()
+        self.removing_region = True
+
+        self.clear_multi_selected_cells()
+
+        # make roi region here in center of view, making ROI half the size of the view
+        roi_width = self.p0.viewRect().width() / 2
+        x_loc = self.p0.viewRect().x() + (roi_width / 2)
+        roi_height = self.p0.viewRect().height() / 2
+        y_loc = self.p0.viewRect().y() + (roi_height / 2)
+
+        pos = [x_loc, y_loc]
+        roi = pg.RectROI(pos, [roi_width, roi_height], pen=pg.mkPen("y", width=2),
+                         removable=True)
+        roi.sigRemoveRequested.connect(self.remove_roi)
+        roi.sigRegionChangeFinished.connect(self.roi_changed)
+        self.p0.addItem(roi)
+        self.remove_roi_obj = roi
+        self.roi_changed(roi)
+
+    def delete_multiple_cells(self):
+        self.unselect_cell()
+        self.disable_buttons_removeROIs()
+        self.DoneDeleteMultipleROIButton.setEnabled(True)
+        self.MakeDeletionRegionButton.setEnabled(True)
+        self.CancelDeleteMultipleROIButton.setEnabled(True)
+        self.deleting_multiple = True
+
+    def done_remove_multiple_cells(self):
+        self.deleting_multiple = False
+        self.removing_region = False
+        self.DoneDeleteMultipleROIButton.setEnabled(False)
+        self.MakeDeletionRegionButton.setEnabled(False)
+        self.CancelDeleteMultipleROIButton.setEnabled(False)
+
+        if self.removing_cells_list:
+            self.removing_cells_list = list(set(self.removing_cells_list))
+            display_remove_list = [i - 1 for i in self.removing_cells_list]
+            print(f"GUI_INFO: removing cells: {display_remove_list}")
+            self.remove_cell(self.removing_cells_list)
+            self.removing_cells_list.clear()
+            self.unselect_cell()
+        self.enable_buttons()
+
+        if self.remove_roi_obj is not None:
+            self.remove_roi(self.remove_roi_obj)
+
+    def merge_cells(self, idx):
+        self.prev_selected = self.selected
+        self.selected = idx
+        if self.selected != self.prev_selected:
+            for z in range(self.NZ):
+                ar0, ac0 = np.nonzero(self.cellpix[z] == self.prev_selected)
+                ar1, ac1 = np.nonzero(self.cellpix[z] == self.selected)
+                touching = np.logical_and((ar0[:, np.newaxis] - ar1) < 3,
+                                          (ac0[:, np.newaxis] - ac1) < 3).sum()
+                ar = np.hstack((ar0, ar1))
+                ac = np.hstack((ac0, ac1))
+                vr0, vc0 = np.nonzero(self.outpix[z] == self.prev_selected)
+                vr1, vc1 = np.nonzero(self.outpix[z] == self.selected)
+                self.outpix[z, vr0, vc0] = 0
+                self.outpix[z, vr1, vc1] = 0
+                if touching > 0:
+                    mask = np.zeros((np.ptp(ar) + 4, np.ptp(ac) + 4), np.uint8)
+                    mask[ar - ar.min() + 2, ac - ac.min() + 2] = 1
+                    contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                                cv2.CHAIN_APPROX_NONE)
+                    pvc, pvr = contours[-2][0].squeeze().T
+                    vr, vc = pvr + ar.min() - 2, pvc + ac.min() - 2
+
+                else:
+                    vr = np.hstack((vr0, vr1))
+                    vc = np.hstack((vc0, vc1))
+                color = self.cellcolors[self.prev_selected]
+                self.draw_mask(z, ar, ac, vr, vc, color, idx=self.prev_selected)
+            self.remove_cell(self.selected)
+            print("GUI_INFO: merged two cells")
+            self.update_layer()
+            io._save_sets_with_check(self)
+            self.undo.setEnabled(False)
+            self.redo.setEnabled(False)
+
+    def undo_remove_cell(self):
+        if len(self.removed_cell) > 0:
+            z = 0
+            ar, ac = self.removed_cell[2]
+            vr, vc = self.removed_cell[3]
+            color = self.removed_cell[1]
+            self.draw_mask(z, ar, ac, vr, vc, color)
+            self.toggle_mask_ops()
+            self.cellcolors = np.append(self.cellcolors, color[np.newaxis, :], axis=0)
+            self.ncells += 1
+            self.ismanual = np.append(self.ismanual, self.removed_cell[0])
+            self.zdraw.append([])
+            print(">>> added back removed cell")
+            self.update_layer()
+            io._save_sets_with_check(self)
+            self.removed_cell = []
+            self.redo.setEnabled(False)
+
+    def remove_stroke(self, delete_points=True, stroke_ind=-1):
+        stroke = np.array(self.strokes[stroke_ind])
+        cZ = self.currentZ
+        inZ = stroke[0, 0] == cZ
+        if inZ:
+            outpix = self.outpix[cZ, stroke[:, 1], stroke[:, 2]] > 0
+            self.layerz[stroke[~outpix, 1], stroke[~outpix, 2]] = np.array([0, 0, 0, 0])
+            cellpix = self.cellpix[cZ, stroke[:, 1], stroke[:, 2]]
+            ccol = self.cellcolors.copy()
+            if self.selected > 0:
+                ccol[self.selected] = np.array([255, 255, 255])
+            col2mask = ccol[cellpix]
+            if self.masksOn:
+                col2mask = np.concatenate(
+                    (col2mask, self.opacity * (cellpix[:, np.newaxis] > 0)), axis=-1)
+            else:
+                col2mask = np.concatenate((col2mask, 0 * (cellpix[:, np.newaxis] > 0)),
+                                          axis=-1)
+            self.layerz[stroke[:, 1], stroke[:, 2], :] = col2mask
+            if self.outlinesOn:
+                self.layerz[stroke[outpix, 1], stroke[outpix,
+                                                      2]] = np.array(self.outcolor)
+            if delete_points:
+                del self.current_point_set[stroke_ind]
+            self.update_layer()
+
+        del self.strokes[stroke_ind]
+
+    def plot_clicked(self, event):
+        if event.button()==QtCore.Qt.LeftButton \
+                and not event.modifiers() & (QtCore.Qt.ShiftModifier | QtCore.Qt.AltModifier)\
+                and not self.removing_region:
+            if event.double():
+                try:
+                    self.p0.setYRange(0, self.Ly + self.pr)
+                except:
+                    self.p0.setYRange(0, self.Ly)
+                self.p0.setXRange(0, self.Lx)
+
+    def cancel_remove_multiple(self):
+        self.clear_multi_selected_cells()
+        self.done_remove_multiple_cells()
+
+    def clear_multi_selected_cells(self):
+        # unselect all previously selected cells:
+        for idx in self.removing_cells_list:
+            self.unselect_cell_multi(idx)
+        self.removing_cells_list.clear()
+
+    def add_roi(self, roi):
+        self.p0.addItem(roi)
+        self.remove_roi_obj = roi
+
+    def remove_roi(self, roi):
+        self.clear_multi_selected_cells()
+        assert roi == self.remove_roi_obj
+        self.remove_roi_obj = None
+        self.p0.removeItem(roi)
+        self.removing_region = False
+
+    def roi_changed(self, roi):
+        # find the overlapping cells and make them selected
+        pos = roi.pos()
+        size = roi.size()
+        x0 = int(pos.x())
+        y0 = int(pos.y())
+        x1 = int(pos.x() + size.x())
+        y1 = int(pos.y() + size.y())
+        if x0 < 0:
+            x0 = 0
+        if y0 < 0:
+            y0 = 0
+        if x1 > self.Lx:
+            x1 = self.Lx
+        if y1 > self.Ly:
+            y1 = self.Ly
+
+        # find cells in that region
+        cell_idxs = np.unique(self.cellpix[self.currentZ, y0:y1, x0:x1])
+        cell_idxs = np.trim_zeros(cell_idxs)
+        # deselect cells not in region by deselecting all and then selecting the ones in the region
+        self.clear_multi_selected_cells()
+
+        for idx in cell_idxs:
+            self.select_cell_multi(idx)
+            self.removing_cells_list.append(idx)
+
+        self.update_layer()
+
+    def mouse_moved(self, pos):
+        items = self.win.scene().items(pos)
+
+    def color_choose(self):
+        self.color = self.RGBDropDown.currentIndex()
+        self.view = 0
+        self.ViewDropDown.setCurrentIndex(self.view)
+        self.update_plot()
+
+    def update_plot(self):
+        self.view = self.ViewDropDown.currentIndex()
+        self.Ly, self.Lx, _ = self.stack[self.currentZ].shape
+
+        if self.view == 0 or self.view == self.ViewDropDown.count() - 1:
+            image = self.stack[
+                self.currentZ] if self.view == 0 else self.stack_filtered[self.currentZ]
+            if self.color == 0:
+                self.img.setImage(image, autoLevels=False, lut=None)
+                if self.nchan > 1:
+                    levels = np.array([
+                        self.saturation[0][self.currentZ],
+                        self.saturation[1][self.currentZ],
+                        self.saturation[2][self.currentZ]
+                    ])
+                    self.img.setLevels(levels)
+                else:
+                    self.img.setLevels(self.saturation[0][self.currentZ])
+            elif self.color > 0 and self.color < 4:
+                if self.nchan > 1:
+                    image = image[:, :, self.color - 1]
+                self.img.setImage(image, autoLevels=False, lut=self.cmap[self.color])
+                if self.nchan > 1:
+                    self.img.setLevels(self.saturation[self.color - 1][self.currentZ])
+                else:
+                    self.img.setLevels(self.saturation[0][self.currentZ])
+            elif self.color == 4:
+                if self.nchan > 1:
+                    image = image.mean(axis=-1)
+                self.img.setImage(image, autoLevels=False, lut=None)
+                self.img.setLevels(self.saturation[0][self.currentZ])
+            elif self.color == 5:
+                if self.nchan > 1:
+                    image = image.mean(axis=-1)
+                self.img.setImage(image, autoLevels=False, lut=self.cmap[0])
+                self.img.setLevels(self.saturation[0][self.currentZ])
+        else:
+            image = np.zeros((self.Ly, self.Lx), np.uint8)
+            if len(self.flows) >= self.view - 1 and len(self.flows[self.view - 1]) > 0:
+                image = self.flows[self.view - 1][self.currentZ]
+            if self.view > 1:
+                self.img.setImage(image, autoLevels=False, lut=self.bwr)
+            else:
+                self.img.setImage(image, autoLevels=False, lut=None)
+            self.img.setLevels([0.0, 255.0])
+
+        for r in range(3):
+            self.sliders[r].setValue([
+                self.saturation[r][self.currentZ][0],
+                self.saturation[r][self.currentZ][1]
+            ])
+        self.win.show()
+        self.show()
+
+
+    def update_layer(self):
+        if self.masksOn or self.outlinesOn:
+            self.layer.setImage(self.layerz, autoLevels=False)
+        self.win.show()
+        self.show()
+
+
+    def add_set(self):
+        if len(self.current_point_set) > 0:
+            while len(self.strokes) > 0:
+                self.remove_stroke(delete_points=False)
+            if len(self.current_point_set[0]) > 8:
+                color = self.colormap[self.ncells.get(), :3]
+                median = self.add_mask(points=self.current_point_set, color=color)
+                if median is not None:
+                    self.removed_cell = []
+                    self.toggle_mask_ops()
+                    self.cellcolors = np.append(self.cellcolors, color[np.newaxis, :],
+                                                axis=0)
+                    self.ncells += 1
+                    self.ismanual = np.append(self.ismanual, True)
+                    if self.NZ == 1:
+                        # only save after each cell if single image
+                        io._save_sets_with_check(self)
+            else:
+                print("GUI_ERROR: cell too small, not drawn")
+            self.current_stroke = []
+            self.strokes = []
+            self.current_point_set = []
+            self.update_layer()
+
+    def add_mask(self, points=None, color=(100, 200, 50), dense=True):
+        # points is list of strokes
+        points_all = np.concatenate(points, axis=0)
+        
+        # loop over z values
+        median = []
+        zdraw = np.unique(points_all[:, 0])
+        z = 0
+        ars, acs, vrs, vcs = np.zeros(0, "int"), np.zeros(0, "int"), np.zeros(
+            0, "int"), np.zeros(0, "int")
+        for stroke in points:
+            stroke = np.concatenate(stroke, axis=0).reshape(-1, 4)
+            vr = stroke[:, 1]
+            vc = stroke[:, 2]
+            # get points inside drawn points
+            mask = np.zeros((np.ptp(vr) + 4, np.ptp(vc) + 4), np.uint8)
+            pts = np.stack((vc - vc.min() + 2, vr - vr.min() + 2),
+                           axis=-1)[:, np.newaxis, :]
+            mask = cv2.fillPoly(mask, [pts], (255, 0, 0))
+            ar, ac = np.nonzero(mask)
+            ar, ac = ar + vr.min() - 2, ac + vc.min() - 2
+            # get dense outline
+            contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+            pvc, pvr = contours[-2][0][:,0].T
+            vr, vc = pvr + vr.min() - 2, pvc + vc.min() - 2
+            # concatenate all points
+            ar, ac = np.hstack((np.vstack((vr, vc)), np.vstack((ar, ac))))
+            # if these pixels are overlapping with another cell, reassign them
+            ioverlap = self.cellpix[z][ar, ac] > 0
+            if (~ioverlap).sum() < 10:
+                print("GUI_ERROR: cell < 10 pixels without overlaps, not drawn")
+                return None
+            elif ioverlap.sum() > 0:
+                ar, ac = ar[~ioverlap], ac[~ioverlap]
+                # compute outline of new mask
+                mask = np.zeros((np.ptp(vr) + 4, np.ptp(vc) + 4), np.uint8)
+                mask[ar - vr.min() + 2, ac - vc.min() + 2] = 1
+                contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                            cv2.CHAIN_APPROX_NONE)
+                pvc, pvr = contours[-2][0][:,0].T
+                vr, vc = pvr + vr.min() - 2, pvc + vc.min() - 2
+            ars = np.concatenate((ars, ar), axis=0)
+            acs = np.concatenate((acs, ac), axis=0)
+            vrs = np.concatenate((vrs, vr), axis=0)
+            vcs = np.concatenate((vcs, vc), axis=0)
+            
+        self.draw_mask(z, ars, acs, vrs, vcs, color)
+        median.append(np.array([np.median(ars), np.median(acs)]))
+
+        self.zdraw.append(zdraw)
+        d = datetime.datetime.now()
+        self.track_changes.append(
+            [d.strftime("%m/%d/%Y, %H:%M:%S"), "added mask", [ar, ac]])
+        return median
+
+    def draw_mask(self, z, ar, ac, vr, vc, color, idx=None):
+        """ draw single mask using outlines and area """
+        if idx is None:
+            idx = self.ncells + 1
+        self.cellpix[z, vr, vc] = idx
+        self.cellpix[z, ar, ac] = idx
+        self.outpix[z, vr, vc] = idx
+        if self.restore and "upsample" in self.restore:
+            if self.resize:
+                self.cellpix_resize[z, vr, vc] = idx
+                self.cellpix_resize[z, ar, ac] = idx
+                self.outpix_resize[z, vr, vc] = idx
+                self.cellpix_orig[z, (vr / self.ratio).astype(int),
+                                  (vc / self.ratio).astype(int)] = idx
+                self.cellpix_orig[z, (ar / self.ratio).astype(int),
+                                  (ac / self.ratio).astype(int)] = idx
+                self.outpix_orig[z, (vr / self.ratio).astype(int),
+                                 (vc / self.ratio).astype(int)] = idx
+            else:
+                self.cellpix_orig[z, vr, vc] = idx
+                self.cellpix_orig[z, ar, ac] = idx
+                self.outpix_orig[z, vr, vc] = idx
+
+                # get upsampled mask
+                vrr = (vr.copy() * self.ratio).astype(int)
+                vcr = (vc.copy() * self.ratio).astype(int)
+                mask = np.zeros((np.ptp(vrr) + 4, np.ptp(vcr) + 4), np.uint8)
+                pts = np.stack((vcr - vcr.min() + 2, vrr - vrr.min() + 2),
+                               axis=-1)[:, np.newaxis, :]
+                mask = cv2.fillPoly(mask, [pts], (255, 0, 0))
+                arr, acr = np.nonzero(mask)
+                arr, acr = arr + vrr.min() - 2, acr + vcr.min() - 2
+                # get dense outline
+                contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                            cv2.CHAIN_APPROX_NONE)
+                pvc, pvr = contours[-2][0].squeeze().T
+                vrr, vcr = pvr + vrr.min() - 2, pvc + vcr.min() - 2
+                # concatenate all points
+                arr, acr = np.hstack((np.vstack((vrr, vcr)), np.vstack((arr, acr))))
+                self.cellpix_resize[z, vrr, vcr] = idx
+                self.cellpix_resize[z, arr, acr] = idx
+                self.outpix_resize[z, vrr, vcr] = idx
+
+        if z == self.currentZ:
+            self.layerz[ar, ac, :3] = color
+            if self.masksOn:
+                self.layerz[ar, ac, -1] = self.opacity
+            if self.outlinesOn:
+                self.layerz[vr, vc] = np.array(self.outcolor)
+
+    def compute_scale(self):
+        # get diameter from gui
+        diameter = self.segmentation_settings.diameter
+        if not diameter:
+            diameter = 30
+
+        self.pr = int(diameter)
+        self.radii_padding = int(self.pr * 1.25)
+        self.radii = np.zeros((self.Ly + self.radii_padding, self.Lx, 4), np.uint8)
+        yy, xx = disk([self.Ly + self.radii_padding / 2 - 1, self.pr / 2 + 1],
+                      self.pr / 2, self.Ly + self.radii_padding, self.Lx)
+        # rgb(150,50,150)
+        self.radii[yy, xx, 0] = 150
+        self.radii[yy, xx, 1] = 50
+        self.radii[yy, xx, 2] = 150
+        self.radii[yy, xx, 3] = 255
+        self.p0.setYRange(0, self.Ly + self.radii_padding)
+        self.p0.setXRange(0, self.Lx)
+
+    def update_scale(self):
+        self.compute_scale()
+        self.scale.setImage(self.radii, autoLevels=False)
+        self.scale.setLevels([0.0, 255.0])
+        self.win.show()
+        self.show()
+
+
+    def draw_layer(self):
+        if self.resize:
+            self.Ly, self.Lx = self.Lyr, self.Lxr
+        else:
+            self.Ly, self.Lx = self.Ly0, self.Lx0
+
+        if self.masksOn or self.outlinesOn:
+            if self.restore and "upsample" in self.restore:
+                if self.resize:
+                    self.cellpix = self.cellpix_resize.copy()
+                    self.outpix = self.outpix_resize.copy()
+                else:
+                    self.cellpix = self.cellpix_orig.copy()
+                    self.outpix = self.outpix_orig.copy()
+
+        self.layerz = np.zeros((self.Ly, self.Lx, 4), np.uint8)
+        if self.masksOn:
+            self.layerz[..., :3] = self.cellcolors[self.cellpix[self.currentZ], :]
+            self.layerz[..., 3] = self.opacity * (self.cellpix[self.currentZ]
+                                                  > 0).astype(np.uint8)
+            if self.selected > 0:
+                self.layerz[self.cellpix[self.currentZ] == self.selected] = np.array(
+                    [255, 255, 255, self.opacity])
+            cZ = self.currentZ
+            stroke_z = np.array([s[0][0] for s in self.strokes])
+            inZ = np.nonzero(stroke_z == cZ)[0]
+            if len(inZ) > 0:
+                for i in inZ:
+                    stroke = np.array(self.strokes[i])
+                    self.layerz[stroke[:, 1], stroke[:,
+                                                     2]] = np.array([255, 0, 255, 100])
+        else:
+            self.layerz[..., 3] = 0
+
+        if self.outlinesOn:
+            self.layerz[self.outpix[self.currentZ] > 0] = np.array(
+                self.outcolor).astype(np.uint8)
+
+
+    def set_normalize_params(self, normalize_params):
+        from cellpose.models import normalize_default
+        if self.restore != "filter":
+            keys = list(normalize_params.keys()).copy()
+            for key in keys:
+                if key != "percentile":
+                    normalize_params[key] = normalize_default[key]
+        normalize_params = {**normalize_default, **normalize_params}
+        out = self.check_filter_params(normalize_params["sharpen_radius"],
+                                       normalize_params["smooth_radius"],
+                                       normalize_params["tile_norm_blocksize"],
+                                       normalize_params["tile_norm_smooth3D"],
+                                       normalize_params["norm3D"],
+                                       normalize_params["invert"])
+
+
+    def check_filter_params(self, sharpen, smooth, tile_norm, smooth3D, norm3D, invert):
+        tile_norm = 0 if tile_norm < 0 else tile_norm
+        sharpen = 0 if sharpen < 0 else sharpen
+        smooth = 0 if smooth < 0 else smooth
+        smooth3D = 0 if smooth3D < 0 else smooth3D
+        norm3D = bool(norm3D)
+        invert = bool(invert)
+        if tile_norm > self.Ly and tile_norm > self.Lx:
+            print(
+                "GUI_ERROR: tile size (tile_norm) bigger than both image dimensions, disabling"
+            )
+            tile_norm = 0
+        self.filt_edits[0].setText(str(sharpen))
+        self.filt_edits[1].setText(str(smooth))
+        self.filt_edits[2].setText(str(tile_norm))
+        self.filt_edits[3].setText(str(smooth3D))
+        self.norm3D_cb.setChecked(norm3D)
+        return sharpen, smooth, tile_norm, smooth3D, norm3D, invert
+
+    def get_normalize_params(self):
+        percentile = [
+            self.segmentation_settings.low_percentile,
+            self.segmentation_settings.high_percentile,
+        ]
+        normalize_params = {"percentile": percentile}
+        norm3D = self.norm3D_cb.isChecked()
+        normalize_params["norm3D"] = norm3D
+        sharpen = float(self.filt_edits[0].text())
+        smooth = float(self.filt_edits[1].text())
+        tile_norm = float(self.filt_edits[2].text())
+        smooth3D = float(self.filt_edits[3].text())
+        invert = False
+        out = self.check_filter_params(sharpen, smooth, tile_norm, smooth3D, norm3D,
+                                        invert)
+        sharpen, smooth, tile_norm, smooth3D, norm3D, invert = out
+        normalize_params["sharpen_radius"] = sharpen
+        normalize_params["smooth_radius"] = smooth
+        normalize_params["tile_norm_blocksize"] = tile_norm
+        normalize_params["tile_norm_smooth3D"] = smooth3D
+        normalize_params["invert"] = invert
+
+        from cellpose.models import normalize_default
+        normalize_params = {**normalize_default, **normalize_params}
+
+        return normalize_params
+    
+    def compute_saturation_if_checked(self):
+        if self.autobtn.isChecked():
+            self.compute_saturation()
+
+    def compute_saturation(self, return_img=False):
+        norm = self.get_normalize_params()
+        print(norm)
+        sharpen, smooth = norm["sharpen_radius"], norm["smooth_radius"]
+        percentile = norm["percentile"]
+        tile_norm = norm["tile_norm_blocksize"]
+        invert = norm["invert"]
+        norm3D = norm["norm3D"]
+        smooth3D = norm["tile_norm_smooth3D"]
+        tile_norm = norm["tile_norm_blocksize"]
+
+        if sharpen > 0 or smooth > 0 or tile_norm > 0:
+            img_norm = self.stack.copy()
+        else:
+            img_norm = self.stack
+
+        if sharpen > 0 or smooth > 0 or tile_norm > 0:
+            self.restore = "filter"
+            print(
+                "GUI_INFO: computing filtered image because sharpen > 0 or tile_norm > 0"
+            )
+            print(
+                "GUI_WARNING: will use memory to create filtered image -- make sure to have RAM for this"
+            )
+            img_norm = self.stack.copy()
+            if sharpen > 0 or smooth > 0:
+                img_norm = smooth_sharpen_img(self.stack, sharpen_radius=sharpen,
+                                              smooth_radius=smooth)
+
+            if tile_norm > 0:
+                img_norm = normalize99_tile(img_norm, blocksize=tile_norm,
+                                            lower=percentile[0], upper=percentile[1],
+                                            smooth3D=smooth3D, norm3D=norm3D)
+            # convert to 0->255
+            img_norm_min = img_norm.min()
+            img_norm_max = img_norm.max()
+            for c in range(img_norm.shape[-1]):
+                if np.ptp(img_norm[..., c]) > 1e-3:
+                    img_norm[..., c] -= img_norm_min
+                    img_norm[..., c] /= (img_norm_max - img_norm_min)
+            img_norm *= 255
+            self.stack_filtered = img_norm
+            self.ViewDropDown.model().item(self.ViewDropDown.count() -
+                                           1).setEnabled(True)
+            self.ViewDropDown.setCurrentIndex(self.ViewDropDown.count() - 1)
+        else:
+            img_norm = self.stack if self.restore is None or self.restore == "filter" else self.stack_filtered
+
+        if self.autobtn.isChecked():
+            self.saturation = []
+            for c in range(img_norm.shape[-1]):
+                self.saturation.append([])
+                if np.ptp(img_norm[..., c]) > 1e-3:
+                    if norm3D:
+                        x01 = np.percentile(img_norm[..., c], percentile[0])
+                        x99 = np.percentile(img_norm[..., c], percentile[1])
+                        if invert:
+                            x01i = 255. - x99
+                            x99i = 255. - x01
+                            x01, x99 = x01i, x99i
+                        for n in range(self.NZ):
+                            self.saturation[-1].append([x01, x99])
+                    else:
+                        for z in range(self.NZ):
+                            if self.NZ > 1:
+                                x01 = np.percentile(img_norm[z, :, :, c], percentile[0])
+                                x99 = np.percentile(img_norm[z, :, :, c], percentile[1])
+                            else:
+                                x01 = np.percentile(img_norm[..., c], percentile[0])
+                                x99 = np.percentile(img_norm[..., c], percentile[1])
+                            if invert:
+                                x01i = 255. - x99
+                                x99i = 255. - x01
+                                x01, x99 = x01i, x99i
+                            self.saturation[-1].append([x01, x99])
+                else:
+                    for n in range(self.NZ):
+                        self.saturation[-1].append([0, 255.])
+            print(self.saturation[2][self.currentZ])
+
+            if img_norm.shape[-1] == 1:
+                self.saturation.append(self.saturation[0])
+                self.saturation.append(self.saturation[0])
+
+        # self.autobtn.setChecked(True)
+        self.update_plot()
+
+
+    def get_model_path(self, custom=False):
+        if custom:
+            self.current_model = self.ModelChooseC.currentText()
+            self.current_model_path = os.fspath(
+                models.MODEL_DIR.joinpath(self.current_model))
+        else:
+            self.current_model = "cpsam"
+            self.current_model_path = models.model_path(self.current_model)
+
+    def initialize_model(self, model_name=None, custom=False):
+        if model_name is None or custom:
+            self.get_model_path(custom=custom)
+            if not os.path.exists(self.current_model_path):
+                raise ValueError("need to specify model (use dropdown)")
+
+        if model_name is None or not isinstance(model_name, str):
+            self.model = models.CellposeModel(gpu=self.useGPU.isChecked(),
+                                              pretrained_model=self.current_model_path)
+        else:
+            self.current_model = model_name
+            self.current_model_path = os.fspath(
+                models.MODEL_DIR.joinpath(self.current_model))
+
+            self.model = models.CellposeModel(gpu=self.useGPU.isChecked(),
+                                             pretrained_model=self.current_model)
+
+    def add_model(self):
+        io._add_model(self)
+        return
+
+    def remove_model(self):
+        io._remove_model(self)
+        return
+
+    def new_model(self):
+        if self.NZ != 1:
+            print("ERROR: cannot train model on 3D data")
+            return
+
+        # train model
+        image_names = self.get_files()[0]
+        self.train_data, self.train_labels, self.train_files, restore, normalize_params = io._get_train_set(
+            image_names)
+        TW = guiparts.TrainWindow(self, models.MODEL_NAMES)
+        train = TW.exec_()
+        if train:
+            self.logger.info(
+                f"training with {[os.path.split(f)[1] for f in self.train_files]}")
+            self.train_model(restore=restore, normalize_params=normalize_params)
+        else:
+            print("GUI_INFO: training cancelled")
+
+    def train_model(self, restore=None, normalize_params=None):
+        from cellpose.models import normalize_default
+        if normalize_params is None:
+            normalize_params = copy.deepcopy(normalize_default)
+        model_type = models.MODEL_NAMES[self.training_params["model_index"]]
+        self.logger.info(f"training new model starting at model {model_type}")
+        self.current_model = model_type
+        
+        self.model = models.CellposeModel(gpu=self.useGPU.isChecked(),
+                                          model_type=model_type)
+        save_path = os.path.dirname(self.filename)
+
+        print("GUI_INFO: name of new model: " + self.training_params["model_name"])
+        self.new_model_path, train_losses = train.train_seg(
+            self.model.net, train_data=self.train_data, train_labels=self.train_labels,
+            normalize=normalize_params, min_train_masks=0,
+            save_path=save_path, nimg_per_epoch=max(2, len(self.train_data)),
+            learning_rate=self.training_params["learning_rate"],
+            weight_decay=self.training_params["weight_decay"],
+            n_epochs=self.training_params["n_epochs"],
+            model_name=self.training_params["model_name"])[:2]
+        # save train losses
+        np.save(str(self.new_model_path) + "_train_losses.npy", train_losses)
+        # run model on next image
+        io._add_model(self, self.new_model_path)
+        diam_labels = self.model.net.diam_labels.item()  #.copy()
+        self.new_model_ind = len(self.model_strings)
+        self.autorun = True
+        self.clear_all()
+        self.restore = restore
+        self.set_normalize_params(normalize_params)
+        self.get_next_image(load_seg=False)
+
+        self.compute_segmentation(custom=True)
+        self.logger.info(
+            f"!!! computed masks for {os.path.split(self.filename)[1]} from new model !!!"
+        )
+
+
+    def compute_cprob(self):
+        if self.recompute_masks:
+            flow_threshold = self.segmentation_settings.flow_threshold
+            cellprob_threshold = self.segmentation_settings.cellprob_threshold
+            niter = self.segmentation_settings.niter
+            min_size = int(self.min_size.text()) if not isinstance(
+                self.min_size, int) else self.min_size
+
+            self.logger.info(
+                    "computing masks with cell prob=%0.3f, flow error threshold=%0.3f" %
+                    (cellprob_threshold, flow_threshold))
+            
+            try:
+                dP = self.flows[2].squeeze()
+                cellprob = self.flows[3].squeeze()
+            except IndexError:
+                self.logger.error("Flows don't exist, try running model again.")
+                return
+            
+            maski = dynamics.resize_and_compute_masks(
+                dP=dP,
+                cellprob=cellprob,
+                niter=niter,
+                do_3D=self.load_3D,
+                min_size=min_size,
+                # max_size_fraction=min_size_fraction, # Leave as default 
+                cellprob_threshold=cellprob_threshold, 
+                flow_threshold=flow_threshold)
+            
+            self.masksOn = True
+            if not self.OCheckBox.isChecked():
+                self.MCheckBox.setChecked(True)
+            if maski.ndim < 3:
+                maski = maski[np.newaxis, ...]
+            self.logger.info("%d cells found" % (len(np.unique(maski)[1:])))
+            io._masks_to_gui(self, maski, outlines=None)
+            self.show()
+
+
+    def compute_segmentation(self, custom=False, model_name=None, load_model=True):
+        self.progress.setValue(0)
+        try:
+            tic = time.time()
+            self.clear_all()
+            self.flows = [[], [], []]
+            if load_model:
+                self.initialize_model(model_name=model_name, custom=custom)
+            self.progress.setValue(10)
+            do_3D = self.load_3D
+            stitch_threshold = float(self.stitch_threshold.text()) if not isinstance(
+                self.stitch_threshold, float) else self.stitch_threshold
+            anisotropy = float(self.anisotropy.text()) if not isinstance(
+                self.anisotropy, float) else self.anisotropy
+            flow3D_smooth = float(self.flow3D_smooth.text()) if not isinstance(
+                self.flow3D_smooth, float) else self.flow3D_smooth
+            min_size = int(self.min_size.text()) if not isinstance(
+                self.min_size, int) else self.min_size
+            
+            do_3D = False if stitch_threshold > 0. else do_3D
+
+            if self.restore == "filter":
+                data = self.stack_filtered.copy().squeeze()
+            else:
+                data = self.stack.copy().squeeze()
+            
+            flow_threshold = self.segmentation_settings.flow_threshold
+            cellprob_threshold = self.segmentation_settings.cellprob_threshold
+            diameter = self.segmentation_settings.diameter
+            niter = self.segmentation_settings.niter
+            
+            normalize_params = self.get_normalize_params()
+            print(normalize_params)
+            try:
+                masks, flows = self.model.eval(
+                    data, 
+                    diameter=diameter,
+                    cellprob_threshold=cellprob_threshold,
+                    flow_threshold=flow_threshold, do_3D=do_3D, niter=niter,
+                    normalize=normalize_params, stitch_threshold=stitch_threshold,
+                    anisotropy=anisotropy, flow3D_smooth=flow3D_smooth,
+                    min_size=min_size, channel_axis=-1,
+                    progress=self.progress, z_axis=0 if self.NZ > 1 else None)[:2]
+            except Exception as e:
+                print("NET ERROR: %s" % e)
+                self.progress.setValue(0)
+                return
+
+            self.progress.setValue(75)
+
+            # convert flows to uint8 and resize to original image size
+            flows_new = []
+            flows_new.append(flows[0].copy())  # RGB flow
+            flows_new.append((np.clip(normalize99(flows[2].copy()), 0, 1) *
+                              255).astype("uint8"))  # cellprob
+            flows_new.append(flows[1].copy()) # XY flows
+            flows_new.append(flows[2].copy()) # original cellprob
+
+            if self.load_3D:
+                if stitch_threshold == 0.:
+                    flows_new.append((flows[1][0] / 10 * 127 + 127).astype("uint8"))
+                else:
+                    flows_new.append(np.zeros(flows[1][0].shape, dtype="uint8"))
+
+            if not self.load_3D:
+                if self.restore and "upsample" in self.restore:
+                    self.Ly, self.Lx = self.Lyr, self.Lxr
+
+                if flows_new[0].shape[-3:-1] != (self.Ly, self.Lx):
+                    self.flows = []
+                    for j in range(len(flows_new)):
+                        self.flows.append(
+                            resize_image(flows_new[j], Ly=self.Ly, Lx=self.Lx,
+                                        interpolation=cv2.INTER_NEAREST))
+                else:
+                    self.flows = flows_new
+            else:
+                self.flows = []
+                Lz, Ly, Lx = self.NZ, self.Ly, self.Lx
+                Lz0, Ly0, Lx0 = flows_new[0].shape[:3]
+                print("GUI_INFO: resizing flows to original image size")
+                for j in range(len(flows_new)):
+                    flow0 = flows_new[j]
+                    if Ly0 != Ly:
+                        flow0 = resize_image(flow0, Ly=Ly, Lx=Lx,
+                                            no_channels=flow0.ndim==3, 
+                                            interpolation=cv2.INTER_NEAREST)
+                    if Lz0 != Lz:
+                        flow0 = np.swapaxes(resize_image(np.swapaxes(flow0, 0, 1),
+                                            Ly=Lz, Lx=Lx,
+                                            no_channels=flow0.ndim==3, 
+                                            interpolation=cv2.INTER_NEAREST), 0, 1)
+                    self.flows.append(flow0)
+
+            # add first axis
+            if self.NZ == 1:
+                masks = masks[np.newaxis, ...]
+                self.flows = [
+                    self.flows[n][np.newaxis, ...] for n in range(len(self.flows))
+                ]
+
+            self.logger.info("%d cells found with model in %0.3f sec" %
+                             (len(np.unique(masks)[1:]), time.time() - tic))
+            self.progress.setValue(80)
+            z = 0
+
+            io._masks_to_gui(self, masks, outlines=None)
+            self.masksOn = True
+            self.MCheckBox.setChecked(True)
+            self.progress.setValue(100)
+            if self.restore != "filter" and self.restore is not None and self.autobtn.isChecked():
+                self.compute_saturation()
+            if not do_3D and not stitch_threshold > 0:
+                self.recompute_masks = True
+            else:
+                self.recompute_masks = False
+        except Exception as e:
+            print("ERROR: %s" % e)

models/seg_post_model/cellpose/gui/gui3d.py

@@ -0,0 +1,667 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer, Michael Rariden and Marius Pachitariu.
+"""
+
+import sys, pathlib, warnings
+
+from qtpy import QtGui, QtCore
+from qtpy.QtWidgets import QApplication, QScrollBar, QCheckBox, QLabel, QLineEdit
+import pyqtgraph as pg
+
+import numpy as np
+from scipy.stats import mode
+import cv2
+
+from . import guiparts, io
+from ..utils import download_url_to_file, masks_to_outlines
+from .gui import MainW
+
+try:
+    import matplotlib.pyplot as plt
+    MATPLOTLIB = True
+except:
+    MATPLOTLIB = False
+
+
+def avg3d(C):
+    """ smooth value of c across nearby points
+        (c is center of grid directly below point)
+        b -- a -- b
+        a -- c -- a
+        b -- a -- b
+    """
+    Ly, Lx = C.shape
+    # pad T by 2
+    T = np.zeros((Ly + 2, Lx + 2), "float32")
+    M = np.zeros((Ly, Lx), "float32")
+    T[1:-1, 1:-1] = C.copy()
+    y, x = np.meshgrid(np.arange(0, Ly, 1, int), np.arange(0, Lx, 1, int),
+                       indexing="ij")
+    y += 1
+    x += 1
+    a = 1. / 2  #/(z**2 + 1)**0.5
+    b = 1. / (1 + 2**0.5)  #(z**2 + 2)**0.5
+    c = 1.
+    M = (b * T[y - 1, x - 1] + a * T[y - 1, x] + b * T[y - 1, x + 1] + a * T[y, x - 1] +
+         c * T[y, x] + a * T[y, x + 1] + b * T[y + 1, x - 1] + a * T[y + 1, x] +
+         b * T[y + 1, x + 1])
+    M /= 4 * a + 4 * b + c
+    return M
+
+
+def interpZ(mask, zdraw):
+    """ find nearby planes and average their values using grid of points
+        zfill is in ascending order
+    """
+    ifill = np.ones(mask.shape[0], "bool")
+    zall = np.arange(0, mask.shape[0], 1, int)
+    ifill[zdraw] = False
+    zfill = zall[ifill]
+    zlower = zdraw[np.searchsorted(zdraw, zfill, side="left") - 1]
+    zupper = zdraw[np.searchsorted(zdraw, zfill, side="right")]
+    for k, z in enumerate(zfill):
+        Z = zupper[k] - zlower[k]
+        zl = (z - zlower[k]) / Z
+        plower = avg3d(mask[zlower[k]]) * (1 - zl)
+        pupper = avg3d(mask[zupper[k]]) * zl
+        mask[z] = (plower + pupper) > 0.33
+    return mask, zfill
+
+
+def run(image=None):
+    from ..io import logger_setup
+    logger, log_file = logger_setup()
+    # Always start by initializing Qt (only once per application)
+    warnings.filterwarnings("ignore")
+    app = QApplication(sys.argv)
+    icon_path = pathlib.Path.home().joinpath(".cellpose", "logo.png")
+    guip_path = pathlib.Path.home().joinpath(".cellpose", "cellpose_gui.png")
+    style_path = pathlib.Path.home().joinpath(".cellpose", "style_choice.npy")
+    if not icon_path.is_file():
+        cp_dir = pathlib.Path.home().joinpath(".cellpose")
+        cp_dir.mkdir(exist_ok=True)
+        print("downloading logo")
+        download_url_to_file(
+            "https://www.cellpose.org/static/images/cellpose_transparent.png",
+            icon_path, progress=True)
+    if not guip_path.is_file():
+        print("downloading help window image")
+        download_url_to_file("https://www.cellpose.org/static/images/cellpose_gui.png",
+                             guip_path, progress=True)
+    icon_path = str(icon_path.resolve())
+    app_icon = QtGui.QIcon()
+    app_icon.addFile(icon_path, QtCore.QSize(16, 16))
+    app_icon.addFile(icon_path, QtCore.QSize(24, 24))
+    app_icon.addFile(icon_path, QtCore.QSize(32, 32))
+    app_icon.addFile(icon_path, QtCore.QSize(48, 48))
+    app_icon.addFile(icon_path, QtCore.QSize(64, 64))
+    app_icon.addFile(icon_path, QtCore.QSize(256, 256))
+    app.setWindowIcon(app_icon)
+    app.setStyle("Fusion")
+    app.setPalette(guiparts.DarkPalette())
+    MainW_3d(image=image, logger=logger)
+    ret = app.exec_()
+    sys.exit(ret)
+
+
+class MainW_3d(MainW):
+
+    def __init__(self, image=None, logger=None):
+        # MainW init
+        MainW.__init__(self, image=image, logger=logger)
+
+        # add gradZ view
+        self.ViewDropDown.insertItem(3, "gradZ")
+
+        # turn off single stroke
+        self.SCheckBox.setChecked(False)
+
+        ### add orthoviews and z-bar
+        # ortho crosshair lines
+        self.vLine = pg.InfiniteLine(angle=90, movable=False)
+        self.hLine = pg.InfiniteLine(angle=0, movable=False)
+        self.vLineOrtho = [
+            pg.InfiniteLine(angle=90, movable=False),
+            pg.InfiniteLine(angle=90, movable=False)
+        ]
+        self.hLineOrtho = [
+            pg.InfiniteLine(angle=0, movable=False),
+            pg.InfiniteLine(angle=0, movable=False)
+        ]
+        self.make_orthoviews()
+
+        # z scrollbar underneath
+        self.scroll = QScrollBar(QtCore.Qt.Horizontal)
+        self.scroll.setMaximum(10)
+        self.scroll.valueChanged.connect(self.move_in_Z)
+        self.lmain.addWidget(self.scroll, 40, 9, 1, 30)
+
+        b = 22
+
+        label = QLabel("stitch\nthreshold:")
+        label.setToolTip(
+            "for 3D volumes, turn on stitch_threshold to stitch masks across planes instead of running cellpose in 3D (see docs for details)"
+        )
+        label.setFont(self.medfont)
+        self.segBoxG.addWidget(label, b, 0, 1, 4)
+        self.stitch_threshold = QLineEdit()
+        self.stitch_threshold.setText("0.0")
+        self.stitch_threshold.setFixedWidth(30)
+        self.stitch_threshold.setFont(self.medfont)
+        self.stitch_threshold.setToolTip(
+            "for 3D volumes, turn on stitch_threshold to stitch masks across planes instead of running cellpose in 3D (see docs for details)"
+        )
+        self.segBoxG.addWidget(self.stitch_threshold, b, 3, 1, 1)
+
+        label = QLabel("flow3D\nsmooth:")
+        label.setToolTip(
+            "for 3D volumes, smooth flows by a Gaussian with standard deviation flow3D_smooth (see docs for details)"
+        )
+        label.setFont(self.medfont)
+        self.segBoxG.addWidget(label, b, 4, 1, 3)
+        self.flow3D_smooth = QLineEdit()
+        self.flow3D_smooth.setText("0.0")
+        self.flow3D_smooth.setFixedWidth(30)
+        self.flow3D_smooth.setFont(self.medfont)
+        self.flow3D_smooth.setToolTip(
+            "for 3D volumes, smooth flows by a Gaussian with standard deviation flow3D_smooth (see docs for details)"
+        )
+        self.segBoxG.addWidget(self.flow3D_smooth, b, 7, 1, 1)
+
+        b+=1
+        label = QLabel("anisotropy:")
+        label.setToolTip(
+            "for 3D volumes, increase in sampling in Z vs XY as a ratio, e.g. set set to 2.0 if Z is sampled half as dense as X or Y (see docs for details)"
+        )
+        label.setFont(self.medfont)
+        self.segBoxG.addWidget(label, b, 0, 1, 3)
+        self.anisotropy = QLineEdit()
+        self.anisotropy.setText("1.0")
+        self.anisotropy.setFixedWidth(30)
+        self.anisotropy.setFont(self.medfont)
+        self.anisotropy.setToolTip(
+            "for 3D volumes, increase in sampling in Z vs XY as a ratio, e.g. set set to 2.0 if Z is sampled half as dense as X or Y (see docs for details)"
+        )
+        self.segBoxG.addWidget(self.anisotropy, b, 3, 1, 1)
+
+        b+=1
+        label = QLabel("min\nsize:")
+        label.setToolTip(
+            "all masks less than this size in pixels (volume) will be removed"
+        )
+        label.setFont(self.medfont)
+        self.segBoxG.addWidget(label, b, 0, 1, 4)
+        self.min_size = QLineEdit()
+        self.min_size.setText("15")
+        self.min_size.setFixedWidth(50)
+        self.min_size.setFont(self.medfont)
+        self.min_size.setToolTip(
+            "all masks less than this size in pixels (volume) will be removed"
+        )
+        self.segBoxG.addWidget(self.min_size, b, 3, 1, 1)
+
+        b += 1
+        self.orthobtn = QCheckBox("ortho")
+        self.orthobtn.setToolTip("activate orthoviews with 3D image")
+        self.orthobtn.setFont(self.medfont)
+        self.orthobtn.setChecked(False)
+        self.l0.addWidget(self.orthobtn, b, 0, 1, 2)
+        self.orthobtn.toggled.connect(self.toggle_ortho)
+
+        label = QLabel("dz:")
+        label.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        label.setFont(self.medfont)
+        self.l0.addWidget(label, b, 2, 1, 1)
+        self.dz = 10
+        self.dzedit = QLineEdit()
+        self.dzedit.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        self.dzedit.setText(str(self.dz))
+        self.dzedit.returnPressed.connect(self.update_ortho)
+        self.dzedit.setFixedWidth(40)
+        self.dzedit.setFont(self.medfont)
+        self.l0.addWidget(self.dzedit, b, 3, 1, 2)
+
+        label = QLabel("z-aspect:")
+        label.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        label.setFont(self.medfont)
+        self.l0.addWidget(label, b, 5, 1, 2)
+        self.zaspect = 1.0
+        self.zaspectedit = QLineEdit()
+        self.zaspectedit.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        self.zaspectedit.setText(str(self.zaspect))
+        self.zaspectedit.returnPressed.connect(self.update_ortho)
+        self.zaspectedit.setFixedWidth(40)
+        self.zaspectedit.setFont(self.medfont)
+        self.l0.addWidget(self.zaspectedit, b, 7, 1, 2)
+
+        b += 1
+        # add z position underneath
+        self.currentZ = 0
+        label = QLabel("Z:")
+        label.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        self.l0.addWidget(label, b, 5, 1, 2)
+        self.zpos = QLineEdit()
+        self.zpos.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        self.zpos.setText(str(self.currentZ))
+        self.zpos.returnPressed.connect(self.update_ztext)
+        self.zpos.setFixedWidth(40)
+        self.zpos.setFont(self.medfont)
+        self.l0.addWidget(self.zpos, b, 7, 1, 2)
+
+        # if called with image, load it
+        if image is not None:
+            self.filename = image
+            io._load_image(self, self.filename, load_3D=True)
+
+        self.load_3D = True
+
+    def add_mask(self, points=None, color=(100, 200, 50), dense=True):
+        # points is list of strokes
+
+        points_all = np.concatenate(points, axis=0)
+
+        # loop over z values
+        median = []
+        zdraw = np.unique(points_all[:, 0])
+        zrange = np.arange(zdraw.min(), zdraw.max() + 1, 1, int)
+        zmin = zdraw.min()
+        pix = np.zeros((2, 0), "uint16")
+        mall = np.zeros((len(zrange), self.Ly, self.Lx), "bool")
+        k = 0
+        for z in zdraw:
+            ars, acs, vrs, vcs = np.zeros(0, "int"), np.zeros(0, "int"), np.zeros(
+                0, "int"), np.zeros(0, "int")
+            for stroke in points:
+                stroke = np.concatenate(stroke, axis=0).reshape(-1, 4)
+                iz = stroke[:, 0] == z
+                vr = stroke[iz, 1]
+                vc = stroke[iz, 2]
+                if iz.sum() > 0:
+                    # get points inside drawn points
+                    mask = np.zeros((np.ptp(vr) + 4, np.ptp(vc) + 4), "uint8")
+                    pts = np.stack((vc - vc.min() + 2, vr - vr.min() + 2),
+                                   axis=-1)[:, np.newaxis, :]
+                    mask = cv2.fillPoly(mask, [pts], (255, 0, 0))
+                    ar, ac = np.nonzero(mask)
+                    ar, ac = ar + vr.min() - 2, ac + vc.min() - 2
+                    # get dense outline
+                    contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                                cv2.CHAIN_APPROX_NONE)
+                    pvc, pvr = contours[-2][0].squeeze().T
+                    vr, vc = pvr + vr.min() - 2, pvc + vc.min() - 2
+                    # concatenate all points
+                    ar, ac = np.hstack((np.vstack((vr, vc)), np.vstack((ar, ac))))
+                    # if these pixels are overlapping with another cell, reassign them
+                    ioverlap = self.cellpix[z][ar, ac] > 0
+                    if (~ioverlap).sum() < 8:
+                        print("ERROR: cell too small without overlaps, not drawn")
+                        return None
+                    elif ioverlap.sum() > 0:
+                        ar, ac = ar[~ioverlap], ac[~ioverlap]
+                        # compute outline of new mask
+                        mask = np.zeros((np.ptp(ar) + 4, np.ptp(ac) + 4), "uint8")
+                        mask[ar - ar.min() + 2, ac - ac.min() + 2] = 1
+                        contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                                    cv2.CHAIN_APPROX_NONE)
+                        pvc, pvr = contours[-2][0].squeeze().T
+                        vr, vc = pvr + ar.min() - 2, pvc + ac.min() - 2
+                    ars = np.concatenate((ars, ar), axis=0)
+                    acs = np.concatenate((acs, ac), axis=0)
+                    vrs = np.concatenate((vrs, vr), axis=0)
+                    vcs = np.concatenate((vcs, vc), axis=0)
+            self.draw_mask(z, ars, acs, vrs, vcs, color)
+
+            median.append(np.array([np.median(ars), np.median(acs)]))
+            mall[z - zmin, ars, acs] = True
+            pix = np.append(pix, np.vstack((ars, acs)), axis=-1)
+
+        mall = mall[:, pix[0].min():pix[0].max() + 1,
+                    pix[1].min():pix[1].max() + 1].astype("float32")
+        ymin, xmin = pix[0].min(), pix[1].min()
+        if len(zdraw) > 1:
+            mall, zfill = interpZ(mall, zdraw - zmin)
+            for z in zfill:
+                mask = mall[z].copy()
+                ar, ac = np.nonzero(mask)
+                ioverlap = self.cellpix[z + zmin][ar + ymin, ac + xmin] > 0
+                if (~ioverlap).sum() < 5:
+                    print("WARNING: stroke on plane %d not included due to overlaps" %
+                          z)
+                elif ioverlap.sum() > 0:
+                    mask[ar[ioverlap], ac[ioverlap]] = 0
+                    ar, ac = ar[~ioverlap], ac[~ioverlap]
+                # compute outline of mask
+                outlines = masks_to_outlines(mask)
+                vr, vc = np.nonzero(outlines)
+                vr, vc = vr + ymin, vc + xmin
+                ar, ac = ar + ymin, ac + xmin
+                self.draw_mask(z + zmin, ar, ac, vr, vc, color)
+
+        self.zdraw.append(zdraw)
+
+        return median
+
+    def move_in_Z(self):
+        if self.loaded:
+            self.currentZ = min(self.NZ, max(0, int(self.scroll.value())))
+            self.zpos.setText(str(self.currentZ))
+            self.update_plot()
+            self.draw_layer()
+            self.update_layer()
+
+    def make_orthoviews(self):
+        self.pOrtho, self.imgOrtho, self.layerOrtho = [], [], []
+        for j in range(2):
+            self.pOrtho.append(
+                pg.ViewBox(lockAspect=True, name=f"plotOrtho{j}",
+                           border=[100, 100, 100], invertY=True, enableMouse=False))
+            self.pOrtho[j].setMenuEnabled(False)
+
+            self.imgOrtho.append(pg.ImageItem(viewbox=self.pOrtho[j], parent=self))
+            self.imgOrtho[j].autoDownsample = False
+
+            self.layerOrtho.append(pg.ImageItem(viewbox=self.pOrtho[j], parent=self))
+            self.layerOrtho[j].setLevels([0., 255.])
+
+            #self.pOrtho[j].scene().contextMenuItem = self.pOrtho[j]
+            self.pOrtho[j].addItem(self.imgOrtho[j])
+            self.pOrtho[j].addItem(self.layerOrtho[j])
+            self.pOrtho[j].addItem(self.vLineOrtho[j], ignoreBounds=False)
+            self.pOrtho[j].addItem(self.hLineOrtho[j], ignoreBounds=False)
+
+        self.pOrtho[0].linkView(self.pOrtho[0].YAxis, self.p0)
+        self.pOrtho[1].linkView(self.pOrtho[1].XAxis, self.p0)
+
+    def add_orthoviews(self):
+        self.yortho = self.Ly // 2
+        self.xortho = self.Lx // 2
+        if self.NZ > 1:
+            self.update_ortho()
+
+        self.win.addItem(self.pOrtho[0], 0, 1, rowspan=1, colspan=1)
+        self.win.addItem(self.pOrtho[1], 1, 0, rowspan=1, colspan=1)
+
+        qGraphicsGridLayout = self.win.ci.layout
+        qGraphicsGridLayout.setColumnStretchFactor(0, 2)
+        qGraphicsGridLayout.setColumnStretchFactor(1, 1)
+        qGraphicsGridLayout.setRowStretchFactor(0, 2)
+        qGraphicsGridLayout.setRowStretchFactor(1, 1)
+
+        self.pOrtho[0].setYRange(0, self.Lx)
+        self.pOrtho[0].setXRange(-self.dz / 3, self.dz * 2 + self.dz / 3)
+        self.pOrtho[1].setYRange(-self.dz / 3, self.dz * 2 + self.dz / 3)
+        self.pOrtho[1].setXRange(0, self.Ly)
+
+        self.p0.addItem(self.vLine, ignoreBounds=False)
+        self.p0.addItem(self.hLine, ignoreBounds=False)
+        self.p0.setYRange(0, self.Lx)
+        self.p0.setXRange(0, self.Ly)
+
+        self.win.show()
+        self.show()
+
+    def remove_orthoviews(self):
+        self.win.removeItem(self.pOrtho[0])
+        self.win.removeItem(self.pOrtho[1])
+        self.p0.removeItem(self.vLine)
+        self.p0.removeItem(self.hLine)
+        self.win.show()
+        self.show()
+
+    def update_crosshairs(self):
+        self.yortho = min(self.Ly - 1, max(0, int(self.yortho)))
+        self.xortho = min(self.Lx - 1, max(0, int(self.xortho)))
+        self.vLine.setPos(self.xortho)
+        self.hLine.setPos(self.yortho)
+        self.vLineOrtho[1].setPos(self.xortho)
+        self.hLineOrtho[1].setPos(self.zc)
+        self.vLineOrtho[0].setPos(self.zc)
+        self.hLineOrtho[0].setPos(self.yortho)
+
+    def update_ortho(self):
+        if self.NZ > 1 and self.orthobtn.isChecked():
+            dzcurrent = self.dz
+            self.dz = min(100, max(3, int(self.dzedit.text())))
+            self.zaspect = max(0.01, min(100., float(self.zaspectedit.text())))
+            self.dzedit.setText(str(self.dz))
+            self.zaspectedit.setText(str(self.zaspect))
+            if self.dz != dzcurrent:
+                self.pOrtho[0].setXRange(-self.dz / 3, self.dz * 2 + self.dz / 3)
+                self.pOrtho[1].setYRange(-self.dz / 3, self.dz * 2 + self.dz / 3)
+            dztot = min(self.NZ, self.dz * 2)
+            y = self.yortho
+            x = self.xortho
+            z = self.currentZ
+            if dztot == self.NZ:
+                zmin, zmax = 0, self.NZ
+            else:
+                if z - self.dz < 0:
+                    zmin = 0
+                    zmax = zmin + self.dz * 2
+                elif z + self.dz >= self.NZ:
+                    zmax = self.NZ
+                    zmin = zmax - self.dz * 2
+                else:
+                    zmin, zmax = z - self.dz, z + self.dz
+            self.zc = z - zmin
+            self.update_crosshairs()
+            if self.view == 0 or self.view == 4:
+                for j in range(2):
+                    if j == 0:
+                        if self.view == 0:
+                            image = self.stack[zmin:zmax, :, x].transpose(1, 0, 2).copy()
+                        else:
+                            image = self.stack_filtered[zmin:zmax, :,
+                                                        x].transpose(1, 0, 2).copy()
+                    else:
+                        image = self.stack[
+                            zmin:zmax,
+                            y, :].copy() if self.view == 0 else self.stack_filtered[zmin:zmax,
+                                                                             y, :].copy()
+                    if self.nchan == 1:
+                        # show single channel
+                        image = image[..., 0]
+                    if self.color == 0:
+                        self.imgOrtho[j].setImage(image, autoLevels=False, lut=None)
+                        if self.nchan > 1:
+                            levels = np.array([
+                                self.saturation[0][self.currentZ],
+                                self.saturation[1][self.currentZ],
+                                self.saturation[2][self.currentZ]
+                            ])
+                            self.imgOrtho[j].setLevels(levels)
+                        else:
+                            self.imgOrtho[j].setLevels(
+                                self.saturation[0][self.currentZ])
+                    elif self.color > 0 and self.color < 4:
+                        if self.nchan > 1:
+                            image = image[..., self.color - 1]
+                        self.imgOrtho[j].setImage(image, autoLevels=False,
+                                                  lut=self.cmap[self.color])
+                        if self.nchan > 1:
+                            self.imgOrtho[j].setLevels(
+                                self.saturation[self.color - 1][self.currentZ])
+                        else:
+                            self.imgOrtho[j].setLevels(
+                                self.saturation[0][self.currentZ])
+                    elif self.color == 4:
+                        if image.ndim > 2:
+                            image = image.astype("float32").mean(axis=2).astype("uint8")
+                        self.imgOrtho[j].setImage(image, autoLevels=False, lut=None)
+                        self.imgOrtho[j].setLevels(self.saturation[0][self.currentZ])
+                    elif self.color == 5:
+                        if image.ndim > 2:
+                            image = image.astype("float32").mean(axis=2).astype("uint8")
+                        self.imgOrtho[j].setImage(image, autoLevels=False,
+                                                  lut=self.cmap[0])
+                        self.imgOrtho[j].setLevels(self.saturation[0][self.currentZ])
+                self.pOrtho[0].setAspectLocked(lock=True, ratio=self.zaspect)
+                self.pOrtho[1].setAspectLocked(lock=True, ratio=1. / self.zaspect)
+
+            else:
+                image = np.zeros((10, 10), "uint8")
+                self.imgOrtho[0].setImage(image, autoLevels=False, lut=None)
+                self.imgOrtho[0].setLevels([0.0, 255.0])
+                self.imgOrtho[1].setImage(image, autoLevels=False, lut=None)
+                self.imgOrtho[1].setLevels([0.0, 255.0])
+
+        zrange = zmax - zmin
+        self.layer_ortho = [
+            np.zeros((self.Ly, zrange, 4), "uint8"),
+            np.zeros((zrange, self.Lx, 4), "uint8")
+        ]
+        if self.masksOn:
+            for j in range(2):
+                if j == 0:
+                    cp = self.cellpix[zmin:zmax, :, x].T
+                else:
+                    cp = self.cellpix[zmin:zmax, y]
+                self.layer_ortho[j][..., :3] = self.cellcolors[cp, :]
+                self.layer_ortho[j][..., 3] = self.opacity * (cp > 0).astype("uint8")
+                if self.selected > 0:
+                    self.layer_ortho[j][cp == self.selected] = np.array(
+                        [255, 255, 255, self.opacity])
+
+        if self.outlinesOn:
+            for j in range(2):
+                if j == 0:
+                    op = self.outpix[zmin:zmax, :, x].T
+                else:
+                    op = self.outpix[zmin:zmax, y]
+                self.layer_ortho[j][op > 0] = np.array(self.outcolor).astype("uint8")
+
+        for j in range(2):
+            self.layerOrtho[j].setImage(self.layer_ortho[j])
+        self.win.show()
+        self.show()
+
+    def toggle_ortho(self):
+        if self.orthobtn.isChecked():
+            self.add_orthoviews()
+        else:
+            self.remove_orthoviews()
+
+    def plot_clicked(self, event):
+        if event.button()==QtCore.Qt.LeftButton \
+                and not event.modifiers() & (QtCore.Qt.ShiftModifier | QtCore.Qt.AltModifier)\
+                and not self.removing_region:
+            if event.double():
+                try:
+                    self.p0.setYRange(0, self.Ly + self.pr)
+                except:
+                    self.p0.setYRange(0, self.Ly)
+                self.p0.setXRange(0, self.Lx)
+            elif self.loaded and not self.in_stroke:
+                if self.orthobtn.isChecked():
+                    items = self.win.scene().items(event.scenePos())
+                    for x in items:
+                        if x == self.p0:
+                            pos = self.p0.mapSceneToView(event.scenePos())
+                            x = int(pos.x())
+                            y = int(pos.y())
+                            if y >= 0 and y < self.Ly and x >= 0 and x < self.Lx:
+                                self.yortho = y
+                                self.xortho = x
+                                self.update_ortho()
+
+    def update_plot(self):
+        super().update_plot()
+        if self.NZ > 1 and self.orthobtn.isChecked():
+            self.update_ortho()
+        self.win.show()
+        self.show()
+
+    def keyPressEvent(self, event):
+        if self.loaded:
+            if not (event.modifiers() &
+                    (QtCore.Qt.ControlModifier | QtCore.Qt.ShiftModifier |
+                     QtCore.Qt.AltModifier) or self.in_stroke):
+                updated = False
+                if len(self.current_point_set) > 0:
+                    if event.key() == QtCore.Qt.Key_Return:
+                        self.add_set()
+                    if self.NZ > 1:
+                        if event.key() == QtCore.Qt.Key_Left:
+                            self.currentZ = max(0, self.currentZ - 1)
+                            self.scroll.setValue(self.currentZ)
+                            updated = True
+                        elif event.key() == QtCore.Qt.Key_Right:
+                            self.currentZ = min(self.NZ - 1, self.currentZ + 1)
+                            self.scroll.setValue(self.currentZ)
+                            updated = True
+                else:
+                    nviews = self.ViewDropDown.count() - 1
+                    nviews += int(
+                        self.ViewDropDown.model().item(self.ViewDropDown.count() -
+                                                       1).isEnabled())
+                    if event.key() == QtCore.Qt.Key_X:
+                        self.MCheckBox.toggle()
+                    if event.key() == QtCore.Qt.Key_Z:
+                        self.OCheckBox.toggle()
+                    if event.key() == QtCore.Qt.Key_Left or event.key(
+                    ) == QtCore.Qt.Key_A:
+                        self.currentZ = max(0, self.currentZ - 1)
+                        self.scroll.setValue(self.currentZ)
+                        updated = True
+                    elif event.key() == QtCore.Qt.Key_Right or event.key(
+                    ) == QtCore.Qt.Key_D:
+                        self.currentZ = min(self.NZ - 1, self.currentZ + 1)
+                        self.scroll.setValue(self.currentZ)
+                        updated = True
+                    elif event.key() == QtCore.Qt.Key_PageDown:
+                        self.view = (self.view + 1) % (nviews)
+                        self.ViewDropDown.setCurrentIndex(self.view)
+                    elif event.key() == QtCore.Qt.Key_PageUp:
+                        self.view = (self.view - 1) % (nviews)
+                        self.ViewDropDown.setCurrentIndex(self.view)
+
+                # can change background or stroke size if cell not finished
+                if event.key() == QtCore.Qt.Key_Up or event.key() == QtCore.Qt.Key_W:
+                    self.color = (self.color - 1) % (6)
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_Down or event.key(
+                ) == QtCore.Qt.Key_S:
+                    self.color = (self.color + 1) % (6)
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_R:
+                    if self.color != 1:
+                        self.color = 1
+                    else:
+                        self.color = 0
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_G:
+                    if self.color != 2:
+                        self.color = 2
+                    else:
+                        self.color = 0
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif event.key() == QtCore.Qt.Key_B:
+                    if self.color != 3:
+                        self.color = 3
+                    else:
+                        self.color = 0
+                    self.RGBDropDown.setCurrentIndex(self.color)
+                elif (event.key() == QtCore.Qt.Key_Comma or
+                      event.key() == QtCore.Qt.Key_Period):
+                    count = self.BrushChoose.count()
+                    gci = self.BrushChoose.currentIndex()
+                    if event.key() == QtCore.Qt.Key_Comma:
+                        gci = max(0, gci - 1)
+                    else:
+                        gci = min(count - 1, gci + 1)
+                    self.BrushChoose.setCurrentIndex(gci)
+                    self.brush_choose()
+                if not updated:
+                    self.update_plot()
+        if event.key() == QtCore.Qt.Key_Minus or event.key() == QtCore.Qt.Key_Equal:
+            self.p0.keyPressEvent(event)
+
+    def update_ztext(self):
+        zpos = self.currentZ
+        try:
+            zpos = int(self.zpos.text())
+        except:
+            print("ERROR: zposition is not a number")
+        self.currentZ = max(0, min(self.NZ - 1, zpos))
+        self.zpos.setText(str(self.currentZ))
+        self.scroll.setValue(self.currentZ)

models/seg_post_model/cellpose/gui/guihelpwindowtext.html

@@ -0,0 +1,143 @@
+<qt>
+    <p class="has-line-data" data-line-start="5" data-line-end="6">
+        <b>Main GUI mouse controls:</b>
+    </p>
+    <ul>
+        <li class="has-line-data" data-line-start="7" data-line-end="8">Pan = left-click + drag</li>
+        <li class="has-line-data" data-line-start="8" data-line-end="9">Zoom = scroll wheel (or +/= and - buttons)</li>
+        <li class="has-line-data" data-line-start="9" data-line-end="10">Full view = double left-click</li>
+        <li class="has-line-data" data-line-start="10" data-line-end="11">Select mask = left-click on mask</li>
+        <li class="has-line-data" data-line-start="11" data-line-end="12">Delete mask = Ctrl (or COMMAND on Mac) +
+            left-click
+        </li>
+        <li class="has-line-data" data-line-start="11" data-line-end="12">Merge masks = Alt + left-click (will merge
+            last two)
+        </li>
+        <li class="has-line-data" data-line-start="12" data-line-end="13">Start draw mask = right-click</li>
+        <li class="has-line-data" data-line-start="13" data-line-end="15">End draw mask = right-click, or return to
+            circle at beginning
+        </li>
+    </ul>
+    <p class="has-line-data" data-line-start="15" data-line-end="16">Overlaps in masks are NOT allowed. If you
+        draw a mask on top of another mask, it is cropped so that it doesn’t overlap with the old mask. Masks in 2D
+        should be single strokes (single stroke is checked). If you want to draw masks in 3D (experimental), then
+        you can turn this option off and draw a stroke on each plane with the cell and then press ENTER. 3D
+        labelling will fill in planes that you have not labelled so that you do not have to as densely label.
+    </p>
+    <p class="has-line-data" data-line-start="17" data-line-end="18"> <b>!NOTE!:</b> The GUI automatically saves after
+        you draw a mask in 2D but NOT after 3D mask drawing and NOT after segmentation. Save in the file menu or
+        with Ctrl+S. The output file is in the same folder as the loaded image with <code>_seg.npy</code> appended.
+    </p>
+
+    <p class="has-line-data" data-line-start="19" data-line-end="20"> <b>Bulk Mask Deletion</b>
+        Clicking the 'delete multiple' button will allow you to select and delete multiple masks at once.
+        Masks can be deselected by clicking on them again. Once you have selected all the masks you want to delete,
+        click the 'done' button to delete them.
+        <br>
+        <br>
+        Alternatively, you can create a rectangular region to delete a regions of masks by clicking the
+        'delete multiple' button, and then moving and/or resizing the region to select the masks you want to delete.
+        Once you have selected the masks you want to delete, click the 'done' button to delete them.
+        <br>
+        <br>
+        At any point in the process, you can click the 'cancel' button to cancel the bulk deletion.
+    </p>
+    <hr>
+    <table class="table table-striped table-bordered">
+        <br>
+        <br>
+        FYI there are tooltips throughout the GUI (hover over text to see)
+        <br>
+        <thead>
+        <tr>
+            <th>Keyboard shortcuts</th>
+            <th>Description</th>
+        </tr>
+        </thead>
+        <tbody>
+        <tr>
+            <td>=/+ button // - button</td>
+            <td>zoom in // zoom out</td>
+        </tr>
+        <tr>
+            <td>CTRL+Z</td>
+            <td>undo previously drawn mask/stroke</td>
+        </tr>
+        <tr>
+            <td>CTRL+Y</td>
+            <td>undo remove mask</td>
+        </tr>
+        <tr>
+            <td>CTRL+0</td>
+            <td>clear all masks</td>
+        </tr>
+        <tr>
+            <td>CTRL+L</td>
+            <td>load image (can alternatively drag and drop image)</td>
+        </tr>
+        <tr>
+            <td>CTRL+S</td>
+            <td>SAVE MASKS IN IMAGE to <code>_seg.npy</code> file</td>
+        </tr>
+        <tr>
+            <td>CTRL+T</td>
+            <td>train model using _seg.npy files in folder
+        </tr>
+        <tr>
+            <td>CTRL+P</td>
+            <td>load <code>_seg.npy</code> file (note: it will load automatically with image if it exists)</td>
+        </tr>
+        <tr>
+            <td>CTRL+M</td>
+            <td>load masks file (must be same size as image with 0 for NO mask, and 1,2,3… for masks)</td>
+        </tr>
+        <tr>
+            <td>CTRL+N</td>
+            <td>save masks as PNG</td>
+        </tr>
+        <tr>
+            <td>CTRL+R</td>
+            <td>save ROIs to native ImageJ ROI format</td>
+        </tr>
+        <tr>
+            <td>CTRL+F</td>
+            <td>save flows to image file</td>
+        </tr>
+        <tr>
+            <td>A/D or LEFT/RIGHT</td>
+            <td>cycle through images in current directory</td>
+        </tr>
+        <tr>
+            <td>W/S or UP/DOWN</td>
+            <td>change color (RGB/gray/red/green/blue)</td>
+        </tr>
+        <tr>
+            <td>R / G / B</td>
+            <td>toggle between RGB and Red or Green or Blue</td>
+        </tr>
+        <tr>
+            <td>PAGE-UP / PAGE-DOWN</td>
+            <td>change to flows and cell prob views (if segmentation computed)</td>
+        </tr>
+        <tr>
+            <td>X</td>
+            <td>turn masks ON or OFF</td>
+        </tr>
+        <tr>
+            <td>Z</td>
+            <td>toggle outlines ON or OFF</td>
+        </tr>
+        <tr>
+            <td>, / .</td>
+            <td>increase / decrease brush size for drawing masks</td>
+        </tr>
+        </tbody>
+    </table>
+    <p class="has-line-data" data-line-start="36" data-line-end="37"><strong>Segmentation options 
+        (2D only) </strong></p>
+    <p class="has-line-data" data-line-start="38" data-line-end="39">use GPU: if you have specially 
+        installed the cuda version of torch, then you can activate this. Due to the size of the 
+        transformer network, it will greatly speed up the processing time.</p>
+    <p class="has-line-data" data-line-start="40" data-line-end="41">There are no channel options
+        in v4.0.1+ since all 3 channels are used for segmentation. </p>
+</qt>

models/seg_post_model/cellpose/gui/guiparts.py

@@ -0,0 +1,793 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+from qtpy import QtGui, QtCore
+from qtpy.QtGui import QPixmap, QDoubleValidator
+from qtpy.QtWidgets import QWidget, QDialog, QGridLayout, QPushButton, QLabel, QLineEdit, QDialogButtonBox, QComboBox, QCheckBox, QVBoxLayout
+import pyqtgraph as pg
+import numpy as np
+import pathlib, os
+
+
+def stylesheet():
+    return """
+        QToolTip { 
+                            background-color: black; 
+                            color: white; 
+                            border: black solid 1px
+                            }
+        QComboBox {color: white;
+                    background-color: rgb(40,40,40);}
+                    QComboBox::item:enabled { color: white;
+                    background-color: rgb(40,40,40);
+                    selection-color: white;
+                    selection-background-color: rgb(50,100,50);}
+                    QComboBox::item:!enabled {
+                            background-color: rgb(40,40,40);
+                            color: rgb(100,100,100);
+                        }
+        QScrollArea > QWidget > QWidget
+                {
+                    background: transparent;
+                    border: none;
+                    margin: 0px 0px 0px 0px;
+                } 
+                           
+        QGroupBox 
+            { border: 1px solid white; color: rgb(255,255,255);
+                           border-radius: 6px;
+                            margin-top: 8px;
+                            padding: 0px 0px;}            
+                           
+        QPushButton:pressed {Text-align: center; 
+                             background-color: rgb(150,50,150); 
+                             border-color: white;
+                             color:white;}
+                            QToolTip { 
+                           background-color: black; 
+                           color: white; 
+                           border: black solid 1px
+                           }
+        QPushButton:!pressed {Text-align: center; 
+                               background-color: rgb(50,50,50);
+                                border-color: white;
+                               color:white;}
+                                QToolTip { 
+                           background-color: black; 
+                           color: white; 
+                           border: black solid 1px
+                           }
+        QPushButton:disabled {Text-align: center; 
+                             background-color: rgb(30,30,30);
+                             border-color: white;
+                              color:rgb(80,80,80);}
+                               QToolTip { 
+                           background-color: black; 
+                           color: white; 
+                           border: black solid 1px
+                           }
+                        
+        """
+
+
+class DarkPalette(QtGui.QPalette):
+    """Class that inherits from pyqtgraph.QtGui.QPalette and renders dark colours for the application.
+    (from pykilosort/kilosort4)
+    """
+
+    def __init__(self):
+        QtGui.QPalette.__init__(self)
+        self.setup()
+
+    def setup(self):
+        self.setColor(QtGui.QPalette.Window, QtGui.QColor(40, 40, 40))
+        self.setColor(QtGui.QPalette.WindowText, QtGui.QColor(255, 255, 255))
+        self.setColor(QtGui.QPalette.Base, QtGui.QColor(34, 27, 24))
+        self.setColor(QtGui.QPalette.AlternateBase, QtGui.QColor(53, 50, 47))
+        self.setColor(QtGui.QPalette.ToolTipBase, QtGui.QColor(255, 255, 255))
+        self.setColor(QtGui.QPalette.ToolTipText, QtGui.QColor(255, 255, 255))
+        self.setColor(QtGui.QPalette.Text, QtGui.QColor(255, 255, 255))
+        self.setColor(QtGui.QPalette.Button, QtGui.QColor(53, 50, 47))
+        self.setColor(QtGui.QPalette.ButtonText, QtGui.QColor(255, 255, 255))
+        self.setColor(QtGui.QPalette.BrightText, QtGui.QColor(255, 0, 0))
+        self.setColor(QtGui.QPalette.Link, QtGui.QColor(42, 130, 218))
+        self.setColor(QtGui.QPalette.Highlight, QtGui.QColor(42, 130, 218))
+        self.setColor(QtGui.QPalette.HighlightedText, QtGui.QColor(0, 0, 0))
+        self.setColor(QtGui.QPalette.Disabled, QtGui.QPalette.Text,
+                      QtGui.QColor(128, 128, 128))
+        self.setColor(
+            QtGui.QPalette.Disabled,
+            QtGui.QPalette.ButtonText,
+            QtGui.QColor(128, 128, 128),
+        )
+        self.setColor(
+            QtGui.QPalette.Disabled,
+            QtGui.QPalette.WindowText,
+            QtGui.QColor(128, 128, 128),
+        )
+
+
+# def create_channel_choose():
+#     # choose channel
+#     ChannelChoose = [QComboBox(), QComboBox()]
+#     ChannelLabels = []
+#     ChannelChoose[0].addItems(["gray", "red", "green", "blue"])
+#     ChannelChoose[1].addItems(["none", "red", "green", "blue"])
+#     cstr = ["chan to segment:", "chan2 (optional): "]
+#     for i in range(2):
+#         ChannelLabels.append(QLabel(cstr[i]))
+#         if i == 0:
+#             ChannelLabels[i].setToolTip(
+#                 "this is the channel in which the cytoplasm or nuclei exist \
+#             that you want to segment")
+#             ChannelChoose[i].setToolTip(
+#                 "this is the channel in which the cytoplasm or nuclei exist \
+#             that you want to segment")
+#         else:
+#             ChannelLabels[i].setToolTip(
+#                 "if <em>cytoplasm</em> model is chosen, and you also have a \
+#             nuclear channel, then choose the nuclear channel for this option")
+#             ChannelChoose[i].setToolTip(
+#                 "if <em>cytoplasm</em> model is chosen, and you also have a \
+#             nuclear channel, then choose the nuclear channel for this option")
+
+#     return ChannelChoose, ChannelLabels
+
+
+class ModelButton(QPushButton):
+
+    def __init__(self, parent, model_name, text):
+        super().__init__()
+        self.setEnabled(False)
+        self.setText(text)
+        self.setFont(parent.boldfont)
+        self.clicked.connect(lambda: self.press(parent))
+        self.model_name = "cpsam"
+
+    def press(self, parent):
+        parent.compute_segmentation(model_name="cpsam")
+
+
+class FilterButton(QPushButton):
+
+    def __init__(self, parent, text):
+        super().__init__()
+        self.setEnabled(False)
+        self.model_type = text
+        self.setText(text)
+        self.setFont(parent.medfont)
+        self.clicked.connect(lambda: self.press(parent))
+
+    def press(self, parent):
+        if self.model_type == "filter":
+            parent.restore = "filter"
+            normalize_params = parent.get_normalize_params()
+            if (normalize_params["sharpen_radius"] == 0 and
+                    normalize_params["smooth_radius"] == 0 and
+                    normalize_params["tile_norm_blocksize"] == 0):
+                print(
+                    "GUI_ERROR: no filtering settings on (use custom filter settings)")
+                parent.restore = None
+                return
+            parent.restore = self.model_type
+            parent.compute_saturation()
+        # elif self.model_type != "none":
+        #     parent.compute_denoise_model(model_type=self.model_type)
+        else:
+            parent.clear_restore()
+        # parent.set_restore_button()
+
+
+class ObservableVariable(QtCore.QObject):
+    valueChanged = QtCore.Signal(object) 
+
+    def __init__(self, initial=None):
+        super().__init__()
+        self._value = initial
+
+    def set(self, new_value):
+        """ Use this method to get emit the value changing and update the ROI count"""
+        if new_value != self._value:
+            self._value = new_value
+            self.valueChanged.emit(new_value)
+
+    def get(self):
+        return self._value
+
+    def __call__(self):
+        return self._value
+    
+    def reset(self):
+        self.set(0)
+
+    def __iadd__(self, amount):
+        if not isinstance(amount, (int, float)):
+            raise TypeError("Value must be numeric.")
+        self.set(self._value + amount)
+        return self
+    
+    def __radd__(self, other):
+        return other + self._value
+
+    def __add__(self, other):
+        return other + self._value
+        
+    def __isub__(self, amount):
+        if not isinstance(amount, (int, float)):
+            raise TypeError("Value must be numeric.")
+        self.set(self._value - amount)
+        return self
+    
+    def __str__(self):
+        return str(self._value)
+    
+    def __lt__(self, x):
+        return self._value < x
+    
+    def __gt__(self, x):
+        return self._value > x
+    
+    def __eq__(self, x):
+        return self._value == x
+
+
+class NormalizationSettings(QWidget):
+    # TODO
+    pass
+
+
+class SegmentationSettings(QWidget):
+    """ Container for gui settings. Validation is done automatically so any attributes can 
+    be acessed without concern.  
+    """
+    def __init__(self, font):
+        super().__init__()
+
+        # Put everything in a grid layout:
+        grid_layout = QGridLayout()
+        widget_container = QWidget()
+        widget_container.setLayout(grid_layout)
+        row = 0
+
+        ########################### Diameter ###########################
+        # TODO: Validate inputs
+        diam_qlabel = QLabel("diameter:")
+        diam_qlabel.setToolTip("diameter of cells in pixels. If not 30, image will be resized to this")
+        diam_qlabel.setFont(font)
+        grid_layout.addWidget(diam_qlabel, row, 0, 1, 2)
+        self.diameter_box = QLineEdit()
+        self.diameter_box.setToolTip("diameter of cells in pixels. If not blank, image will be resized relative to 30 pixel cell diameters")
+        self.diameter_box.setFont(font)
+        self.diameter_box.setFixedWidth(40)
+        self.diameter_box.setText(' ')
+        grid_layout.addWidget(self.diameter_box, row, 2, 1, 2)
+
+        row += 1
+
+        ########################### Flow threshold ###########################
+        # TODO: Validate inputs
+        flow_threshold_qlabel = QLabel("flow\nthreshold:")
+        flow_threshold_qlabel.setToolTip("threshold on flow error to accept a mask (set higher to get more cells, e.g. in range from (0.1, 3.0), OR set to 0.0 to turn off so no cells discarded);\n press enter to recompute if model already run")
+        flow_threshold_qlabel.setFont(font)
+        grid_layout.addWidget(flow_threshold_qlabel, row, 0, 1, 2)
+        self.flow_threshold_box = QLineEdit()
+        self.flow_threshold_box.setText("0.4")
+        self.flow_threshold_box.setFixedWidth(40)
+        self.flow_threshold_box.setFont(font)
+        grid_layout.addWidget(self.flow_threshold_box, row, 2, 1, 2)
+        self.flow_threshold_box.setToolTip("threshold on flow error to accept a mask (set higher to get more cells, e.g. in range from (0.1, 3.0), OR set to 0.0 to turn off so no cells discarded);\n press enter to recompute if model already run")
+        
+        ########################### Cellprob threshold ###########################
+        # TODO: Validate inputs
+        cellprob_qlabel = QLabel("cellprob\nthreshold:")
+        cellprob_qlabel.setToolTip("threshold on cellprob output to seed cell masks (set lower to include more pixels or higher to include fewer, e.g. in range from (-6, 6)); \n press enter to recompute if model already run")
+        cellprob_qlabel.setFont(font)
+        grid_layout.addWidget(cellprob_qlabel, row, 4, 1, 2)
+        self.cellprob_threshold_box = QLineEdit()
+        self.cellprob_threshold_box.setText("0.0")
+        self.cellprob_threshold_box.setFixedWidth(40)
+        self.cellprob_threshold_box.setFont(font)
+        self.cellprob_threshold_box.setToolTip("threshold on cellprob output to seed cell masks (set lower to include more pixels or higher to include fewer, e.g. in range from (-6, 6)); \n press enter to recompute if model already run")
+        grid_layout.addWidget(self.cellprob_threshold_box, row, 6, 1, 2)
+
+        row += 1
+
+        ########################### Norm percentiles ###########################
+        norm_percentiles_qlabel = QLabel("norm percentiles:")
+        norm_percentiles_qlabel.setToolTip("sets normalization percentiles for segmentation and denoising\n(pixels at lower percentile set to 0.0 and at upper set to 1.0 for network)")
+        norm_percentiles_qlabel.setFont(font)
+        grid_layout.addWidget(norm_percentiles_qlabel, row, 0, 1, 8)
+
+        row += 1
+        validator = QDoubleValidator(0.0, 100.0, 2)
+        validator.setNotation(QDoubleValidator.StandardNotation)
+
+        low_norm_qlabel = QLabel('lower:')
+        low_norm_qlabel.setToolTip("pixels at this percentile set to 0 (default 1.0)")
+        low_norm_qlabel.setFont(font)
+        grid_layout.addWidget(low_norm_qlabel, row, 0, 1, 2)
+        self.norm_percentile_low_box = QLineEdit()
+        self.norm_percentile_low_box.setText("1.0")
+        self.norm_percentile_low_box.setFont(font)
+        self.norm_percentile_low_box.setFixedWidth(40)
+        self.norm_percentile_low_box.setToolTip("pixels at this percentile set to 0 (default 1.0)")
+        self.norm_percentile_low_box.setValidator(validator)
+        self.norm_percentile_low_box.editingFinished.connect(self.validate_normalization_range)
+        grid_layout.addWidget(self.norm_percentile_low_box, row, 2, 1, 1)
+
+        high_norm_qlabel = QLabel('upper:')
+        high_norm_qlabel.setToolTip("pixels at this percentile set to 1 (default 99.0)")
+        high_norm_qlabel.setFont(font)
+        grid_layout.addWidget(high_norm_qlabel, row, 4, 1, 2)
+        self.norm_percentile_high_box = QLineEdit()
+        self.norm_percentile_high_box.setText("99.0")
+        self.norm_percentile_high_box.setFont(font)
+        self.norm_percentile_high_box.setFixedWidth(40)
+        self.norm_percentile_high_box.setToolTip("pixels at this percentile set to 1 (default 99.0)")
+        self.norm_percentile_high_box.setValidator(validator)
+        self.norm_percentile_high_box.editingFinished.connect(self.validate_normalization_range)
+        grid_layout.addWidget(self.norm_percentile_high_box, row, 6, 1, 2)
+
+        row += 1
+
+        ########################### niter ###########################
+        # TODO: change this to follow the same default logic as 'diameter' above
+        # TODO: input validation
+        niter_qlabel = QLabel("niter dynamics:")
+        niter_qlabel.setFont(font)
+        niter_qlabel.setToolTip("number of iterations for dynamics (0 uses default based on diameter); use 2000 for bacteria")
+        grid_layout.addWidget(niter_qlabel, row, 0, 1, 4)
+        self.niter_box = QLineEdit()
+        self.niter_box.setText("0")
+        self.niter_box.setFixedWidth(40)
+        self.niter_box.setFont(font)
+        self.niter_box.setToolTip("number of iterations for dynamics (0 uses default based on diameter); use 2000 for bacteria")
+        grid_layout.addWidget(self.niter_box, row, 4, 1, 2)
+
+        self.setLayout(grid_layout)
+
+    def validate_normalization_range(self):
+        low_text = self.norm_percentile_low_box.text()
+        high_text = self.norm_percentile_high_box.text()
+        
+        if not low_text or low_text.isspace():
+            self.norm_percentile_low_box.setText('1.0')
+            low_text = '1.0'
+        elif not high_text or high_text.isspace():
+            self.norm_percentile_high_box.setText('1.0')
+            high_text = '99.0'
+
+        low = float(low_text)
+        high = float(high_text)
+
+        if low >= high:
+            # Invalid: show error and mark fields
+            self.norm_percentile_low_box.setStyleSheet("border: 1px solid red;")
+            self.norm_percentile_high_box.setStyleSheet("border: 1px solid red;")
+        else:
+            # Valid: clear style
+            self.norm_percentile_low_box.setStyleSheet("")
+            self.norm_percentile_high_box.setStyleSheet("")
+
+    @property
+    def low_percentile(self):
+        """ Also validate the low input by returning 1.0 if text doesn't work """
+        low_text = self.norm_percentile_low_box.text()
+        if not low_text or low_text.isspace():
+            self.norm_percentile_low_box.setText('1.0')
+            low_text = '1.0'
+        return float(self.norm_percentile_low_box.text())
+    
+    @property
+    def high_percentile(self):
+        """ Also validate the high input by returning 99.0 if text doesn't work """
+        high_text = self.norm_percentile_high_box.text()
+        if not high_text or high_text.isspace():
+            self.norm_percentile_high_box.setText('99.0')
+            high_text = '99.0'
+        return float(self.norm_percentile_high_box.text())
+    
+    @property
+    def diameter(self):
+        """ Get the diameter from the diameter box, if box isn't a number return None"""
+        try:
+            d = float(self.diameter_box.text())
+        except ValueError:
+            d = None
+        return d 
+    
+    @property
+    def flow_threshold(self):
+        return float(self.flow_threshold_box.text())
+    
+    @property
+    def cellprob_threshold(self):
+        return float(self.cellprob_threshold_box.text())
+    
+    @property
+    def niter(self):
+        num = int(self.niter_box.text())
+        if num < 1:
+            self.niter_box.setText('200')
+            return 200
+        else:
+            return num
+
+
+
+class TrainWindow(QDialog):
+
+    def __init__(self, parent, model_strings):
+        super().__init__(parent)
+        self.setGeometry(100, 100, 900, 550)
+        self.setWindowTitle("train settings")
+        self.win = QWidget(self)
+        self.l0 = QGridLayout()
+        self.win.setLayout(self.l0)
+
+        yoff = 0
+        qlabel = QLabel("train model w/ images + _seg.npy in current folder >>")
+        qlabel.setFont(QtGui.QFont("Arial", 10, QtGui.QFont.Bold))
+
+        qlabel.setAlignment(QtCore.Qt.AlignVCenter)
+        self.l0.addWidget(qlabel, yoff, 0, 1, 2)
+
+        # choose initial model
+        yoff += 1
+        self.ModelChoose = QComboBox()
+        self.ModelChoose.addItems(model_strings)
+        self.ModelChoose.setFixedWidth(150)
+        self.ModelChoose.setCurrentIndex(parent.training_params["model_index"])
+        self.l0.addWidget(self.ModelChoose, yoff, 1, 1, 1)
+        qlabel = QLabel("initial model: ")
+        qlabel.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+        self.l0.addWidget(qlabel, yoff, 0, 1, 1)
+
+        # choose parameters
+        labels = ["learning_rate", "weight_decay", "n_epochs", "model_name"]
+        self.edits = []
+        yoff += 1
+        for i, label in enumerate(labels):
+            qlabel = QLabel(label)
+            qlabel.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+            self.l0.addWidget(qlabel, i + yoff, 0, 1, 1)
+            self.edits.append(QLineEdit())
+            self.edits[-1].setText(str(parent.training_params[label]))
+            self.edits[-1].setFixedWidth(200)
+            self.l0.addWidget(self.edits[-1], i + yoff, 1, 1, 1)
+
+        yoff += len(labels)
+
+        yoff += 1
+        self.use_norm = QCheckBox(f"use restored/filtered image")
+        self.use_norm.setChecked(True)
+
+        yoff += 2
+        qlabel = QLabel(
+            "(to remove files, click cancel then remove \nfrom folder and reopen train window)"
+        )
+        self.l0.addWidget(qlabel, yoff, 0, 2, 4)
+
+        # click button
+        yoff += 3
+        QBtn = QDialogButtonBox.Ok | QDialogButtonBox.Cancel
+        self.buttonBox = QDialogButtonBox(QBtn)
+        self.buttonBox.accepted.connect(lambda: self.accept(parent))
+        self.buttonBox.rejected.connect(self.reject)
+        self.l0.addWidget(self.buttonBox, yoff, 0, 1, 4)
+
+        # list files in folder
+        qlabel = QLabel("filenames")
+        qlabel.setFont(QtGui.QFont("Arial", 8, QtGui.QFont.Bold))
+        self.l0.addWidget(qlabel, 0, 4, 1, 1)
+        qlabel = QLabel("# of masks")
+        qlabel.setFont(QtGui.QFont("Arial", 8, QtGui.QFont.Bold))
+        self.l0.addWidget(qlabel, 0, 5, 1, 1)
+
+        for i in range(10):
+            if i > len(parent.train_files) - 1:
+                break
+            elif i == 9 and len(parent.train_files) > 10:
+                label = "..."
+                nmasks = "..."
+            else:
+                label = os.path.split(parent.train_files[i])[-1]
+                nmasks = str(parent.train_labels[i].max())
+            qlabel = QLabel(label)
+            self.l0.addWidget(qlabel, i + 1, 4, 1, 1)
+            qlabel = QLabel(nmasks)
+            qlabel.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
+            self.l0.addWidget(qlabel, i + 1, 5, 1, 1)
+
+    def accept(self, parent):
+        # set training params
+        parent.training_params = {
+            "model_index": self.ModelChoose.currentIndex(),
+            "learning_rate": float(self.edits[0].text()),
+            "weight_decay": float(self.edits[1].text()),
+            "n_epochs": int(self.edits[2].text()),
+            "model_name": self.edits[3].text(),
+            #"use_norm": True if self.use_norm.isChecked() else False,
+        }
+        self.done(1)
+
+
+class ExampleGUI(QDialog):
+
+    def __init__(self, parent=None):
+        super(ExampleGUI, self).__init__(parent)
+        self.setGeometry(100, 100, 1300, 900)
+        self.setWindowTitle("GUI layout")
+        self.win = QWidget(self)
+        layout = QGridLayout()
+        self.win.setLayout(layout)
+        guip_path = pathlib.Path.home().joinpath(".cellpose", "cellposeSAM_gui.png")
+        guip_path = str(guip_path.resolve())
+        pixmap = QPixmap(guip_path)
+        label = QLabel(self)
+        label.setPixmap(pixmap)
+        pixmap.scaled
+        layout.addWidget(label, 0, 0, 1, 1)
+
+
+class HelpWindow(QDialog):
+
+    def __init__(self, parent=None):
+        super(HelpWindow, self).__init__(parent)
+        self.setGeometry(100, 50, 700, 1000)
+        self.setWindowTitle("cellpose help")
+        self.win = QWidget(self)
+        layout = QGridLayout()
+        self.win.setLayout(layout)
+
+        text_file = pathlib.Path(__file__).parent.joinpath("guihelpwindowtext.html")
+        with open(str(text_file.resolve()), "r") as f:
+            text = f.read()
+
+        label = QLabel(text)
+        label.setFont(QtGui.QFont("Arial", 8))
+        label.setWordWrap(True)
+        layout.addWidget(label, 0, 0, 1, 1)
+        self.show()
+
+
+class TrainHelpWindow(QDialog):
+
+    def __init__(self, parent=None):
+        super(TrainHelpWindow, self).__init__(parent)
+        self.setGeometry(100, 50, 700, 300)
+        self.setWindowTitle("training instructions")
+        self.win = QWidget(self)
+        layout = QGridLayout()
+        self.win.setLayout(layout)
+
+        text_file = pathlib.Path(__file__).parent.joinpath(
+            "guitrainhelpwindowtext.html")
+        with open(str(text_file.resolve()), "r") as f:
+            text = f.read()
+
+        label = QLabel(text)
+        label.setFont(QtGui.QFont("Arial", 8))
+        label.setWordWrap(True)
+        layout.addWidget(label, 0, 0, 1, 1)
+        self.show()
+
+
+class ViewBoxNoRightDrag(pg.ViewBox):
+
+    def __init__(self, parent=None, border=None, lockAspect=False, enableMouse=True,
+                 invertY=False, enableMenu=True, name=None, invertX=False):
+        pg.ViewBox.__init__(self, None, border, lockAspect, enableMouse, invertY,
+                            enableMenu, name, invertX)
+        self.parent = parent
+        self.axHistoryPointer = -1
+
+    def keyPressEvent(self, ev):
+        """
+        This routine should capture key presses in the current view box.
+        The following events are implemented:
+        +/= : moves forward in the zooming stack (if it exists)
+        - : moves backward in the zooming stack (if it exists)
+
+        """
+        ev.accept()
+        if ev.text() == "-":
+            self.scaleBy([1.1, 1.1])
+        elif ev.text() in ["+", "="]:
+            self.scaleBy([0.9, 0.9])
+        else:
+            ev.ignore()
+
+
+class ImageDraw(pg.ImageItem):
+    """
+    **Bases:** :class:`GraphicsObject <pyqtgraph.GraphicsObject>`
+    GraphicsObject displaying an image. Optimized for rapid update (ie video display).
+    This item displays either a 2D numpy array (height, width) or
+    a 3D array (height, width, RGBa). This array is optionally scaled (see
+    :func:`setLevels <pyqtgraph.ImageItem.setLevels>`) and/or colored
+    with a lookup table (see :func:`setLookupTable <pyqtgraph.ImageItem.setLookupTable>`)
+    before being displayed.
+    ImageItem is frequently used in conjunction with
+    :class:`HistogramLUTItem <pyqtgraph.HistogramLUTItem>` or
+    :class:`HistogramLUTWidget <pyqtgraph.HistogramLUTWidget>` to provide a GUI
+    for controlling the levels and lookup table used to display the image.
+    """
+
+    sigImageChanged = QtCore.Signal()
+
+    def __init__(self, image=None, viewbox=None, parent=None, **kargs):
+        super(ImageDraw, self).__init__()
+        self.levels = np.array([0, 255])
+        self.lut = None
+        self.autoDownsample = False
+        self.axisOrder = "row-major"
+        self.removable = False
+
+        self.parent = parent
+        self.setDrawKernel(kernel_size=self.parent.brush_size)
+        self.parent.current_stroke = []
+        self.parent.in_stroke = False
+
+    def mouseClickEvent(self, ev):
+        if (self.parent.masksOn or
+                self.parent.outlinesOn) and not self.parent.removing_region:
+            is_right_click = ev.button() == QtCore.Qt.RightButton
+            if self.parent.loaded \
+                    and (is_right_click or ev.modifiers() & QtCore.Qt.ShiftModifier and not ev.double())\
+                    and not self.parent.deleting_multiple:
+                if not self.parent.in_stroke:
+                    ev.accept()
+                    self.create_start(ev.pos())
+                    self.parent.stroke_appended = False
+                    self.parent.in_stroke = True
+                    self.drawAt(ev.pos(), ev)
+                else:
+                    ev.accept()
+                    self.end_stroke()
+                    self.parent.in_stroke = False
+            elif not self.parent.in_stroke:
+                y, x = int(ev.pos().y()), int(ev.pos().x())
+                if y >= 0 and y < self.parent.Ly and x >= 0 and x < self.parent.Lx:
+                    if ev.button() == QtCore.Qt.LeftButton and not ev.double():
+                        idx = self.parent.cellpix[self.parent.currentZ][y, x]
+                        if idx > 0:
+                            if ev.modifiers() & QtCore.Qt.ControlModifier:
+                                # delete mask selected
+                                self.parent.remove_cell(idx)
+                            elif ev.modifiers() & QtCore.Qt.AltModifier:
+                                self.parent.merge_cells(idx)
+                            elif self.parent.masksOn and not self.parent.deleting_multiple:
+                                self.parent.unselect_cell()
+                                self.parent.select_cell(idx)
+                            elif self.parent.deleting_multiple:
+                                if idx in self.parent.removing_cells_list:
+                                    self.parent.unselect_cell_multi(idx)
+                                    self.parent.removing_cells_list.remove(idx)
+                                else:
+                                    self.parent.select_cell_multi(idx)
+                                    self.parent.removing_cells_list.append(idx)
+
+                        elif self.parent.masksOn and not self.parent.deleting_multiple:
+                            self.parent.unselect_cell()
+
+    def mouseDragEvent(self, ev):
+        ev.ignore()
+        return
+
+    def hoverEvent(self, ev):
+        if self.parent.in_stroke:
+            if self.parent.in_stroke:
+                # continue stroke if not at start
+                self.drawAt(ev.pos())
+                if self.is_at_start(ev.pos()):
+                    self.end_stroke()
+        else:
+            ev.acceptClicks(QtCore.Qt.RightButton)
+
+    def create_start(self, pos):
+        self.scatter = pg.ScatterPlotItem([pos.x()], [pos.y()], pxMode=False,
+                                          pen=pg.mkPen(color=(255, 0, 0),
+                                                       width=self.parent.brush_size),
+                                          size=max(3 * 2,
+                                                   self.parent.brush_size * 1.8 * 2),
+                                          brush=None)
+        self.parent.p0.addItem(self.scatter)
+
+    def is_at_start(self, pos):
+        thresh_out = max(6, self.parent.brush_size * 3)
+        thresh_in = max(3, self.parent.brush_size * 1.8)
+        # first check if you ever left the start
+        if len(self.parent.current_stroke) > 3:
+            stroke = np.array(self.parent.current_stroke)
+            dist = (((stroke[1:, 1:] -
+                      stroke[:1, 1:][np.newaxis, :, :])**2).sum(axis=-1))**0.5
+            dist = dist.flatten()
+            has_left = (dist > thresh_out).nonzero()[0]
+            if len(has_left) > 0:
+                first_left = np.sort(has_left)[0]
+                has_returned = (dist[max(4, first_left + 1):] < thresh_in).sum()
+                if has_returned > 0:
+                    return True
+                else:
+                    return False
+            else:
+                return False
+
+    def end_stroke(self):
+        self.parent.p0.removeItem(self.scatter)
+        if not self.parent.stroke_appended:
+            self.parent.strokes.append(self.parent.current_stroke)
+            self.parent.stroke_appended = True
+            self.parent.current_stroke = np.array(self.parent.current_stroke)
+            ioutline = self.parent.current_stroke[:, 3] == 1
+            self.parent.current_point_set.append(
+                list(self.parent.current_stroke[ioutline]))
+            self.parent.current_stroke = []
+            if self.parent.autosave:
+                self.parent.add_set()
+        if len(self.parent.current_point_set) and len(
+                self.parent.current_point_set[0]) > 0 and self.parent.autosave:
+            self.parent.add_set()
+        self.parent.in_stroke = False
+
+    def tabletEvent(self, ev):
+        pass
+
+    def drawAt(self, pos, ev=None):
+        mask = self.strokemask
+        stroke = self.parent.current_stroke
+        pos = [int(pos.y()), int(pos.x())]
+        dk = self.drawKernel
+        kc = self.drawKernelCenter
+        sx = [0, dk.shape[0]]
+        sy = [0, dk.shape[1]]
+        tx = [pos[0] - kc[0], pos[0] - kc[0] + dk.shape[0]]
+        ty = [pos[1] - kc[1], pos[1] - kc[1] + dk.shape[1]]
+        kcent = kc.copy()
+        if tx[0] <= 0:
+            sx[0] = 0
+            sx[1] = kc[0] + 1
+            tx = sx
+            kcent[0] = 0
+        if ty[0] <= 0:
+            sy[0] = 0
+            sy[1] = kc[1] + 1
+            ty = sy
+            kcent[1] = 0
+        if tx[1] >= self.parent.Ly - 1:
+            sx[0] = dk.shape[0] - kc[0] - 1
+            sx[1] = dk.shape[0]
+            tx[0] = self.parent.Ly - kc[0] - 1
+            tx[1] = self.parent.Ly
+            kcent[0] = tx[1] - tx[0] - 1
+        if ty[1] >= self.parent.Lx - 1:
+            sy[0] = dk.shape[1] - kc[1] - 1
+            sy[1] = dk.shape[1]
+            ty[0] = self.parent.Lx - kc[1] - 1
+            ty[1] = self.parent.Lx
+            kcent[1] = ty[1] - ty[0] - 1
+
+        ts = (slice(tx[0], tx[1]), slice(ty[0], ty[1]))
+        ss = (slice(sx[0], sx[1]), slice(sy[0], sy[1]))
+        self.image[ts] = mask[ss]
+
+        for ky, y in enumerate(np.arange(ty[0], ty[1], 1, int)):
+            for kx, x in enumerate(np.arange(tx[0], tx[1], 1, int)):
+                iscent = np.logical_and(kx == kcent[0], ky == kcent[1])
+                stroke.append([self.parent.currentZ, x, y, iscent])
+        self.updateImage()
+
+    def setDrawKernel(self, kernel_size=3):
+        bs = kernel_size
+        kernel = np.ones((bs, bs), np.uint8)
+        self.drawKernel = kernel
+        self.drawKernelCenter = [
+            int(np.floor(kernel.shape[0] / 2)),
+            int(np.floor(kernel.shape[1] / 2))
+        ]
+        onmask = 255 * kernel[:, :, np.newaxis]
+        offmask = np.zeros((bs, bs, 1))
+        opamask = 100 * kernel[:, :, np.newaxis]
+        self.redmask = np.concatenate((onmask, offmask, offmask, onmask), axis=-1)
+        self.strokemask = np.concatenate((onmask, offmask, onmask, opamask), axis=-1)

models/seg_post_model/cellpose/gui/guitrainhelpwindowtext.html

@@ -0,0 +1,25 @@
+<qt>
+    Check out this <a href="https://youtu.be/3Y1VKcxjNy4">video</a> to learn the process.
+    <ol>
+        <li>Drag and drop an image from a folder of images with a similar style (like similar cell types).</li>
+        <li>Run the built-in models on one of the images using the "model zoo" and find the one that works best for your
+            data. Make sure that if you have a nuclear channel you have selected it for CHAN2.
+        </li>
+        <li>Fix the labelling by drawing new ROIs (right-click) and deleting incorrect ones (CTRL+click). The GUI
+            autosaves any manual changes (but does not autosave after running the model, for that click CTRL+S). The
+            segmentation is saved in a "_seg.npy" file.
+        </li>
+        <li> Go to the "Models" menu in the File bar at the top and click "Train new model..." or use shortcut CTRL+T.
+        </li>
+        <li> Choose the pretrained model to start the training from (the model you used in #2), and type in the model
+            name that you want to use. The other parameters should work well in general for most data types. Then click
+            OK.
+        </li>
+        <li> The model will train (much faster if you have a GPU) and then auto-run on the next image in the folder.
+            Next you can repeat #3-#5 as many times as is necessary.
+        </li>
+        <li> The trained model is available to use in the future in the GUI in the "custom model" section and is saved
+            in your image folder.
+        </li>
+    </ol>
+</qt>

models/seg_post_model/cellpose/gui/io.py

@@ -0,0 +1,634 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os, gc
+import numpy as np
+import cv2
+import fastremap
+
+from ..io import imread, imread_2D, imread_3D, imsave, outlines_to_text, add_model, remove_model, save_rois
+from ..models import normalize_default, MODEL_DIR, MODEL_LIST_PATH, get_user_models
+from ..utils import masks_to_outlines, outlines_list
+
+try:
+    import qtpy
+    from qtpy.QtWidgets import QFileDialog
+    GUI = True
+except:
+    GUI = False
+
+try:
+    import matplotlib.pyplot as plt
+    MATPLOTLIB = True
+except:
+    MATPLOTLIB = False
+
+
+def _init_model_list(parent):
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    parent.model_list_path = MODEL_LIST_PATH
+    parent.model_strings = get_user_models()
+
+
+def _add_model(parent, filename=None, load_model=True):
+    if filename is None:
+        name = QFileDialog.getOpenFileName(parent, "Add model to GUI")
+        filename = name[0]
+    add_model(filename)
+    fname = os.path.split(filename)[-1]
+    parent.ModelChooseC.addItems([fname])
+    parent.model_strings.append(fname)
+
+    for ind, model_string in enumerate(parent.model_strings[:-1]):
+        if model_string == fname:
+            _remove_model(parent, ind=ind + 1, verbose=False)
+
+    parent.ModelChooseC.setCurrentIndex(len(parent.model_strings))
+    if load_model:
+        parent.model_choose(custom=True)
+
+
+def _remove_model(parent, ind=None, verbose=True):
+    if ind is None:
+        ind = parent.ModelChooseC.currentIndex()
+    if ind > 0:
+        ind -= 1
+        parent.ModelChooseC.removeItem(ind + 1)
+        del parent.model_strings[ind]
+        # remove model from txt path
+        modelstr = parent.ModelChooseC.currentText()
+        remove_model(modelstr)
+        if len(parent.model_strings) > 0:
+            parent.ModelChooseC.setCurrentIndex(len(parent.model_strings))
+        else:
+            parent.ModelChooseC.setCurrentIndex(0)
+    else:
+        print("ERROR: no model selected to delete")
+
+
+def _get_train_set(image_names):
+    """ get training data and labels for images in current folder image_names"""
+    train_data, train_labels, train_files = [], [], []
+    restore = None
+    normalize_params = normalize_default
+    for image_name_full in image_names:
+        image_name = os.path.splitext(image_name_full)[0]
+        label_name = None
+        if os.path.exists(image_name + "_seg.npy"):
+            dat = np.load(image_name + "_seg.npy", allow_pickle=True).item()
+            masks = dat["masks"].squeeze()
+            if masks.ndim == 2:
+                fastremap.renumber(masks, in_place=True)
+                label_name = image_name + "_seg.npy"
+            else:
+                print(f"GUI_INFO: _seg.npy found for {image_name} but masks.ndim!=2")
+            if "img_restore" in dat:
+                data = dat["img_restore"].squeeze()
+                restore = dat["restore"]
+            else:
+                data = imread(image_name_full)
+            normalize_params = dat[
+                "normalize_params"] if "normalize_params" in dat else normalize_default
+        if label_name is not None:
+            train_files.append(image_name_full)
+            train_data.append(data)
+            train_labels.append(masks)
+    if restore:
+        print(f"GUI_INFO: using {restore} images (dat['img_restore'])")
+    return train_data, train_labels, train_files, restore, normalize_params
+
+
+def _load_image(parent, filename=None, load_seg=True, load_3D=False):
+    """ load image with filename; if None, open QFileDialog
+    if image is grey change view to default to grey scale 
+    """
+
+    if parent.load_3D:
+        load_3D = True
+
+    if filename is None:
+        name = QFileDialog.getOpenFileName(parent, "Load image")
+        filename = name[0]
+        if filename == "":
+            return
+    manual_file = os.path.splitext(filename)[0] + "_seg.npy"
+    load_mask = False
+    if load_seg:
+        if os.path.isfile(manual_file) and not parent.autoloadMasks.isChecked():
+            if filename is not None:
+                image = (imread_2D(filename) if not load_3D else 
+                         imread_3D(filename))
+            else:
+                image = None
+            _load_seg(parent, manual_file, image=image, image_file=filename,
+                      load_3D=load_3D)
+            return
+        elif parent.autoloadMasks.isChecked():
+            mask_file = os.path.splitext(filename)[0] + "_masks" + os.path.splitext(
+                filename)[-1]
+            mask_file = os.path.splitext(filename)[
+                0] + "_masks.tif" if not os.path.isfile(mask_file) else mask_file
+            load_mask = True if os.path.isfile(mask_file) else False
+    try:
+        print(f"GUI_INFO: loading image: {filename}")
+        if not load_3D:
+            image = imread_2D(filename)
+        else:
+            image = imread_3D(filename)
+        parent.loaded = True
+    except Exception as e:
+        print("ERROR: images not compatible")
+        print(f"ERROR: {e}")
+
+    if parent.loaded:
+        parent.reset()
+        parent.filename = filename
+        filename = os.path.split(parent.filename)[-1]
+        _initialize_images(parent, image, load_3D=load_3D)
+        parent.loaded = True
+        parent.enable_buttons()
+        if load_mask:
+            _load_masks(parent, filename=mask_file)
+
+    # check if gray and adjust viewer:
+    if len(np.unique(image[..., 1:])) == 1:
+        parent.color = 4
+        parent.RGBDropDown.setCurrentIndex(4) # gray
+        parent.update_plot()
+
+        
+def _initialize_images(parent, image, load_3D=False):
+    """ format image for GUI
+
+    assumes image is Z x W x H x C
+
+    """
+    load_3D = parent.load_3D if load_3D is False else load_3D
+
+    parent.stack = image
+    print(f"GUI_INFO: image shape: {image.shape}")
+    if load_3D:
+        parent.NZ = len(parent.stack)
+        parent.scroll.setMaximum(parent.NZ - 1)
+    else:
+        parent.NZ = 1
+        parent.stack = parent.stack[np.newaxis, ...]
+
+    img_min = image.min()
+    img_max = image.max()
+    parent.stack = parent.stack.astype(np.float32)
+    parent.stack -= img_min
+    if img_max > img_min + 1e-3:
+        parent.stack /= (img_max - img_min)
+    parent.stack *= 255
+
+    if load_3D:
+        print("GUI_INFO: converted to float and normalized values to 0.0->255.0")
+
+    del image
+    gc.collect()
+
+    parent.imask = 0
+    parent.Ly, parent.Lx = parent.stack.shape[-3:-1]
+    parent.Ly0, parent.Lx0 = parent.stack.shape[-3:-1]
+    parent.layerz = 255 * np.ones((parent.Ly, parent.Lx, 4), "uint8")
+    if hasattr(parent, "stack_filtered"):
+        parent.Lyr, parent.Lxr = parent.stack_filtered.shape[-3:-1]
+    elif parent.restore and "upsample" in parent.restore:
+        parent.Lyr, parent.Lxr = int(parent.Ly * parent.ratio), int(parent.Lx *
+                                                                    parent.ratio)
+    else:
+        parent.Lyr, parent.Lxr = parent.Ly, parent.Lx
+    parent.clear_all()
+
+    if not hasattr(parent, "stack_filtered") and parent.restore:
+        print("GUI_INFO: no 'img_restore' found, applying current settings")
+        parent.compute_restore()
+
+    if parent.autobtn.isChecked():
+        if parent.restore is None or parent.restore != "filter":
+            print(
+                "GUI_INFO: normalization checked: computing saturation levels (and optionally filtered image)"
+            )
+            parent.compute_saturation()
+    # elif len(parent.saturation) != parent.NZ:
+    #     parent.saturation = []
+    #     for r in range(3):
+    #         parent.saturation.append([])
+    #         for n in range(parent.NZ):
+    #             parent.saturation[-1].append([0, 255])
+    #         parent.sliders[r].setValue([0, 255])
+    parent.compute_scale()
+    parent.track_changes = []
+
+    if load_3D:
+        parent.currentZ = int(np.floor(parent.NZ / 2))
+        parent.scroll.setValue(parent.currentZ)
+        parent.zpos.setText(str(parent.currentZ))
+    else:
+        parent.currentZ = 0
+
+
+def _load_seg(parent, filename=None, image=None, image_file=None, load_3D=False):
+    """ load *_seg.npy with filename; if None, open QFileDialog """
+    if filename is None:
+        name = QFileDialog.getOpenFileName(parent, "Load labelled data", filter="*.npy")
+        filename = name[0]
+    try:
+        dat = np.load(filename, allow_pickle=True).item()
+        # check if there are keys in filename
+        dat["outlines"]
+        parent.loaded = True
+    except:
+        parent.loaded = False
+        print("ERROR: not NPY")
+        return
+
+    parent.reset()
+    if image is None:
+        found_image = False
+        if "filename" in dat:
+            parent.filename = dat["filename"]
+            if os.path.isfile(parent.filename):
+                parent.filename = dat["filename"]
+                found_image = True
+            else:
+                imgname = os.path.split(parent.filename)[1]
+                root = os.path.split(filename)[0]
+                parent.filename = root + "/" + imgname
+                if os.path.isfile(parent.filename):
+                    found_image = True
+        if found_image:
+            try:
+                print(parent.filename)
+                image = (imread_2D(parent.filename) if not load_3D else 
+                         imread_3D(parent.filename))
+            except:
+                parent.loaded = False
+                found_image = False
+                print("ERROR: cannot find image file, loading from npy")
+        if not found_image:
+            parent.filename = filename[:-8]
+            print(parent.filename)
+            if "img" in dat:
+                image = dat["img"]
+            else:
+                print("ERROR: no image file found and no image in npy")
+                return
+    else:
+        parent.filename = image_file
+
+    parent.restore = None
+    parent.ratio = 1.
+
+    if "normalize_params" in dat:
+        parent.set_normalize_params(dat["normalize_params"])
+
+    _initialize_images(parent, image, load_3D=load_3D)
+    print(parent.stack.shape)
+
+    if "outlines" in dat:
+        if isinstance(dat["outlines"], list):
+            # old way of saving files
+            dat["outlines"] = dat["outlines"][::-1]
+            for k, outline in enumerate(dat["outlines"]):
+                if "colors" in dat:
+                    color = dat["colors"][k]
+                else:
+                    col_rand = np.random.randint(1000)
+                    color = parent.colormap[col_rand, :3]
+                median = parent.add_mask(points=outline, color=color)
+                if median is not None:
+                    parent.cellcolors = np.append(parent.cellcolors,
+                                                  color[np.newaxis, :], axis=0)
+                    parent.ncells += 1
+        else:
+            if dat["masks"].min() == -1:
+                dat["masks"] += 1
+                dat["outlines"] += 1
+            parent.ncells.set(dat["masks"].max())
+            if "colors" in dat and len(dat["colors"]) == dat["masks"].max():
+                colors = dat["colors"]
+            else:
+                colors = parent.colormap[:parent.ncells.get(), :3]
+
+            _masks_to_gui(parent, dat["masks"], outlines=dat["outlines"], colors=colors)
+
+            parent.draw_layer()
+
+        if "manual_changes" in dat:
+            parent.track_changes = dat["manual_changes"]
+            print("GUI_INFO: loaded in previous changes")
+        if "zdraw" in dat:
+            parent.zdraw = dat["zdraw"]
+        else:
+            parent.zdraw = [None for n in range(parent.ncells.get())]
+        parent.loaded = True
+    else:
+        parent.clear_all()
+
+    parent.ismanual = np.zeros(parent.ncells.get(), bool)
+    if "ismanual" in dat:
+        if len(dat["ismanual"]) == parent.ncells:
+            parent.ismanual = dat["ismanual"]
+
+    if "current_channel" in dat:
+        parent.color = (dat["current_channel"] + 2) % 5
+        parent.RGBDropDown.setCurrentIndex(parent.color)
+
+    if "flows" in dat:
+        parent.flows = dat["flows"]
+        try:
+            if parent.flows[0].shape[-3] != dat["masks"].shape[-2]:
+                Ly, Lx = dat["masks"].shape[-2:]
+                for i in range(len(parent.flows)):
+                    parent.flows[i] = cv2.resize(
+                        parent.flows[i].squeeze(), (Lx, Ly),
+                        interpolation=cv2.INTER_NEAREST)[np.newaxis, ...]
+            if parent.NZ == 1:
+                parent.recompute_masks = True
+            else:
+                parent.recompute_masks = False
+
+        except:
+            try:
+                if len(parent.flows[0]) > 0:
+                    parent.flows = parent.flows[0]
+            except:
+                parent.flows = [[], [], [], [], [[]]]
+            parent.recompute_masks = False
+
+    parent.enable_buttons()
+    parent.update_layer()
+    del dat
+    gc.collect()
+
+
+def _load_masks(parent, filename=None):
+    """ load zeros-based masks (0=no cell, 1=cell 1, ...) """
+    if filename is None:
+        name = QFileDialog.getOpenFileName(parent, "Load masks (PNG or TIFF)")
+        filename = name[0]
+    print(f"GUI_INFO: loading masks: {filename}")
+    masks = imread(filename)
+    outlines = None
+    if masks.ndim > 3:
+        # Z x nchannels x Ly x Lx
+        if masks.shape[-1] > 5:
+            parent.flows = list(np.transpose(masks[:, :, :, 2:], (3, 0, 1, 2)))
+            outlines = masks[..., 1]
+            masks = masks[..., 0]
+        else:
+            parent.flows = list(np.transpose(masks[:, :, :, 1:], (3, 0, 1, 2)))
+            masks = masks[..., 0]
+    elif masks.ndim == 3:
+        if masks.shape[-1] < 5:
+            masks = masks[np.newaxis, :, :, 0]
+    elif masks.ndim < 3:
+        masks = masks[np.newaxis, :, :]
+    # masks should be Z x Ly x Lx
+    if masks.shape[0] != parent.NZ:
+        print("ERROR: masks are not same depth (number of planes) as image stack")
+        return
+
+    _masks_to_gui(parent, masks, outlines)
+    if parent.ncells > 0:
+        parent.draw_layer()
+        parent.toggle_mask_ops()
+    del masks
+    gc.collect()
+    parent.update_layer()
+    parent.update_plot()
+
+
+def _masks_to_gui(parent, masks, outlines=None, colors=None):
+    """ masks loaded into GUI """
+    # get unique values
+    shape = masks.shape
+    if len(fastremap.unique(masks)) != masks.max() + 1:
+        print("GUI_INFO: renumbering masks")
+        fastremap.renumber(masks, in_place=True)
+        outlines = None
+        masks = masks.reshape(shape)
+    if masks.ndim == 2:
+        outlines = None
+    masks = masks.astype(np.uint16) if masks.max() < 2**16 - 1 else masks.astype(
+        np.uint32)
+    if parent.restore and "upsample" in parent.restore:
+        parent.cellpix_resize = masks.copy()
+        parent.cellpix = parent.cellpix_resize.copy()
+        parent.cellpix_orig = cv2.resize(
+            masks.squeeze(), (parent.Lx0, parent.Ly0),
+            interpolation=cv2.INTER_NEAREST)[np.newaxis, :, :]
+        parent.resize = True
+    else:
+        parent.cellpix = masks
+    if parent.cellpix.ndim == 2:
+        parent.cellpix = parent.cellpix[np.newaxis, :, :]
+        if parent.restore and "upsample" in parent.restore:
+            if parent.cellpix_resize.ndim == 2:
+                parent.cellpix_resize = parent.cellpix_resize[np.newaxis, :, :]
+            if parent.cellpix_orig.ndim == 2:
+                parent.cellpix_orig = parent.cellpix_orig[np.newaxis, :, :]
+
+    print(f"GUI_INFO: {masks.max()} masks found")
+
+    # get outlines
+    if outlines is None:  # parent.outlinesOn
+        parent.outpix = np.zeros_like(parent.cellpix)
+        if parent.restore and "upsample" in parent.restore:
+            parent.outpix_orig = np.zeros_like(parent.cellpix_orig)
+        for z in range(parent.NZ):
+            outlines = masks_to_outlines(parent.cellpix[z])
+            parent.outpix[z] = outlines * parent.cellpix[z]
+            if parent.restore and "upsample" in parent.restore:
+                outlines = masks_to_outlines(parent.cellpix_orig[z])
+                parent.outpix_orig[z] = outlines * parent.cellpix_orig[z]
+            if z % 50 == 0 and parent.NZ > 1:
+                print("GUI_INFO: plane %d outlines processed" % z)
+        if parent.restore and "upsample" in parent.restore:
+            parent.outpix_resize = parent.outpix.copy()
+    else:
+        parent.outpix = outlines
+        if parent.restore and "upsample" in parent.restore:
+            parent.outpix_resize = parent.outpix.copy()
+            parent.outpix_orig = np.zeros_like(parent.cellpix_orig)
+            for z in range(parent.NZ):
+                outlines = masks_to_outlines(parent.cellpix_orig[z])
+                parent.outpix_orig[z] = outlines * parent.cellpix_orig[z]
+                if z % 50 == 0 and parent.NZ > 1:
+                    print("GUI_INFO: plane %d outlines processed" % z)
+
+    if parent.outpix.ndim == 2:
+        parent.outpix = parent.outpix[np.newaxis, :, :]
+        if parent.restore and "upsample" in parent.restore:
+            if parent.outpix_resize.ndim == 2:
+                parent.outpix_resize = parent.outpix_resize[np.newaxis, :, :]
+            if parent.outpix_orig.ndim == 2:
+                parent.outpix_orig = parent.outpix_orig[np.newaxis, :, :]
+
+    parent.ncells.set(parent.cellpix.max())
+    colors = parent.colormap[:parent.ncells.get(), :3] if colors is None else colors
+    print("GUI_INFO: creating cellcolors and drawing masks")
+    parent.cellcolors = np.concatenate((np.array([[255, 255, 255]]), colors),
+                                       axis=0).astype(np.uint8)
+    if parent.ncells > 0:
+        parent.draw_layer()
+        parent.toggle_mask_ops()
+    parent.ismanual = np.zeros(parent.ncells.get(), bool)
+    parent.zdraw = list(-1 * np.ones(parent.ncells.get(), np.int16))
+
+    if hasattr(parent, "stack_filtered"):
+        parent.ViewDropDown.setCurrentIndex(parent.ViewDropDown.count() - 1)
+        print("set denoised/filtered view")
+    else:
+        parent.ViewDropDown.setCurrentIndex(0)
+
+
+def _save_png(parent):
+    """ save masks to png or tiff (if 3D) """
+    filename = parent.filename
+    base = os.path.splitext(filename)[0]
+    if parent.NZ == 1:
+        if parent.cellpix[0].max() > 65534:
+            print("GUI_INFO: saving 2D masks to tif (too many masks for PNG)")
+            imsave(base + "_cp_masks.tif", parent.cellpix[0])
+        else:
+            print("GUI_INFO: saving 2D masks to png")
+            imsave(base + "_cp_masks.png", parent.cellpix[0].astype(np.uint16))
+    else:
+        print("GUI_INFO: saving 3D masks to tiff")
+        imsave(base + "_cp_masks.tif", parent.cellpix)
+
+
+def _save_flows(parent):
+    """ save flows and cellprob to tiff """
+    filename = parent.filename
+    base = os.path.splitext(filename)[0]
+    print("GUI_INFO: saving flows and cellprob to tiff")
+    if len(parent.flows) > 0:
+        imsave(base + "_cp_cellprob.tif", parent.flows[1])
+        for i in range(3):
+            imsave(base + f"_cp_flows_{i}.tif", parent.flows[0][..., i])
+        if len(parent.flows) > 2:
+            imsave(base + "_cp_flows.tif", parent.flows[2])
+        print("GUI_INFO: saved flows and cellprob")
+    else:
+        print("ERROR: no flows or cellprob found")
+
+
+def _save_rois(parent):
+    """ save masks as rois in .zip file for ImageJ """
+    filename = parent.filename
+    if parent.NZ == 1:
+        print(
+            f"GUI_INFO: saving {parent.cellpix[0].max()} ImageJ ROIs to .zip archive.")
+        save_rois(parent.cellpix[0], parent.filename)
+    else:
+        print("ERROR: cannot save 3D outlines")
+
+
+def _save_outlines(parent):
+    filename = parent.filename
+    base = os.path.splitext(filename)[0]
+    if parent.NZ == 1:
+        print(
+            "GUI_INFO: saving 2D outlines to text file, see docs for info to load into ImageJ"
+        )
+        outlines = outlines_list(parent.cellpix[0])
+        outlines_to_text(base, outlines)
+    else:
+        print("ERROR: cannot save 3D outlines")
+
+
+def _save_sets_with_check(parent):
+    """ Save masks and update *_seg.npy file. Use this function when saving should be optional
+     based on the disableAutosave checkbox. Otherwise, use _save_sets """
+    if not parent.disableAutosave.isChecked():
+        _save_sets(parent)
+
+
+def _save_sets(parent):
+    """ save masks to *_seg.npy. This function should be used when saving
+    is forced, e.g. when clicking the save button. Otherwise, use _save_sets_with_check
+    """
+    filename = parent.filename
+    base = os.path.splitext(filename)[0]
+    flow_threshold = parent.segmentation_settings.flow_threshold
+    cellprob_threshold = parent.segmentation_settings.cellprob_threshold
+
+    if parent.NZ > 1:
+        dat = {
+            "outlines":
+                parent.outpix,
+            "colors":
+                parent.cellcolors[1:],
+            "masks":
+                parent.cellpix,
+            "current_channel": (parent.color - 2) % 5,
+            "filename":
+                parent.filename,
+            "flows":
+                parent.flows,
+            "zdraw":
+                parent.zdraw,
+            "model_path":
+                parent.current_model_path
+                if hasattr(parent, "current_model_path") else 0,
+            "flow_threshold":
+                flow_threshold,
+            "cellprob_threshold":
+                cellprob_threshold,
+            "normalize_params":
+                parent.get_normalize_params(),
+            "restore":
+                parent.restore,
+            "ratio":
+                parent.ratio,
+            "diameter":
+                parent.segmentation_settings.diameter
+        }
+        if parent.restore is not None:
+            dat["img_restore"] = parent.stack_filtered
+    else:
+        dat = {
+            "outlines":
+                parent.outpix.squeeze() if parent.restore is None or
+                not "upsample" in parent.restore else parent.outpix_resize.squeeze(),
+            "colors":
+                parent.cellcolors[1:],
+            "masks":
+                parent.cellpix.squeeze() if parent.restore is None or
+                not "upsample" in parent.restore else parent.cellpix_resize.squeeze(),
+            "filename":
+                parent.filename,
+            "flows":
+                parent.flows,
+            "ismanual":
+                parent.ismanual,
+            "manual_changes":
+                parent.track_changes,
+            "model_path":
+                parent.current_model_path
+                if hasattr(parent, "current_model_path") else 0,
+            "flow_threshold":
+                flow_threshold,
+            "cellprob_threshold":
+                cellprob_threshold,
+            "normalize_params":
+                parent.get_normalize_params(),
+            "restore":
+                parent.restore,
+            "ratio":
+                parent.ratio,
+            "diameter":
+                parent.segmentation_settings.diameter
+        }
+        if parent.restore is not None:
+            dat["img_restore"] = parent.stack_filtered
+    try:
+        np.save(base + "_seg.npy", dat)
+        print("GUI_INFO: %d ROIs saved to %s" % (parent.ncells.get(), base + "_seg.npy"))
+    except Exception as e:
+        print(f"ERROR: {e}")
+    del dat

models/seg_post_model/cellpose/gui/make_train.py

@@ -0,0 +1,107 @@
+import os, argparse
+import numpy as np
+from cellpose import io, transforms
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Make slices of XYZ image data for training. Assumes image is ZXYC unless specified otherwise using --channel_axis and --z_axis')
+
+    input_img_args = parser.add_argument_group("input image arguments")
+    input_img_args.add_argument('--dir', default=[], type=str,
+                                help='folder containing data to run or train on.')
+    input_img_args.add_argument(
+        '--image_path', default=[], type=str, help=
+        'if given and --dir not given, run on single image instead of folder (cannot train with this option)'
+    )
+    input_img_args.add_argument(
+        '--look_one_level_down', action='store_true',
+        help='run processing on all subdirectories of current folder')
+    input_img_args.add_argument('--img_filter', default=[], type=str,
+                                help='end string for images to run on')
+    input_img_args.add_argument(
+        '--channel_axis', default=-1, type=int,
+        help='axis of image which corresponds to image channels')
+    input_img_args.add_argument('--z_axis', default=0, type=int,
+                                help='axis of image which corresponds to Z dimension')
+    input_img_args.add_argument(
+        '--chan', default=0, type=int, help=
+        'Deprecated')
+    input_img_args.add_argument(
+        '--chan2', default=0, type=int, help=
+        'Deprecated'
+    )
+    input_img_args.add_argument('--invert', action='store_true',
+                                help='invert grayscale channel')
+    input_img_args.add_argument(
+        '--all_channels', action='store_true', help=
+        'deprecated')
+    input_img_args.add_argument("--anisotropy", required=False, default=1.0, type=float,
+                                help="anisotropy of volume in 3D")
+    
+
+    # algorithm settings
+    algorithm_args = parser.add_argument_group("algorithm arguments")
+    algorithm_args.add_argument('--sharpen_radius', required=False, default=0.0,
+                                type=float, help='high-pass filtering radius. Default: %(default)s')
+    algorithm_args.add_argument('--tile_norm', required=False, default=0, type=int,
+                                help='tile normalization block size. Default: %(default)s')
+    algorithm_args.add_argument('--nimg_per_tif', required=False, default=10, type=int,
+                                help='number of crops in XY to save per tiff. Default: %(default)s')
+    algorithm_args.add_argument('--crop_size', required=False, default=512, type=int,
+                                help='size of random crop to save. Default: %(default)s')
+
+    args = parser.parse_args()
+
+    # find images
+    if len(args.img_filter) > 0:
+        imf = args.img_filter
+    else:
+        imf = None
+
+    if len(args.dir) > 0:
+        image_names = io.get_image_files(args.dir, "_masks", imf=imf,
+                                     look_one_level_down=args.look_one_level_down)
+        dirname = args.dir
+    else:
+        if os.path.exists(args.image_path):
+            image_names = [args.image_path]
+            dirname = os.path.split(args.image_path)[0]
+        else:
+            raise ValueError(f"ERROR: no file found at {args.image_path}")
+        
+    np.random.seed(0)
+    nimg_per_tif = args.nimg_per_tif
+    crop_size = args.crop_size
+    os.makedirs(os.path.join(dirname, 'train/'), exist_ok=True)
+    pm = [(0, 1, 2, 3), (2, 0, 1, 3), (1, 0, 2, 3)]
+    npm = ["YX", "ZY", "ZX"]
+    for name in image_names:
+        name0 = os.path.splitext(os.path.split(name)[-1])[0]
+        img0 = io.imread_3D(name)
+        try: 
+            img0 = transforms.convert_image(img0, channel_axis=args.channel_axis, 
+                                            z_axis=args.z_axis, do_3D=True)
+        except ValueError:
+            print('Error converting image. Did you provide the correct --channel_axis and --z_axis ?') 
+            
+        for p in range(3):
+            img = img0.transpose(pm[p]).copy()
+            print(npm[p], img[0].shape)
+            Ly, Lx = img.shape[1:3]
+            imgs = img[np.random.permutation(img.shape[0])[:args.nimg_per_tif]]
+            if args.anisotropy > 1.0 and p > 0:
+                imgs = transforms.resize_image(imgs, Ly=int(args.anisotropy * Ly), Lx=Lx)
+            for k, img in enumerate(imgs):
+                if args.tile_norm:
+                    img = transforms.normalize99_tile(img, blocksize=args.tile_norm)
+                if args.sharpen_radius:
+                    img = transforms.smooth_sharpen_img(img,
+                                                        sharpen_radius=args.sharpen_radius)
+                ly = 0 if Ly - crop_size <= 0 else np.random.randint(0, Ly - crop_size)
+                lx = 0 if Lx - crop_size <= 0 else np.random.randint(0, Lx - crop_size)
+                io.imsave(os.path.join(dirname, f'train/{name0}_{npm[p]}_{k}.tif'),
+                        img[ly:ly + args.crop_size, lx:lx + args.crop_size].squeeze())
+
+
+if __name__ == '__main__':
+    main()

models/seg_post_model/cellpose/gui/menus.py

@@ -0,0 +1,145 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+from qtpy.QtWidgets import QAction
+from . import io
+
+
+def mainmenu(parent):
+    main_menu = parent.menuBar()
+    file_menu = main_menu.addMenu("&File")
+    # load processed data
+    loadImg = QAction("&Load image (*.tif, *.png, *.jpg)", parent)
+    loadImg.setShortcut("Ctrl+L")
+    loadImg.triggered.connect(lambda: io._load_image(parent))
+    file_menu.addAction(loadImg)
+
+    parent.autoloadMasks = QAction("Autoload masks from _masks.tif file", parent,
+                                   checkable=True)
+    parent.autoloadMasks.setChecked(False)
+    file_menu.addAction(parent.autoloadMasks)
+
+    parent.disableAutosave = QAction("Disable autosave _seg.npy file", parent,
+                                     checkable=True)
+    parent.disableAutosave.setChecked(False)
+    file_menu.addAction(parent.disableAutosave)
+
+    parent.loadMasks = QAction("Load &masks (*.tif, *.png, *.jpg)", parent)
+    parent.loadMasks.setShortcut("Ctrl+M")
+    parent.loadMasks.triggered.connect(lambda: io._load_masks(parent))
+    file_menu.addAction(parent.loadMasks)
+    parent.loadMasks.setEnabled(False)
+
+    loadManual = QAction("Load &processed/labelled image (*_seg.npy)", parent)
+    loadManual.setShortcut("Ctrl+P")
+    loadManual.triggered.connect(lambda: io._load_seg(parent))
+    file_menu.addAction(loadManual)
+
+    parent.saveSet = QAction("&Save masks and image (as *_seg.npy)", parent)
+    parent.saveSet.setShortcut("Ctrl+S")
+    parent.saveSet.triggered.connect(lambda: io._save_sets(parent))
+    file_menu.addAction(parent.saveSet)
+    parent.saveSet.setEnabled(False)
+
+    parent.savePNG = QAction("Save masks as P&NG/tif", parent)
+    parent.savePNG.setShortcut("Ctrl+N")
+    parent.savePNG.triggered.connect(lambda: io._save_png(parent))
+    file_menu.addAction(parent.savePNG)
+    parent.savePNG.setEnabled(False)
+
+    parent.saveOutlines = QAction("Save &Outlines as text for imageJ", parent)
+    parent.saveOutlines.setShortcut("Ctrl+O")
+    parent.saveOutlines.triggered.connect(lambda: io._save_outlines(parent))
+    file_menu.addAction(parent.saveOutlines)
+    parent.saveOutlines.setEnabled(False)
+
+    parent.saveROIs = QAction("Save outlines as .zip archive of &ROI files for ImageJ",
+                              parent)
+    parent.saveROIs.setShortcut("Ctrl+R")
+    parent.saveROIs.triggered.connect(lambda: io._save_rois(parent))
+    file_menu.addAction(parent.saveROIs)
+    parent.saveROIs.setEnabled(False)
+
+    parent.saveFlows = QAction("Save &Flows and cellprob as tif", parent)
+    parent.saveFlows.setShortcut("Ctrl+F")
+    parent.saveFlows.triggered.connect(lambda: io._save_flows(parent))
+    file_menu.addAction(parent.saveFlows)
+    parent.saveFlows.setEnabled(False)
+
+
+def editmenu(parent):
+    main_menu = parent.menuBar()
+    edit_menu = main_menu.addMenu("&Edit")
+    parent.undo = QAction("Undo previous mask/trace", parent)
+    parent.undo.setShortcut("Ctrl+Z")
+    parent.undo.triggered.connect(parent.undo_action)
+    parent.undo.setEnabled(False)
+    edit_menu.addAction(parent.undo)
+
+    parent.redo = QAction("Undo remove mask", parent)
+    parent.redo.setShortcut("Ctrl+Y")
+    parent.redo.triggered.connect(parent.undo_remove_action)
+    parent.redo.setEnabled(False)
+    edit_menu.addAction(parent.redo)
+
+    parent.ClearButton = QAction("Clear all masks", parent)
+    parent.ClearButton.setShortcut("Ctrl+0")
+    parent.ClearButton.triggered.connect(parent.clear_all)
+    parent.ClearButton.setEnabled(False)
+    edit_menu.addAction(parent.ClearButton)
+
+    parent.remcell = QAction("Remove selected cell (Ctrl+CLICK)", parent)
+    parent.remcell.setShortcut("Ctrl+Click")
+    parent.remcell.triggered.connect(parent.remove_action)
+    parent.remcell.setEnabled(False)
+    edit_menu.addAction(parent.remcell)
+
+    parent.mergecell = QAction("FYI: Merge cells by Alt+Click", parent)
+    parent.mergecell.setEnabled(False)
+    edit_menu.addAction(parent.mergecell)
+
+
+def modelmenu(parent):
+    main_menu = parent.menuBar()
+    io._init_model_list(parent)
+    model_menu = main_menu.addMenu("&Models")
+    parent.addmodel = QAction("Add custom torch model to GUI", parent)
+    #parent.addmodel.setShortcut("Ctrl+A")
+    parent.addmodel.triggered.connect(parent.add_model)
+    parent.addmodel.setEnabled(True)
+    model_menu.addAction(parent.addmodel)
+
+    parent.removemodel = QAction("Remove selected custom model from GUI", parent)
+    #parent.removemodel.setShortcut("Ctrl+R")
+    parent.removemodel.triggered.connect(parent.remove_model)
+    parent.removemodel.setEnabled(True)
+    model_menu.addAction(parent.removemodel)
+
+    parent.newmodel = QAction("&Train new model with image+masks in folder", parent)
+    parent.newmodel.setShortcut("Ctrl+T")
+    parent.newmodel.triggered.connect(parent.new_model)
+    parent.newmodel.setEnabled(False)
+    model_menu.addAction(parent.newmodel)
+
+    openTrainHelp = QAction("Training instructions", parent)
+    openTrainHelp.triggered.connect(parent.train_help_window)
+    model_menu.addAction(openTrainHelp)
+
+
+def helpmenu(parent):
+    main_menu = parent.menuBar()
+    help_menu = main_menu.addMenu("&Help")
+
+    openHelp = QAction("&Help with GUI", parent)
+    openHelp.setShortcut("Ctrl+H")
+    openHelp.triggered.connect(parent.help_window)
+    help_menu.addAction(openHelp)
+
+    openGUI = QAction("&GUI layout", parent)
+    openGUI.setShortcut("Ctrl+G")
+    openGUI.triggered.connect(parent.gui_window)
+    help_menu.addAction(openGUI)
+
+    openTrainHelp = QAction("Training instructions", parent)
+    openTrainHelp.triggered.connect(parent.train_help_window)
+    help_menu.addAction(openTrainHelp)

models/seg_post_model/cellpose/io.py

@@ -0,0 +1,816 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os, warnings, glob, shutil
+from natsort import natsorted
+import numpy as np
+import cv2
+import tifffile
+import logging, pathlib, sys
+from tqdm import tqdm
+from pathlib import Path
+import re
+from .version import version_str
+from roifile import ImagejRoi, roiwrite
+
+try:
+    from qtpy import QtGui, QtCore, Qt, QtWidgets
+    from qtpy.QtWidgets import QMessageBox
+    GUI = True
+except:
+    GUI = False
+
+try:
+    import matplotlib.pyplot as plt
+    MATPLOTLIB = True
+except:
+    MATPLOTLIB = False
+
+try:
+    import nd2
+    ND2 = True
+except:
+    ND2 = False
+
+try:
+    import nrrd
+    NRRD = True
+except:
+    NRRD = False
+
+try:
+    from google.cloud import storage
+    SERVER_UPLOAD = True
+except:
+    SERVER_UPLOAD = False
+
+io_logger = logging.getLogger(__name__)
+
+def logger_setup(cp_path=".cellpose", logfile_name="run.log", stdout_file_replacement=None):
+    cp_dir = pathlib.Path.home().joinpath(cp_path)
+    cp_dir.mkdir(exist_ok=True)
+    log_file = cp_dir.joinpath(logfile_name)
+    try:
+        log_file.unlink()
+    except:
+        print('creating new log file')
+    handlers = [logging.FileHandler(log_file),]
+    if stdout_file_replacement is not None:
+        handlers.append(logging.FileHandler(stdout_file_replacement))
+    else:
+        handlers.append(logging.StreamHandler(sys.stdout))
+    logging.basicConfig(
+                    level=logging.INFO,
+                    format="%(asctime)s [%(levelname)s] %(message)s",
+                    handlers=handlers,
+                    force=True
+    )
+    logger = logging.getLogger(__name__)
+    logger.info(f"WRITING LOG OUTPUT TO {log_file}")
+    logger.info(version_str)
+
+    return logger, log_file
+
+
+from . import utils, plot, transforms
+
+# helper function to check for a path; if it doesn't exist, make it
+def check_dir(path):
+    if not os.path.isdir(path):
+        os.mkdir(path)
+
+
+def outlines_to_text(base, outlines):
+    with open(base + "_cp_outlines.txt", "w") as f:
+        for o in outlines:
+            xy = list(o.flatten())
+            xy_str = ",".join(map(str, xy))
+            f.write(xy_str)
+            f.write("\n")
+
+
+def load_dax(filename):
+    ### modified from ZhuangLab github:
+    ### https://github.com/ZhuangLab/storm-analysis/blob/71ae493cbd17ddb97938d0ae2032d97a0eaa76b2/storm_analysis/sa_library/datareader.py#L156
+
+    inf_filename = os.path.splitext(filename)[0] + ".inf"
+    if not os.path.exists(inf_filename):
+        io_logger.critical(
+            f"ERROR: no inf file found for dax file {filename}, cannot load dax without it"
+        )
+        return None
+
+    ### get metadata
+    image_height, image_width = None, None
+    # extract the movie information from the associated inf file
+    size_re = re.compile(r"frame dimensions = ([\d]+) x ([\d]+)")
+    length_re = re.compile(r"number of frames = ([\d]+)")
+    endian_re = re.compile(r" (big|little) endian")
+
+    with open(inf_filename, "r") as inf_file:
+        lines = inf_file.read().split("\n")
+        for line in lines:
+            m = size_re.match(line)
+            if m:
+                image_height = int(m.group(2))
+                image_width = int(m.group(1))
+            m = length_re.match(line)
+            if m:
+                number_frames = int(m.group(1))
+            m = endian_re.search(line)
+            if m:
+                if m.group(1) == "big":
+                    bigendian = 1
+                else:
+                    bigendian = 0
+    # set defaults, warn the user that they couldn"t be determined from the inf file.
+    if not image_height:
+        io_logger.warning("could not determine dax image size, assuming 256x256")
+        image_height = 256
+        image_width = 256
+
+    ### load image
+    img = np.memmap(filename, dtype="uint16",
+                    shape=(number_frames, image_height, image_width))
+    if bigendian:
+        img = img.byteswap()
+    img = np.array(img)
+
+    return img
+
+
+def imread(filename):
+    """
+    Read in an image file with tif or image file type supported by cv2.
+
+    Args:
+        filename (str): The path to the image file.
+
+    Returns:
+        numpy.ndarray: The image data as a NumPy array.
+
+    Raises:
+        None
+
+    Raises an error if the image file format is not supported.
+
+    Examples:
+        >>> img = imread("image.tif")
+    """
+    # ensure that extension check is not case sensitive
+    ext = os.path.splitext(filename)[-1].lower()
+    if ext == ".tif" or ext == ".tiff" or ext == ".flex":
+        with tifffile.TiffFile(filename) as tif:
+            ltif = len(tif.pages)
+            try:
+                full_shape = tif.shaped_metadata[0]["shape"]
+            except:
+                try:
+                    page = tif.series[0][0]
+                    full_shape = tif.series[0].shape
+                except:
+                    ltif = 0
+            if ltif < 10:
+                img = tif.asarray()
+            else:
+                page = tif.series[0][0]
+                shape, dtype = page.shape, page.dtype
+                ltif = int(np.prod(full_shape) / np.prod(shape))
+                io_logger.info(f"reading tiff with {ltif} planes")
+                img = np.zeros((ltif, *shape), dtype=dtype)
+                for i, page in enumerate(tqdm(tif.series[0])):
+                    img[i] = page.asarray()
+                img = img.reshape(full_shape)
+        return img
+    elif ext == ".dax":
+        img = load_dax(filename)
+        return img
+    elif ext == ".nd2":
+        if not ND2:
+            io_logger.critical("ERROR: need to 'pip install nd2' to load in .nd2 file")
+            return None
+    elif ext == ".nrrd":
+        if not NRRD:
+            io_logger.critical(
+                "ERROR: need to 'pip install pynrrd' to load in .nrrd file")
+            return None
+        else:
+            img, metadata = nrrd.read(filename)
+            if img.ndim == 3:
+                img = img.transpose(2, 0, 1)
+            return img
+    elif ext != ".npy":
+        try:
+            img = cv2.imread(filename, -1)  #cv2.LOAD_IMAGE_ANYDEPTH)
+            if img.ndim > 2:
+                img = img[..., [2, 1, 0]]
+            return img
+        except Exception as e:
+            io_logger.critical("ERROR: could not read file, %s" % e)
+            return None
+    else:
+        try:
+            dat = np.load(filename, allow_pickle=True).item()
+            masks = dat["masks"]
+            return masks
+        except Exception as e:
+            io_logger.critical("ERROR: could not read masks from file, %s" % e)
+            return None
+
+
+def imread_2D(img_file):
+    """
+    Read in a 2D image file and convert it to a 3-channel image. Attempts to do this for multi-channel and grayscale images.
+    If the image has more than 3 channels, only the first 3 channels are kept.
+    
+    Args:
+        img_file (str): The path to the image file.
+
+    Returns:
+        img_out (numpy.ndarray): The 3-channel image data as a NumPy array.
+    """
+    img = imread(img_file)
+    return transforms.convert_image(img, do_3D=False)
+
+
+def imread_3D(img_file):
+    """
+    Read in a 3D image file and convert it to have a channel axis last automatically. Attempts to do this for multi-channel and grayscale images.
+
+    If multichannel image, the channel axis is assumed to be the smallest dimension, and the z axis is the next smallest dimension. 
+    Use `cellpose.io.imread()` to load the full image without selecting the z and channel axes. 
+    
+    Args:
+        img_file (str): The path to the image file.
+
+    Returns:
+        img_out (numpy.ndarray): The image data as a NumPy array.
+    """
+    img = imread(img_file)
+
+    dimension_lengths = list(img.shape)
+
+    # grayscale images:
+    if img.ndim == 3:
+        channel_axis = None
+        # guess at z axis:
+        z_axis = np.argmin(dimension_lengths)
+
+    elif img.ndim == 4:
+        # guess at channel axis:
+        channel_axis = np.argmin(dimension_lengths)
+
+        # guess at z axis: 
+        # set channel axis to max so argmin works:
+        dimension_lengths[channel_axis] = max(dimension_lengths)
+        z_axis = np.argmin(dimension_lengths)
+
+    else: 
+        raise ValueError(f'image shape error, 3D image must 3 or 4 dimensional. Number of dimensions: {img.ndim}')
+    
+    try:
+        return transforms.convert_image(img, channel_axis=channel_axis, z_axis=z_axis, do_3D=True)
+    except Exception as e:
+        io_logger.critical("ERROR: could not read file, %s" % e)
+        io_logger.critical("ERROR: Guessed z_axis: %s, channel_axis: %s" % (z_axis, channel_axis))
+        return None
+
+def remove_model(filename, delete=False):
+    """ remove model from .cellpose custom model list """
+    filename = os.path.split(filename)[-1]
+    from . import models
+    model_strings = models.get_user_models()
+    if len(model_strings) > 0:
+        with open(models.MODEL_LIST_PATH, "w") as textfile:
+            for fname in model_strings:
+                textfile.write(fname + "\n")
+    else:
+        # write empty file
+        textfile = open(models.MODEL_LIST_PATH, "w")
+        textfile.close()
+    print(f"{filename} removed from custom model list")
+    if delete:
+        os.remove(os.fspath(models.MODEL_DIR.joinpath(fname)))
+        print("model deleted")
+
+
+def add_model(filename):
+    """ add model to .cellpose models folder to use with GUI or CLI """
+    from . import models
+    fname = os.path.split(filename)[-1]
+    try:
+        shutil.copyfile(filename, os.fspath(models.MODEL_DIR.joinpath(fname)))
+    except shutil.SameFileError:
+        pass
+    print(f"{filename} copied to models folder {os.fspath(models.MODEL_DIR)}")
+    if fname not in models.get_user_models():
+        with open(models.MODEL_LIST_PATH, "a") as textfile:
+            textfile.write(fname + "\n")
+
+
+def imsave(filename, arr):
+    """
+    Saves an image array to a file.
+
+    Args:
+        filename (str): The name of the file to save the image to.
+        arr (numpy.ndarray): The image array to be saved.
+
+    Returns:
+        None
+    """
+    ext = os.path.splitext(filename)[-1].lower()
+    if ext == ".tif" or ext == ".tiff":
+        tifffile.imwrite(filename, data=arr, compression="zlib")
+    else:
+        if len(arr.shape) > 2:
+            arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
+        cv2.imwrite(filename, arr)
+
+
+def get_image_files(folder, mask_filter, imf=None, look_one_level_down=False):
+    """
+    Finds all images in a folder and its subfolders (if specified) with the given file extensions.
+
+    Args:
+        folder (str): The path to the folder to search for images.
+        mask_filter (str): The filter for mask files.
+        imf (str, optional): The additional filter for image files. Defaults to None.
+        look_one_level_down (bool, optional): Whether to search for images in subfolders. Defaults to False.
+
+    Returns:
+        list: A list of image file paths.
+
+    Raises:
+        ValueError: If no files are found in the specified folder.
+        ValueError: If no images are found in the specified folder with the supported file extensions.
+        ValueError: If no images are found in the specified folder without the mask or flow file endings.
+    """
+    mask_filters = ["_cp_output", "_flows", "_flows_0", "_flows_1",
+                    "_flows_2", "_cellprob", "_masks", mask_filter]
+    image_names = []
+    if imf is None:
+        imf = ""
+
+    folders = []
+    if look_one_level_down:
+        folders = natsorted(glob.glob(os.path.join(folder, "*/")))
+    folders.append(folder)
+    exts = [".png", ".jpg", ".jpeg", ".tif", ".tiff", ".flex", ".dax", ".nd2", ".nrrd"]
+    l0 = 0
+    al = 0
+    for folder in folders:
+        all_files = glob.glob(folder + "/*")
+        al += len(all_files)
+        for ext in exts:
+            image_names.extend(glob.glob(folder + f"/*{imf}{ext}"))
+            image_names.extend(glob.glob(folder + f"/*{imf}{ext.upper()}"))
+        l0 += len(image_names)
+
+    # return error if no files found
+    if al == 0:
+        raise ValueError("ERROR: no files in --dir folder ")
+    elif l0 == 0:
+        raise ValueError(
+            "ERROR: no images in --dir folder with extensions .png, .jpg, .jpeg, .tif, .tiff, .flex"
+        )
+
+    image_names = natsorted(image_names)
+    imn = []
+    for im in image_names:
+        imfile = os.path.splitext(im)[0]
+        igood = all([(len(imfile) > len(mask_filter) and
+                      imfile[-len(mask_filter):] != mask_filter) or
+                     len(imfile) <= len(mask_filter) for mask_filter in mask_filters])
+        if len(imf) > 0:
+            igood &= imfile[-len(imf):] == imf
+        if igood:
+            imn.append(im)
+
+    image_names = imn
+
+    # remove duplicates
+    image_names = [*set(image_names)]
+    image_names = natsorted(image_names)
+
+    if len(image_names) == 0:
+        raise ValueError(
+            "ERROR: no images in --dir folder without _masks or _flows or _cellprob ending")
+
+    return image_names
+
+def get_label_files(image_names, mask_filter, imf=None):
+    """
+    Get the label files corresponding to the given image names and mask filter.
+
+    Args:
+        image_names (list): List of image names.
+        mask_filter (str): Mask filter to be applied.
+        imf (str, optional): Image file extension. Defaults to None.
+
+    Returns:
+        tuple: A tuple containing the label file names and flow file names (if present).
+    """
+    nimg = len(image_names)
+    label_names0 = [os.path.splitext(image_names[n])[0] for n in range(nimg)]
+
+    if imf is not None and len(imf) > 0:
+        label_names = [label_names0[n][:-len(imf)] for n in range(nimg)]
+    else:
+        label_names = label_names0
+
+    # check for flows
+    if os.path.exists(label_names0[0] + "_flows.tif"):
+        flow_names = [label_names0[n] + "_flows.tif" for n in range(nimg)]
+    else:
+        flow_names = [label_names[n] + "_flows.tif" for n in range(nimg)]
+    if not all([os.path.exists(flow) for flow in flow_names]):
+        io_logger.info(
+            "not all flows are present, running flow generation for all images")
+        flow_names = None
+
+    # check for masks
+    if mask_filter == "_seg.npy":
+        label_names = [label_names[n] + mask_filter for n in range(nimg)]
+        return label_names, None
+
+    if os.path.exists(label_names[0] + mask_filter + ".tif"):
+        label_names = [label_names[n] + mask_filter + ".tif" for n in range(nimg)]
+    elif os.path.exists(label_names[0] + mask_filter + ".tiff"):
+        label_names = [label_names[n] + mask_filter + ".tiff" for n in range(nimg)]
+    elif os.path.exists(label_names[0] + mask_filter + ".png"):
+        label_names = [label_names[n] + mask_filter + ".png" for n in range(nimg)]
+    # TODO, allow _seg.npy
+    #elif os.path.exists(label_names[0] + "_seg.npy"):
+    #    io_logger.info("labels found as _seg.npy files, converting to tif")
+    else:
+        if not flow_names:
+            raise ValueError("labels not provided with correct --mask_filter")
+        else:
+            label_names = None
+    if not all([os.path.exists(label) for label in label_names]):
+        if not flow_names:
+            raise ValueError(
+                "labels not provided for all images in train and/or test set")
+        else:
+            label_names = None
+
+    return label_names, flow_names
+
+
+def load_images_labels(tdir, mask_filter="_masks", image_filter=None,
+                       look_one_level_down=False):
+    """
+    Loads images and corresponding labels from a directory.
+
+    Args:
+        tdir (str): The directory path.
+        mask_filter (str, optional): The filter for mask files. Defaults to "_masks".
+        image_filter (str, optional): The filter for image files. Defaults to None.
+        look_one_level_down (bool, optional): Whether to look for files one level down. Defaults to False.
+
+    Returns:
+        tuple: A tuple containing a list of images, a list of labels, and a list of image names.
+    """
+    image_names = get_image_files(tdir, mask_filter, image_filter, look_one_level_down)
+    nimg = len(image_names)
+
+    # training data
+    label_names, flow_names = get_label_files(image_names, mask_filter,
+                                              imf=image_filter)
+
+    images = []
+    labels = []
+    k = 0
+    for n in range(nimg):
+        if (os.path.isfile(label_names[n]) or
+            (flow_names is not None and os.path.isfile(flow_names[0]))):
+            image = imread(image_names[n])
+            if label_names is not None:
+                label = imread(label_names[n])
+            if flow_names is not None:
+                flow = imread(flow_names[n])
+                if flow.shape[0] < 4:
+                    label = np.concatenate((label[np.newaxis, :, :], flow), axis=0)
+                else:
+                    label = flow
+            images.append(image)
+            labels.append(label)
+            k += 1
+    io_logger.info(f"{k} / {nimg} images in {tdir} folder have labels")
+    return images, labels, image_names
+
+def load_train_test_data(train_dir, test_dir=None, image_filter=None,
+                         mask_filter="_masks", look_one_level_down=False):
+    """
+    Loads training and testing data for a Cellpose model.
+
+    Args:
+        train_dir (str): The directory path containing the training data.
+        test_dir (str, optional): The directory path containing the testing data. Defaults to None.
+        image_filter (str, optional): The filter for selecting image files. Defaults to None.
+        mask_filter (str, optional): The filter for selecting mask files. Defaults to "_masks".
+        look_one_level_down (bool, optional): Whether to look for data in subdirectories of train_dir and test_dir. Defaults to False.
+
+    Returns:
+        images, labels, image_names, test_images, test_labels, test_image_names
+
+    """
+    images, labels, image_names = load_images_labels(train_dir, mask_filter,
+                                                     image_filter, look_one_level_down)
+    # testing data
+    test_images, test_labels, test_image_names = None, None, None
+    if test_dir is not None:
+        test_images, test_labels, test_image_names = load_images_labels(
+            test_dir, mask_filter, image_filter, look_one_level_down)
+
+    return images, labels, image_names, test_images, test_labels, test_image_names
+
+
+def masks_flows_to_seg(images, masks, flows, file_names, 
+                       channels=None,
+                       imgs_restore=None, restore_type=None, ratio=1.):
+    """Save output of model eval to be loaded in GUI.
+
+    Can be list output (run on multiple images) or single output (run on single image).
+
+    Saved to file_names[k]+"_seg.npy".
+
+    Args:
+        images (list): Images input into cellpose.
+        masks (list): Masks output from Cellpose.eval, where 0=NO masks; 1,2,...=mask labels.
+        flows (list): Flows output from Cellpose.eval.
+        file_names (list, str): Names of files of images.
+        diams (float array): Diameters used to run Cellpose. Defaults to 30. TODO: remove this
+        channels (list, int, optional): Channels used to run Cellpose. Defaults to None.
+
+    Returns:
+        None
+    """
+
+    if channels is None:
+        channels = [0, 0]
+
+    if isinstance(masks, list):
+        if imgs_restore is None:
+            imgs_restore = [None] * len(masks)
+        if isinstance(file_names, str):
+            file_names = [file_names] * len(masks)
+        for k, [image, mask, flow, 
+                # diam, 
+                file_name, img_restore
+               ] in enumerate(zip(images, masks, flows, 
+                                #   diams, 
+                                  file_names,
+                                  imgs_restore)):
+            channels_img = channels
+            if channels_img is not None and len(channels) > 2:
+                channels_img = channels[k]
+            masks_flows_to_seg(image, mask, flow, file_name, 
+                            #    diams=diam,
+                               channels=channels_img, imgs_restore=img_restore,
+                               restore_type=restore_type, ratio=ratio)
+        return
+
+    if len(channels) == 1:
+        channels = channels[0]
+
+    flowi = []
+    if flows[0].ndim == 3:
+        Ly, Lx = masks.shape[-2:]
+        flowi.append(
+            cv2.resize(flows[0], (Lx, Ly), interpolation=cv2.INTER_NEAREST)[np.newaxis,
+                                                                            ...])
+    else:
+        flowi.append(flows[0])
+
+    if flows[0].ndim == 3:
+        cellprob = (np.clip(transforms.normalize99(flows[2]), 0, 1) * 255).astype(
+            np.uint8)
+        cellprob = cv2.resize(cellprob, (Lx, Ly), interpolation=cv2.INTER_NEAREST)
+        flowi.append(cellprob[np.newaxis, ...])
+        flowi.append(np.zeros(flows[0].shape, dtype=np.uint8))
+        flowi[-1] = flowi[-1][np.newaxis, ...]
+    else:
+        flowi.append(
+            (np.clip(transforms.normalize99(flows[2]), 0, 1) * 255).astype(np.uint8))
+        flowi.append((flows[1][0] / 10 * 127 + 127).astype(np.uint8))
+    if len(flows) > 2:
+        if len(flows) > 3:
+            flowi.append(flows[3])
+        else:
+            flowi.append([])
+        flowi.append(np.concatenate((flows[1], flows[2][np.newaxis, ...]), axis=0))
+    outlines = masks * utils.masks_to_outlines(masks)
+    base = os.path.splitext(file_names)[0]
+
+    dat = {
+        "outlines":
+            outlines.astype(np.uint16) if outlines.max() < 2**16 -
+            1 else outlines.astype(np.uint32),
+        "masks":
+            masks.astype(np.uint16) if outlines.max() < 2**16 -
+            1 else masks.astype(np.uint32),
+        "chan_choose":
+            channels,
+        "ismanual":
+            np.zeros(masks.max(), bool),
+        "filename":
+            file_names,
+        "flows":
+            flowi,
+        "diameter":
+            np.nan
+    }
+    if restore_type is not None and imgs_restore is not None:
+        dat["restore"] = restore_type
+        dat["ratio"] = ratio
+        dat["img_restore"] = imgs_restore
+
+    np.save(base + "_seg.npy", dat)
+
+def save_to_png(images, masks, flows, file_names):
+    """ deprecated (runs io.save_masks with png=True)
+
+        does not work for 3D images
+
+    """
+    save_masks(images, masks, flows, file_names, png=True)
+
+
+def save_rois(masks, file_name, multiprocessing=None):
+    """ save masks to .roi files in .zip archive for ImageJ/Fiji
+
+    Args:
+        masks (np.ndarray): masks output from Cellpose.eval, where 0=NO masks; 1,2,...=mask labels
+        file_name (str): name to save the .zip file to
+
+    Returns:
+        None
+    """
+    outlines = utils.outlines_list(masks, multiprocessing=multiprocessing)
+    nonempty_outlines = [outline for outline in outlines if len(outline)!=0]
+    if len(outlines)!=len(nonempty_outlines):
+        print(f"empty outlines found, saving {len(nonempty_outlines)} ImageJ ROIs to .zip archive.")
+    rois = [ImagejRoi.frompoints(outline) for outline in nonempty_outlines]
+    file_name = os.path.splitext(file_name)[0] + '_rois.zip'
+
+
+    # Delete file if it exists; the roifile lib appends to existing zip files.
+    # If the user removed a mask it will still be in the zip file
+    if os.path.exists(file_name):
+        os.remove(file_name)
+
+    roiwrite(file_name, rois)
+
+
+def save_masks(images, masks, flows, file_names, png=True, tif=False, channels=[0, 0],
+               suffix="_cp_masks", save_flows=False, save_outlines=False, dir_above=False,
+               in_folders=False, savedir=None, save_txt=False, save_mpl=False):
+    """ Save masks + nicely plotted segmentation image to png and/or tiff.
+
+    Can save masks, flows to different directories, if in_folders is True.
+
+    If png, masks[k] for images[k] are saved to file_names[k]+"_cp_masks.png".
+
+    If tif, masks[k] for images[k] are saved to file_names[k]+"_cp_masks.tif".
+
+    If png and matplotlib installed, full segmentation figure is saved to file_names[k]+"_cp.png".
+
+    Only tif option works for 3D data, and only tif option works for empty masks.
+
+    Args:
+        images (list): Images input into cellpose.
+        masks (list): Masks output from Cellpose.eval, where 0=NO masks; 1,2,...=mask labels.
+        flows (list): Flows output from Cellpose.eval.
+        file_names (list, str): Names of files of images.
+        png (bool, optional): Save masks to PNG. Defaults to True.
+        tif (bool, optional): Save masks to TIF. Defaults to False.
+        channels (list, int, optional): Channels used to run Cellpose. Defaults to [0,0].
+        suffix (str, optional): Add name to saved masks. Defaults to "_cp_masks".
+        save_flows (bool, optional): Save flows output from Cellpose.eval. Defaults to False.
+        save_outlines (bool, optional): Save outlines of masks. Defaults to False.
+        dir_above (bool, optional): Save masks/flows in directory above. Defaults to False.
+        in_folders (bool, optional): Save masks/flows in separate folders. Defaults to False.
+        savedir (str, optional): Absolute path where images will be saved. If None, saves to image directory. Defaults to None.
+        save_txt (bool, optional): Save masks as list of outlines for ImageJ. Defaults to False.
+        save_mpl (bool, optional): If True, saves a matplotlib figure of the original image/segmentation/flows. Does not work for 3D.
+                This takes a long time for large images. Defaults to False.
+
+    Returns:
+        None
+    """
+
+    if isinstance(masks, list):
+        for image, mask, flow, file_name in zip(images, masks, flows, file_names):
+            save_masks(image, mask, flow, file_name, png=png, tif=tif, suffix=suffix,
+                       dir_above=dir_above, save_flows=save_flows,
+                       save_outlines=save_outlines, savedir=savedir, save_txt=save_txt,
+                       in_folders=in_folders, save_mpl=save_mpl)
+        return
+
+    if masks.ndim > 2 and not tif:
+        raise ValueError("cannot save 3D outputs as PNG, use tif option instead")
+
+    if masks.max() == 0:
+        io_logger.warning("no masks found, will not save PNG or outlines")
+        if not tif:
+            return
+        else:
+            png = False
+            save_outlines = False
+            save_flows = False
+            save_txt = False
+
+    if savedir is None:
+        if dir_above:
+            savedir = Path(file_names).parent.parent.absolute(
+            )  #go up a level to save in its own folder
+        else:
+            savedir = Path(file_names).parent.absolute()
+
+    check_dir(savedir)
+
+    basename = os.path.splitext(os.path.basename(file_names))[0]
+    if in_folders:
+        maskdir = os.path.join(savedir, "masks")
+        outlinedir = os.path.join(savedir, "outlines")
+        txtdir = os.path.join(savedir, "txt_outlines")
+        flowdir = os.path.join(savedir, "flows")
+    else:
+        maskdir = savedir
+        outlinedir = savedir
+        txtdir = savedir
+        flowdir = savedir
+
+    check_dir(maskdir)
+
+    exts = []
+    if masks.ndim > 2:
+        png = False
+        tif = True
+    if png:
+        if masks.max() < 2**16:
+            masks = masks.astype(np.uint16)
+            exts.append(".png")
+        else:
+            png = False
+            tif = True
+            io_logger.warning(
+                "found more than 65535 masks in each image, cannot save PNG, saving as TIF"
+            )
+    if tif:
+        exts.append(".tif")
+
+    # save masks
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        for ext in exts:
+            imsave(os.path.join(maskdir, basename + suffix + ext), masks)
+
+    if save_mpl and png and MATPLOTLIB and not min(images.shape) > 3:
+        # Make and save original/segmentation/flows image
+
+        img = images.copy()
+        if img.ndim < 3:
+            img = img[:, :, np.newaxis]
+        elif img.shape[0] < 8:
+            np.transpose(img, (1, 2, 0))
+
+        fig = plt.figure(figsize=(12, 3))
+        plot.show_segmentation(fig, img, masks, flows[0])
+        fig.savefig(os.path.join(savedir, basename + "_cp_output" + suffix + ".png"),
+                    dpi=300)
+        plt.close(fig)
+
+    # ImageJ txt outline files
+    if masks.ndim < 3 and save_txt:
+        check_dir(txtdir)
+        outlines = utils.outlines_list(masks)
+        outlines_to_text(os.path.join(txtdir, basename), outlines)
+
+    # RGB outline images
+    if masks.ndim < 3 and save_outlines:
+        check_dir(outlinedir)
+        outlines = utils.masks_to_outlines(masks)
+        outX, outY = np.nonzero(outlines)
+        img0 = transforms.normalize99(images)
+        if img0.shape[0] < 4:
+            img0 = np.transpose(img0, (1, 2, 0))
+        if img0.shape[-1] < 3 or img0.ndim < 3:
+            img0 = plot.image_to_rgb(img0, channels=channels)
+        else:
+            if img0.max() <= 50.0:
+                img0 = np.uint8(np.clip(img0 * 255, 0, 1))
+        imgout = img0.copy()
+        imgout[outX, outY] = np.array([255, 0, 0])  #pure red
+        imsave(os.path.join(outlinedir, basename + "_outlines" + suffix + ".png"),
+               imgout)
+
+    # save RGB flow picture
+    if masks.ndim < 3 and save_flows:
+        check_dir(flowdir)
+        imsave(os.path.join(flowdir, basename + "_flows" + suffix + ".tif"),
+               (flows[0] * (2**16 - 1)).astype(np.uint16))
+        #save full flow data
+        imsave(os.path.join(flowdir, basename + '_dP' + suffix + '.tif'), flows[1])

models/seg_post_model/cellpose/metrics.py

@@ -0,0 +1,205 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import numpy as np
+from . import utils
+from scipy.optimize import linear_sum_assignment
+from scipy.ndimage import convolve
+from scipy.sparse import csr_matrix 
+
+
+def mask_ious(masks_true, masks_pred):
+    """Return best-matched masks."""
+    iou = _intersection_over_union(masks_true, masks_pred)[1:, 1:]
+    n_min = min(iou.shape[0], iou.shape[1])
+    costs = -(iou >= 0.5).astype(float) - iou / (2 * n_min)
+    true_ind, pred_ind = linear_sum_assignment(costs)
+    iout = np.zeros(masks_true.max())
+    iout[true_ind] = iou[true_ind, pred_ind]
+    preds = np.zeros(masks_true.max(), "int")
+    preds[true_ind] = pred_ind + 1
+    return iout, preds
+
+
+def boundary_scores(masks_true, masks_pred, scales):
+    """
+    Calculate boundary precision, recall, and F-score.
+
+    Args:
+        masks_true (list): List of true masks.
+        masks_pred (list): List of predicted masks.
+        scales (list): List of scales.
+
+    Returns:
+        tuple: A tuple containing precision, recall, and F-score arrays.
+    """
+    diams = [utils.diameters(lbl)[0] for lbl in masks_true]
+    precision = np.zeros((len(scales), len(masks_true)))
+    recall = np.zeros((len(scales), len(masks_true)))
+    fscore = np.zeros((len(scales), len(masks_true)))
+    for j, scale in enumerate(scales):
+        for n in range(len(masks_true)):
+            diam = max(1, scale * diams[n])
+            rs, ys, xs = utils.circleMask([int(np.ceil(diam)), int(np.ceil(diam))])
+            filt = (rs <= diam).astype(np.float32)
+            otrue = utils.masks_to_outlines(masks_true[n])
+            otrue = convolve(otrue, filt)
+            opred = utils.masks_to_outlines(masks_pred[n])
+            opred = convolve(opred, filt)
+            tp = np.logical_and(otrue == 1, opred == 1).sum()
+            fp = np.logical_and(otrue == 0, opred == 1).sum()
+            fn = np.logical_and(otrue == 1, opred == 0).sum()
+            precision[j, n] = tp / (tp + fp)
+            recall[j, n] = tp / (tp + fn)
+        fscore[j] = 2 * precision[j] * recall[j] / (precision[j] + recall[j])
+    return precision, recall, fscore
+
+
+def aggregated_jaccard_index(masks_true, masks_pred):
+    """ 
+    AJI = intersection of all matched masks / union of all masks 
+    
+    Args:
+        masks_true (list of np.ndarrays (int) or np.ndarray (int)): 
+            where 0=NO masks; 1,2... are mask labels
+        masks_pred (list of np.ndarrays (int) or np.ndarray (int)): 
+            np.ndarray (int) where 0=NO masks; 1,2... are mask labels
+
+    Returns:
+        aji (float): aggregated jaccard index for each set of masks
+    """
+    aji = np.zeros(len(masks_true))
+    for n in range(len(masks_true)):
+        iout, preds = mask_ious(masks_true[n], masks_pred[n])
+        inds = np.arange(0, masks_true[n].max(), 1, int)
+        overlap = _label_overlap(masks_true[n], masks_pred[n])
+        union = np.logical_or(masks_true[n] > 0, masks_pred[n] > 0).sum()
+        overlap = overlap[inds[preds > 0] + 1, preds[preds > 0].astype(int)]
+        aji[n] = overlap.sum() / union
+    return aji
+
+
+def average_precision(masks_true, masks_pred, threshold=[0.5, 0.75, 0.9]):
+    """ 
+    Average precision estimation: AP = TP / (TP + FP + FN)
+
+    This function is based heavily on the *fast* stardist matching functions
+    (https://github.com/mpicbg-csbd/stardist/blob/master/stardist/matching.py)
+
+    Args:
+        masks_true (list of np.ndarrays (int) or np.ndarray (int)): 
+            where 0=NO masks; 1,2... are mask labels
+        masks_pred (list of np.ndarrays (int) or np.ndarray (int)): 
+            np.ndarray (int) where 0=NO masks; 1,2... are mask labels
+
+    Returns:
+        ap (array [len(masks_true) x len(threshold)]): 
+            average precision at thresholds
+        tp (array [len(masks_true) x len(threshold)]): 
+            number of true positives at thresholds
+        fp (array [len(masks_true) x len(threshold)]): 
+            number of false positives at thresholds
+        fn (array [len(masks_true) x len(threshold)]): 
+            number of false negatives at thresholds
+    """
+    not_list = False
+    if not isinstance(masks_true, list):
+        masks_true = [masks_true]
+        masks_pred = [masks_pred]
+        not_list = True
+    if not isinstance(threshold, list) and not isinstance(threshold, np.ndarray):
+        threshold = [threshold]
+
+    if len(masks_true) != len(masks_pred):
+        raise ValueError(
+            "metrics.average_precision requires len(masks_true)==len(masks_pred)")
+
+    ap = np.zeros((len(masks_true), len(threshold)), np.float32)
+    tp = np.zeros((len(masks_true), len(threshold)), np.float32)
+    fp = np.zeros((len(masks_true), len(threshold)), np.float32)
+    fn = np.zeros((len(masks_true), len(threshold)), np.float32)
+    n_true = np.array([len(np.unique(mt)) - 1 for mt in masks_true])
+    n_pred = np.array([len(np.unique(mp)) - 1 for mp in masks_pred])
+
+    for n in range(len(masks_true)):
+        #_,mt = np.reshape(np.unique(masks_true[n], return_index=True), masks_pred[n].shape)
+        if n_pred[n] > 0:
+            iou = _intersection_over_union(masks_true[n], masks_pred[n])[1:, 1:]
+            for k, th in enumerate(threshold):
+                tp[n, k] = _true_positive(iou, th)
+        fp[n] = n_pred[n] - tp[n]
+        fn[n] = n_true[n] - tp[n]
+        ap[n] = tp[n] / (tp[n] + fp[n] + fn[n])
+
+    if not_list:
+        ap, tp, fp, fn = ap[0], tp[0], fp[0], fn[0]
+    return ap, tp, fp, fn
+
+
+def _intersection_over_union(masks_true, masks_pred):
+    """Calculate the intersection over union of all mask pairs.
+
+    Parameters:
+        masks_true (np.ndarray, int): Ground truth masks, where 0=NO masks; 1,2... are mask labels.
+        masks_pred (np.ndarray, int): Predicted masks, where 0=NO masks; 1,2... are mask labels.
+
+    Returns:
+        iou (np.ndarray, float): Matrix of IOU pairs of size [x.max()+1, y.max()+1].
+
+    How it works:
+        The overlap matrix is a lookup table of the area of intersection
+        between each set of labels (true and predicted). The true labels
+        are taken to be along axis 0, and the predicted labels are taken 
+        to be along axis 1. The sum of the overlaps along axis 0 is thus
+        an array giving the total overlap of the true labels with each of
+        the predicted labels, and likewise the sum over axis 1 is the
+        total overlap of the predicted labels with each of the true labels.
+        Because the label 0 (background) is included, this sum is guaranteed
+        to reconstruct the total area of each label. Adding this row and
+        column vectors gives a 2D array with the areas of every label pair
+        added together. This is equivalent to the union of the label areas
+        except for the duplicated overlap area, so the overlap matrix is
+        subtracted to find the union matrix. 
+    """
+    if masks_true.size != masks_pred.size:
+        raise ValueError(f"masks_true.size {masks_true.shape} != masks_pred.size {masks_pred.shape}")
+    overlap = csr_matrix((np.ones((masks_true.size,), "int"), 
+                         (masks_true.flatten(), masks_pred.flatten())),
+                         shape=(masks_true.max()+1, masks_pred.max()+1))
+    overlap = overlap.toarray()
+    n_pixels_pred = np.sum(overlap, axis=0, keepdims=True)
+    n_pixels_true = np.sum(overlap, axis=1, keepdims=True)
+    iou = overlap / (n_pixels_pred + n_pixels_true - overlap)
+    iou[np.isnan(iou)] = 0.0
+    return iou
+
+
+def _true_positive(iou, th):
+    """Calculate the true positive at threshold th.
+
+    Args:
+        iou (float, np.ndarray): Array of IOU pairs.
+        th (float): Threshold on IOU for positive label.
+
+    Returns:
+        tp (float): Number of true positives at threshold.
+
+    How it works:
+        (1) Find minimum number of masks.
+        (2) Define cost matrix; for a given threshold, each element is negative
+            the higher the IoU is (perfect IoU is 1, worst is 0). The second term
+            gets more negative with higher IoU, but less negative with greater
+            n_min (but that's a constant...).
+        (3) Solve the linear sum assignment problem. The costs array defines the cost
+            of matching a true label with a predicted label, so the problem is to 
+            find the set of pairings that minimizes this cost. The scipy.optimize
+            function gives the ordered lists of corresponding true and predicted labels. 
+        (4) Extract the IoUs from these pairings and then threshold to get a boolean array
+            whose sum is the number of true positives that is returned. 
+    """
+    n_min = min(iou.shape[0], iou.shape[1])
+    costs = -(iou >= th).astype(float) - iou / (2 * n_min)
+    true_ind, pred_ind = linear_sum_assignment(costs)
+    match_ok = iou[true_ind, pred_ind] >= th
+    tp = match_ok.sum()
+    return tp

models/seg_post_model/cellpose/models.py

@@ -0,0 +1,524 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer, Michael Rariden and Marius Pachitariu.
+"""
+
+import os, time
+from pathlib import Path
+import numpy as np
+from tqdm import trange
+import torch
+from scipy.ndimage import gaussian_filter
+import gc
+import cv2
+
+import logging
+
+models_logger = logging.getLogger(__name__)
+
+from . import transforms, dynamics, utils, plot
+from .vit_sam import Transformer
+from .core import assign_device, run_net, run_3D
+
+_CPSAM_MODEL_URL = "https://huggingface.co/mouseland/cellpose-sam/resolve/main/cpsam"
+_MODEL_DIR_ENV = os.environ.get("CELLPOSE_LOCAL_MODELS_PATH")
+# _MODEL_DIR_DEFAULT = Path.home().joinpath(".cellpose", "models")
+_MODEL_DIR_DEFAULT = Path("/media/data1/huix/seg/cellpose_models")
+MODEL_DIR = Path(_MODEL_DIR_ENV) if _MODEL_DIR_ENV else _MODEL_DIR_DEFAULT
+
+MODEL_NAMES = ["cpsam"]
+
+MODEL_LIST_PATH = os.fspath(MODEL_DIR.joinpath("gui_models.txt"))
+
+normalize_default = {
+    "lowhigh": None,
+    "percentile": None,
+    "normalize": True,
+    "norm3D": True,
+    "sharpen_radius": 0,
+    "smooth_radius": 0,
+    "tile_norm_blocksize": 0,
+    "tile_norm_smooth3D": 1,
+    "invert": False
+}
+
+
+# def model_path(model_type, model_index=0):
+#     return cache_CPSAM_model_path()
+
+
+# def cache_CPSAM_model_path():
+#     MODEL_DIR.mkdir(parents=True, exist_ok=True)
+#     cached_file = os.fspath(MODEL_DIR.joinpath('cpsam'))
+#     if not os.path.exists(cached_file):
+#         models_logger.info('Downloading: "{}" to {}\n'.format(_CPSAM_MODEL_URL, cached_file))
+#         utils.download_url_to_file(_CPSAM_MODEL_URL, cached_file, progress=True)
+#     return cached_file
+
+
+def get_user_models():
+    model_strings = []
+    if os.path.exists(MODEL_LIST_PATH):
+        with open(MODEL_LIST_PATH, "r") as textfile:
+            lines = [line.rstrip() for line in textfile]
+            if len(lines) > 0:
+                model_strings.extend(lines)
+    return model_strings
+
+
+class CellposeModel():
+    """
+    Class representing a Cellpose model.
+
+    Attributes:
+        diam_mean (float): Mean "diameter" value for the model.
+        builtin (bool): Whether the model is a built-in model or not.
+        device (torch device): Device used for model running / training.
+        nclasses (int): Number of classes in the model.
+        nbase (list): List of base values for the model.
+        net (CPnet): Cellpose network.
+        pretrained_model (str): Path to pretrained cellpose model.
+        pretrained_model_ortho (str): Path or model_name for pretrained cellpose model for ortho views in 3D.
+        backbone (str): Type of network ("default" is the standard res-unet, "transformer" for the segformer).
+
+    Methods:
+        __init__(self, gpu=False, pretrained_model=False, model_type=None, diam_mean=30., device=None):
+            Initialize the CellposeModel.
+        
+        eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None, z_axis=None, normalize=True, invert=False, rescale=None, diameter=None, flow_threshold=0.4, cellprob_threshold=0.0, do_3D=False, anisotropy=None, stitch_threshold=0.0, min_size=15, niter=None, augment=False, tile_overlap=0.1, bsize=224, interp=True, compute_masks=True, progress=None):
+            Segment list of images x, or 4D array - Z x C x Y x X.
+
+    """
+
+    def __init__(self, gpu=False, pretrained_model="", model_type=None,
+                 diam_mean=None, device=None, nchan=None, use_bfloat16=True, vit_checkpoint=None):
+        """
+        Initialize the CellposeModel.
+
+        Parameters:
+            gpu (bool, optional): Whether or not to save model to GPU, will check if GPU available.
+            pretrained_model (str or list of strings, optional): Full path to pretrained cellpose model(s), if None or False, no model loaded.
+            model_type (str, optional): Any model that is available in the GUI, use name in GUI e.g. "livecell" (can be user-trained or model zoo).
+            diam_mean (float, optional): Mean "diameter", 30. is built-in value for "cyto" model; 17. is built-in value for "nuclei" model; if saved in custom model file (cellpose>=2.0) then it will be loaded automatically and overwrite this value.
+            device (torch device, optional): Device used for model running / training (torch.device("cuda") or torch.device("cpu")), overrides gpu input, recommended if you want to use a specific GPU (e.g. torch.device("cuda:1")).
+            use_bfloat16 (bool, optional): Use 16bit float precision instead of 32bit for model weights. Default to 16bit (True).
+        """
+        # if diam_mean is not None:
+        #     models_logger.warning(
+        #         "diam_mean argument are not used in v4.0.1+. Ignoring this argument..."
+        #     )
+        # if model_type is not None:
+        #     models_logger.warning(
+        #         "model_type argument is not used in v4.0.1+. Ignoring this argument..."
+        #     )
+        # if nchan is not None:
+        #     models_logger.warning("nchan argument is deprecated in v4.0.1+. Ignoring this argument")
+
+        ### assign model device
+        self.device = assign_device(gpu=gpu)[0] if device is None else device
+        if torch.cuda.is_available():
+            device_gpu = self.device.type == "cuda"
+        elif torch.backends.mps.is_available():
+            device_gpu = self.device.type == "mps"
+        else:
+            device_gpu = False
+        self.gpu = device_gpu
+
+        if pretrained_model is None:
+            # raise ValueError("Must specify a pretrained model, training from scratch is not implemented")
+            pretrained_model = ""
+
+        ### create neural network
+        if pretrained_model and not os.path.exists(pretrained_model):
+            # check if pretrained model is in the models directory
+            model_strings = get_user_models()
+            all_models = MODEL_NAMES.copy()
+            all_models.extend(model_strings)
+            if pretrained_model in all_models:
+                pretrained_model = os.path.join(MODEL_DIR, pretrained_model)
+            else:
+                pretrained_model = os.path.join(MODEL_DIR, "cpsam")
+                models_logger.warning(
+                    f"pretrained model {pretrained_model} not found, using default model"
+                )
+
+        self.pretrained_model = pretrained_model
+        dtype = torch.bfloat16 if use_bfloat16 else torch.float32
+        self.net = Transformer(dtype=dtype, checkpoint=vit_checkpoint).to(self.device)
+
+        if os.path.exists(self.pretrained_model):
+            models_logger.info(f">>>> loading model {self.pretrained_model}")
+            self.net.load_model(self.pretrained_model, device=self.device)
+        # else:
+            # try:
+            #     if os.path.split(self.pretrained_model)[-1] != 'cpsam':
+            #         raise FileNotFoundError('model file not recognized')
+            #     cache_CPSAM_model_path()
+            #     self.net.load_model(self.pretrained_model, device=self.device)
+            # except:
+                # print("ViT not initialized")
+        
+        
+    def eval(self, x, feat=None, batch_size=8, resample=True, channels=None, channel_axis=None,
+             z_axis=None, normalize=True, invert=False, rescale=None, diameter=None,
+             flow_threshold=0.4, cellprob_threshold=0.0, do_3D=False, anisotropy=None,
+             flow3D_smooth=0, stitch_threshold=0.0, 
+             min_size=15, max_size_fraction=0.4, niter=None, 
+             augment=False, tile_overlap=0.1, bsize=256, 
+             compute_masks=True, progress=None):
+        
+
+        # if rescale is not None:
+        #     models_logger.warning("rescaling deprecated in v4.0.1+") 
+        # if channels is not None:
+        #     models_logger.warning("channels deprecated in v4.0.1+. If data contain more than 3 channels, only the first 3 channels will be used")
+
+        if isinstance(x, list) or x.squeeze().ndim == 5:
+            self.timing = []
+            masks, styles, flows = [], [], []
+            tqdm_out = utils.TqdmToLogger(models_logger, level=logging.INFO)
+            nimg = len(x)
+            iterator = trange(nimg, file=tqdm_out,
+                              mininterval=30) if nimg > 1 else range(nimg)
+            for i in iterator:
+                tic = time.time()
+                maski, flowi, stylei = self.eval(
+                    x[i], 
+                    feat=None if feat is None else feat[i],
+                    batch_size=batch_size,
+                    channel_axis=channel_axis, 
+                    z_axis=z_axis,
+                    normalize=normalize, 
+                    invert=invert,
+                    diameter=diameter[i] if isinstance(diameter, list) or
+                        isinstance(diameter, np.ndarray) else diameter, 
+                    do_3D=do_3D,
+                    anisotropy=anisotropy, 
+                    augment=augment, 
+                    tile_overlap=tile_overlap, 
+                    bsize=bsize, 
+                    resample=resample,
+                    flow_threshold=flow_threshold,
+                    cellprob_threshold=cellprob_threshold, 
+                    compute_masks=compute_masks,
+                    min_size=min_size, 
+                    max_size_fraction=max_size_fraction, 
+                    stitch_threshold=stitch_threshold, 
+                    flow3D_smooth=flow3D_smooth,
+                    progress=progress, 
+                    niter=niter)
+                masks.append(maski)
+                flows.append(flowi)
+                styles.append(stylei)
+                self.timing.append(time.time() - tic)
+            return masks, flows, styles
+
+        ############# actual eval code ############
+        # reshape image
+        x = transforms.convert_image(x, channel_axis=channel_axis,
+                                        z_axis=z_axis, 
+                                        do_3D=(do_3D or stitch_threshold > 0))
+        
+        # Add batch dimension if not present
+        if x.ndim < 4:
+            x = x[np.newaxis, ...]
+        if feat is not None:
+            if feat.ndim < 4:
+                feat = feat[np.newaxis, ...]
+        nimg = x.shape[0]
+        
+        image_scaling = None
+        Ly_0 = x.shape[1]
+        Lx_0 = x.shape[2]
+        Lz_0 = None
+        if do_3D or stitch_threshold > 0:
+            Lz_0 = x.shape[0]
+        if diameter is not None:
+            image_scaling = 30. / diameter
+            x = transforms.resize_image(x,
+                                        Ly=int(x.shape[1] * image_scaling),
+                                        Lx=int(x.shape[2] * image_scaling))
+            if feat is not None:
+                feat = transforms.resize_image(feat,
+                                        Ly=int(feat.shape[1] * image_scaling),
+                                        Lx=int(feat.shape[2] * image_scaling))
+
+
+        # normalize image
+        normalize_params = normalize_default
+        if isinstance(normalize, dict):
+            normalize_params = {**normalize_params, **normalize}
+        elif not isinstance(normalize, bool):
+            raise ValueError("normalize parameter must be a bool or a dict")
+        else:
+            normalize_params["normalize"] = normalize
+            normalize_params["invert"] = invert
+
+        # pre-normalize if 3D stack for stitching or do_3D
+        do_normalization = True if normalize_params["normalize"] else False
+        if nimg > 1 and do_normalization and (stitch_threshold or do_3D):
+            normalize_params["norm3D"] = True if do_3D else normalize_params["norm3D"]
+            x = transforms.normalize_img(x, **normalize_params)
+            do_normalization = False # do not normalize again
+        else:
+            if normalize_params["norm3D"] and nimg > 1 and do_normalization:
+                models_logger.warning(
+                    "normalize_params['norm3D'] is True but do_3D is False and stitch_threshold=0, so setting to False"
+                )
+                normalize_params["norm3D"] = False
+        if do_normalization:
+            x = transforms.normalize_img(x, **normalize_params)
+
+        if feat is not None:
+            if feat.shape[-1] > feat.shape[1]:
+                # transpose feat to have channels last
+                feat = np.moveaxis(feat, 1, -1)
+
+        # ajust the anisotropy when diameter is specified and images are resized:
+        if isinstance(anisotropy, (float, int)) and image_scaling:
+            anisotropy = image_scaling * anisotropy
+
+        dP, cellprob, styles = self._run_net(
+            x, 
+            feat=feat,
+            augment=augment, 
+            batch_size=batch_size, 
+            tile_overlap=tile_overlap, 
+            bsize=bsize,
+            do_3D=do_3D, 
+            anisotropy=anisotropy)
+
+        if do_3D:    
+            if flow3D_smooth > 0:
+                models_logger.info(f"smoothing flows with sigma={flow3D_smooth}")
+                dP = gaussian_filter(dP, (0, flow3D_smooth, flow3D_smooth, flow3D_smooth))
+            torch.cuda.empty_cache()
+            gc.collect()
+
+        if resample:
+            # upsample flows before computing them: 
+            dP = self._resize_gradients(dP, to_y_size=Ly_0, to_x_size=Lx_0, to_z_size=Lz_0)
+            cellprob = self._resize_cellprob(cellprob, to_x_size=Lx_0, to_y_size=Ly_0, to_z_size=Lz_0)
+
+
+        if compute_masks:
+            niter0 = 200
+            niter = niter0 if niter is None or niter == 0 else niter
+            masks = self._compute_masks(x.shape, dP, cellprob, flow_threshold=flow_threshold,
+                            cellprob_threshold=cellprob_threshold, min_size=min_size,
+                        max_size_fraction=max_size_fraction, niter=niter,
+                        stitch_threshold=stitch_threshold, do_3D=do_3D)
+        else:
+            masks = np.zeros(0) #pass back zeros if not compute_masks
+        
+        masks, dP, cellprob = masks.squeeze(), dP.squeeze(), cellprob.squeeze()
+
+        # undo resizing:
+        if image_scaling is not None or anisotropy is not None:
+
+            dP = self._resize_gradients(dP, to_y_size=Ly_0, to_x_size=Lx_0, to_z_size=Lz_0) # works for 2 or 3D: 
+            cellprob = self._resize_cellprob(cellprob, to_x_size=Lx_0, to_y_size=Ly_0, to_z_size=Lz_0)
+
+            if do_3D:
+                if compute_masks:
+                    # Rescale xy then xz:
+                    masks = transforms.resize_image(masks, Ly=Ly_0, Lx=Lx_0, no_channels=True, interpolation=cv2.INTER_NEAREST)
+                    masks = masks.transpose(1, 0, 2)
+                    masks = transforms.resize_image(masks, Ly=Lz_0, Lx=Lx_0, no_channels=True, interpolation=cv2.INTER_NEAREST)
+                    masks = masks.transpose(1, 0, 2)
+
+            else:
+                # 2D or 3D stitching case:
+                if compute_masks:
+                    masks = transforms.resize_image(masks, Ly=Ly_0, Lx=Lx_0, no_channels=True, interpolation=cv2.INTER_NEAREST)
+
+        return masks, [plot.dx_to_circ(dP), dP, cellprob], styles
+    
+
+    def _resize_cellprob(self, prob: np.ndarray, to_y_size: int, to_x_size: int, to_z_size: int = None) -> np.ndarray:
+        """
+        Resize cellprob array to specified dimensions for either 2D or 3D.
+
+        Parameters:
+            prob (numpy.ndarray): The cellprobs to resize, either in 2D or 3D. Returns the same ndim as provided.
+            to_y_size (int): The target size along the Y-axis.
+            to_x_size (int): The target size along the X-axis.
+            to_z_size (int, optional): The target size along the Z-axis. Required
+                for 3D cellprobs.
+
+        Returns:
+            numpy.ndarray: The resized cellprobs array with the same number of dimensions
+            as the input.
+
+        Raises:
+            ValueError: If the input cellprobs array does not have 3 or 4 dimensions.
+        """
+        prob_shape = prob.shape
+        prob = prob.squeeze()
+        squeeze_happened = prob.shape != prob_shape
+        prob_shape = np.array(prob_shape)
+
+        if prob.ndim == 2:
+            # 2D case:
+            prob = transforms.resize_image(prob, Ly=to_y_size, Lx=to_x_size, no_channels=True)
+            if squeeze_happened:
+                prob = np.expand_dims(prob, int(np.argwhere(prob_shape == 1))) # add back empty axis for compatibility
+        elif prob.ndim == 3:
+            # 3D case: 
+            prob = transforms.resize_image(prob, Ly=to_y_size, Lx=to_x_size, no_channels=True)
+            prob = prob.transpose(1, 0, 2)
+            prob = transforms.resize_image(prob, Ly=to_z_size, Lx=to_x_size, no_channels=True)
+            prob = prob.transpose(1, 0, 2)
+        else:
+            raise ValueError(f'gradients have incorrect dimension after squeezing. Should be 2 or 3, prob shape: {prob.shape}')
+        
+        return prob
+
+
+    def _resize_gradients(self, grads: np.ndarray, to_y_size: int, to_x_size: int, to_z_size: int = None) -> np.ndarray:
+        """
+        Resize gradient arrays to specified dimensions for either 2D or 3D gradients.
+
+        Parameters:
+            grads (np.ndarray): The gradients to resize, either in 2D or 3D. Returns the same ndim as provided.
+            to_y_size (int): The target size along the Y-axis.
+            to_x_size (int): The target size along the X-axis.
+            to_z_size (int, optional): The target size along the Z-axis. Required
+                for 3D gradients.
+
+        Returns:
+            numpy.ndarray: The resized gradient array with the same number of dimensions
+            as the input.
+
+        Raises:
+            ValueError: If the input gradient array does not have 3 or 4 dimensions.
+        """
+        grads_shape = grads.shape
+        grads = grads.squeeze()
+        squeeze_happened = grads.shape != grads_shape
+        grads_shape = np.array(grads_shape)
+
+        if grads.ndim == 3:
+            # 2D case, with XY flows in 2 channels:
+            grads = np.moveaxis(grads, 0, -1) # Put gradients last
+            grads = transforms.resize_image(grads, Ly=to_y_size, Lx=to_x_size, no_channels=False)
+            grads = np.moveaxis(grads, -1, 0) # Put gradients first
+
+            if squeeze_happened:
+                grads = np.expand_dims(grads, int(np.argwhere(grads_shape == 1))) # add back empty axis for compatibility
+        elif grads.ndim == 4:
+            # dP has gradients that can be treated as channels:
+            grads = grads.transpose(1, 2, 3, 0) # move gradients last:
+            grads = transforms.resize_image(grads, Ly=to_y_size, Lx=to_x_size, no_channels=False)
+            grads = grads.transpose(1, 0, 2, 3) # switch axes to resize again
+            grads = transforms.resize_image(grads, Ly=to_z_size, Lx=to_x_size, no_channels=False)
+            grads = grads.transpose(3, 1, 0, 2) # undo transposition
+        else:
+            raise ValueError(f'gradients have incorrect dimension after squeezing. Should be 3 or 4, grads shape: {grads.shape}')
+        
+        return grads
+
+
+    def _run_net(self, x, feat=None,
+                augment=False, 
+                batch_size=8, tile_overlap=0.1,
+                bsize=224, anisotropy=1.0, do_3D=False):
+        """ run network on image x """
+        tic = time.time()
+        shape = x.shape
+        nimg = shape[0]
+
+
+        if do_3D:
+            Lz, Ly, Lx = shape[:-1]
+            if anisotropy is not None and anisotropy != 1.0:
+                models_logger.info(f"resizing 3D image with anisotropy={anisotropy}")
+                x = transforms.resize_image(x.transpose(1,0,2,3),
+                                        Ly=int(Lz*anisotropy), 
+                                        Lx=int(Lx)).transpose(1,0,2,3)
+            yf, styles = run_3D(self.net, x,
+                                batch_size=batch_size, augment=augment,  
+                                tile_overlap=tile_overlap, 
+                                bsize=bsize
+                                )
+            cellprob = yf[..., -1]
+            dP = yf[..., :-1].transpose((3, 0, 1, 2))
+        else:
+            yf, styles = run_net(self.net, x, feat=feat, bsize=bsize, augment=augment,
+                                batch_size=batch_size,  
+                                tile_overlap=tile_overlap, 
+                                )
+            cellprob = yf[..., -1]
+            dP = yf[..., -3:-1].transpose((3, 0, 1, 2))
+            if yf.shape[-1] > 3:
+                styles = yf[..., :-3]
+        
+        styles = styles.squeeze()
+
+        net_time = time.time() - tic
+        if nimg > 1:
+            models_logger.info("network run in %2.2fs" % (net_time))
+
+        return dP, cellprob, styles
+    
+    def _compute_masks(self, shape, dP, cellprob, flow_threshold=0.4, cellprob_threshold=0.0,
+                       min_size=15, max_size_fraction=0.4, niter=None,
+                       do_3D=False, stitch_threshold=0.0):
+        """ compute masks from flows and cell probability """
+        changed_device_from = None
+        if self.device.type == "mps" and do_3D:
+            models_logger.warning("MPS does not support 3D post-processing, switching to CPU")
+            self.device = torch.device("cpu")
+            changed_device_from = "mps"
+        Lz, Ly, Lx = shape[:3]
+        tic = time.time()
+        if do_3D:
+            masks = dynamics.resize_and_compute_masks(
+                dP, cellprob, niter=niter, cellprob_threshold=cellprob_threshold,
+                flow_threshold=flow_threshold, do_3D=do_3D,
+                min_size=min_size, max_size_fraction=max_size_fraction, 
+                resize=shape[:3] if (np.array(dP.shape[-3:])!=np.array(shape[:3])).sum() 
+                        else None,
+                device=self.device)
+        else:
+            nimg = shape[0]
+            Ly0, Lx0 = cellprob[0].shape 
+            resize = None if Ly0==Ly and Lx0==Lx else [Ly, Lx]
+            tqdm_out = utils.TqdmToLogger(models_logger, level=logging.INFO)
+            iterator = trange(nimg, file=tqdm_out,
+                            mininterval=30) if nimg > 1 else range(nimg)
+            for i in iterator:
+                # turn off min_size for 3D stitching
+                min_size0 = min_size if stitch_threshold == 0 or nimg == 1 else -1
+                outputs = dynamics.resize_and_compute_masks(
+                    dP[:, i], cellprob[i],
+                    niter=niter, cellprob_threshold=cellprob_threshold,
+                    flow_threshold=flow_threshold, resize=resize,
+                    min_size=min_size0, max_size_fraction=max_size_fraction,
+                    device=self.device)
+                if i==0 and nimg > 1:
+                    masks = np.zeros((nimg, shape[1], shape[2]), outputs.dtype)
+                if nimg > 1:
+                    masks[i] = outputs
+                else:
+                    masks = outputs
+
+            if stitch_threshold > 0 and nimg > 1:
+                models_logger.info(
+                    f"stitching {nimg} planes using stitch_threshold={stitch_threshold:0.3f} to make 3D masks"
+                )
+                masks = utils.stitch3D(masks, stitch_threshold=stitch_threshold)
+                masks = utils.fill_holes_and_remove_small_masks(
+                    masks, min_size=min_size)
+            elif nimg > 1:
+                models_logger.warning(
+                    "3D stack used, but stitch_threshold=0 and do_3D=False, so masks are made per plane only"
+                )
+
+        flow_time = time.time() - tic
+        if shape[0] > 1:
+            models_logger.info("masks created in %2.2fs" % (flow_time))
+        
+        if changed_device_from is not None:
+            models_logger.info("switching back to device %s" % self.device)
+            self.device = torch.device(changed_device_from)
+        return masks

models/seg_post_model/cellpose/plot.py

@@ -0,0 +1,281 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os
+import numpy as np
+import cv2
+from scipy.ndimage import gaussian_filter
+from . import utils, io, transforms
+
+try:
+    import matplotlib
+    MATPLOTLIB_ENABLED = True
+except:
+    MATPLOTLIB_ENABLED = False
+
+try:
+    from skimage import color
+    from skimage.segmentation import find_boundaries
+    SKIMAGE_ENABLED = True
+except:
+    SKIMAGE_ENABLED = False
+
+
+# modified to use sinebow color
+def dx_to_circ(dP):
+    """Converts the optic flow representation to a circular color representation.
+
+    Args:
+        dP (ndarray): Flow field components [dy, dx].
+        
+    Returns:
+        ndarray: The circular color representation of the optic flow.
+
+    """
+    mag = 255 * np.clip(transforms.normalize99(np.sqrt(np.sum(dP**2, axis=0))), 0, 1.)
+    angles = np.arctan2(dP[1], dP[0]) + np.pi
+    a = 2
+    mag /= a
+    rgb = np.zeros((*dP.shape[1:], 3), "uint8")
+    rgb[..., 0] = np.clip(mag * (np.cos(angles) + 1), 0, 255).astype("uint8")
+    rgb[..., 1] = np.clip(mag * (np.cos(angles + 2 * np.pi / 3) + 1), 0, 255).astype("uint8")
+    rgb[..., 2] = np.clip(mag * (np.cos(angles + 4 * np.pi / 3) + 1), 0, 255).astype("uint8")
+    
+    return rgb
+
+
+def show_segmentation(fig, img, maski, flowi, channels=[0, 0], file_name=None):
+    """Plot segmentation results (like on website).
+
+    Can save each panel of figure with file_name option. Use channels option if
+    img input is not an RGB image with 3 channels.
+
+    Args:
+        fig (matplotlib.pyplot.figure): Figure in which to make plot.
+        img (ndarray): 2D or 3D array. Image input into cellpose.
+        maski (int, ndarray): For image k, masks[k] output from Cellpose.eval, where 0=NO masks; 1,2,...=mask labels.
+        flowi (int, ndarray): For image k, flows[k][0] output from Cellpose.eval (RGB of flows).
+        channels (list of int, optional): Channels used to run Cellpose, no need to use if image is RGB. Defaults to [0, 0].
+        file_name (str, optional): File name of image. If file_name is not None, figure panels are saved. Defaults to None.
+        seg_norm (bool, optional): Improve cell visibility under labels. Defaults to False.
+    """
+    if not MATPLOTLIB_ENABLED:
+        raise ImportError(
+            "matplotlib not installed, install with 'pip install matplotlib'")
+    ax = fig.add_subplot(1, 4, 1)
+    img0 = img.copy()
+
+    if img0.shape[0] < 4:
+        img0 = np.transpose(img0, (1, 2, 0))
+    if img0.shape[-1] < 3 or img0.ndim < 3:
+        img0 = image_to_rgb(img0, channels=channels)
+    else:
+        if img0.max() <= 50.0:
+            img0 = np.uint8(np.clip(img0, 0, 1) * 255)
+    ax.imshow(img0)
+    ax.set_title("original image")
+    ax.axis("off")
+
+    outlines = utils.masks_to_outlines(maski)
+
+    overlay = mask_overlay(img0, maski)
+
+    ax = fig.add_subplot(1, 4, 2)
+    outX, outY = np.nonzero(outlines)
+    imgout = img0.copy()
+    imgout[outX, outY] = np.array([255, 0, 0])  # pure red
+
+    ax.imshow(imgout)
+    ax.set_title("predicted outlines")
+    ax.axis("off")
+
+    ax = fig.add_subplot(1, 4, 3)
+    ax.imshow(overlay)
+    ax.set_title("predicted masks")
+    ax.axis("off")
+
+    ax = fig.add_subplot(1, 4, 4)
+    ax.imshow(flowi)
+    ax.set_title("predicted cell pose")
+    ax.axis("off")
+
+    if file_name is not None:
+        save_path = os.path.splitext(file_name)[0]
+        io.imsave(save_path + "_overlay.jpg", overlay)
+        io.imsave(save_path + "_outlines.jpg", imgout)
+        io.imsave(save_path + "_flows.jpg", flowi)
+
+
+def mask_rgb(masks, colors=None):
+    """Masks in random RGB colors.
+
+    Args:
+        masks (int, 2D array): Masks where 0=NO masks; 1,2,...=mask labels.
+        colors (int, 2D array, optional): Size [nmasks x 3], each entry is a color in 0-255 range.
+
+    Returns:
+        RGB (uint8, 3D array): Array of masks overlaid on grayscale image.
+    """
+    if colors is not None:
+        if colors.max() > 1:
+            colors = np.float32(colors)
+            colors /= 255
+        colors = utils.rgb_to_hsv(colors)
+
+    HSV = np.zeros((masks.shape[0], masks.shape[1], 3), np.float32)
+    HSV[:, :, 2] = 1.0
+    for n in range(int(masks.max())):
+        ipix = (masks == n + 1).nonzero()
+        if colors is None:
+            HSV[ipix[0], ipix[1], 0] = np.random.rand()
+        else:
+            HSV[ipix[0], ipix[1], 0] = colors[n, 0]
+        HSV[ipix[0], ipix[1], 1] = np.random.rand() * 0.5 + 0.5
+        HSV[ipix[0], ipix[1], 2] = np.random.rand() * 0.5 + 0.5
+    RGB = (utils.hsv_to_rgb(HSV) * 255).astype(np.uint8)
+    return RGB
+
+
+def mask_overlay(img, masks, colors=None):
+    """Overlay masks on image (set image to grayscale).
+
+    Args:
+        img (int or float, 2D or 3D array): Image of size [Ly x Lx (x nchan)].
+        masks (int, 2D array): Masks where 0=NO masks; 1,2,...=mask labels.
+        colors (int, 2D array, optional): Size [nmasks x 3], each entry is a color in 0-255 range.
+
+    Returns:
+        RGB (uint8, 3D array): Array of masks overlaid on grayscale image.
+    """
+    if colors is not None:
+        if colors.max() > 1:
+            colors = np.float32(colors)
+            colors /= 255
+        colors = utils.rgb_to_hsv(colors)
+    if img.ndim > 2:
+        img = img.astype(np.float32).mean(axis=-1)
+    else:
+        img = img.astype(np.float32)
+
+    HSV = np.zeros((img.shape[0], img.shape[1], 3), np.float32)
+    HSV[:, :, 2] = np.clip((img / 255. if img.max() > 1 else img) * 1.5, 0, 1)
+    hues = np.linspace(0, 1, masks.max() + 1)[np.random.permutation(masks.max())]
+    for n in range(int(masks.max())):
+        ipix = (masks == n + 1).nonzero()
+        if colors is None:
+            HSV[ipix[0], ipix[1], 0] = hues[n]
+        else:
+            HSV[ipix[0], ipix[1], 0] = colors[n, 0]
+        HSV[ipix[0], ipix[1], 1] = 1.0
+    RGB = (utils.hsv_to_rgb(HSV) * 255).astype(np.uint8)
+    return RGB
+
+
+def image_to_rgb(img0, channels=[0, 0]):
+    """Converts image from 2 x Ly x Lx or Ly x Lx x 2 to RGB Ly x Lx x 3.
+
+    Args:
+        img0 (ndarray): Input image of shape 2 x Ly x Lx or Ly x Lx x 2.
+
+    Returns:
+        ndarray: RGB image of shape Ly x Lx x 3.
+
+    """
+    img = img0.copy()
+    img = img.astype(np.float32)
+    if img.ndim < 3:
+        img = img[:, :, np.newaxis]
+    if img.shape[0] < 5:
+        img = np.transpose(img, (1, 2, 0))
+    if channels[0] == 0:
+        img = img.mean(axis=-1)[:, :, np.newaxis]
+    for i in range(img.shape[-1]):
+        if np.ptp(img[:, :, i]) > 0:
+            img[:, :, i] = np.clip(transforms.normalize99(img[:, :, i]), 0, 1)
+            img[:, :, i] = np.clip(img[:, :, i], 0, 1)
+    img *= 255
+    img = np.uint8(img)
+    RGB = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
+    if img.shape[-1] == 1:
+        RGB = np.tile(img, (1, 1, 3))
+    else:
+        RGB[:, :, channels[0] - 1] = img[:, :, 0]
+        if channels[1] > 0:
+            RGB[:, :, channels[1] - 1] = img[:, :, 1]
+    return RGB
+
+
+def interesting_patch(mask, bsize=130):
+    """
+    Get patch of size bsize x bsize with most masks.
+
+    Args:
+        mask (ndarray): Input mask.
+        bsize (int): Size of the patch.
+
+    Returns:
+        tuple: Patch coordinates (y, x).
+
+    """
+    Ly, Lx = mask.shape
+    m = np.float32(mask > 0)
+    m = gaussian_filter(m, bsize / 2)
+    y, x = np.unravel_index(np.argmax(m), m.shape)
+    ycent = max(bsize // 2, min(y, Ly - bsize // 2))
+    xcent = max(bsize // 2, min(x, Lx - bsize // 2))
+    patch = [
+        np.arange(ycent - bsize // 2, ycent + bsize // 2, 1, int),
+        np.arange(xcent - bsize // 2, xcent + bsize // 2, 1, int)
+    ]
+    return patch
+
+
+def disk(med, r, Ly, Lx):
+    """Returns the pixels of a disk with a given radius and center.
+
+    Args:
+        med (tuple): The center coordinates of the disk.
+        r (float): The radius of the disk.
+        Ly (int): The height of the image.
+        Lx (int): The width of the image.
+
+    Returns:
+        tuple: A tuple containing the y and x coordinates of the pixels within the disk.
+
+    """
+    yy, xx = np.meshgrid(np.arange(0, Ly, 1, int), np.arange(0, Lx, 1, int),
+                         indexing="ij")
+    inds = ((yy - med[0])**2 + (xx - med[1])**2)**0.5 <= r
+    y = yy[inds].flatten()
+    x = xx[inds].flatten()
+    return y, x
+
+
+def outline_view(img0, maski, color=[1, 0, 0], mode="inner"):
+    """
+    Generates a red outline overlay onto the image.
+
+    Args:
+        img0 (numpy.ndarray): The input image.
+        maski (numpy.ndarray): The mask representing the region of interest.
+        color (list, optional): The color of the outline overlay. Defaults to [1, 0, 0] (red).
+        mode (str, optional): The mode for generating the outline. Defaults to "inner".
+
+    Returns:
+        numpy.ndarray: The image with the red outline overlay.
+
+    """
+    if img0.ndim == 2:
+        img0 = np.stack([img0] * 3, axis=-1)
+    elif img0.ndim != 3:
+        raise ValueError("img0 not right size (must have ndim 2 or 3)")
+
+    if SKIMAGE_ENABLED:
+        outlines = find_boundaries(maski, mode=mode)
+    else:
+        outlines = utils.masks_to_outlines(maski, mode=mode)
+    outY, outX = np.nonzero(outlines)
+    imgout = img0.copy()
+    imgout[outY, outX] = np.array(color)
+
+    return imgout

models/seg_post_model/cellpose/transforms.py

@@ -0,0 +1,1261 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import logging
+
+import cv2
+import numpy as np
+import torch
+from scipy.ndimage import gaussian_filter1d
+from torch.fft import fft2, fftshift, ifft2
+
+transforms_logger = logging.getLogger(__name__)
+
+
+def _taper_mask(ly=224, lx=224, sig=7.5):
+    """
+    Generate a taper mask.
+
+    Args:
+        ly (int): The height of the mask. Default is 224.
+        lx (int): The width of the mask. Default is 224.
+        sig (float): The sigma value for the tapering function. Default is 7.5.
+
+    Returns:
+        numpy.ndarray: The taper mask.
+
+    """
+    bsize = max(224, max(ly, lx))
+    xm = np.arange(bsize)
+    xm = np.abs(xm - xm.mean())
+    mask = 1 / (1 + np.exp((xm - (bsize / 2 - 20)) / sig))
+    mask = mask * mask[:, np.newaxis]
+    mask = mask[bsize // 2 - ly // 2:bsize // 2 + ly // 2 + ly % 2,
+                bsize // 2 - lx // 2:bsize // 2 + lx // 2 + lx % 2]
+    return mask
+
+
+def unaugment_tiles(y):
+    """Reverse test-time augmentations for averaging (includes flipping of flowsY and flowsX).
+
+    Args:
+        y (float32): Array of shape (ntiles_y, ntiles_x, chan, Ly, Lx) where chan = (flowsY, flowsX, cell prob).
+
+    Returns:
+        float32: Array of shape (ntiles_y, ntiles_x, chan, Ly, Lx).
+
+    """
+    for j in range(y.shape[0]):
+        for i in range(y.shape[1]):
+            if j % 2 == 0 and i % 2 == 1:
+                y[j, i] = y[j, i, :, ::-1, :]
+                y[j, i, 0] *= -1
+            elif j % 2 == 1 and i % 2 == 0:
+                y[j, i] = y[j, i, :, :, ::-1]
+                y[j, i, 1] *= -1
+            elif j % 2 == 1 and i % 2 == 1:
+                y[j, i] = y[j, i, :, ::-1, ::-1]
+                y[j, i, 0] *= -1
+                y[j, i, 1] *= -1
+    return y
+
+
+def average_tiles(y, ysub, xsub, Ly, Lx):
+    """
+    Average the results of the network over tiles.
+
+    Args:
+        y (float): Output of cellpose network for each tile. Shape: [ntiles x nclasses x bsize x bsize]
+        ysub (list): List of arrays with start and end of tiles in Y of length ntiles
+        xsub (list): List of arrays with start and end of tiles in X of length ntiles
+        Ly (int): Size of pre-tiled image in Y (may be larger than original image if image size is less than bsize)
+        Lx (int): Size of pre-tiled image in X (may be larger than original image if image size is less than bsize)
+
+    Returns:
+        yf (float32): Network output averaged over tiles. Shape: [nclasses x Ly x Lx]
+    """
+    Navg = np.zeros((Ly, Lx))
+    yf = np.zeros((y.shape[1], Ly, Lx), np.float32)
+    # taper edges of tiles
+    mask = _taper_mask(ly=y.shape[-2], lx=y.shape[-1])
+    for j in range(len(ysub)):
+        yf[:, ysub[j][0]:ysub[j][1], xsub[j][0]:xsub[j][1]] += y[j] * mask
+        Navg[ysub[j][0]:ysub[j][1], xsub[j][0]:xsub[j][1]] += mask
+    yf /= Navg
+    return yf
+
+
+def make_tiles(imgi, bsize=224, augment=False, tile_overlap=0.1):
+    """Make tiles of image to run at test-time.
+
+    Args:
+        imgi (np.ndarray): Array of shape (nchan, Ly, Lx) representing the input image.
+        bsize (int, optional): Size of tiles. Defaults to 224.
+        augment (bool, optional): Whether to flip tiles and set tile_overlap=2. Defaults to False.
+        tile_overlap (float, optional): Fraction of overlap of tiles. Defaults to 0.1.
+
+    Returns:
+        A tuple containing (IMG, ysub, xsub, Ly, Lx):
+        IMG (np.ndarray): Array of shape (ntiles, nchan, bsize, bsize) representing the tiles.
+        ysub (list): List of arrays with start and end of tiles in Y of length ntiles.
+        xsub (list): List of arrays with start and end of tiles in X of length ntiles.
+        Ly (int): Height of the input image.
+        Lx (int): Width of the input image.
+    """
+    nchan, Ly, Lx = imgi.shape
+    if augment:
+        bsize = np.int32(bsize)
+        # pad if image smaller than bsize
+        if Ly < bsize:
+            imgi = np.concatenate((imgi, np.zeros((nchan, bsize - Ly, Lx))), axis=1)
+            Ly = bsize
+        if Lx < bsize:
+            imgi = np.concatenate((imgi, np.zeros((nchan, Ly, bsize - Lx))), axis=2)
+        Ly, Lx = imgi.shape[-2:]
+        
+        # tiles overlap by half of tile size
+        ny = max(2, int(np.ceil(2. * Ly / bsize)))
+        nx = max(2, int(np.ceil(2. * Lx / bsize)))
+        ystart = np.linspace(0, Ly - bsize, ny).astype(int)
+        xstart = np.linspace(0, Lx - bsize, nx).astype(int)
+
+        ysub = []
+        xsub = []
+
+        # flip tiles so that overlapping segments are processed in rotation
+        IMG = np.zeros((len(ystart), len(xstart), nchan, bsize, bsize), np.float32)
+        for j in range(len(ystart)):
+            for i in range(len(xstart)):
+                ysub.append([ystart[j], ystart[j] + bsize])
+                xsub.append([xstart[i], xstart[i] + bsize])
+                IMG[j, i] = imgi[:, ysub[-1][0]:ysub[-1][1], xsub[-1][0]:xsub[-1][1]]
+                # flip tiles to allow for augmentation of overlapping segments
+                if j % 2 == 0 and i % 2 == 1:
+                    IMG[j, i] = IMG[j, i, :, ::-1, :]
+                elif j % 2 == 1 and i % 2 == 0:
+                    IMG[j, i] = IMG[j, i, :, :, ::-1]
+                elif j % 2 == 1 and i % 2 == 1:
+                    IMG[j, i] = IMG[j, i, :, ::-1, ::-1]
+    else:
+        tile_overlap = min(0.5, max(0.05, tile_overlap))
+        bsizeY, bsizeX = min(bsize, Ly), min(bsize, Lx)
+        bsizeY = np.int32(bsizeY)
+        bsizeX = np.int32(bsizeX)
+        # tiles overlap by 10% tile size
+        ny = 1 if Ly <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Ly / bsize))
+        nx = 1 if Lx <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Lx / bsize))
+        ystart = np.linspace(0, Ly - bsizeY, ny).astype(int)
+        xstart = np.linspace(0, Lx - bsizeX, nx).astype(int)
+
+        ysub = []
+        xsub = []
+        IMG = np.zeros((len(ystart), len(xstart), nchan, bsizeY, bsizeX), np.float32)
+        for j in range(len(ystart)):
+            for i in range(len(xstart)):
+                ysub.append([ystart[j], ystart[j] + bsizeY])
+                xsub.append([xstart[i], xstart[i] + bsizeX])
+                IMG[j, i] = imgi[:, ysub[-1][0]:ysub[-1][1], xsub[-1][0]:xsub[-1][1]]
+
+    return IMG, ysub, xsub, Ly, Lx
+
+
+def normalize99(Y, lower=1, upper=99, copy=True, downsample=False):
+    """
+    Normalize the image so that 0.0 corresponds to the 1st percentile and 1.0 corresponds to the 99th percentile.
+
+    Args:
+        Y (ndarray): The input image (for downsample, use [Ly x Lx] or [Lz x Ly x Lx]).
+        lower (int, optional): The lower percentile. Defaults to 1.
+        upper (int, optional): The upper percentile. Defaults to 99.
+        copy (bool, optional): Whether to create a copy of the input image. Defaults to True.
+        downsample (bool, optional): Whether to downsample image to compute percentiles. Defaults to False.
+
+    Returns:
+        ndarray: The normalized image.
+    """
+    X = Y.copy() if copy else Y
+    X = X.astype("float32") if X.dtype!="float64" and X.dtype!="float32" else X
+    if downsample and X.size > 224**3:
+        nskip = [max(1, X.shape[i] // 224) for i in range(X.ndim)]
+        nskip[0] = max(1, X.shape[0] // 50) if X.ndim == 3 else nskip[0]
+        slc = tuple([slice(0, X.shape[i], nskip[i]) for i in range(X.ndim)])
+        x01 = np.percentile(X[slc], lower)
+        x99 = np.percentile(X[slc], upper)
+    else:
+        x01 = np.percentile(X, lower)
+        x99 = np.percentile(X, upper)
+    if x99 - x01 > 1e-3:
+        X -= x01 
+        X /= (x99 - x01)
+    else:
+        X[:] = 0
+    return X
+
+
+def normalize99_tile(img, blocksize=100, lower=1., upper=99., tile_overlap=0.1,
+                     norm3D=False, smooth3D=1, is3D=False):
+    """Compute normalization like normalize99 function but in tiles.
+
+    Args:
+        img (numpy.ndarray): Array of shape (Lz x) Ly x Lx (x nchan) containing the image.
+        blocksize (float, optional): Size of tiles. Defaults to 100.
+        lower (float, optional): Lower percentile for normalization. Defaults to 1.0.
+        upper (float, optional): Upper percentile for normalization. Defaults to 99.0.
+        tile_overlap (float, optional): Fraction of overlap of tiles. Defaults to 0.1.
+        norm3D (bool, optional): Use same tiled normalization for each z-plane. Defaults to False.
+        smooth3D (int, optional): Smoothing factor for 3D normalization. Defaults to 1.
+        is3D (bool, optional): Set to True if image is a 3D stack. Defaults to False.
+
+    Returns:
+        numpy.ndarray: Normalized image array of shape (Lz x) Ly x Lx (x nchan).
+    """
+    is1c = True if img.ndim == 2 or (is3D and img.ndim == 3) else False
+    is3D = True if img.ndim > 3 or (is3D and img.ndim == 3) else False
+    img = img[..., np.newaxis] if is1c else img
+    img = img[np.newaxis, ...] if img.ndim == 3 else img
+    Lz, Ly, Lx, nchan = img.shape
+
+    tile_overlap = min(0.5, max(0.05, tile_overlap))
+    blocksizeY, blocksizeX = min(blocksize, Ly), min(blocksize, Lx)
+    blocksizeY = np.int32(blocksizeY)
+    blocksizeX = np.int32(blocksizeX)
+    # tiles overlap by 10% tile size
+    ny = 1 if Ly <= blocksize else int(np.ceil(
+        (1. + 2 * tile_overlap) * Ly / blocksize))
+    nx = 1 if Lx <= blocksize else int(np.ceil(
+        (1. + 2 * tile_overlap) * Lx / blocksize))
+    ystart = np.linspace(0, Ly - blocksizeY, ny).astype(int)
+    xstart = np.linspace(0, Lx - blocksizeX, nx).astype(int)
+    ysub = []
+    xsub = []
+    for j in range(len(ystart)):
+        for i in range(len(xstart)):
+            ysub.append([ystart[j], ystart[j] + blocksizeY])
+            xsub.append([xstart[i], xstart[i] + blocksizeX])
+
+    x01_tiles_z = []
+    x99_tiles_z = []
+    for z in range(Lz):
+        IMG = np.zeros((len(ystart), len(xstart), blocksizeY, blocksizeX, nchan),
+                       "float32")
+        k = 0
+        for j in range(len(ystart)):
+            for i in range(len(xstart)):
+                IMG[j, i] = img[z, ysub[k][0]:ysub[k][1], xsub[k][0]:xsub[k][1], :]
+                k += 1
+        x01_tiles = np.percentile(IMG, lower, axis=(-3, -2))
+        x99_tiles = np.percentile(IMG, upper, axis=(-3, -2))
+
+        # fill areas with small differences with neighboring squares
+        to_fill = np.zeros(x01_tiles.shape[:2], "bool")
+        for c in range(nchan):
+            to_fill = x99_tiles[:, :, c] - x01_tiles[:, :, c] < +1e-3
+            if to_fill.sum() > 0 and to_fill.sum() < x99_tiles[:, :, c].size:
+                fill_vals = np.nonzero(to_fill)
+                fill_neigh = np.nonzero(~to_fill)
+                nearest_neigh = (
+                    (fill_vals[0] - fill_neigh[0][:, np.newaxis])**2 +
+                    (fill_vals[1] - fill_neigh[1][:, np.newaxis])**2).argmin(axis=0)
+                x01_tiles[fill_vals[0], fill_vals[1],
+                          c] = x01_tiles[fill_neigh[0][nearest_neigh],
+                                         fill_neigh[1][nearest_neigh], c]
+                x99_tiles[fill_vals[0], fill_vals[1],
+                          c] = x99_tiles[fill_neigh[0][nearest_neigh],
+                                         fill_neigh[1][nearest_neigh], c]
+            elif to_fill.sum() > 0 and to_fill.sum() == x99_tiles[:, :, c].size:
+                x01_tiles[:, :, c] = 0
+                x99_tiles[:, :, c] = 1
+        x01_tiles_z.append(x01_tiles)
+        x99_tiles_z.append(x99_tiles)
+
+    x01_tiles_z = np.array(x01_tiles_z)
+    x99_tiles_z = np.array(x99_tiles_z)
+    # do not smooth over z-axis if not normalizing separately per plane
+    for a in range(2):
+        x01_tiles_z = gaussian_filter1d(x01_tiles_z, 1, axis=a)
+        x99_tiles_z = gaussian_filter1d(x99_tiles_z, 1, axis=a)
+    if norm3D:
+        smooth3D = 1 if smooth3D == 0 else smooth3D
+        x01_tiles_z = gaussian_filter1d(x01_tiles_z, smooth3D, axis=a)
+        x99_tiles_z = gaussian_filter1d(x99_tiles_z, smooth3D, axis=a)
+
+    if not norm3D and Lz > 1:
+        x01 = np.zeros((len(x01_tiles_z), Ly, Lx, nchan), "float32")
+        x99 = np.zeros((len(x01_tiles_z), Ly, Lx, nchan), "float32")
+        for z in range(Lz):
+            x01_rsz = cv2.resize(x01_tiles_z[z], (Lx, Ly),
+                                 interpolation=cv2.INTER_LINEAR)
+            x01[z] = x01_rsz[..., np.newaxis] if nchan == 1 else x01_rsz
+            x99_rsz = cv2.resize(x99_tiles_z[z], (Lx, Ly),
+                                 interpolation=cv2.INTER_LINEAR)
+            x99[z] = x99_rsz[..., np.newaxis] if nchan == 1 else x01_rsz
+        if (x99 - x01).min() < 1e-3:
+            raise ZeroDivisionError(
+                "cannot use norm3D=False with tile_norm, sample is too sparse; set norm3D=True or tile_norm=0"
+            )
+    else:
+        x01 = cv2.resize(x01_tiles_z.mean(axis=0), (Lx, Ly),
+                         interpolation=cv2.INTER_LINEAR)
+        x99 = cv2.resize(x99_tiles_z.mean(axis=0), (Lx, Ly),
+                         interpolation=cv2.INTER_LINEAR)
+        if x01.ndim < 3:
+            x01 = x01[..., np.newaxis]
+            x99 = x99[..., np.newaxis]
+
+    if is1c:
+        img, x01, x99 = img.squeeze(), x01.squeeze(), x99.squeeze()
+    elif not is3D:
+        img, x01, x99 = img[0], x01[0], x99[0]
+
+    # normalize
+    img -= x01 
+    img /= (x99 - x01)
+
+    return img
+
+
+def gaussian_kernel(sigma, Ly, Lx, device=torch.device("cpu")):
+    """
+    Generates a 2D Gaussian kernel.
+
+    Args:
+        sigma (float): Standard deviation of the Gaussian distribution.
+        Ly (int): Number of pixels in the y-axis.
+        Lx (int): Number of pixels in the x-axis.
+        device (torch.device, optional): Device to store the kernel tensor. Defaults to torch.device("cpu").
+
+    Returns:
+        torch.Tensor: 2D Gaussian kernel tensor.
+
+    """
+    y = torch.linspace(-Ly / 2, Ly / 2 + 1, Ly, device=device)
+    x = torch.linspace(-Ly / 2, Ly / 2 + 1, Lx, device=device)
+    y, x = torch.meshgrid(y, x, indexing="ij")
+    kernel = torch.exp(-(y**2 + x**2) / (2 * sigma**2))
+    kernel /= kernel.sum()
+    return kernel
+
+
+def smooth_sharpen_img(img, smooth_radius=6, sharpen_radius=12,
+                       device=torch.device("cpu"), is3D=False):
+    """Sharpen blurry images with surround subtraction and/or smooth noisy images.
+
+    Args:
+        img (float32): Array that's (Lz x) Ly x Lx (x nchan).
+        smooth_radius (float, optional): Size of gaussian smoothing filter, recommended to be 1/10-1/4 of cell diameter
+            (if also sharpening, should be 2-3x smaller than sharpen_radius). Defaults to 6.
+        sharpen_radius (float, optional): Size of gaussian surround filter, recommended to be 1/8-1/2 of cell diameter
+            (if also smoothing, should be 2-3x larger than smooth_radius). Defaults to 12.
+        device (torch.device, optional): Device on which to perform sharpening.
+            Will be faster on GPU but need to ensure GPU has RAM for image. Defaults to torch.device("cpu").
+        is3D (bool, optional): If image is 3D stack (only necessary to set if img.ndim==3). Defaults to False.
+
+    Returns:
+        img_sharpen (float32): Array that's (Lz x) Ly x Lx (x nchan).
+    """
+    img_sharpen = torch.from_numpy(img.astype("float32")).to(device)
+    shape = img_sharpen.shape
+
+    is1c = True if img_sharpen.ndim == 2 or (is3D and img_sharpen.ndim == 3) else False
+    is3D = True if img_sharpen.ndim > 3 or (is3D and img_sharpen.ndim == 3) else False
+    img_sharpen = img_sharpen.unsqueeze(-1) if is1c else img_sharpen
+    img_sharpen = img_sharpen.unsqueeze(0) if img_sharpen.ndim == 3 else img_sharpen
+    Lz, Ly, Lx, nchan = img_sharpen.shape
+
+    if smooth_radius > 0:
+        kernel = gaussian_kernel(smooth_radius, Ly, Lx, device=device)
+        if sharpen_radius > 0:
+            kernel += -1 * gaussian_kernel(sharpen_radius, Ly, Lx, device=device)
+    elif sharpen_radius > 0:
+        kernel = -1 * gaussian_kernel(sharpen_radius, Ly, Lx, device=device)
+        kernel[Ly // 2, Lx // 2] = 1
+
+    fhp = fft2(kernel)
+    for z in range(Lz):
+        for c in range(nchan):
+            img_filt = torch.real(ifft2(
+                fft2(img_sharpen[z, :, :, c]) * torch.conj(fhp)))
+            img_filt = fftshift(img_filt)
+            img_sharpen[z, :, :, c] = img_filt
+
+    img_sharpen = img_sharpen.reshape(shape)
+    return img_sharpen.cpu().numpy()
+
+
+def move_axis(img, m_axis=-1, first=True):
+    """ move axis m_axis to first or last position """
+    if m_axis == -1:
+        m_axis = img.ndim - 1
+    m_axis = min(img.ndim - 1, m_axis)
+    axes = np.arange(0, img.ndim)
+    if first:
+        axes[1:m_axis + 1] = axes[:m_axis]
+        axes[0] = m_axis
+    else:
+        axes[m_axis:-1] = axes[m_axis + 1:]
+        axes[-1] = m_axis
+    img = img.transpose(tuple(axes))
+    return img
+
+
+def move_min_dim(img, force=False):
+    """Move the minimum dimension last as channels if it is less than 10 or force is True.
+
+    Args:
+        img (ndarray): The input image.
+        force (bool, optional): If True, the minimum dimension will always be moved.
+            Defaults to False.
+
+    Returns:
+        ndarray: The image with the minimum dimension moved to the last axis as channels.
+    """
+    if len(img.shape) > 2:
+        min_dim = min(img.shape)
+        if min_dim < 10 or force:
+            if img.shape[-1] == min_dim:
+                channel_axis = -1
+            else:
+                channel_axis = (img.shape).index(min_dim)
+            img = move_axis(img, m_axis=channel_axis, first=False)
+    return img
+
+
+def update_axis(m_axis, to_squeeze, ndim):
+    """
+    Squeeze the axis value based on the given parameters.
+
+    Args:
+        m_axis (int): The current axis value.
+        to_squeeze (numpy.ndarray): An array of indices to squeeze.
+        ndim (int): The number of dimensions.
+
+    Returns:
+        m_axis (int or None): The updated axis value.
+    """
+    if m_axis == -1:
+        m_axis = ndim - 1
+    if (to_squeeze == m_axis).sum() == 1:
+        m_axis = None
+    else:
+        inds = np.ones(ndim, bool)
+        inds[to_squeeze] = False
+        m_axis = np.nonzero(np.arange(0, ndim)[inds] == m_axis)[0]
+        if len(m_axis) > 0:
+            m_axis = m_axis[0]
+        else:
+            m_axis = None
+    return m_axis
+
+
+def _convert_image_3d(x, channel_axis=None, z_axis=None):
+    """
+    Convert a 3D or 4D image array to have dimensions ordered as (Z, X, Y, C).
+
+    Arrays of ndim=3 are assumed to be grayscale and must be specified with z_axis. 
+    Arrays of ndim=4 must have both `channel_axis` and `z_axis` specified.
+    
+    Args:
+        x (numpy.ndarray): Input image array. Must be either 3D (assumed to be grayscale 3D) or 4D. 
+        channel_axis (int): The axis index corresponding to the channel dimension in the input array. \
+            Must be specified for 4D images.
+        z_axis (int): The axis index corresponding to the depth (Z) dimension in the input array. \
+            Must be specified for both 3D and 4D images.
+
+    Returns:
+        numpy.ndarray: A 4D image array with dimensions ordered as (Z, X, Y, C), where C is the channel 
+        dimension. If the input has fewer than 3 channels, the output will be padded with zeros to \
+            have 3 channels. If the input has more than 3 channels, only the first 3 channels will be retained.
+    
+    Raises:
+        ValueError: If `z_axis` is not specified for 3D images. If either `channel_axis` or `z_axis` \
+            is not specified for 4D images. If the input image does not have 3 or 4 dimensions.
+
+    Notes:
+        - For 3D images (ndim=3), the function assumes the input is grayscale and adds a singleton channel dimension.
+        - The function reorders the dimensions of the input array to ensure the output has the desired (Z, X, Y, C) order.
+        - If the number of channels is not equal to 3, the function either truncates or pads the \
+            channels to ensure the output has exactly 3 channels.
+    """
+
+    if x.ndim < 3:
+        raise ValueError(f"Input image must have at least 3 dimensions, input shape: {x.shape}, ndim={x.ndim}")
+    
+    if z_axis is not None and z_axis < 0:
+        z_axis += x.ndim
+
+    # if image is ndim==3, assume it is greyscale 3D and use provided z_axis
+    if x.ndim == 3 and z_axis is not None:
+        # add in channel axis
+        x = x[..., np.newaxis]
+        channel_axis = 3
+    elif x.ndim == 3 and z_axis is None:
+        raise ValueError("z_axis must be specified when segmenting 3D images of ndim=3")
+
+
+    if channel_axis is None or z_axis is None:
+        raise ValueError("For 4D images, both `channel_axis` and `z_axis` must be explicitly specified. Please provide values for both parameters.")
+    if channel_axis is not None and channel_axis < 0:
+        channel_axis += x.ndim
+    if channel_axis is None or channel_axis >= x.ndim:
+        raise IndexError(f"channel_axis {channel_axis} is out of bounds for input array with {x.ndim} dimensions")
+    assert x.ndim == 4, f"input image must have ndim == 4, ndim={x.ndim}"
+    
+    x_dim_shapes = list(x.shape)
+    num_z_layers = x_dim_shapes[z_axis]
+    num_channels = x_dim_shapes[channel_axis]
+    x_xy_axes = [i for i in range(x.ndim)]
+    
+    # need to remove the z and channels from the shapes:
+    # delete the one with the bigger index first 
+    if z_axis > channel_axis:
+        del x_dim_shapes[z_axis]
+        del x_dim_shapes[channel_axis]
+
+        del x_xy_axes[z_axis]
+        del x_xy_axes[channel_axis]
+
+    else: 
+        del x_dim_shapes[channel_axis]
+        del x_dim_shapes[z_axis]
+
+        del x_xy_axes[channel_axis]
+        del x_xy_axes[z_axis]
+
+    x = x.transpose((z_axis, x_xy_axes[0], x_xy_axes[1], channel_axis))
+
+    # Handle cases with not 3 channels:
+    if num_channels != 3:
+        x_chans_to_copy = min(3, num_channels)
+
+        if num_channels > 3:
+            transforms_logger.warning("more than 3 channels provided, only segmenting on first 3 channels")
+            x = x[..., :x_chans_to_copy]
+        else: 
+            # less than 3 channels: pad up to 
+            pad_width = [(0, 0), (0, 0), (0, 0), (0, 3 - x_chans_to_copy)]
+            x = np.pad(x, pad_width, mode='constant', constant_values=0)
+
+    return x
+
+
+def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
+    """Converts the image to have the z-axis first, channels last. Image will be converted to 3 channels if it is not already.
+    If more than 3 channels are provided, only the first 3 channels will be used. 
+
+    Accepts: 
+        - 2D images with no channel dimension: `z_axis` and `channel_axis` must be `None`
+        - 2D images with channel dimension: `channel_axis` will be guessed between first or last axis, can also specify `channel_axis`. `z_axis` must be `None`
+        - 3D images with or without channels: 
+
+    Args:
+        x (numpy.ndarray or torch.Tensor): The input image.
+        channel_axis (int or None): The axis of the channels in the input image. If None, the axis is determined automatically.
+        z_axis (int or None): The axis of the z-dimension in the input image. If None, the axis is determined automatically.
+        do_3D (bool): Whether to process the image in 3D mode. Defaults to False.
+
+    Returns:
+        numpy.ndarray: The converted image.
+
+    Raises:
+        ValueError: If the input image is 2D and do_3D is True.
+        ValueError: If the input image is 4D and do_3D is False.
+    """
+
+    # check if image is a torch array instead of numpy array, convert to numpy
+    ndim = x.ndim
+    if torch.is_tensor(x):
+        transforms_logger.warning("torch array used as input, converting to numpy")
+        x = x.cpu().numpy()
+
+    # should be 2D
+    if z_axis is not None and not do_3D:
+        raise ValueError("2D image provided, but z_axis is not None. Set z_axis=None to process 2D images of ndim=2 or 3.")
+
+    # make sure that channel_axis and z_axis are specified if 3D
+    if ndim == 4 and not do_3D:
+        raise ValueError("3D input image provided, but do_3D is False. Set do_3D=True to process 3D images. ndims=4")
+
+    # make sure that channel_axis and z_axis are specified if 3D
+    if do_3D:
+        return _convert_image_3d(x, channel_axis=channel_axis, z_axis=z_axis)
+    
+    ######################## 2D reshaping ########################
+    # if user specifies channel axis, return early
+    if channel_axis is not None:
+        if ndim == 2:
+            raise ValueError("2D image provided, but channel_axis is not None. Set channel_axis=None to process 2D images of ndim=2.")
+        
+        # Put channel axis last:
+        # Find the indices of the dims that need to be put in dim 0 and 1
+        n_channels = x.shape[channel_axis]
+        x_shape_dims = list(x.shape)
+        del x_shape_dims[channel_axis]
+        dimension_indicies = [i for i in range(x.ndim)]
+        del dimension_indicies[channel_axis]
+
+        x = x.transpose((dimension_indicies[0], dimension_indicies[1], channel_axis))
+
+        if n_channels != 3:
+            x_chans_to_copy = min(3, n_channels)
+
+            if n_channels > 3: 
+                transforms_logger.warning("more than 3 channels provided, only segmenting on first 3 channels")
+                x = x[..., :x_chans_to_copy]
+            else: 
+                x_out = np.zeros((x_shape_dims[0], x_shape_dims[1], 3), dtype=x.dtype)
+                x_out[..., :x_chans_to_copy] = x[...]
+                x = x_out
+                del x_out
+
+        return x
+
+    # do image padding and channel conversion
+    if ndim == 2:
+        # grayscale image, make 3 channels
+        x_out = np.zeros((x.shape[0], x.shape[1], 3), dtype=x.dtype)
+        x_out[..., 0] = x
+        x = x_out
+        del x_out
+    elif ndim == 3:
+        # assume 2d with channels
+        # find dim with smaller size between first and last dims
+        move_channel_axis = x.shape[0] < x.shape[2]
+        if move_channel_axis:
+            x = x.transpose((1, 2, 0))
+
+        # zero padding up to 3 channels: 
+        num_channels = x.shape[-1]
+        if num_channels > 3: 
+            transforms_logger.warning("Found more than 3 channels, only using first 3")
+            num_channels = 3
+        x_out = np.zeros((x.shape[0], x.shape[1], 3), dtype=x.dtype)
+        x_out[..., :num_channels] = x[..., :num_channels]
+        x = x_out
+        del x_out
+    else:
+        # something is wrong: yell
+        expected_shapes = "2D (H, W), 3D (H, W, C), or 4D (Z, H, W, C)"
+        transforms_logger.critical(f"ERROR: Unexpected image shape: {str(x.shape)}. Expected shapes: {expected_shapes}")
+        raise ValueError(f"ERROR: Unexpected image shape: {str(x.shape)}. Expected shapes: {expected_shapes}")
+
+    return x
+    
+
+def normalize_img(img, normalize=True, norm3D=True, invert=False, lowhigh=None,
+                  percentile=(1., 99.), sharpen_radius=0, smooth_radius=0,
+                  tile_norm_blocksize=0, tile_norm_smooth3D=1, axis=-1):
+    """Normalize each channel of the image with optional inversion, smoothing, and sharpening.
+
+    Args:
+        img (ndarray): The input image. It should have at least 3 dimensions.
+            If it is 4-dimensional, it assumes the first non-channel axis is the Z dimension.
+        normalize (bool, optional): Whether to perform normalization. Defaults to True.
+        norm3D (bool, optional): Whether to normalize in 3D. If True, the entire 3D stack will
+            be normalized per channel. If False, normalization is applied per Z-slice. Defaults to False.
+        invert (bool, optional): Whether to invert the image. Useful if cells are dark instead of bright.
+            Defaults to False.
+        lowhigh (tuple or ndarray, optional): The lower and upper bounds for normalization.
+            Can be a tuple of two values (applied to all channels) or an array of shape (nchan, 2)
+            for per-channel normalization. Incompatible with smoothing and sharpening.
+            Defaults to None.
+        percentile (tuple, optional): The lower and upper percentiles for normalization. If provided, it should be
+            a tuple of two values. Each value should be between 0 and 100. Defaults to (1.0, 99.0).
+        sharpen_radius (int, optional): The radius for sharpening the image. Defaults to 0.
+        smooth_radius (int, optional): The radius for smoothing the image. Defaults to 0.
+        tile_norm_blocksize (int, optional): The block size for tile-based normalization. Defaults to 0.
+        tile_norm_smooth3D (int, optional): The smoothness factor for tile-based normalization in 3D. Defaults to 1.
+        axis (int, optional): The channel axis to loop over for normalization. Defaults to -1.
+
+    Returns:
+        ndarray: The normalized image of the same size.
+
+    Raises:
+        ValueError: If the image has less than 3 dimensions.
+        ValueError: If the provided lowhigh or percentile values are invalid.
+        ValueError: If the image is inverted without normalization.
+
+    """
+    if img.ndim < 3:
+        error_message = "Image needs to have at least 3 dimensions"
+        transforms_logger.critical(error_message)
+        raise ValueError(error_message)
+
+    img_norm = img if img.dtype=="float32" else img.astype(np.float32)
+    if axis != -1 and axis != img_norm.ndim - 1:
+        img_norm = np.moveaxis(img_norm, axis, -1)  # Move channel axis to last
+
+    nchan = img_norm.shape[-1]
+
+    # Validate and handle lowhigh bounds
+    if lowhigh is not None:
+        lowhigh = np.array(lowhigh)
+        if lowhigh.shape == (2,):
+            lowhigh = np.tile(lowhigh, (nchan, 1))  # Expand to per-channel bounds
+        elif lowhigh.shape != (nchan, 2):
+            error_message = "`lowhigh` must have shape (2,) or (nchan, 2)"
+            transforms_logger.critical(error_message)
+            raise ValueError(error_message)
+
+    # Validate percentile
+    if percentile is None:
+        percentile = (1.0, 99.0)
+    elif not (0 <= percentile[0] < percentile[1] <= 100):
+        error_message = "Invalid percentile range, should be between 0 and 100"
+        transforms_logger.critical(error_message)
+        raise ValueError(error_message)
+
+    # Apply normalization based on lowhigh or percentile
+    cgood = np.zeros(nchan, "bool")
+    if lowhigh is not None:
+        for c in range(nchan):
+            lower = lowhigh[c, 0]
+            upper = lowhigh[c, 1]
+            img_norm[..., c] -= lower 
+            img_norm[..., c] /= (upper - lower)
+            cgood[c] = True
+    else:
+        # Apply sharpening and smoothing if specified
+        if sharpen_radius > 0 or smooth_radius > 0:
+            img_norm = smooth_sharpen_img(
+                img_norm, sharpen_radius=sharpen_radius, smooth_radius=smooth_radius
+            )
+
+        # Apply tile-based normalization or standard normalization
+        if tile_norm_blocksize > 0:
+            img_norm = normalize99_tile(
+                img_norm,
+                blocksize=tile_norm_blocksize,
+                lower=percentile[0],
+                upper=percentile[1],
+                smooth3D=tile_norm_smooth3D,
+                norm3D=norm3D,
+            )
+            cgood[:] = True
+        elif normalize:
+            if img_norm.ndim == 3 or norm3D:  # i.e. if YXC, or ZYXC with norm3D=True
+                for c in range(nchan):
+                    if np.ptp(img_norm[..., c]) > 0.:
+                        img_norm[..., c] = normalize99(
+                            img_norm[..., c],
+                            lower=percentile[0],
+                            upper=percentile[1],
+                            copy=False, downsample=True,
+                        )
+                        cgood[c] = True
+            else:  # i.e. if ZYXC with norm3D=False then per Z-slice
+                for z in range(img_norm.shape[0]):
+                    for c in range(nchan):
+                        if np.ptp(img_norm[z, ..., c]) > 0.:
+                            img_norm[z, ..., c] = normalize99(
+                                img_norm[z, ..., c],
+                                lower=percentile[0],
+                                upper=percentile[1],
+                                copy=False, downsample=True,
+                            )
+                            cgood[c] = True
+
+
+    if invert:
+        if lowhigh is not None or tile_norm_blocksize > 0 or normalize:
+            for c in range(nchan):
+                if cgood[c]:
+                    img_norm[..., c] = 1 - img_norm[..., c]
+        else:
+            error_message = "Cannot invert image without normalization"
+            transforms_logger.critical(error_message)
+            raise ValueError(error_message)
+
+    # Move channel axis back to the original position
+    if axis != -1 and axis != img_norm.ndim - 1:
+        img_norm = np.moveaxis(img_norm, -1, axis)
+
+    # The transformer can get confused if a channel is all 1's instead of all 0's:
+    for i, chan_did_normalize in enumerate(cgood):
+        if not chan_did_normalize:
+            if img_norm.ndim == 3:
+                img_norm[:, :, i] = 0
+            if img_norm.ndim == 4:
+                img_norm[:, :, :, i] = 0
+
+    return img_norm
+
+def resize_safe(img, Ly, Lx, interpolation=cv2.INTER_LINEAR):
+    """OpenCV resize function does not support uint32.
+
+    This function converts the image to float32 before resizing and then converts it back to uint32. Not safe!
+    References issue: https://github.com/MouseLand/cellpose/issues/937
+
+    Implications:
+    * Runtime: Runtime increases by 5x-50x due to type casting. However, with resizing being very efficient, this is not
+    a big issue. A 10,000x10,000 image takes 0.47s instead of 0.016s to cast and resize on 32 cores on GPU.
+    * Memory: However, memory usage increases. Not tested by how much.
+
+    Args:
+        img (ndarray): Image of size [Ly x Lx].
+        Ly (int): Desired height of the resized image.
+        Lx (int): Desired width of the resized image.
+        interpolation (int, optional): OpenCV interpolation method. Defaults to cv2.INTER_LINEAR.
+
+    Returns:
+        ndarray: Resized image of size [Ly x Lx].
+
+    """
+
+    # cast image
+    cast = img.dtype == np.uint32
+    if cast:
+        img = img.astype(np.float32)
+
+    # resize
+    img = cv2.resize(img, (Lx, Ly), interpolation=interpolation)
+
+    # cast back
+    if cast:
+        img = img.round().astype(np.uint32)
+
+    return img
+
+
+def resize_image(img0, Ly=None, Lx=None, rsz=None, interpolation=cv2.INTER_LINEAR,
+                 no_channels=False):
+    """Resize image for computing flows / unresize for computing dynamics.
+
+    Args:
+        img0 (ndarray): Image of size [Y x X x nchan] or [Lz x Y x X x nchan] or [Lz x Y x X].
+        Ly (int, optional): Desired height of the resized image. Defaults to None.
+        Lx (int, optional): Desired width of the resized image. Defaults to None.
+        rsz (float, optional): Resize coefficient(s) for the image. If Ly is None, rsz is used. Defaults to None.
+        interpolation (int, optional): OpenCV interpolation method. Defaults to cv2.INTER_LINEAR.
+        no_channels (bool, optional): Flag indicating whether to treat the third dimension as a channel.
+            Defaults to False.
+
+    Returns:
+        ndarray: Resized image of size [Ly x Lx x nchan] or [Lz x Ly x Lx x nchan].
+
+    Raises:
+        ValueError: If Ly is None and rsz is None.
+
+    """
+    if Ly is None and rsz is None:
+        error_message = "must give size to resize to or factor to use for resizing"
+        transforms_logger.critical(error_message)
+        raise ValueError(error_message)
+
+    if Ly is None:
+        # determine Ly and Lx using rsz
+        if not isinstance(rsz, list) and not isinstance(rsz, np.ndarray):
+            rsz = [rsz, rsz]
+        if no_channels:
+            Ly = int(img0.shape[-2] * rsz[-2])
+            Lx = int(img0.shape[-1] * rsz[-1])
+        else:
+            Ly = int(img0.shape[-3] * rsz[-2])
+            Lx = int(img0.shape[-2] * rsz[-1])
+
+    # no_channels useful for z-stacks, so the third dimension is not treated as a channel
+    # but if this is called for grayscale images, they first become [Ly,Lx,2] so ndim=3 but
+    if (img0.ndim > 2 and no_channels) or (img0.ndim == 4 and not no_channels):
+        if Ly == 0 or Lx == 0:
+            raise ValueError(
+                "anisotropy too high / low -- not enough pixels to resize to ratio")
+        for i, img in enumerate(img0):
+            imgi = resize_safe(img, Ly, Lx, interpolation=interpolation)
+            if i==0:
+                if no_channels:
+                    imgs = np.zeros((img0.shape[0], Ly, Lx), imgi.dtype)
+                else:
+                    imgs = np.zeros((img0.shape[0], Ly, Lx, img0.shape[-1]), imgi.dtype)
+            imgs[i] = imgi if imgi.ndim > 2 or no_channels else imgi[..., np.newaxis]
+    else:
+        imgs = resize_safe(img0, Ly, Lx, interpolation=interpolation)
+    return imgs
+
+def get_pad_yx(Ly, Lx, div=16, extra=1, min_size=None):
+    if min_size is None or Ly >= min_size[-2]:
+        Lpad = int(div * np.ceil(Ly / div) - Ly)
+    else:
+        Lpad = min_size[-2] - Ly
+    ypad1 = extra * div // 2 + Lpad // 2
+    ypad2 = extra * div // 2 + Lpad - Lpad // 2
+    if min_size is None or Lx >= min_size[-1]:
+        Lpad = int(div * np.ceil(Lx / div) - Lx)
+    else:
+        Lpad = min_size[-1] - Lx
+    xpad1 = extra * div // 2 + Lpad // 2
+    xpad2 = extra * div // 2 + Lpad - Lpad // 2
+
+    return ypad1, ypad2, xpad1, xpad2
+
+
+def pad_image_ND(img0, div=16, extra=1, min_size=None, zpad=False):
+    """Pad image for test-time so that its dimensions are a multiple of 16 (2D or 3D).
+
+    Args:
+        img0 (ndarray): Image of size [nchan (x Lz) x Ly x Lx].
+        div (int, optional): Divisor for padding. Defaults to 16.
+        extra (int, optional): Extra padding. Defaults to 1.
+        min_size (tuple, optional): Minimum size of the image. Defaults to None.
+
+    Returns:
+        A tuple containing (I, ysub, xsub) or (I, ysub, xsub, zsub), I is padded image, -sub are ranges of pixels in the padded image corresponding to img0.
+            
+    """
+    Ly, Lx = img0.shape[-2:]
+    ypad1, ypad2, xpad1, xpad2 = get_pad_yx(Ly, Lx, div=div, extra=extra, min_size=min_size)
+
+    if img0.ndim > 3:
+        if zpad:
+            Lpad = int(div * np.ceil(img0.shape[-3] / div) - img0.shape[-3])
+            zpad1 = extra * div // 2 + Lpad // 2
+            zpad2 = extra * div // 2 + Lpad - Lpad // 2
+        else:
+            zpad1, zpad2 = 0, 0
+        pads = np.array([[0, 0], [zpad1, zpad2], [ypad1, ypad2], [xpad1, xpad2]])
+    else:
+        pads = np.array([[0, 0], [ypad1, ypad2], [xpad1, xpad2]])
+
+    I = np.pad(img0, pads, mode="constant")
+
+    ysub = np.arange(ypad1, ypad1 + Ly)
+    xsub = np.arange(xpad1, xpad1 + Lx)
+    if zpad:
+        zsub = np.arange(zpad1, zpad1 + img0.shape[-3])
+        return I, ysub, xsub, zsub
+    else:
+        return I, ysub, xsub
+
+
+def random_rotate_and_resize(X, Y=None, scale_range=1., xy=(224, 224), do_3D=False,
+                             zcrop=48, do_flip=True, rotate=True, rescale=None, unet=False,
+                             random_per_image=True):
+    """Augmentation by random rotation and resizing.
+
+    Args:
+        X (list of ND-arrays, float): List of image arrays of size [nchan x Ly x Lx] or [Ly x Lx].
+        Y (list of ND-arrays, float, optional): List of image labels of size [nlabels x Ly x Lx] or [Ly x Lx].
+            The 1st channel of Y is always nearest-neighbor interpolated (assumed to be masks or 0-1 representation).
+            If Y.shape[0]==3 and not unet, then the labels are assumed to be [cell probability, Y flow, X flow].
+            If unet, second channel is dist_to_bound. Defaults to None.
+        scale_range (float, optional): Range of resizing of images for augmentation.
+            Images are resized by (1-scale_range/2) + scale_range * np.random.rand(). Defaults to 1.0.
+        xy (tuple, int, optional): Size of transformed images to return. Defaults to (224,224).
+        do_flip (bool, optional): Whether or not to flip images horizontally. Defaults to True.
+        rotate (bool, optional): Whether or not to rotate images. Defaults to True.
+        rescale (array, float, optional): How much to resize images by before performing augmentations. Defaults to None.
+        unet (bool, optional): Whether or not to use unet. Defaults to False.
+        random_per_image (bool, optional): Different random rotate and resize per image. Defaults to True.
+
+    Returns:
+        A tuple containing (imgi, lbl, scale): imgi (ND-array, float): Transformed images in array [nimg x nchan x xy[0] x xy[1]]; 
+        lbl (ND-array, float): Transformed labels in array [nimg x nchan x xy[0] x xy[1]]; 
+        scale (array, float): Amount each image was resized by.
+    """
+    scale_range = max(0, min(2, float(scale_range))) if scale_range is not None else scale_range
+    nimg = len(X)
+    if X[0].ndim > 2:
+        nchan = X[0].shape[0]
+    else:
+        nchan = 1
+    if do_3D and X[0].ndim > 3:
+        shape = (zcrop, xy[0], xy[1])
+    else:
+        shape = (xy[0], xy[1])
+    imgi = np.zeros((nimg, nchan, *shape), "float32")
+
+    lbl = []
+    if Y is not None:
+        if Y[0].ndim > 2:
+            nt = Y[0].shape[0]
+        else:
+            nt = 1
+        lbl = np.zeros((nimg, nt, *shape), np.float32)
+
+    scale = np.ones(nimg, np.float32)
+
+    for n in range(nimg):
+
+        if random_per_image or n == 0:
+            Ly, Lx = X[n].shape[-2:]
+            # generate random augmentation parameters
+            flip = np.random.rand() > .5
+            theta = np.random.rand() * np.pi * 2 if rotate else 0.
+            if scale_range is None:
+                scale[n] = 2 ** (4 * np.random.rand() - 2)
+            else:
+                scale[n] =  (1 - scale_range / 2) + scale_range * np.random.rand()
+            if rescale is not None:
+                scale[n] *= 1. / rescale[n]
+            dxy = np.maximum(0, np.array([Lx * scale[n] - xy[1],
+                                          Ly * scale[n] - xy[0]]))
+            dxy = (np.random.rand(2,) - .5) * dxy
+
+            # create affine transform
+            cc = np.array([Lx / 2, Ly / 2])
+            cc1 = cc - np.array([Lx - xy[1], Ly - xy[0]]) / 2 + dxy
+            pts1 = np.float32([cc, cc + np.array([1, 0]), cc + np.array([0, 1])])
+            pts2 = np.float32([
+                cc1,
+                cc1 + scale[n] * np.array([np.cos(theta), np.sin(theta)]),
+                cc1 + scale[n] *
+                np.array([np.cos(np.pi / 2 + theta),
+                          np.sin(np.pi / 2 + theta)])
+            ])
+            M = cv2.getAffineTransform(pts1, pts2)
+
+        img = X[n].copy()
+        if Y is not None:
+            labels = Y[n].copy()
+            if labels.ndim < 3:
+                labels = labels[np.newaxis, :, :]
+
+        if do_3D:
+            Lz = X[n].shape[-3]
+            flip_z = np.random.rand() > .5
+            lz = int(np.round(zcrop / scale[n]))
+            iz = np.random.randint(0, Lz - lz)
+            img = img[:,iz:iz + lz,:,:]
+            if Y is not None:
+                labels = labels[:,iz:iz + lz,:,:]
+        
+        if do_flip:
+            if flip:
+                img = img[..., ::-1]
+                if Y is not None:
+                    labels = labels[..., ::-1]
+                    if nt > 1 and not unet:
+                        labels[-1] = -labels[-1]
+            if do_3D and flip_z:
+                img = img[:, ::-1]
+                if Y is not None:
+                    labels = labels[:,::-1]
+                    if nt > 1 and not unet:
+                        labels[-3] = -labels[-3]
+
+        for k in range(nchan):
+            if do_3D:
+                img0 = np.zeros((lz, xy[0], xy[1]), "float32")
+                for z in range(lz):
+                    I = cv2.warpAffine(img[k, z], M, (xy[1], xy[0]),
+                                       flags=cv2.INTER_LINEAR)
+                    img0[z] = I
+                if scale[n] != 1.0:
+                    for y in range(imgi.shape[-2]):
+                        imgi[n, k, :, y] = cv2.resize(img0[:, y], (xy[1], zcrop),
+                                                      interpolation=cv2.INTER_LINEAR)
+                else:
+                    imgi[n, k] = img0
+            else:
+                I = cv2.warpAffine(img[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)
+                imgi[n, k] = I
+
+        if Y is not None:
+            for k in range(nt):
+                flag = cv2.INTER_NEAREST if k < nt-2 else cv2.INTER_LINEAR
+                if do_3D:
+                    lbl0 = np.zeros((lz, xy[0], xy[1]), "float32")
+                    for z in range(lz):
+                        I = cv2.warpAffine(labels[k, z], M, (xy[1], xy[0]),
+                                                      flags=flag)
+                        lbl0[z] = I
+                    if scale[n] != 1.0:
+                        for y in range(lbl.shape[-2]):
+                            lbl[n, k, :, y] = cv2.resize(lbl0[:, y], (xy[1], zcrop),
+                                                          interpolation=flag)
+                    else:
+                        lbl[n, k] = lbl0
+                else:
+                    lbl[n, k] = cv2.warpAffine(labels[k], M, (xy[1], xy[0]), flags=flag)
+
+            if nt > 1 and not unet:
+                v1 = lbl[n, -1].copy()
+                v2 = lbl[n, -2].copy()
+                lbl[n, -2] = (-v1 * np.sin(-theta) + v2 * np.cos(-theta))
+                lbl[n, -1] = (v1 * np.cos(-theta) + v2 * np.sin(-theta))
+
+    return imgi, lbl, scale
+
+
+def random_rotate_and_resize_with_feat(X, Y=None, feat=None, scale_range=1., xy=(224, 224), do_3D=False,
+                             zcrop=48, do_flip=True, rotate=True, rescale=None, unet=False,
+                             random_per_image=True):
+    """Augmentation by random rotation and resizing.
+
+    Args:
+        X (list of ND-arrays, float): List of image arrays of size [nchan x Ly x Lx] or [Ly x Lx].
+        Y (list of ND-arrays, float, optional): List of image labels of size [nlabels x Ly x Lx] or [Ly x Lx].
+            The 1st channel of Y is always nearest-neighbor interpolated (assumed to be masks or 0-1 representation).
+            If Y.shape[0]==3 and not unet, then the labels are assumed to be [cell probability, Y flow, X flow].
+            If unet, second channel is dist_to_bound. Defaults to None.
+        scale_range (float, optional): Range of resizing of images for augmentation.
+            Images are resized by (1-scale_range/2) + scale_range * np.random.rand(). Defaults to 1.0.
+        xy (tuple, int, optional): Size of transformed images to return. Defaults to (224,224).
+        do_flip (bool, optional): Whether or not to flip images horizontally. Defaults to True.
+        rotate (bool, optional): Whether or not to rotate images. Defaults to True.
+        rescale (array, float, optional): How much to resize images by before performing augmentations. Defaults to None.
+        unet (bool, optional): Whether or not to use unet. Defaults to False.
+        random_per_image (bool, optional): Different random rotate and resize per image. Defaults to True.
+
+    Returns:
+        A tuple containing (imgi, lbl, scale): imgi (ND-array, float): Transformed images in array [nimg x nchan x xy[0] x xy[1]]; 
+        lbl (ND-array, float): Transformed labels in array [nimg x nchan x xy[0] x xy[1]]; 
+        scale (array, float): Amount each image was resized by.
+    """
+    scale_range = max(0, min(2, float(scale_range))) if scale_range is not None else scale_range
+    nimg = len(X)
+    if X[0].ndim > 2:
+        nchan = X[0].shape[0]
+    else:
+        nchan = 1
+    if do_3D and X[0].ndim > 3:
+        shape = (zcrop, xy[0], xy[1])
+    else:
+        shape = (xy[0], xy[1])
+    imgi = np.zeros((nimg, nchan, *shape), "float32")
+
+    lbl = []
+    if Y is not None:
+        if Y[0].ndim > 2:
+            nt = Y[0].shape[0]
+        else:
+            nt = 1
+        lbl = np.zeros((nimg, nt, *shape), np.float32)
+    
+    if feat is not None:
+        if feat[0].ndim > 2:
+            nf = feat[0].shape[0]
+        else:
+            nf = 1
+        feat_out = np.zeros((nimg, nf, *shape), "float32")
+
+    scale = np.ones(nimg, np.float32)
+
+    for n in range(nimg):
+
+        if random_per_image or n == 0:
+            Ly, Lx = X[n].shape[-2:]
+            # generate random augmentation parameters
+            flip = np.random.rand() > .5
+            theta = np.random.rand() * np.pi * 2 if rotate else 0.
+            if scale_range is None:
+                scale[n] = 2 ** (4 * np.random.rand() - 2)
+            else:
+                scale[n] =  (1 - scale_range / 2) + scale_range * np.random.rand()
+            if rescale is not None:
+                scale[n] *= 1. / rescale[n]
+            dxy = np.maximum(0, np.array([Lx * scale[n] - xy[1],
+                                          Ly * scale[n] - xy[0]]))
+            dxy = (np.random.rand(2,) - .5) * dxy
+
+            # create affine transform
+            cc = np.array([Lx / 2, Ly / 2])
+            cc1 = cc - np.array([Lx - xy[1], Ly - xy[0]]) / 2 + dxy
+            pts1 = np.float32([cc, cc + np.array([1, 0]), cc + np.array([0, 1])])
+            pts2 = np.float32([
+                cc1,
+                cc1 + scale[n] * np.array([np.cos(theta), np.sin(theta)]),
+                cc1 + scale[n] *
+                np.array([np.cos(np.pi / 2 + theta),
+                          np.sin(np.pi / 2 + theta)])
+            ])
+            M = cv2.getAffineTransform(pts1, pts2)
+
+        img = X[n].copy()
+        if Y is not None:
+            labels = Y[n].copy()
+            if labels.ndim < 3:
+                labels = labels[np.newaxis, :, :]
+        if feat is not None:
+            feats = feat[n].copy()
+            if feats.ndim < 3:
+                feats = feats[np.newaxis, :, :]
+
+        if do_3D:
+            Lz = X[n].shape[-3]
+            flip_z = np.random.rand() > .5
+            lz = int(np.round(zcrop / scale[n]))
+            iz = np.random.randint(0, Lz - lz)
+            img = img[:,iz:iz + lz,:,:]
+            if Y is not None:
+                labels = labels[:,iz:iz + lz,:,:]
+            if feat is not None:
+                feats = feats[:,iz:iz + lz,:,:]
+        
+        if do_flip:
+            if flip:
+                img = img[..., ::-1]
+                if Y is not None:
+                    labels = labels[..., ::-1]
+                    if nt > 1 and not unet:
+                        labels[-1] = -labels[-1]
+                if feat is not None:
+                    feats = feats[..., ::-1]
+            if do_3D and flip_z:
+                img = img[:, ::-1]
+                if Y is not None:
+                    labels = labels[:,::-1]
+                    if nt > 1 and not unet:
+                        labels[-3] = -labels[-3]
+                if feat is not None:
+                    feats = feats[:, ::-1]
+
+        for k in range(nchan):
+            if do_3D:
+                img0 = np.zeros((lz, xy[0], xy[1]), "float32")
+                for z in range(lz):
+                    I = cv2.warpAffine(img[k, z], M, (xy[1], xy[0]),
+                                       flags=cv2.INTER_LINEAR)
+                    img0[z] = I
+                if scale[n] != 1.0:
+                    for y in range(imgi.shape[-2]):
+                        imgi[n, k, :, y] = cv2.resize(img0[:, y], (xy[1], zcrop),
+                                                      interpolation=cv2.INTER_LINEAR)
+                else:
+                    imgi[n, k] = img0
+            else:
+                I = cv2.warpAffine(img[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)
+                imgi[n, k] = I
+
+        if Y is not None:
+            for k in range(nt):
+                flag = cv2.INTER_NEAREST if k < nt-2 else cv2.INTER_LINEAR
+                if do_3D:
+                    lbl0 = np.zeros((lz, xy[0], xy[1]), "float32")
+                    for z in range(lz):
+                        I = cv2.warpAffine(labels[k, z], M, (xy[1], xy[0]),
+                                                      flags=flag)
+                        lbl0[z] = I
+                    if scale[n] != 1.0:
+                        for y in range(lbl.shape[-2]):
+                            lbl[n, k, :, y] = cv2.resize(lbl0[:, y], (xy[1], zcrop),
+                                                          interpolation=flag)
+                    else:
+                        lbl[n, k] = lbl0
+                else:
+                    lbl[n, k] = cv2.warpAffine(labels[k], M, (xy[1], xy[0]), flags=flag)
+
+            if nt > 1 and not unet:
+                v1 = lbl[n, -1].copy()
+                v2 = lbl[n, -2].copy()
+                lbl[n, -2] = (-v1 * np.sin(-theta) + v2 * np.cos(-theta))
+                lbl[n, -1] = (v1 * np.cos(-theta) + v2 * np.sin(-theta))
+        
+        if feat is not None:
+            for k in range(nf):
+                if do_3D:
+                    feat0 = np.zeros((lz, xy[0], xy[1]), "float32")
+                    for z in range(lz):
+                        I = cv2.warpAffine(feats[k, z], M, (xy[1], xy[0]),
+                                                      flags=cv2.INTER_LINEAR)
+                        feat0[z] = I
+                    if scale[n] != 1.0:
+                        for y in range(feat_out.shape[-2]):
+                            feat_out[n, k, :, y] = cv2.resize(feat0[:, y], (xy[1], zcrop),
+                                                          interpolation=cv2.INTER_LINEAR)
+                    else:
+                        feat_out[n, k] = feat0
+                else:
+                    feat_out[n, k] = cv2.warpAffine(feats[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)
+
+
+
+    return imgi, lbl, feat_out, scale

models/seg_post_model/cellpose/utils.py

@@ -0,0 +1,667 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import logging
+import os, tempfile, shutil, io
+from tqdm import tqdm, trange
+from urllib.request import urlopen
+import cv2
+from scipy.ndimage import find_objects, gaussian_filter, generate_binary_structure, label
+from scipy.spatial import ConvexHull
+import numpy as np
+import colorsys
+import fastremap
+import fill_voids
+from multiprocessing import Pool, cpu_count
+# try:
+#     from cellpose import metrics
+# except:
+#     import metrics as metrics
+from models.seg_post_model.cellpose import metrics
+
+try:
+    from skimage.morphology import remove_small_holes
+    SKIMAGE_ENABLED = True
+except:
+    SKIMAGE_ENABLED = False
+
+
+class TqdmToLogger(io.StringIO):
+    """
+        Output stream for TQDM which will output to logger module instead of
+        the StdOut.
+    """
+    logger = None
+    level = None
+    buf = ""
+
+    def __init__(self, logger, level=None):
+        super(TqdmToLogger, self).__init__()
+        self.logger = logger
+        self.level = level or logging.INFO
+
+    def write(self, buf):
+        self.buf = buf.strip("\r\n\t ")
+
+    def flush(self):
+        self.logger.log(self.level, self.buf)
+
+
+def rgb_to_hsv(arr):
+    rgb_to_hsv_channels = np.vectorize(colorsys.rgb_to_hsv)
+    r, g, b = np.rollaxis(arr, axis=-1)
+    h, s, v = rgb_to_hsv_channels(r, g, b)
+    hsv = np.stack((h, s, v), axis=-1)
+    return hsv
+
+
+def hsv_to_rgb(arr):
+    hsv_to_rgb_channels = np.vectorize(colorsys.hsv_to_rgb)
+    h, s, v = np.rollaxis(arr, axis=-1)
+    r, g, b = hsv_to_rgb_channels(h, s, v)
+    rgb = np.stack((r, g, b), axis=-1)
+    return rgb
+
+
+def download_url_to_file(url, dst, progress=True):
+    r"""Download object at the given URL to a local path.
+            Thanks to torch, slightly modified
+    Args:
+        url (string): URL of the object to download
+        dst (string): Full path where object will be saved, e.g. `/tmp/temporary_file`
+        progress (bool, optional): whether or not to display a progress bar to stderr
+            Default: True
+    """
+    file_size = None
+    import ssl
+    ssl._create_default_https_context = ssl._create_unverified_context
+    u = urlopen(url)
+    meta = u.info()
+    if hasattr(meta, "getheaders"):
+        content_length = meta.getheaders("Content-Length")
+    else:
+        content_length = meta.get_all("Content-Length")
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+    # We deliberately save it in a temp file and move it after
+    dst = os.path.expanduser(dst)
+    dst_dir = os.path.dirname(dst)
+    f = tempfile.NamedTemporaryFile(delete=False, dir=dst_dir)
+    try:
+        with tqdm(total=file_size, disable=not progress, unit="B", unit_scale=True,
+                  unit_divisor=1024) as pbar:
+            while True:
+                buffer = u.read(8192)
+                if len(buffer) == 0:
+                    break
+                f.write(buffer)
+                pbar.update(len(buffer))
+        f.close()
+        shutil.move(f.name, dst)
+    finally:
+        f.close()
+        if os.path.exists(f.name):
+            os.remove(f.name)
+
+
+def distance_to_boundary(masks):
+    """Get the distance to the boundary of mask pixels.
+
+    Args:
+        masks (int, 2D or 3D array): The masks array. Size [Ly x Lx] or [Lz x Ly x Lx], where 0 represents no mask and 1, 2, ... represent mask labels.
+
+    Returns:
+        dist_to_bound (2D or 3D array): The distance to the boundary. Size [Ly x Lx] or [Lz x Ly x Lx].
+
+    Raises:
+        ValueError: If the masks array is not 2D or 3D.
+
+    """
+    if masks.ndim > 3 or masks.ndim < 2:
+        raise ValueError("distance_to_boundary takes 2D or 3D array, not %dD array" %
+                         masks.ndim)
+    dist_to_bound = np.zeros(masks.shape, np.float64)
+
+    if masks.ndim == 3:
+        for i in range(masks.shape[0]):
+            dist_to_bound[i] = distance_to_boundary(masks[i])
+        return dist_to_bound
+    else:
+        slices = find_objects(masks)
+        for i, si in enumerate(slices):
+            if si is not None:
+                sr, sc = si
+                mask = (masks[sr, sc] == (i + 1)).astype(np.uint8)
+                contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                            cv2.CHAIN_APPROX_NONE)
+                pvc, pvr = np.concatenate(contours[-2], axis=0).squeeze().T
+                ypix, xpix = np.nonzero(mask)
+                min_dist = ((ypix[:, np.newaxis] - pvr)**2 +
+                            (xpix[:, np.newaxis] - pvc)**2).min(axis=1)
+                dist_to_bound[ypix + sr.start, xpix + sc.start] = min_dist
+        return dist_to_bound
+
+
+def masks_to_edges(masks, threshold=1.0):
+    """Get edges of masks as a 0-1 array.
+
+    Args:
+        masks (int, 2D or 3D array): Size [Ly x Lx] or [Lz x Ly x Lx], where 0=NO masks and 1,2,...=mask labels.
+        threshold (float, optional): Threshold value for distance to boundary. Defaults to 1.0.
+
+    Returns:
+        edges (2D or 3D array): Size [Ly x Lx] or [Lz x Ly x Lx], where True pixels are edge pixels.
+    """
+    dist_to_bound = distance_to_boundary(masks)
+    edges = (dist_to_bound < threshold) * (masks > 0)
+    return edges
+
+
+def remove_edge_masks(masks, change_index=True):
+    """Removes masks with pixels on the edge of the image.
+
+    Args:
+        masks (int, 2D or 3D array): The masks to be processed. Size [Ly x Lx] or [Lz x Ly x Lx], where 0 represents no mask and 1, 2, ... represent mask labels.
+        change_index (bool, optional): If True, after removing masks, changes the indexing so that there are no missing label numbers. Defaults to True.
+
+    Returns:
+        outlines (2D or 3D array): The processed masks. Size [Ly x Lx] or [Lz x Ly x Lx], where 0 represents no mask and 1, 2, ... represent mask labels.
+    """
+    slices = find_objects(masks.astype(int))
+    for i, si in enumerate(slices):
+        remove = False
+        if si is not None:
+            for d, sid in enumerate(si):
+                if sid.start == 0 or sid.stop == masks.shape[d]:
+                    remove = True
+                    break
+            if remove:
+                masks[si][masks[si] == i + 1] = 0
+    shape = masks.shape
+    if change_index:
+        _, masks = np.unique(masks, return_inverse=True)
+        masks = np.reshape(masks, shape).astype(np.int32)
+
+    return masks
+
+
+def masks_to_outlines(masks):
+    """Get outlines of masks as a 0-1 array.
+
+    Args:
+        masks (int, 2D or 3D array): Size [Ly x Lx] or [Lz x Ly x Lx], where 0=NO masks and 1,2,...=mask labels.
+
+    Returns:
+        outlines (2D or 3D array): Size [Ly x Lx] or [Lz x Ly x Lx], where True pixels are outlines.
+    """
+    if masks.ndim > 3 or masks.ndim < 2:
+        raise ValueError("masks_to_outlines takes 2D or 3D array, not %dD array" %
+                         masks.ndim)
+    outlines = np.zeros(masks.shape, bool)
+
+    if masks.ndim == 3:
+        for i in range(masks.shape[0]):
+            outlines[i] = masks_to_outlines(masks[i])
+        return outlines
+    else:
+        slices = find_objects(masks.astype(int))
+        for i, si in enumerate(slices):
+            if si is not None:
+                sr, sc = si
+                mask = (masks[sr, sc] == (i + 1)).astype(np.uint8)
+                contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                            cv2.CHAIN_APPROX_NONE)
+                pvc, pvr = np.concatenate(contours[-2], axis=0).squeeze().T
+                vr, vc = pvr + sr.start, pvc + sc.start
+                outlines[vr, vc] = 1
+        return outlines
+
+
+def outlines_list(masks, multiprocessing_threshold=1000, multiprocessing=None):
+    """Get outlines of masks as a list to loop over for plotting.
+
+    Args:
+        masks (ndarray): Array of masks.
+        multiprocessing_threshold (int, optional): Threshold for enabling multiprocessing. Defaults to 1000.
+        multiprocessing (bool, optional): Flag to enable multiprocessing. Defaults to None.
+
+    Returns:
+        list: List of outlines.
+
+    Raises:
+        None
+
+    Notes:
+        - This function is a wrapper for outlines_list_single and outlines_list_multi.
+        - Multiprocessing is disabled for Windows.
+    """
+    # default to use multiprocessing if not few_masks, but allow user to override
+    if multiprocessing is None:
+        few_masks = np.max(masks) < multiprocessing_threshold
+        multiprocessing = not few_masks
+
+    # disable multiprocessing for Windows
+    if os.name == "nt":
+        if multiprocessing:
+            logging.getLogger(__name__).warning(
+                "Multiprocessing is disabled for Windows")
+        multiprocessing = False
+
+    if multiprocessing:
+        return outlines_list_multi(masks)
+    else:
+        return outlines_list_single(masks)
+
+
+def outlines_list_single(masks):
+    """Get outlines of masks as a list to loop over for plotting.
+
+    Args:
+        masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)
+
+    Returns:
+        list: List of outlines as pixel coordinates.
+
+    """
+    outpix = []
+    for n in np.unique(masks)[1:]:
+        mn = masks == n
+        if mn.sum() > 0:
+            contours = cv2.findContours(mn.astype(np.uint8), mode=cv2.RETR_EXTERNAL,
+                                        method=cv2.CHAIN_APPROX_NONE)
+            contours = contours[-2]
+            cmax = np.argmax([c.shape[0] for c in contours])
+            pix = contours[cmax].astype(int).squeeze()
+            if len(pix) > 4:
+                outpix.append(pix)
+            else:
+                outpix.append(np.zeros((0, 2)))
+    return outpix
+
+
+def outlines_list_multi(masks, num_processes=None):
+    """
+    Get outlines of masks as a list to loop over for plotting.
+
+    Args:
+        masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)
+
+    Returns:
+        list: List of outlines as pixel coordinates.
+    """
+    if num_processes is None:
+        num_processes = cpu_count()
+
+    unique_masks = np.unique(masks)[1:]
+    with Pool(processes=num_processes) as pool:
+        outpix = pool.map(get_outline_multi, [(masks, n) for n in unique_masks])
+    return outpix
+
+
+def get_outline_multi(args):
+    """Get the outline of a specific mask in a multi-mask image.
+
+    Args:
+        args (tuple): A tuple containing the masks and the mask number.
+
+    Returns:
+        numpy.ndarray: The outline of the specified mask as an array of coordinates.
+
+    """
+    masks, n = args
+    mn = masks == n
+    if mn.sum() > 0:
+        contours = cv2.findContours(mn.astype(np.uint8), mode=cv2.RETR_EXTERNAL,
+                                    method=cv2.CHAIN_APPROX_NONE)
+        contours = contours[-2]
+        cmax = np.argmax([c.shape[0] for c in contours])
+        pix = contours[cmax].astype(int).squeeze()
+        return pix if len(pix) > 4 else np.zeros((0, 2))
+    return np.zeros((0, 2))
+
+
+def dilate_masks(masks, n_iter=5):
+    """Dilate masks by n_iter pixels.
+
+    Args:
+        masks (ndarray): Array of masks.
+        n_iter (int, optional): Number of pixels to dilate the masks. Defaults to 5.
+
+    Returns:
+        ndarray: Dilated masks.
+    """
+    dilated_masks = masks.copy()
+    for n in range(n_iter):
+        # define the structuring element to use for dilation
+        kernel = np.ones((3, 3), "uint8")
+        # find the distance to each mask (distances are zero within masks)
+        dist_transform = cv2.distanceTransform((dilated_masks == 0).astype("uint8"),
+                                               cv2.DIST_L2, 5)
+        # dilate each mask and assign to it the pixels along the border of the mask
+        # (does not allow dilation into other masks since dist_transform is zero there)
+        for i in range(1, np.max(masks) + 1):
+            mask = (dilated_masks == i).astype("uint8")
+            dilated_mask = cv2.dilate(mask, kernel, iterations=1)
+            dilated_mask = np.logical_and(dist_transform < 2, dilated_mask)
+            dilated_masks[dilated_mask > 0] = i
+    return dilated_masks
+
+
+def get_perimeter(points):
+    """
+    Calculate the perimeter of a set of points.
+
+    Parameters:
+        points (ndarray): An array of points with shape (npoints, ndim).
+
+    Returns:
+        float: The perimeter of the points.
+
+    """
+    if points.shape[0] > 4:
+        points = np.append(points, points[:1], axis=0)
+        return ((np.diff(points, axis=0)**2).sum(axis=1)**0.5).sum()
+    else:
+        return 0
+
+
+def get_mask_compactness(masks):
+    """
+    Calculate the compactness of masks.
+    
+    Parameters:
+        masks (ndarray): Binary masks representing objects.
+        
+    Returns:
+        ndarray: Array of compactness values for each mask.
+    """
+    perimeters = get_mask_perimeters(masks)
+    npoints = np.unique(masks, return_counts=True)[1][1:]
+    areas = npoints
+    compactness = 4 * np.pi * areas / perimeters**2
+    compactness[perimeters == 0] = 0
+    compactness[compactness > 1.0] = 1.0
+    return compactness
+
+
+def get_mask_perimeters(masks):
+    """
+    Calculate the perimeters of the given masks.
+
+    Parameters:
+        masks (numpy.ndarray): Binary masks representing objects.
+
+    Returns:
+        numpy.ndarray: Array containing the perimeters of each mask.
+    """
+    perimeters = np.zeros(masks.max())
+    for n in range(masks.max()):
+        mn = masks == (n + 1)
+        if mn.sum() > 0:
+            contours = cv2.findContours(mn.astype(np.uint8), mode=cv2.RETR_EXTERNAL,
+                                        method=cv2.CHAIN_APPROX_NONE)[-2]
+            perimeters[n] = np.array(
+                [get_perimeter(c.astype(int).squeeze()) for c in contours]).sum()
+
+    return perimeters
+
+
+def circleMask(d0):
+    """
+    Creates an array with indices which are the radius of that x,y point.
+
+    Args:
+        d0 (tuple): Patch of (-d0, d0+1) over which radius is computed.
+
+    Returns:
+        tuple: A tuple containing:
+            - rs (ndarray): Array of radii with shape (2*d0[0]+1, 2*d0[1]+1).
+            - dx (ndarray): Indices of the patch along the x-axis.
+            - dy (ndarray): Indices of the patch along the y-axis.
+    """
+    dx = np.tile(np.arange(-d0[1], d0[1] + 1), (2 * d0[0] + 1, 1))
+    dy = np.tile(np.arange(-d0[0], d0[0] + 1), (2 * d0[1] + 1, 1))
+    dy = dy.transpose()
+
+    rs = (dy**2 + dx**2)**0.5
+    return rs, dx, dy
+
+
+def get_mask_stats(masks_true):
+    """
+    Calculate various statistics for the given binary masks.
+
+    Parameters:
+        masks_true (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)
+
+    Returns:
+        convexity (ndarray): Convexity values for each mask.
+        solidity (ndarray): Solidity values for each mask.
+        compactness (ndarray): Compactness values for each mask.
+    """
+    mask_perimeters = get_mask_perimeters(masks_true)
+
+    # disk for compactness
+    rs, dy, dx = circleMask(np.array([100, 100]))
+    rsort = np.sort(rs.flatten())
+
+    # area for solidity
+    npoints = np.unique(masks_true, return_counts=True)[1][1:]
+    areas = npoints - mask_perimeters / 2 - 1
+
+    compactness = np.zeros(masks_true.max())
+    convexity = np.zeros(masks_true.max())
+    solidity = np.zeros(masks_true.max())
+    convex_perimeters = np.zeros(masks_true.max())
+    convex_areas = np.zeros(masks_true.max())
+    for ic in range(masks_true.max()):
+        points = np.array(np.nonzero(masks_true == (ic + 1))).T
+        if len(points) > 15 and mask_perimeters[ic] > 0:
+            med = np.median(points, axis=0)
+            # compute compactness of ROI
+            r2 = ((points - med)**2).sum(axis=1)**0.5
+            compactness[ic] = (rsort[:r2.size].mean() + 1e-10) / r2.mean()
+            try:
+                hull = ConvexHull(points)
+                convex_perimeters[ic] = hull.area
+                convex_areas[ic] = hull.volume
+            except:
+                convex_perimeters[ic] = 0
+
+    convexity[mask_perimeters > 0.0] = (convex_perimeters[mask_perimeters > 0.0] /
+                                        mask_perimeters[mask_perimeters > 0.0])
+    solidity[convex_areas > 0.0] = (areas[convex_areas > 0.0] /
+                                    convex_areas[convex_areas > 0.0])
+    convexity = np.clip(convexity, 0.0, 1.0)
+    solidity = np.clip(solidity, 0.0, 1.0)
+    compactness = np.clip(compactness, 0.0, 1.0)
+    return convexity, solidity, compactness
+
+
+def get_masks_unet(output, cell_threshold=0, boundary_threshold=0):
+    """Create masks using cell probability and cell boundary.
+
+    Args:
+        output (ndarray): The output array containing cell probability and cell boundary.
+        cell_threshold (float, optional): The threshold value for cell probability. Defaults to 0.
+        boundary_threshold (float, optional): The threshold value for cell boundary. Defaults to 0.
+
+    Returns:
+        ndarray: The masks representing the segmented cells.
+
+    """
+    cells = (output[..., 1] - output[..., 0]) > cell_threshold
+    selem = generate_binary_structure(cells.ndim, connectivity=1)
+    labels, nlabels = label(cells, selem)
+
+    if output.shape[-1] > 2:
+        slices = find_objects(labels)
+        dists = 10000 * np.ones(labels.shape, np.float32)
+        mins = np.zeros(labels.shape, np.int32)
+        borders = np.logical_and(~(labels > 0), output[..., 2] > boundary_threshold)
+        pad = 10
+        for i, slc in enumerate(slices):
+            if slc is not None:
+                slc_pad = tuple([
+                    slice(max(0, sli.start - pad), min(labels.shape[j], sli.stop + pad))
+                    for j, sli in enumerate(slc)
+                ])
+                msk = (labels[slc_pad] == (i + 1)).astype(np.float32)
+                msk = 1 - gaussian_filter(msk, 5)
+                dists[slc_pad] = np.minimum(dists[slc_pad], msk)
+                mins[slc_pad][dists[slc_pad] == msk] = (i + 1)
+        labels[labels == 0] = borders[labels == 0] * mins[labels == 0]
+
+    masks = labels
+    shape0 = masks.shape
+    _, masks = np.unique(masks, return_inverse=True)
+    masks = np.reshape(masks, shape0)
+    return masks
+
+
+def stitch3D(masks, stitch_threshold=0.25):
+    """
+    Stitch 2D masks into a 3D volume using a stitch_threshold on IOU.
+
+    Args:
+        masks (list or ndarray): List of 2D masks.
+        stitch_threshold (float, optional): Threshold value for stitching. Defaults to 0.25.
+
+    Returns:
+        list: List of stitched 3D masks.
+    """
+    mmax = masks[0].max()
+    empty = 0
+    for i in trange(len(masks) - 1):
+        iou = metrics._intersection_over_union(masks[i + 1], masks[i])[1:, 1:]
+        if not iou.size and empty == 0:
+            masks[i + 1] = masks[i + 1]
+            mmax = masks[i + 1].max()
+        elif not iou.size and not empty == 0:
+            icount = masks[i + 1].max()
+            istitch = np.arange(mmax + 1, mmax + icount + 1, 1, masks.dtype)
+            mmax += icount
+            istitch = np.append(np.array(0), istitch)
+            masks[i + 1] = istitch[masks[i + 1]]
+        else:
+            iou[iou < stitch_threshold] = 0.0
+            iou[iou < iou.max(axis=0)] = 0.0
+            istitch = iou.argmax(axis=1) + 1
+            ino = np.nonzero(iou.max(axis=1) == 0.0)[0]
+            istitch[ino] = np.arange(mmax + 1, mmax + len(ino) + 1, 1, masks.dtype)
+            mmax += len(ino)
+            istitch = np.append(np.array(0), istitch)
+            masks[i + 1] = istitch[masks[i + 1]]
+            empty = 1
+
+    return masks
+
+
+def diameters(masks):
+    """
+    Calculate the diameters of the objects in the given masks.
+
+    Parameters:
+    masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)
+
+    Returns:
+        tuple: A tuple containing the median diameter and an array of diameters for each object.
+
+    Examples:
+    >>> masks = np.array([[0, 1, 1], [1, 0, 0], [1, 1, 0]])
+    >>> diameters(masks)
+    (1.0, array([1.41421356, 1.0, 1.0]))
+    """
+    uniq, counts = fastremap.unique(masks.astype("int32"), return_counts=True)
+    counts = counts[1:]
+    md = np.median(counts**0.5)
+    if np.isnan(md):
+        md = 0
+    md /= (np.pi**0.5) / 2
+    return md, counts**0.5
+
+
+def radius_distribution(masks, bins):
+    """
+    Calculate the radius distribution of masks.
+
+    Args:
+        masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)
+        bins (int): Number of bins for the histogram.
+
+    Returns:
+        A tuple containing a normalized histogram of radii, median radius, array of radii.
+
+    """
+    unique, counts = np.unique(masks, return_counts=True)
+    counts = counts[unique != 0]
+    nb, _ = np.histogram((counts**0.5) * 0.5, bins)
+    nb = nb.astype(np.float32)
+    if nb.sum() > 0:
+        nb = nb / nb.sum()
+    md = np.median(counts**0.5) * 0.5
+    if np.isnan(md):
+        md = 0
+    md /= (np.pi**0.5) / 2
+    return nb, md, (counts**0.5) / 2
+
+
+def size_distribution(masks):
+    """
+    Calculates the size distribution of masks.
+
+    Args:
+        masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)
+
+    Returns:
+        float: The ratio of the 25th percentile of mask sizes to the 75th percentile of mask sizes.
+    """
+    counts = np.unique(masks, return_counts=True)[1][1:]
+    return np.percentile(counts, 25) / np.percentile(counts, 75)
+
+
+def fill_holes_and_remove_small_masks(masks, min_size=15):
+    """ Fills holes in masks (2D/3D) and discards masks smaller than min_size.
+
+    This function fills holes in each mask using fill_voids.fill.
+    It also removes masks that are smaller than the specified min_size.
+
+    Parameters:
+    masks (ndarray): Int, 2D or 3D array of labelled masks.
+        0 represents no mask, while positive integers represent mask labels.
+        The size can be [Ly x Lx] or [Lz x Ly x Lx].
+    min_size (int, optional): Minimum number of pixels per mask.
+        Masks smaller than min_size will be removed.
+        Set to -1 to turn off this functionality. Default is 15.
+
+    Returns:
+        ndarray: Int, 2D or 3D array of masks with holes filled and small masks removed.
+            0 represents no mask, while positive integers represent mask labels.
+            The size is [Ly x Lx] or [Lz x Ly x Lx].
+    """
+
+    if masks.ndim > 3 or masks.ndim < 2:
+        raise ValueError("masks_to_outlines takes 2D or 3D array, not %dD array" %
+                         masks.ndim)
+
+    # Filter small masks
+    if min_size > 0:
+        counts = fastremap.unique(masks, return_counts=True)[1][1:]
+        masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)
+        fastremap.renumber(masks, in_place=True)
+        
+    slices = find_objects(masks)
+    j = 0
+    for i, slc in enumerate(slices):
+        if slc is not None:
+            msk = masks[slc] == (i + 1)
+            msk = fill_voids.fill(msk)
+            masks[slc][msk] = (j + 1)
+            j += 1
+
+    if min_size > 0:
+        counts = fastremap.unique(masks, return_counts=True)[1][1:]
+        masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)
+        fastremap.renumber(masks, in_place=True)
+    
+    return masks

models/seg_post_model/cellpose/version.py

@@ -0,0 +1,18 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+from importlib.metadata import PackageNotFoundError, version
+import sys
+from platform import python_version
+import torch
+
+try:
+    version = version("cellpose")
+except PackageNotFoundError:
+    version = "unknown"
+
+version_str = f"""
+cellpose version: \t{version} 
+platform:       \t{sys.platform} 
+python version: \t{python_version()} 
+torch version:  \t{torch.__version__}"""

models/seg_post_model/cellpose/vit_sam.py

@@ -0,0 +1,195 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.
+"""
+
+import torch
+from segment_anything import sam_model_registry
+torch.backends.cuda.matmul.allow_tf32 = True
+from torch import nn 
+import torch.nn.functional as F
+
+class Transformer(nn.Module):
+    def __init__(self, backbone="vit_l", ps=8, nout=3, bsize=256, rdrop=0.4,
+                  checkpoint=None, dtype=torch.float32):
+        super(Transformer, self).__init__()
+        """
+        print(self.encoder.patch_embed)
+            PatchEmbed(
+            (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))
+            )
+        print(self.encoder.neck)
+            Sequential(
+            (0): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
+            (1): LayerNorm2d()
+            (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (3): LayerNorm2d()
+            )
+        """
+        # instantiate the vit model, default to not loading SAM
+        # checkpoint = sam_vit_l_0b3195.pth is standard pretrained SAM
+        self.encoder = sam_model_registry[backbone](checkpoint).image_encoder
+        w = self.encoder.patch_embed.proj.weight.detach()
+        nchan = w.shape[0]
+        
+        # change token size to ps x ps
+        self.ps = ps
+        self.encoder.patch_embed.proj = nn.Conv2d(3, nchan, stride=ps, kernel_size=ps)
+        self.encoder.patch_embed.proj.weight.data = w[:,:,::16//ps,::16//ps]
+        
+        # adjust position embeddings for new bsize and new token size
+        ds = (1024 // 16) // (bsize // ps)
+        self.encoder.pos_embed = nn.Parameter(self.encoder.pos_embed[:,::ds,::ds], requires_grad=True)
+
+        # readout weights for nout output channels
+        # if nout is changed, weights will not load correctly from pretrained Cellpose-SAM
+        self.nout = nout
+        self.out = nn.Conv2d(256, self.nout * ps**2, kernel_size=1)
+
+        # W2 reshapes token space to pixel space, not trainable
+        self.W2 = nn.Parameter(torch.eye(self.nout * ps**2).reshape(self.nout*ps**2, self.nout, ps, ps), 
+                               requires_grad=False)
+        
+        # fraction of layers to drop at random during training
+        self.rdrop = rdrop
+
+        # average diameter of ROIs from training images from fine-tuning 
+        self.diam_labels = nn.Parameter(torch.tensor([30.]), requires_grad=False)
+        # average diameter of ROIs during main training
+        self.diam_mean = nn.Parameter(torch.tensor([30.]), requires_grad=False)
+        
+        # set attention to global in every layer
+        for blk in self.encoder.blocks:
+            blk.window_size = 0
+
+        self.dtype = dtype
+
+    def forward(self, x, feat=None):      
+        # same progression as SAM until readout
+        x = self.encoder.patch_embed(x)
+        if feat is not None:
+            feat = self.encoder.patch_embed(feat)
+            x = x + x * feat * 0.5
+        
+        if self.encoder.pos_embed is not None:
+            x = x + self.encoder.pos_embed
+        
+        if self.training and self.rdrop > 0:
+            nlay = len(self.encoder.blocks)
+            rdrop = (torch.rand((len(x), nlay), device=x.device) < 
+                     torch.linspace(0, self.rdrop, nlay, device=x.device)).to(x.dtype)
+            for i, blk in enumerate(self.encoder.blocks):            
+                mask = rdrop[:,i].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+                x = x * mask + blk(x) * (1-mask)
+        else:
+            for blk in self.encoder.blocks:
+                x = blk(x)
+
+        x = self.encoder.neck(x.permute(0, 3, 1, 2))
+
+        # readout is changed here
+        x1 = self.out(x)
+        x1 = F.conv_transpose2d(x1, self.W2, stride = self.ps, padding = 0)
+        
+        # maintain the second output of feature size 256 for backwards compatibility
+           
+        return x1, torch.randn((x.shape[0], 256), device=x.device)
+    
+    def load_model(self, PATH, device, strict = False):        
+        state_dict = torch.load(PATH, map_location = device, weights_only=True)
+        keys = [k for k in state_dict.keys()]
+        if keys[0][:7] == "module.":
+            from collections import OrderedDict
+            new_state_dict = OrderedDict()
+            for k, v in state_dict.items():
+                name = k[7:] # remove 'module.' of DataParallel/DistributedDataParallel
+                new_state_dict[name] = v
+            self.load_state_dict(new_state_dict, strict = strict)
+        else:
+            self.load_state_dict(state_dict, strict = strict)
+
+        if self.dtype != torch.float32:
+            self = self.to(self.dtype)
+
+    
+    @property
+    def device(self):
+        """
+        Get the device of the model.
+
+        Returns:
+            torch.device: The device of the model.
+        """
+        return next(self.parameters()).device
+
+    def save_model(self, filename):
+        """
+        Save the model to a file.
+
+        Args:
+            filename (str): The path to the file where the model will be saved.
+        """
+        torch.save(self.state_dict(), filename)
+
+
+
+class CPnetBioImageIO(Transformer):
+    """
+    A subclass of the CP-SAM model compatible with the BioImage.IO Spec.
+
+    This subclass addresses the limitation of CPnet's incompatibility with the BioImage.IO Spec,
+    allowing the CPnet model to use the weights uploaded to the BioImage.IO Model Zoo.
+    """
+
+    def forward(self, x):
+        """
+        Perform a forward pass of the CPnet model and return unpacked tensors.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            tuple: A tuple containing the output tensor, style tensor, and downsampled tensors.
+        """
+        output_tensor, style_tensor, downsampled_tensors = super().forward(x)
+        return output_tensor, style_tensor, *downsampled_tensors
+    
+
+    def load_model(self, filename, device=None):
+        """
+        Load the model from a file.
+
+        Args:
+            filename (str): The path to the file where the model is saved.
+            device (torch.device, optional): The device to load the model on. Defaults to None.
+        """
+        if (device is not None) and (device.type != "cpu"):
+            state_dict = torch.load(filename, map_location=device, weights_only=True)
+        else:
+            self.__init__(self.nout)
+            state_dict = torch.load(filename, map_location=torch.device("cpu"), 
+                                    weights_only=True)
+
+        self.load_state_dict(state_dict)
+
+    def load_state_dict(self, state_dict):
+        """
+        Load the state dictionary into the model.
+
+        This method overrides the default `load_state_dict` to handle Cellpose's custom
+        loading mechanism and ensures compatibility with BioImage.IO Core.
+
+        Args:
+            state_dict (Mapping[str, Any]): A state dictionary to load into the model
+        """
+        if state_dict["output.2.weight"].shape[0] != self.nout:
+            for name in self.state_dict():
+                if "output" not in name:
+                    self.state_dict()[name].copy_(state_dict[name])
+        else:
+            super().load_state_dict(
+                {name: param for name, param in state_dict.items()},
+                strict=False)
+
+
+    
+

models/seg_post_model/cellpose/vit_sam_new.py

@@ -0,0 +1,197 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.
+"""
+
+import torch
+from segment_anything import sam_model_registry
+torch.backends.cuda.matmul.allow_tf32 = True
+from torch import nn 
+import torch.nn.functional as F
+
+class Transformer(nn.Module):
+    def __init__(self, backbone="vit_l", ps=16, nout=3, bsize=256, rdrop=0.4,
+                  checkpoint=None, dtype=torch.float32):
+        super(Transformer, self).__init__()
+        """
+        print(self.encoder.patch_embed)
+            PatchEmbed(
+            (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))
+            )
+        print(self.encoder.neck)
+            Sequential(
+            (0): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
+            (1): LayerNorm2d()
+            (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (3): LayerNorm2d()
+            )
+        """
+        # instantiate the vit model, default to not loading SAM
+        # checkpoint = sam_vit_l_0b3195.pth is standard pretrained SAM
+        if checkpoint is None:
+            checkpoint = "sam_vit_l_0b3195.pth"
+        self.encoder = sam_model_registry[backbone](checkpoint).image_encoder
+        w = self.encoder.patch_embed.proj.weight.detach()
+        nchan = w.shape[0]
+        
+        # change token size to ps x ps
+        self.ps = ps
+        # self.encoder.patch_embed.proj = nn.Conv2d(3, nchan, stride=ps, kernel_size=ps)
+        # self.encoder.patch_embed.proj.weight.data = w[:,:,::16//ps,::16//ps]
+        
+        # adjust position embeddings for new bsize and new token size
+        ds = (1024 // 16) // (bsize // ps)
+        self.encoder.pos_embed = nn.Parameter(self.encoder.pos_embed[:,::ds,::ds], requires_grad=True)
+
+        # readout weights for nout output channels
+        # if nout is changed, weights will not load correctly from pretrained Cellpose-SAM
+        self.nout = nout
+        self.out = nn.Conv2d(256, self.nout * ps**2, kernel_size=1)
+
+        # W2 reshapes token space to pixel space, not trainable
+        self.W2 = nn.Parameter(torch.eye(self.nout * ps**2).reshape(self.nout*ps**2, self.nout, ps, ps), 
+                               requires_grad=False)
+        
+        # fraction of layers to drop at random during training
+        self.rdrop = rdrop
+
+        # average diameter of ROIs from training images from fine-tuning 
+        self.diam_labels = nn.Parameter(torch.tensor([30.]), requires_grad=False)
+        # average diameter of ROIs during main training
+        self.diam_mean = nn.Parameter(torch.tensor([30.]), requires_grad=False)
+        
+        # set attention to global in every layer
+        for blk in self.encoder.blocks:
+            blk.window_size = 0
+
+        self.dtype = dtype
+
+    def forward(self, x, feat=None):      
+        # same progression as SAM until readout
+        x = self.encoder.patch_embed(x)
+        if feat is not None:
+            feat = self.encoder.patch_embed(feat)
+            x = x + x * feat * 0.5
+        
+        if self.encoder.pos_embed is not None:
+            x = x + self.encoder.pos_embed
+        
+        if self.training and self.rdrop > 0:
+            nlay = len(self.encoder.blocks)
+            rdrop = (torch.rand((len(x), nlay), device=x.device) < 
+                     torch.linspace(0, self.rdrop, nlay, device=x.device)).to(x.dtype)
+            for i, blk in enumerate(self.encoder.blocks):            
+                mask = rdrop[:,i].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+                x = x * mask + blk(x) * (1-mask)
+        else:
+            for blk in self.encoder.blocks:
+                x = blk(x)
+
+        x = self.encoder.neck(x.permute(0, 3, 1, 2))
+
+        # readout is changed here
+        x1 = self.out(x)
+        x1 = F.conv_transpose2d(x1, self.W2, stride = self.ps, padding = 0)
+        
+        # maintain the second output of feature size 256 for backwards compatibility
+           
+        return x1, torch.randn((x.shape[0], 256), device=x.device)
+    
+    def load_model(self, PATH, device, strict = False):        
+        state_dict = torch.load(PATH, map_location = device, weights_only=True)
+        keys = [k for k in state_dict.keys()]
+        if keys[0][:7] == "module.":
+            from collections import OrderedDict
+            new_state_dict = OrderedDict()
+            for k, v in state_dict.items():
+                name = k[7:] # remove 'module.' of DataParallel/DistributedDataParallel
+                new_state_dict[name] = v
+            self.load_state_dict(new_state_dict, strict = strict)
+        else:
+            self.load_state_dict(state_dict, strict = strict)
+
+        if self.dtype != torch.float32:
+            self = self.to(self.dtype)
+
+    
+    @property
+    def device(self):
+        """
+        Get the device of the model.
+
+        Returns:
+            torch.device: The device of the model.
+        """
+        return next(self.parameters()).device
+
+    def save_model(self, filename):
+        """
+        Save the model to a file.
+
+        Args:
+            filename (str): The path to the file where the model will be saved.
+        """
+        torch.save(self.state_dict(), filename)
+
+
+
+class CPnetBioImageIO(Transformer):
+    """
+    A subclass of the CP-SAM model compatible with the BioImage.IO Spec.
+
+    This subclass addresses the limitation of CPnet's incompatibility with the BioImage.IO Spec,
+    allowing the CPnet model to use the weights uploaded to the BioImage.IO Model Zoo.
+    """
+
+    def forward(self, x):
+        """
+        Perform a forward pass of the CPnet model and return unpacked tensors.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            tuple: A tuple containing the output tensor, style tensor, and downsampled tensors.
+        """
+        output_tensor, style_tensor, downsampled_tensors = super().forward(x)
+        return output_tensor, style_tensor, *downsampled_tensors
+    
+
+    def load_model(self, filename, device=None):
+        """
+        Load the model from a file.
+
+        Args:
+            filename (str): The path to the file where the model is saved.
+            device (torch.device, optional): The device to load the model on. Defaults to None.
+        """
+        if (device is not None) and (device.type != "cpu"):
+            state_dict = torch.load(filename, map_location=device, weights_only=True)
+        else:
+            self.__init__(self.nout)
+            state_dict = torch.load(filename, map_location=torch.device("cpu"), 
+                                    weights_only=True)
+
+        self.load_state_dict(state_dict)
+
+    def load_state_dict(self, state_dict):
+        """
+        Load the state dictionary into the model.
+
+        This method overrides the default `load_state_dict` to handle Cellpose's custom
+        loading mechanism and ensures compatibility with BioImage.IO Core.
+
+        Args:
+            state_dict (Mapping[str, Any]): A state dictionary to load into the model
+        """
+        if state_dict["output.2.weight"].shape[0] != self.nout:
+            for name in self.state_dict():
+                if "output" not in name:
+                    self.state_dict()[name].copy_(state_dict[name])
+        else:
+            super().load_state_dict(
+                {name: param for name, param in state_dict.items()},
+                strict=False)
+
+
+    
+