clean up

app.py

@@ -314,7 +314,7 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
 
     try:
         mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
-        print("📏 mask shape:", mask.shape, "dtype:", mask.dtype, "unique:", np.unique(mask))
+        print("📏 mask shape:", mask.shape, "dtype:", mask.dtype)
     except Exception as e:
         print(f"❌ Inference failed: {str(e)}")
         return None, None, {}
@@ -344,8 +344,6 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
     inst_mask = mask_np.astype(np.int32)
     unique_ids = np.unique(inst_mask)
     num_instances = len(unique_ids[unique_ids != 0])
-    print(f"✅ Instance IDs found: {unique_ids}, Total instances: {num_instances}")
-
     if num_instances == 0:
         print("⚠️ No instance found, returning dummy red image")
         return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None, {}

models/model.py

@@ -5,35 +5,10 @@ import os
 import clip
 import sys
 import numpy as np
-from models.seg_post_model.cellpose.models import CellposeModel
+from models.seg_post_model.models import SegModel
 
 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):
@@ -55,60 +30,6 @@ class Counting_with_SD_features_track(nn.Module):
         self.regressor = regressor_with_SD_features_tra()
 
 
-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__()
@@ -188,69 +109,6 @@ class adapter_roi_loca(nn.Module):
 
 
 
-
-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__()
@@ -301,57 +159,6 @@ class regressor_with_SD_features(nn.Module):
                 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 *
 
 
@@ -363,7 +170,7 @@ class regressor_with_SD_features_seg_vit_c3(nn.Module):
         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_model = SegModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)
         self.vit = self.vit_model.net
 
     def forward(self, img, attn_stack, feature_list):
@@ -380,7 +187,7 @@ class regressor_with_SD_features_seg_vit_c3(nn.Module):
 
 
         
-        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())[0]
+        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())
         if out.dtype == np.uint16:
             out = out.astype(np.int16)
         out = torch.from_numpy(out).unsqueeze(0).to(x.device)
@@ -403,12 +210,11 @@ class regressor_with_SD_features_tra(nn.Module):
 
         # 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_model = SegModel(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], ...]
@@ -422,7 +228,7 @@ class regressor_with_SD_features_tra(nn.Module):
         x = self.inc_0(x)
         feat = x
         
-        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())[0]
+        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())
         if out.dtype == np.uint16:
             out = out.astype(np.int16)
         out = torch.from_numpy(out).unsqueeze(0).to(x.device)
@@ -450,204 +256,3 @@ class regressor_with_SD_features_tra(nn.Module):
                 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/__init__.py

@@ -0,0 +1 @@
+

models/seg_post_model/cellpose/__init__.py

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

models/seg_post_model/cellpose/io.py

@@ -1,816 +0,0 @@
-"""
-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/plot.py

@@ -1,281 +0,0 @@
-"""
-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/utils.py

@@ -1,667 +0,0 @@
-"""
-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

@@ -1,18 +0,0 @@
-"""
-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/core.py → core.py}

@@ -109,39 +109,6 @@ def assign_device(use_torch=True, gpu=False, device=0):
     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.
 
@@ -152,15 +119,19 @@ def _forward(net, x, feat=None):
     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 not isinstance(x, torch.Tensor):
+        X = torch.from_numpy(x).to(net.device, dtype=net.dtype)
+    else:
+        X = x
     if feat is not None:
-        feat = _to_device(feat, device=net.device, dtype=net.dtype)
+        if not isinstance(feat, torch.Tensor):
+            feat = torch.from_numpy(feat).to(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)
+    y = y.detach().cpu().to(torch.float32).numpy()
+    style = style.detach().cpu().to(torch.float32).numpy()
     return y, style
 
 
@@ -269,54 +240,3 @@ def run_net(net, imgi, feat=None, batch_size=8, augment=False, tile_overlap=0.1,
     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/dynamics.py → dynamics.py}

@@ -151,126 +151,6 @@ def masks_to_flows_gpu(masks, device=torch.device("cpu"), niter=None):
         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):
@@ -372,29 +252,6 @@ def steps_interp(dP, inds, niter, device=torch.device("cpu")):
         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")):
@@ -551,7 +408,6 @@ def get_masks_torch(pt, inds, shape0, rpad=20, max_size_fraction=0.4):
         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 
@@ -580,7 +436,6 @@ def get_masks_torch(pt, inds, shape0, rpad=20, max_size_fraction=0.4):
     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
 
 
@@ -652,7 +507,7 @@ def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
             mask = np.zeros(shape, "uint16")
             return mask
 
-        p_final = follow_flows(dP * (cellprob > cellprob_threshold) / 5., 
+        p_final = steps_interp(dP * (cellprob > cellprob_threshold) / 5., 
                                inds=inds, niter=niter, 
                                 device=device)
         if not torch.is_tensor(p_final):

models/seg_post_model/io.py

@@ -0,0 +1,174 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import os
+import numpy as np
+import cv2
+import tifffile
+import logging
+from tqdm import tqdm
+import re
+
+try:
+    import nd2
+    ND2 = True
+except:
+    ND2 = False
+
+try:
+    import nrrd
+    NRRD = True
+except:
+    NRRD = False
+
+
+io_logger = logging.getLogger(__name__)
+
+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 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)
+

models/seg_post_model/{cellpose/metrics.py → metrics.py}

@@ -2,83 +2,10 @@
 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)

models/seg_post_model/{cellpose/models.py → models.py}

@@ -3,31 +3,29 @@ Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer,
 """
 
 import os, time
-from pathlib import Path
+# from pathlib import Path
 import numpy as np
 from tqdm import trange
 import torch
 from scipy.ndimage import gaussian_filter
-import gc
+# import gc
 import cv2
 
 import logging
 
 models_logger = logging.getLogger(__name__)
 
-from . import transforms, dynamics, utils, plot
+from . import transforms, dynamics, utils
 from .vit_sam import Transformer
-from .core import assign_device, run_net, run_3D
+from .core import assign_device, run_net
 
-_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_DIR_ENV = os.environ.get("CELLPOSE_LOCAL_MODELS_PATH")
+# _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_NAMES = ["cpsam"]
 
-MODEL_LIST_PATH = os.fspath(MODEL_DIR.joinpath("gui_models.txt"))
+# MODEL_LIST_PATH = os.fspath(MODEL_DIR.joinpath("gui_models.txt"))
 
 normalize_default = {
     "lowhigh": None,
@@ -42,30 +40,17 @@ normalize_default = {
 }
 
 
-# def model_path(model_type, model_index=0):
-#     return cache_CPSAM_model_path()
+# 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
 
 
-# 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 SegModel():
     """
     Class representing a Cellpose model.
 
@@ -102,17 +87,6 @@ class CellposeModel():
             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():
@@ -127,36 +101,10 @@ class CellposeModel():
             # 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,
@@ -166,12 +114,6 @@ class CellposeModel():
              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 = [], [], []
@@ -230,7 +172,7 @@ class CellposeModel():
         Ly_0 = x.shape[1]
         Lx_0 = x.shape[2]
         Lz_0 = None
-        if do_3D or stitch_threshold > 0:
+        if stitch_threshold > 0:
             Lz_0 = x.shape[0]
         if diameter is not None:
             image_scaling = 30. / diameter
@@ -273,7 +215,7 @@ class CellposeModel():
                 # transpose feat to have channels last
                 feat = np.moveaxis(feat, 1, -1)
 
-        # ajust the anisotropy when diameter is specified and images are resized:
+        # adjust the anisotropy when diameter is specified and images are resized:
         if isinstance(anisotropy, (float, int)) and image_scaling:
             anisotropy = image_scaling * anisotropy
 
@@ -287,12 +229,6 @@ class CellposeModel():
             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: 
@@ -310,28 +246,16 @@ class CellposeModel():
         else:
             masks = np.zeros(0) #pass back zeros if not compute_masks
         
-        masks, dP, cellprob = masks.squeeze(), dP.squeeze(), cellprob.squeeze()
+        masks = masks.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)
+            if compute_masks:
+                masks = transforms.resize_image(masks, Ly=Ly_0, Lx=Lx_0, no_channels=True, interpolation=cv2.INTER_NEAREST)
 
-            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
 
-        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:
@@ -428,29 +352,15 @@ class CellposeModel():
         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]
+        
+        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()
 
@@ -471,48 +381,48 @@ class CellposeModel():
             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,
+        # 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)
-        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"
-                )
+            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/seg_post_model/{cellpose/transforms.py → transforms.py}

@@ -398,145 +398,6 @@ def move_axis(img, m_axis=-1, first=True):
     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.
@@ -571,13 +432,9 @@ def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
     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
@@ -956,10 +813,10 @@ def random_rotate_and_resize(X, Y=None, scale_range=1., xy=(224, 224), do_3D=Fal
         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])
+    # 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 = []
@@ -1008,15 +865,6 @@ def random_rotate_and_resize(X, Y=None, scale_range=1., xy=(224, 224), do_3D=Fal
             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]
@@ -1024,47 +872,15 @@ def random_rotate_and_resize(X, Y=None, scale_range=1., xy=(224, 224), do_3D=Fal
                     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
+            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)
+                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()
@@ -1106,10 +922,7 @@ def random_rotate_and_resize_with_feat(X, Y=None, feat=None, scale_range=1., xy=
         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])
+    shape = (xy[0], xy[1])
     imgi = np.zeros((nimg, nchan, *shape), "float32")
 
     lbl = []
@@ -1169,17 +982,6 @@ def random_rotate_and_resize_with_feat(X, Y=None, feat=None, scale_range=1., xy=
             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]
@@ -1189,49 +991,16 @@ def random_rotate_and_resize_with_feat(X, Y=None, feat=None, scale_range=1., xy=
                         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
+            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)
+                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()
@@ -1241,20 +1010,7 @@ def random_rotate_and_resize_with_feat(X, Y=None, feat=None, scale_range=1., xy=
         
         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)
+                feat_out[n, k] = cv2.warpAffine(feats[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)
 
 
 

models/seg_post_model/utils.py

@@ -0,0 +1,214 @@
+"""
+Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.
+"""
+import logging
+import io
+from tqdm import tqdm, trange
+import cv2
+from scipy.ndimage import find_objects
+import numpy as np
+import fastremap
+import fill_voids
+from models.seg_post_model import metrics
+
+
+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 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 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/vit_sam.py → vit_sam.py}

@@ -130,66 +130,3 @@ class Transformer(nn.Module):
         """
         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)
-
-
-    
-

segmentation.py

@@ -1,9 +1,6 @@
 import os
-import pprint
 from typing import Any, List, Optional
-import argparse
 from huggingface_hub import hf_hub_download
-import pyrallis
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 import torch
 import os
@@ -27,7 +24,7 @@ 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
 from _utils.seg_eval import *
-from models.seg_post_model.cellpose import metrics
+from models.seg_post_model import metrics
 from datetime import datetime
 import json
 import logging
@@ -35,7 +32,7 @@ from PIL import Image
 import torchvision.transforms as T
 import cv2
 from skimage import io, measure
-logging.getLogger('models.seg_post_model.cellpose.models').setLevel(logging.ERROR)
+logging.getLogger('models.seg_post_model.models').setLevel(logging.ERROR)
 
 SCALE = 1
 
@@ -76,8 +73,6 @@ class SegmentationModule(pl.LightningModule):
                 " `placeholder_token` that is not already in the tokenizer."
             )
         try:
-            # print("loading pretrained task embedding from {}".format("pretrained/task_embed.pth"))
-            # task_embed_from_pretrain = torch.load("pretrained/task_embed.pth")
             task_embed_from_pretrain = hf_hub_download(
                 repo_id="phoebe777777/111",
                 filename="task_embed.pth",
@@ -190,11 +185,6 @@ class SegmentationModule(pl.LightningModule):
         exemplar_attention_maps2 = []
         exemplar_attention_maps3 = []
 
-        cross_self_task_attn_maps = []
-        cross_self_exe_attn_maps1 = []
-        cross_self_exe_attn_maps2 = []
-        cross_self_exe_attn_maps3 = []
-
         # only use 64x64 self-attention
         self_attn_aggregate = attn_utils.aggregate_attention( # [res, res, 4096]
                 prompts=[self.config.prompt for i in range(bsz)],        # 这里要改么
@@ -323,7 +313,7 @@ def inference(data_path, box=None, save_path="./example_imgs", visualize=False):
         ax[0].axis("off")
         ax[1].imshow(img_show)
         for inst_id in np.unique(mask_show):
-            if inst_id == 0:  # 0 通常是背景
+            if inst_id == 0:  # 0 background
                 continue
             # 生成二值 mask
             binary_mask = (mask_show == inst_id).astype(np.uint8)
@@ -352,12 +342,6 @@ def main():
 from matplotlib import cm
 
 def overlay_instances(img, mask, alpha=0.5, cmap_name="tab20"):
-    """
-    img: 原图 (H, W, 3)，范围 [0,255] 或 [0,1]
-    mask: 实例分割结果 (H, W)，背景=0，实例=1,2,...
-    alpha: 透明度
-    cmap_name: 颜色映射表
-    """
     img = img.astype(np.float32)
     if len(img.shape) == 2:
         img = np.stack([img]*3, axis=-1)
@@ -369,7 +353,7 @@ def overlay_instances(img, mask, alpha=0.5, cmap_name="tab20"):
     cmap = cm.get_cmap(cmap_name, np.max(mask)+1)
 
     for inst_id in np.unique(mask):
-        if inst_id == 0:  # 背景跳过
+        if inst_id == 0:  
             continue
         color = np.array(cmap(inst_id)[:3])  # RGB
         overlay[mask == inst_id] = (1 - alpha) * overlay[mask == inst_id] + alpha * color