Upload example images and assets using LFS

README.md

@@ -1,14 +1,14 @@
 ---
-title: FluoGen
-emoji: 🚀
-colorFrom: indigo
-colorTo: pink
+title: FluoGen Demo
+emoji: 📉
+colorFrom: red
+colorTo: red
 sdk: gradio
 sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: mit
-short_description: TMP
+short_description: 'Demo space of FluoGen: An Open-Source Generative Foundation '
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,664 @@
+import torch
+import numpy as np
+import gradio as gr
+from PIL import Image
+import os
+import json
+import glob
+import random
+import tifffile
+import re
+import imageio
+from torchvision import transforms, models
+import accelerate
+import shutil
+import time
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+
+# --- Imports from both scripts ---
+from diffusers import DDPMScheduler, DDIMScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+from accelerate.state import AcceleratorState
+from transformers.utils import ContextManagers
+
+# --- Custom Model Imports ---
+from models.pipeline_ddpm_text_encoder import DDPMPipeline
+from models.unet_2d import UNet2DModel
+from models.controlnet import ControlNetModel
+from models.unet_2d_condition import UNet2DConditionModel
+from models.pipeline_controlnet import DDPMControlnetPipeline
+
+# --- New Import for Segmentation ---
+from cellpose import models as cellpose_models
+from cellpose import plot as cellpose_plot
+from huggingface_hub import hf_hub_download
+
+# --- 0. Configuration & Constants ---
+# --- General ---
+MODEL_TITLE = "🔬 FluoGen: AI-Powered Fluorescence Microscopy Suite"
+MODEL_DESCRIPTION = """
+**Paper**: *Generative AI empowering fluorescence microscopy imaging and analysis*
+<br>
+Select a task from the tabs below: generate new images from text, enhance existing images using super-resolution, denoise them, generate training data from segmentation masks, perform cell segmentation, or classify cell images.
+"""
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+WEIGHT_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
+LOGO_PATH = "utils/logo_0801_2.png"
+
+# --- Global switch to control example saving ---
+SAVE_EXAMPLES = False
+
+# --- Base directory for all models ---
+# NOTE: All model paths are now relative.
+# Run the `copy_weights.py` script once to copy all necessary model files into this local directory.
+REPO_ID = "rayquaza384mega/FluoGen-demo-test-ckpts"
+MODELS_ROOT_DIR = hf_hub_download(repo_id=REPO_ID) #"models_collection"
+
+
+# --- Tab 1: Mask-to-Image Config (Formerly Segmentation-to-Image) ---
+M2I_CONTROLNET_PATH = f"{MODELS_ROOT_DIR}/ControlNet_M2I/checkpoint-30000"
+M2I_EXAMPLE_IMG_DIR = "example_images_m2i"
+
+# --- Tab 2: Text-to-Image Config ---
+T2I_PROMPTS = ["F-actin of COS-7", "ER of COS-7", "Mitochondria of BPAE", "Nucleus of BPAE", "ER of HeLa", "Microtubules of HeLa"]
+T2I_EXAMPLE_IMG_DIR = "example_images"
+T2I_CHECKPOINT = 285000
+T2I_PRETRAINED_MODEL_PATH = f"{MODELS_ROOT_DIR}/stable-diffusion-v1-5"
+T2I_UNET_PATH = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-{T2I_CHECKPOINT}"
+
+# --- Tab 3, 4: ControlNet-based Tasks Config ---
+CONTROLNET_CLIP_PATH = f"{MODELS_ROOT_DIR}/stable-diffusion-v1-5"
+CONTROLNET_UNET_PATH = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"
+
+# Super-Resolution Models
+SR_CONTROLNET_MODELS = {
+    "Checkpoint CCPs": f"{MODELS_ROOT_DIR}/ControlNet_SR/CCPs/checkpoint-100000",
+    "Checkpoint F-actin": f"{MODELS_ROOT_DIR}/ControlNet_SR/F-actin/checkpoint-35000",
+}
+SR_EXAMPLE_IMG_DIR = "example_images_sr"
+
+# Denoising Model
+DN_CONTROLNET_PATH = f"{MODELS_ROOT_DIR}/ControlNet_DN/checkpoint-10000"
+DN_PROMPT_RULES = {'MICE': 'mouse brain tissues', 'FISH': 'zebrafish embryos', 'BPAE_B': 'nucleus of BPAE', 'BPAE_R': 'mitochondria of BPAE', 'BPAE_G': 'F-actin of BPAE'}
+DN_EXAMPLE_IMG_DIR = "example_images_dn"
+
+# --- Tab 5: Cell Segmentation Config ---
+SEG_MODELS = {
+    "DynamicNet Model": f"{MODELS_ROOT_DIR}/Cellpose/DynamicNet_baseline/CP_dynamic_ten_epoch_0100",
+    "DynamicNet Model + FluoGen": f"{MODELS_ROOT_DIR}/Cellpose/DynamicNet_FluoGen/CP_dynamic_epoch_0300",
+    "DSB Model": f"{MODELS_ROOT_DIR}/Cellpose/DSB_baseline/CP_dsb_baseline_ratio_1_epoch_0135",
+    "DSB Model + FluoGen": f"{MODELS_ROOT_DIR}/Cellpose/DSB_FluoGen/CP_dsb_ten_epoch_0135",
+}
+SEG_EXAMPLE_IMG_DIR = "example_images_seg"
+
+
+# --- Tab 6: Classification Config ---
+CLS_MODEL_PATHS = OrderedDict({
+    "5shot":        f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_5_shot_re",
+    #"10shot":       f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_10_shot_re",
+    #"15shot":       f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_15_shot_re",
+    #"20shot":       f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_20_shot_re",
+    "5shot+FluoGen":    f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_5_shot_aug_re",
+    #"10shot_aug":   f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_10_shot_aug_re",
+    #"15shot_aug":   f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_15_shot_aug_re",
+    #"20shot_aug":   f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_20_shot_aug_re",
+})
+CLS_CLASS_NAMES = ['dap', 'erdak', 'giant', 'gpp130', 'h4b4', 'mc151', 'nucle', 'phal', 'tfr', 'tubul']
+CLS_EXAMPLE_IMG_DIR = "example_images_cls"
+
+
+# --- Helper Functions ---
+def sanitize_prompt_for_filename(prompt):
+    prompt = prompt.lower(); prompt = re.sub(r'\s+of\s+', '_', prompt); prompt = re.sub(r'[^a-z0-9-_]+', '', prompt)
+    return f"{prompt}.png"
+
+def min_max_norm(x):
+    x = x.astype(np.float32); min_val, max_val = np.min(x), np.max(x)
+    if max_val - min_val < 1e-8: return np.zeros_like(x)
+    return (x - min_val) / (max_val - min_val)
+
+def numpy_to_pil(image_np, target_mode="RGB"):
+    # If the input is already a PIL image, just ensure mode and return
+    if isinstance(image_np, Image.Image):
+        if target_mode == "RGB" and image_np.mode != "RGB": return image_np.convert("RGB")
+        if target_mode == "L" and image_np.mode != "L": return image_np.convert("L")
+        return image_np
+    
+    # Handle numpy array conversion
+    squeezed_np = np.squeeze(image_np); 
+    if squeezed_np.dtype == np.uint8:
+        # If it's already uint8, it's likely in the 0-255 range.
+        image_8bit = squeezed_np
+    else:
+        # Normalize and scale for other types
+        normalized_np = min_max_norm(squeezed_np)
+        image_8bit = (normalized_np * 255).astype(np.uint8)
+    
+    pil_image = Image.fromarray(image_8bit)
+    
+    if target_mode == "RGB" and pil_image.mode != "RGB": pil_image = pil_image.convert("RGB")
+    elif target_mode == "L" and pil_image.mode != "L": pil_image = pil_image.convert("L")
+    return pil_image
+
+# --- 1. Model Loading ---
+print("--- Initializing FluoGen Application ---")
+t2i_pipe, controlnet_pipe = None, None
+try:
+    print("Loading Text-to-Image model...")
+    t2i_noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear", prediction_type="v_prediction", rescale_betas_zero_snr=True, timestep_spacing="trailing")
+    t2i_unet = UNet2DModel.from_pretrained(T2I_UNET_PATH, subfolder="unet")
+    t2i_text_encoder = CLIPTextModel.from_pretrained(T2I_PRETRAINED_MODEL_PATH, subfolder="text_encoder").to(DEVICE)
+    t2i_tokenizer = CLIPTokenizer.from_pretrained(T2I_PRETRAINED_MODEL_PATH, subfolder="tokenizer")
+    t2i_pipe = DDPMPipeline(unet=t2i_unet, scheduler=t2i_noise_scheduler, text_encoder=t2i_text_encoder, tokenizer=t2i_tokenizer)
+    t2i_pipe.to(DEVICE)
+    print("✓ Text-to-Image model loaded successfully!")
+except Exception as e:
+    print(f"!!!!!! FATAL: Text-to-Image Model Loading Failed !!!!!!\nError: {e}")
+
+try:
+    print("Loading shared ControlNet pipeline components...")
+    controlnet_unet = UNet2DConditionModel.from_pretrained(CONTROLNET_UNET_PATH, subfolder="unet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    default_controlnet_path = M2I_CONTROLNET_PATH # Start with the first tab's model
+    controlnet_controlnet = ControlNetModel.from_pretrained(default_controlnet_path, subfolder="controlnet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    controlnet_scheduler = DDIMScheduler(num_train_timesteps=1000, beta_schedule="linear", prediction_type="v_prediction", rescale_betas_zero_snr=False, timestep_spacing="trailing")
+    controlnet_tokenizer = CLIPTokenizer.from_pretrained(CONTROLNET_CLIP_PATH, subfolder="tokenizer")
+    with ContextManagers([]):
+        controlnet_text_encoder = CLIPTextModel.from_pretrained(CONTROLNET_CLIP_PATH, subfolder="text_encoder").to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    controlnet_pipe = DDPMControlnetPipeline(unet=controlnet_unet, controlnet=controlnet_controlnet, scheduler=controlnet_scheduler, text_encoder=controlnet_text_encoder, tokenizer=controlnet_tokenizer)
+    controlnet_pipe.to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    controlnet_pipe.current_controlnet_path = default_controlnet_path
+    print("✓ Shared ControlNet pipeline loaded successfully!")
+except Exception as e:
+    print(f"!!!!!! FATAL: ControlNet Pipeline Loading Failed !!!!!!\nError: {e}")
+
+# --- 2. Core Logic Functions ---
+def swap_controlnet(pipe, target_path):
+    if os.path.normpath(getattr(pipe, 'current_controlnet_path', '')) != os.path.normpath(target_path):
+        print(f"Swapping ControlNet model to: {target_path}")
+        try:
+            pipe.controlnet = ControlNetModel.from_pretrained(target_path, subfolder="controlnet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
+            pipe.current_controlnet_path = target_path
+        except Exception as e:
+            raise gr.Error(f"Failed to load ControlNet model '{target_path}'. Error: {e}")
+    return pipe
+
+def generate_t2i(prompt, num_inference_steps):
+    if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
+    print(f"\nTask started... | Prompt: '{prompt}' | Steps: {num_inference_steps}")
+    image_np = t2i_pipe(prompt.lower(), generator=None, num_inference_steps=int(num_inference_steps), output_type="np").images
+    generated_image = numpy_to_pil(image_np)
+    print("✓ Image generated")
+    if SAVE_EXAMPLES:
+        example_filepath = os.path.join(T2I_EXAMPLE_IMG_DIR, sanitize_prompt_for_filename(prompt))
+        if not os.path.exists(example_filepath):
+            generated_image.save(example_filepath); print(f"✓ New T2I example saved: {example_filepath}")
+    return generated_image
+
+def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed):
+    if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
+    if mask_file_obj is None: raise gr.Error("Please upload a segmentation mask TIF file.")
+    if not cell_type or not cell_type.strip(): raise gr.Error("Please enter a cell type.")
+    
+    if SAVE_EXAMPLES:
+        input_path = mask_file_obj.name
+        filename = os.path.basename(input_path)
+        dest_path = os.path.join(M2I_EXAMPLE_IMG_DIR, filename)
+        if not os.path.exists(dest_path):
+            shutil.copy(input_path, dest_path)
+            print(f"✓ New Mask-to-Image example saved: {dest_path}")
+            
+    pipe = swap_controlnet(controlnet_pipe, M2I_CONTROLNET_PATH)
+    try:
+        mask_np = tifffile.imread(mask_file_obj.name)
+    except Exception as e:
+        raise gr.Error(f"Failed to read the TIF file. Error: {e}")
+        
+    input_display_image = numpy_to_pil(mask_np, "L")
+    mask_normalized = min_max_norm(mask_np)
+    image_tensor = torch.from_numpy(mask_normalized.astype(np.float32))
+    image_tensor = image_tensor.unsqueeze(0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    image_tensor = transforms.Resize((512, 512), antialias=True)(image_tensor)
+    
+    prompt = f"nuclei of {cell_type.strip()}"
+    print(f"\nTask started... | Task: Mask-to-Image | Prompt: '{prompt}' | Steps: {steps} | Images: {num_images}")
+    
+    generated_images_pil = []
+    for i in range(int(num_images)):
+        current_seed = int(seed) + i
+        generator = torch.Generator(device=DEVICE).manual_seed(current_seed)
+        with torch.autocast("cuda"):
+            output_np = pipe(prompt=prompt, image_cond=image_tensor, num_inference_steps=int(steps), generator=generator, output_type="np").images
+        pil_image = numpy_to_pil(output_np)
+        generated_images_pil.append(pil_image)
+        print(f"✓ Generated image {i+1}/{int(num_images)}")
+        
+    return input_display_image, generated_images_pil
+
+def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps, seed):
+    if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
+    if low_res_file_obj is None: raise gr.Error("Please upload a low-resolution TIF file.")
+    
+    if SAVE_EXAMPLES:
+        input_path = low_res_file_obj.name
+        filename = os.path.basename(input_path)
+        dest_path = os.path.join(SR_EXAMPLE_IMG_DIR, filename)
+        if not os.path.exists(dest_path):
+            shutil.copy(input_path, dest_path)
+            print(f"✓ New SR example saved: {dest_path}")
+            
+    target_path = SR_CONTROLNET_MODELS.get(controlnet_model_name)
+    if not target_path: raise gr.Error(f"ControlNet model '{controlnet_model_name}' not found.")
+    
+    pipe = swap_controlnet(controlnet_pipe, target_path)
+    try:
+        image_stack_np = tifffile.imread(low_res_file_obj.name)
+    except Exception as e:
+        raise gr.Error(f"Failed to read the TIF file. Error: {e}")
+        
+    if image_stack_np.ndim != 3 or image_stack_np.shape[-3] != 9:
+        raise gr.Error(f"Invalid TIF shape. Expected 9 channels (shape 9, H, W), but got {image_stack_np.shape}.")
+        
+    average_projection_np = np.mean(image_stack_np, axis=0)
+    input_display_image = numpy_to_pil(average_projection_np, "L")
+    
+    image_tensor = torch.from_numpy(image_stack_np.astype(np.float32) / 65535.0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    image_tensor = transforms.Resize((512, 512), antialias=True)(image_tensor)
+    
+    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
+    with torch.autocast("cuda"):
+        output_np = pipe(prompt=prompt, image_cond=image_tensor, num_inference_steps=int(steps), generator=generator, output_type="np").images
+        
+    return input_display_image, numpy_to_pil(output_np)
+
+def run_denoising(noisy_image_np, image_type, steps, seed):
+    if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
+    if noisy_image_np is None: raise gr.Error("Please upload a noisy image.")
+    
+    if SAVE_EXAMPLES:
+        timestamp = int(time.time() * 1000)
+        filename = f"dn_input_{image_type}_{timestamp}.tif"
+        dest_path = os.path.join(DN_EXAMPLE_IMG_DIR, filename)
+        try:
+            img_to_save = noisy_image_np.astype(np.uint8) if noisy_image_np.dtype != np.uint8 else noisy_image_np
+            tifffile.imwrite(dest_path, img_to_save)
+            print(f"✓ New Denoising example saved: {dest_path}")
+        except Exception as e:
+            print(f"✗ Failed to save denoising example: {e}")
+            
+    pipe = swap_controlnet(controlnet_pipe, DN_CONTROLNET_PATH)
+    prompt = DN_PROMPT_RULES.get(image_type, 'microscopy image')
+    print(f"\nTask started... | Task: Denoising | Prompt: '{prompt}' | Steps: {steps}")
+    
+    image_tensor = torch.from_numpy(noisy_image_np.astype(np.float32) / 255.0)
+    image_tensor = image_tensor.unsqueeze(0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    image_tensor = transforms.Resize((512, 512), antialias=True)(image_tensor)
+    
+    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
+    with torch.autocast("cuda"):
+        output_np = pipe(prompt=prompt, image_cond=image_tensor, num_inference_steps=int(steps), generator=generator, output_type="np").images
+        
+    return numpy_to_pil(noisy_image_np, "L"), numpy_to_pil(output_np)
+
+def run_segmentation(input_image_np, model_name, diameter, flow_threshold, cellprob_threshold):
+    """
+    Runs cell segmentation and creates a dark red overlay.
+    """
+    if input_image_np is None:
+        raise gr.Error("Please upload an image to segment.")
+        
+    model_path = SEG_MODELS.get(model_name)
+    if not model_path:
+        raise gr.Error(f"Segmentation model '{model_name}' not found.")
+    
+    if not os.path.exists(model_path):
+        raise gr.Error(f"Model file not found at path: {model_path}. Please check the configuration.")
+
+    print(f"\nTask started... | Task: Cell Segmentation | Model: '{model_name}'")
+    
+    # 1. Load Cellpose Model
+    try:
+        use_gpu = torch.cuda.is_available()
+        model = cellpose_models.CellposeModel(gpu=use_gpu, pretrained_model=model_path)
+    except Exception as e:
+        raise gr.Error(f"Failed to load Cellpose model. Error: {e}")
+
+    diameter_to_use = model.diam_labels if diameter == 0 else float(diameter)
+    print(f"Using Diameter: {diameter_to_use}")
+
+    # 2. Run model evaluation
+    try:
+        masks, _, _ = model.eval(
+            [input_image_np], 
+            channels=[0, 0],
+            diameter=diameter_to_use,
+            flow_threshold=flow_threshold,
+            cellprob_threshold=cellprob_threshold
+        )
+        mask_output = masks[0]
+    except Exception as e:
+        raise gr.Error(f"Cellpose model evaluation failed. Error: {e}")
+
+    # 3. Create custom dark red overlay
+    # Ensure input image is uint8 and 3-channel for blending
+    original_rgb = numpy_to_pil(input_image_np, "RGB")
+    original_rgb_np = np.array(original_rgb)
+
+    # Create a blank layer for the red mask
+    red_mask_layer = np.zeros_like(original_rgb_np)
+    dark_red_color = [139, 0, 0]
+    
+    # Apply the red color where the mask is present
+    is_mask_pixels = mask_output > 0
+    red_mask_layer[is_mask_pixels] = dark_red_color
+    
+    # Blend the original image with the red mask layer
+    alpha = 0.4  # Opacity of the mask
+    blended_image_np = ((1 - alpha) * original_rgb_np + alpha * red_mask_layer).astype(np.uint8)
+    
+    # 4. Save example if enabled
+    if SAVE_EXAMPLES:
+        timestamp = int(time.time() * 1000)
+        filename = f"seg_input_{timestamp}.tif"
+        dest_path = os.path.join(SEG_EXAMPLE_IMG_DIR, filename)
+        try:
+            img_to_save = input_image_np.astype(np.uint8) if input_image_np.dtype != np.uint8 else input_image_np
+            tifffile.imwrite(dest_path, img_to_save)
+            print(f"✓ New Segmentation example saved: {dest_path}")
+        except Exception as e:
+            print(f"✗ Failed to save segmentation example: {e}")
+            
+    print("✓ Segmentation complete")
+    
+    return numpy_to_pil(input_image_np, "L"), numpy_to_pil(blended_image_np, "RGB")
+
+def run_classification(input_image_np, model_name):
+    """
+    Runs classification on a single image using a pre-trained ResNet50 model.
+    """
+    if input_image_np is None:
+        raise gr.Error("Please upload an image to classify.")
+
+    model_dir = CLS_MODEL_PATHS.get(model_name)
+    if not model_dir:
+        raise gr.Error(f"Classification model '{model_name}' not found.")
+        
+    model_path = os.path.join(model_dir, "best_resnet50.pth")
+    if not os.path.exists(model_path):
+        raise gr.Error(f"Model file not found at {model_path}. Please check the configuration.")
+
+    print(f"\nTask started... | Task: Classification | Model: '{model_name}'")
+
+    # 1. Load Model
+    try:
+        model = models.resnet50(weights=None)
+        num_features = model.fc.in_features
+        model.fc = nn.Linear(num_features, len(CLS_CLASS_NAMES))
+        model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+        model.to(DEVICE)
+        model.eval()
+    except Exception as e:
+        raise gr.Error(f"Failed to load classification model. Error: {e}")
+
+    # 2. Preprocess Image
+    # Grayscale numpy -> RGB PIL -> transform -> tensor
+    input_pil = numpy_to_pil(input_image_np, "RGB")
+    
+    transform_test = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # ResNet needs 3-channel norm
+    ])
+    input_tensor = transform_test(input_pil).unsqueeze(0).to(DEVICE)
+
+    # 3. Perform Inference
+    with torch.no_grad():
+        outputs = model(input_tensor)
+        probabilities = F.softmax(outputs, dim=1).squeeze().cpu().numpy()
+
+    # 4. Format output for Gradio Label component
+    confidences = {name: float(prob) for name, prob in zip(CLS_CLASS_NAMES, probabilities)}
+    
+    # 5. Save example
+    if SAVE_EXAMPLES:
+        timestamp = int(time.time() * 1000)
+        filename = f"cls_input_{timestamp}.png" # Save as png for compatibility
+        dest_path = os.path.join(CLS_EXAMPLE_IMG_DIR, filename)
+        try:
+            input_pil.save(dest_path)
+            print(f"✓ New Classification example saved: {dest_path}")
+        except Exception as e:
+            print(f"✗ Failed to save classification example: {e}")
+            
+    print("✓ Classification complete")
+
+    return numpy_to_pil(input_image_np, "L"), confidences
+
+
+# --- 3. Gradio UI Layout ---
+print("Building Gradio interface...")
+# Create directories for all example types
+os.makedirs(M2I_EXAMPLE_IMG_DIR, exist_ok=True)
+os.makedirs(T2I_EXAMPLE_IMG_DIR, exist_ok=True)
+os.makedirs(SR_EXAMPLE_IMG_DIR, exist_ok=True)
+os.makedirs(DN_EXAMPLE_IMG_DIR, exist_ok=True)
+os.makedirs(SEG_EXAMPLE_IMG_DIR, exist_ok=True)
+os.makedirs(CLS_EXAMPLE_IMG_DIR, exist_ok=True)
+
+# --- Load examples ---
+filename_to_prompt_map = { sanitize_prompt_for_filename(prompt): prompt for prompt in T2I_PROMPTS }
+t2i_gallery_examples = []
+for filename in os.listdir(T2I_EXAMPLE_IMG_DIR):
+    if filename in filename_to_prompt_map:
+        filepath = os.path.join(T2I_EXAMPLE_IMG_DIR, filename)
+        prompt = filename_to_prompt_map[filename]
+        t2i_gallery_examples.append((filepath, prompt))
+
+def load_image_examples(example_dir, is_stack=False):
+    examples = []
+    if not os.path.exists(example_dir): return examples
+    for f in sorted(os.listdir(example_dir)):
+        if f.lower().endswith(('.tif', '.tiff', '.png', '.jpg', '.jpeg')):
+            filepath = os.path.join(example_dir, f)
+            try:
+                if f.lower().endswith(('.tif', '.tiff')):
+                    img_np = tifffile.imread(filepath)
+                else:
+                    img_np = np.array(Image.open(filepath).convert("L"))
+
+                if is_stack and img_np.ndim == 3:
+                    img_np = np.mean(img_np, axis=0)
+                
+                display_img = numpy_to_pil(img_np, "L")
+                examples.append((display_img, filepath))
+            except Exception as e:
+                print(f"Warning: Could not load gallery image {filepath}. Error: {e}")
+    return examples
+
+m2i_gallery_examples = load_image_examples(M2I_EXAMPLE_IMG_DIR)
+sr_gallery_examples = load_image_examples(SR_EXAMPLE_IMG_DIR, is_stack=True)
+dn_gallery_examples = load_image_examples(DN_EXAMPLE_IMG_DIR)
+seg_gallery_examples = load_image_examples(SEG_EXAMPLE_IMG_DIR)
+cls_gallery_examples = load_image_examples(CLS_EXAMPLE_IMG_DIR)
+
+# --- Universal event handlers ---
+def select_example_prompt(evt: gr.SelectData):
+    return evt.value['caption']
+    
+def select_example_input_file(evt: gr.SelectData):
+    return evt.value['caption']
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    with gr.Row():
+        gr.Image(value=LOGO_PATH, width=300, height=200, container=False, interactive=False, show_download_button=False, show_fullscreen_button=False)
+        gr.Markdown(f"# {MODEL_TITLE}\n{MODEL_DESCRIPTION}")
+
+    with gr.Tabs():
+        # --- TAB 1: Mask-to-Image ---
+        with gr.Tab("Mask-to-Image", id="mask2img"):
+            gr.Markdown("""
+            ### Instructions
+            1.  Upload a single-channel segmentation mask (`.tif` file), or select one from the examples gallery below.
+            2.  Enter the corresponding 'Cell Type' (e.g., 'CoNSS', 'HeLa') to create the prompt.
+            3.  Select how many sample images you want to generate.
+            4.  Adjust 'Inference Steps' and 'Seed' as needed.
+            5.  Click 'Generate Training Samples' to start the process.
+            6.  The 'Generated Samples' will appear in the main gallery, with the 'Input Mask' shown below for reference.
+            """) # Content hidden for brevity
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1, min_width=350):
+                    m2i_input_file = gr.File(label="Upload Segmentation Mask (.tif)", file_types=['.tif', '.tiff'])
+                    m2i_cell_type_input = gr.Textbox(label="Cell Type (for prompt)", placeholder="e.g., CoNSS, HeLa, MCF-7")
+                    m2i_num_images_slider = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of Images to Generate")
+                    m2i_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")
+                    m2i_seed_input = gr.Number(label="Seed", value=42)
+                    m2i_generate_button = gr.Button("Generate Training Samples", variant="primary")
+                with gr.Column(scale=2):
+                    m2i_output_gallery = gr.Gallery(label="Generated Samples", columns=5, object_fit="contain", height="auto")
+                    m2i_input_display = gr.Image(label="Input Mask", type="pil", interactive=False)
+            m2i_gallery = gr.Gallery(value=m2i_gallery_examples, label="Input Examples (Click an image to use it as input)", columns=6, object_fit="contain", height="auto")
+
+        # --- TAB 2: Text-to-Image ---
+        with gr.Tab("Text-to-Image Generation", id="txt2img"):
+            gr.Markdown("""
+            ### Instructions
+            1. Select a desired prompt from the dropdown menu.
+            2. Adjust the 'Inference Steps' slider to control generation quality.
+            3. Click the 'Generate' button to create a new image.
+            4. Explore the 'Examples' gallery; clicking an image will load its prompt.
+            
+            **Notice:** This model currently supports 3566 prompt categories. However, data for many cell structures and lines is still lacking. We welcome data source contributions to improve the model.
+            """) # Content hidden for brevity
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1, min_width=350):
+                    t2i_prompt_input = gr.Dropdown(choices=T2I_PROMPTS, value=T2I_PROMPTS[0], label="Select a Prompt")
+                    t2i_steps_slider = gr.Slider(minimum=10, maximum=200, step=1, value=50, label="Inference Steps")
+                    t2i_generate_button = gr.Button("Generate", variant="primary")
+                with gr.Column(scale=2):
+                    t2i_generated_output = gr.Image(label="Generated Image", type="pil", interactive=False)
+            t2i_gallery = gr.Gallery(value=t2i_gallery_examples, label="Examples (Click an image to use its prompt)", columns=6, object_fit="contain", height="auto")
+
+        # --- TAB 3: Super-Resolution ---
+        with gr.Tab("Super-Resolution", id="super_res"):
+            gr.Markdown("""
+            ### Instructions
+            1. Upload a low-resolution 9-channel TIF stack, or select one from the examples.
+            2. Select a 'Super-Resolution Model' from the dropdown.
+            3. Enter a descriptive 'Prompt' related to the image content (e.g., 'CCPs of COS-7').
+            4. Adjust 'Inference Steps' and 'Seed' as needed.
+            5. Click 'Generate Super-Resolution' to process the image.
+            
+            **Notice:** This model was trained on the **BioSR** dataset. If your data's characteristics differ significantly, please consider fine-tuning the model using our project on GitHub for optimal results.
+            """) # Content hidden for brevity
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1, min_width=350):
+                    sr_input_file = gr.File(label="Upload 9-Channel TIF Stack", file_types=['.tif', '.tiff'])
+                    sr_model_selector = gr.Dropdown(choices=list(SR_CONTROLNET_MODELS.keys()), value=list(SR_CONTROLNET_MODELS.keys())[-1], label="Select Super-Resolution Model")
+                    sr_prompt_input = gr.Textbox(label="Prompt (e.g., structure name)", value="CCPs of COS-7")
+                    sr_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")
+                    sr_seed_input = gr.Number(label="Seed", value=42)
+                    sr_generate_button = gr.Button("Generate Super-Resolution", variant="primary")
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        sr_input_display = gr.Image(label="Input (Average Projection)", type="pil", interactive=False)
+                        sr_output_image = gr.Image(label="Super-Resolved Image", type="pil", interactive=False)
+            sr_gallery = gr.Gallery(value=sr_gallery_examples, label="Input Examples (Click an image to use it as input)", columns=6, object_fit="contain", height="auto")
+
+        # --- TAB 4: Denoising ---
+        with gr.Tab("Denoising", id="denoising"):
+            gr.Markdown("""
+            ### Instructions
+            1. Upload a noisy single-channel image, or select one from the examples.
+            2. Select the 'Image Type' from the dropdown to provide context for the model.
+            3. Adjust 'Inference Steps' and 'Seed' as needed.
+            4. Click 'Denoise Image' to reduce the noise.
+            
+            **Notice:** This model was trained on the **FMD** dataset. If your data's characteristics differ significantly, please consider fine-tuning the model using our project on GitHub for optimal results.
+            """) # Content hidden for brevity
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1, min_width=350):
+                    dn_input_image = gr.Image(type="numpy", label="Upload Noisy Image", image_mode="L")
+                    dn_image_type_selector = gr.Dropdown(choices=list(DN_PROMPT_RULES.keys()), value='MICE', label="Select Image Type (for Prompt)")
+                    dn_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")
+                    dn_seed_input = gr.Number(label="Seed", value=42)
+                    dn_generate_button = gr.Button("Denoise Image", variant="primary")
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        dn_original_display = gr.Image(label="Original Noisy Image", type="pil", interactive=False)
+                        dn_output_image = gr.Image(label="Denoised Image", type="pil", interactive=False)
+            dn_gallery = gr.Gallery(value=dn_gallery_examples, label="Input Examples (Click an image to use it as input)", columns=6, object_fit="contain", height="auto")
+
+        # --- TAB 5: Cell Segmentation ---
+        with gr.Tab("Cell Segmentation", id="segmentation"):
+            gr.Markdown("""
+            ### Instructions
+            1. Upload a single-channel image for segmentation, or select one from the examples.
+            2. Select a 'Segmentation Model' from the dropdown menu.
+            3. Set the expected 'Diameter' of the cells in pixels. Set to 0 to let the model automatically estimate it.
+            4. Adjust 'Flow Threshold' and 'Cell Probability Threshold' for finer control.
+            5. Click 'Segment Cells'. The result will be shown as a dark red overlay on the original image.
+            """)
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1, min_width=350):
+                    gr.Markdown("### 1. Inputs & Controls")
+                    seg_input_image = gr.Image(type="numpy", label="Upload Image for Segmentation", image_mode="L")
+                    seg_model_selector = gr.Dropdown(choices=list(SEG_MODELS.keys()), value=list(SEG_MODELS.keys())[0], label="Select Segmentation Model")
+                    seg_diameter_input = gr.Number(label="Cell Diameter (pixels, 0=auto)", value=30)
+                    seg_flow_slider = gr.Slider(minimum=0.0, maximum=3.0, step=0.1, value=0.4, label="Flow Threshold")
+                    seg_cellprob_slider = gr.Slider(minimum=-6.0, maximum=6.0, step=0.5, value=0.0, label="Cell Probability Threshold")
+                    seg_generate_button = gr.Button("Segment Cells", variant="primary")
+                with gr.Column(scale=2):
+                    gr.Markdown("### 2. Results")
+                    with gr.Row():
+                        seg_original_display = gr.Image(label="Original Image", type="pil", interactive=False)
+                        seg_output_image = gr.Image(label="Segmented Image (Overlay)", type="pil", interactive=False)
+            seg_gallery = gr.Gallery(value=seg_gallery_examples, label="Input Examples (Click an image to use it as input)", columns=6, object_fit="contain", height="auto")
+            
+        # --- NEW TAB 6: Classification ---
+        with gr.Tab("Classification", id="classification"):
+            gr.Markdown("""
+            ### Instructions
+            1.  Upload a single-channel image for classification, or select an example.
+            2.  Select a pre-trained 'Classification Model' from the dropdown menu.
+            3.  Click 'Classify Image' to view the prediction probabilities for each class.
+            
+            **Note:** The models provided are ResNet50 trained on the 2D HeLa dataset.
+            """)
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1, min_width=350):
+                    gr.Markdown("### 1. Inputs & Controls")
+                    cls_input_image = gr.Image(type="numpy", label="Upload Image for Classification", image_mode="L")
+                    cls_model_selector = gr.Dropdown(choices=list(CLS_MODEL_PATHS.keys()), value=list(CLS_MODEL_PATHS.keys())[0], label="Select Classification Model")
+                    cls_generate_button = gr.Button("Classify Image", variant="primary")
+                with gr.Column(scale=2):
+                    gr.Markdown("### 2. Results")
+                    cls_original_display = gr.Image(label="Input Image", type="pil", interactive=False)
+                    cls_output_label = gr.Label(label="Classification Results", num_top_classes=len(CLS_CLASS_NAMES))
+            cls_gallery = gr.Gallery(value=cls_gallery_examples, label="Input Examples (Click an image to use it as input)", columns=6, object_fit="contain", height="auto")
+
+
+    # --- Event Handlers ---
+    m2i_generate_button.click(fn=run_mask_to_image_generation, inputs=[m2i_input_file, m2i_cell_type_input, m2i_num_images_slider, m2i_steps_slider, m2i_seed_input], outputs=[m2i_input_display, m2i_output_gallery])
+    m2i_gallery.select(fn=select_example_input_file, outputs=m2i_input_file)
+    
+    t2i_generate_button.click(fn=generate_t2i, inputs=[t2i_prompt_input, t2i_steps_slider], outputs=[t2i_generated_output])
+    t2i_gallery.select(fn=select_example_prompt, outputs=t2i_prompt_input)
+    
+    sr_generate_button.click(fn=run_super_resolution, inputs=[sr_input_file, sr_model_selector, sr_prompt_input, sr_steps_slider, sr_seed_input], outputs=[sr_input_display, sr_output_image])
+    sr_gallery.select(fn=select_example_input_file, outputs=sr_input_file)
+    
+    dn_generate_button.click(fn=run_denoising, inputs=[dn_input_image, dn_image_type_selector, dn_steps_slider, dn_seed_input], outputs=[dn_original_display, dn_output_image])
+    dn_gallery.select(fn=select_example_input_file, outputs=dn_input_image)
+
+    seg_generate_button.click(fn=run_segmentation, inputs=[seg_input_image, seg_model_selector, seg_diameter_input, seg_flow_slider, seg_cellprob_slider], outputs=[seg_original_display, seg_output_image])
+    seg_gallery.select(fn=select_example_input_file, outputs=seg_input_image)
+
+    cls_generate_button.click(fn=run_classification, inputs=[cls_input_image, cls_model_selector], outputs=[cls_original_display, cls_output_label])
+    cls_gallery.select(fn=select_example_input_file, outputs=cls_input_image)
+
+
+# --- 4. Launch Application ---
+if __name__ == "__main__":
+    print("Interface built. Launching server...")
+    demo.launch()

cellpose/__init__.py

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

cellpose/__main__.py

@@ -0,0 +1,380 @@
+"""
+Copyright © 2023 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.
+"""
+
+import sys, os, glob, pathlib, time
+import numpy as np
+from natsort import natsorted
+from tqdm import tqdm
+from cellpose import utils, models, io, version_str, train, denoise
+from cellpose.cli import get_arg_parser
+
+try:
+    from cellpose.gui import gui3d, gui
+    GUI_ENABLED = True
+except ImportError as err:
+    GUI_ERROR = err
+    GUI_ENABLED = False
+    GUI_IMPORT = True
+except Exception as err:
+    GUI_ENABLED = False
+    GUI_ERROR = err
+    GUI_IMPORT = False
+    raise
+
+import logging
+
+
+# settings re-grouped a bit
+def main():
+    """ Run cellpose from command line
+    """
+
+    args = get_arg_parser().parse_args(
+    )  # this has to be in a separate file for autodoc to work
+
+    if args.version:
+        print(version_str)
+        return
+
+    if args.check_mkl:
+        mkl_enabled = models.check_mkl()
+    else:
+        mkl_enabled = True
+
+    if len(args.dir) == 0 and len(args.image_path) == 0:
+        if args.add_model:
+            io.add_model(args.add_model)
+        else:
+            if not GUI_ENABLED:
+                print("GUI ERROR: %s" % GUI_ERROR)
+                if GUI_IMPORT:
+                    print(
+                        "GUI FAILED: GUI dependencies may not be installed, to install, run"
+                    )
+                    print("     pip install 'cellpose[gui]'")
+            else:
+                if args.Zstack:
+                    gui3d.run()
+                else:
+                    gui.run()
+
+    else:
+        if args.verbose:
+            from .io import logger_setup
+            logger, log_file = logger_setup()
+        else:
+            print(
+                ">>>> !LOGGING OFF BY DEFAULT! To see cellpose progress, set --verbose")
+            print("No --verbose => no progress or info printed")
+            logger = logging.getLogger(__name__)
+
+        use_gpu = False
+        channels = [args.chan, args.chan2]
+
+        # find images
+        if len(args.img_filter) > 0:
+            imf = args.img_filter
+        else:
+            imf = None
+
+        # Check with user if they REALLY mean to run without saving anything
+        if not (args.train or args.train_size):
+            saving_something = args.save_png or args.save_tif or args.save_flows or args.save_txt
+
+        device, gpu = models.assign_device(use_torch=True, gpu=args.use_gpu,
+                                           device=args.gpu_device)
+
+        if args.pretrained_model is None or args.pretrained_model == "None" or args.pretrained_model == "False" or args.pretrained_model == "0":
+            pretrained_model = False
+        else:
+            pretrained_model = args.pretrained_model
+
+        restore_type = args.restore_type
+        if restore_type is not None:
+            try:
+                denoise.model_path(restore_type)
+            except Exception as e:
+                raise ValueError("restore_type invalid")
+            if args.train or args.train_size:
+                raise ValueError("restore_type cannot be used with training on CLI yet")
+
+        if args.transformer and (restore_type is None):
+            default_model = "transformer_cp3"
+            backbone = "transformer"
+        elif args.transformer and restore_type is not None:
+            raise ValueError("no transformer based restoration")
+        else:
+            default_model = "cyto3"
+            backbone = "default"
+
+        if args.norm_percentile is not None:
+            value1, value2 = args.norm_percentile
+            normalize = {'percentile': (float(value1), float(value2))} 
+        else:
+            normalize = (not args.no_norm)
+        
+
+        model_type = None
+        if pretrained_model and not os.path.exists(pretrained_model):
+            model_type = pretrained_model if pretrained_model is not None else "cyto3"
+            model_strings = models.get_user_models()
+            all_models = models.MODEL_NAMES.copy()
+            all_models.extend(model_strings)
+            if ~np.any([model_type == s for s in all_models]):
+                model_type = default_model
+                logger.warning(
+                    f"pretrained model has incorrect path, using {default_model}")
+            if model_type == "nuclei":
+                szmean = 17.
+            else:
+                szmean = 30.
+        builtin_size = (model_type == "cyto" or model_type == "cyto2" or
+                        model_type == "nuclei" or model_type == "cyto3")
+
+        if len(args.image_path) > 0 and (args.train or args.train_size):
+            raise ValueError("ERROR: cannot train model with single image input")
+
+        if not args.train and not args.train_size:
+            tic = time.time()
+            if len(args.dir) > 0:
+                image_names = io.get_image_files(
+                    args.dir, args.mask_filter, imf=imf,
+                    look_one_level_down=args.look_one_level_down)
+            else:
+                if os.path.exists(args.image_path):
+                    image_names = [args.image_path]
+                else:
+                    raise ValueError(f"ERROR: no file found at {args.image_path}")
+            nimg = len(image_names)
+
+            if args.savedir:
+                if not os.path.exists(args.savedir):
+                    raise FileExistsError("--savedir {args.savedir} does not exist")
+
+            cstr0 = ["GRAY", "RED", "GREEN", "BLUE"]
+            cstr1 = ["NONE", "RED", "GREEN", "BLUE"]
+            logger.info(
+                ">>>> running cellpose on %d images using chan_to_seg %s and chan (opt) %s"
+                % (nimg, cstr0[channels[0]], cstr1[channels[1]]))
+
+            # handle built-in model exceptions
+            if builtin_size and restore_type is None and not args.pretrained_model_ortho:
+                model = models.Cellpose(gpu=gpu, device=device, model_type=model_type,
+                                        backbone=backbone)
+            else:
+                builtin_size = False
+                if args.all_channels:
+                    channels = None
+                    img = io.imread(image_names[0])
+                    if img.ndim == 3:
+                        nchan = min(img.shape)
+                    elif img.ndim == 2:
+                        nchan = 1
+                    channels = None
+                else:
+                    nchan = 2
+
+                pretrained_model = None if model_type is not None else pretrained_model
+                if restore_type is None:
+                    pretrained_model_ortho = None if args.pretrained_model_ortho is None else args.pretrained_model_ortho
+                    model = models.CellposeModel(gpu=gpu, device=device,
+                                                 pretrained_model=pretrained_model,
+                                                 model_type=model_type,
+                                                 nchan=nchan,
+                                                 backbone=backbone,
+                                                 pretrained_model_ortho=pretrained_model_ortho)
+                else:
+                    model = denoise.CellposeDenoiseModel(
+                        gpu=gpu, device=device, pretrained_model=pretrained_model,
+                        model_type=model_type, restore_type=restore_type, nchan=nchan,
+                        chan2_restore=args.chan2_restore)
+
+            # handle diameters
+            if args.diameter == 0:
+                if builtin_size:
+                    diameter = None
+                    logger.info(">>>> estimating diameter for each image")
+                else:
+                    if restore_type is None:
+                        logger.info(
+                            ">>>> not using cyto3, cyto, cyto2, or nuclei model, cannot auto-estimate diameter"
+                        )
+                    else:
+                        logger.info(
+                            ">>>> cannot auto-estimate diameter for image restoration")
+                    diameter = model.diam_labels
+                    logger.info(">>>> using diameter %0.3f for all images" % diameter)
+            else:
+                diameter = args.diameter
+                logger.info(">>>> using diameter %0.3f for all images" % diameter)
+
+            tqdm_out = utils.TqdmToLogger(logger, level=logging.INFO)
+
+            for image_name in tqdm(image_names, file=tqdm_out):
+                image = io.imread(image_name)
+                out = model.eval(
+                    image, channels=channels, diameter=diameter, do_3D=args.do_3D,
+                    augment=args.augment, resample=(not args.no_resample),
+                    flow_threshold=args.flow_threshold,
+                    cellprob_threshold=args.cellprob_threshold,
+                    stitch_threshold=args.stitch_threshold, min_size=args.min_size,
+                    invert=args.invert, batch_size=args.batch_size,
+                    interp=(not args.no_interp), normalize=normalize,
+                    channel_axis=args.channel_axis, z_axis=args.z_axis,
+                    anisotropy=args.anisotropy, niter=args.niter,
+                    flow3D_smooth=args.flow3D_smooth)
+                masks, flows = out[:2]
+                if len(out) > 3 and restore_type is None:
+                    diams = out[-1]
+                else:
+                    diams = diameter
+                ratio = 1.
+                if restore_type is not None:
+                    imgs_dn = out[-1]
+                    ratio = diams / model.dn.diam_mean if "upsample" in restore_type else 1.
+                    diams = model.dn.diam_mean if "upsample" in restore_type and model.dn.diam_mean > diams else diams
+                else:
+                    imgs_dn = None
+                if args.exclude_on_edges:
+                    masks = utils.remove_edge_masks(masks)
+                if not args.no_npy:
+                    io.masks_flows_to_seg(image, masks, flows, image_name,
+                                          imgs_restore=imgs_dn, channels=channels,
+                                          diams=diams, restore_type=restore_type,
+                                          ratio=1.)
+                if saving_something:
+                    suffix = "_cp_masks"
+                    if args.output_name is not None: 
+                        # (1) If `savedir` is not defined, then must have a non-zero `suffix`
+                        if args.savedir is None and len(args.output_name) > 0:
+                            suffix = args.output_name
+                        elif args.savedir is not None and not os.path.samefile(args.savedir, args.dir):
+                            # (2) If `savedir` is defined, and different from `dir` then                              
+                            # takes the value passed as a param. (which can be empty string)
+                            suffix = args.output_name
+
+                    io.save_masks(image, masks, flows, image_name,
+                                  suffix=suffix, png=args.save_png,
+                                  tif=args.save_tif, save_flows=args.save_flows,
+                                  save_outlines=args.save_outlines,
+                                  dir_above=args.dir_above, savedir=args.savedir,
+                                  save_txt=args.save_txt, in_folders=args.in_folders,
+                                  save_mpl=args.save_mpl)
+                if args.save_rois:
+                    io.save_rois(masks, image_name)
+            logger.info(">>>> completed in %0.3f sec" % (time.time() - tic))
+        else:
+            test_dir = None if len(args.test_dir) == 0 else args.test_dir
+            images, labels, image_names, train_probs = None, None, None, None
+            test_images, test_labels, image_names_test, test_probs = None, None, None, None
+            compute_flows = False
+            if len(args.file_list) > 0:
+                if os.path.exists(args.file_list):
+                    dat = np.load(args.file_list, allow_pickle=True).item()
+                    image_names = dat["train_files"]
+                    image_names_test = dat.get("test_files", None)
+                    train_probs = dat.get("train_probs", None)
+                    test_probs = dat.get("test_probs", None)
+                    compute_flows = dat.get("compute_flows", False)
+                    load_files = False
+                else:
+                    logger.critical(f"ERROR: {args.file_list} does not exist")
+            else:
+                output = io.load_train_test_data(args.dir, test_dir, imf,
+                                                 args.mask_filter,
+                                                 args.look_one_level_down)
+                images, labels, image_names, test_images, test_labels, image_names_test = output
+                load_files = True
+
+            # training with all channels
+            if args.all_channels:
+                img = images[0] if images is not None else io.imread(image_names[0])
+                if img.ndim == 3:
+                    nchan = min(img.shape)
+                elif img.ndim == 2:
+                    nchan = 1
+                channels = None
+            else:
+                nchan = 2
+
+            # model path
+            szmean = args.diam_mean
+            if not os.path.exists(pretrained_model) and model_type is None:
+                if not args.train:
+                    error_message = "ERROR: model path missing or incorrect - cannot train size model"
+                    logger.critical(error_message)
+                    raise ValueError(error_message)
+                pretrained_model = False
+                logger.info(">>>> training from scratch")
+            if args.train:
+                logger.info(
+                    ">>>> during training rescaling images to fixed diameter of %0.1f pixels"
+                    % args.diam_mean)
+
+            # initialize model
+            model = models.CellposeModel(
+                device=device, model_type=model_type, diam_mean=szmean, nchan=nchan,
+                pretrained_model=pretrained_model if model_type is None else None,
+                backbone=backbone)
+
+            # train segmentation model
+            if args.train:
+                cpmodel_path = train.train_seg(
+                    model.net, images, labels, train_files=image_names,
+                    test_data=test_images, test_labels=test_labels,
+                    test_files=image_names_test, train_probs=train_probs,
+                    test_probs=test_probs, compute_flows=compute_flows,
+                    load_files=load_files, normalize=normalize,
+                    channels=channels, channel_axis=args.channel_axis, rgb=(nchan == 3),
+                    learning_rate=args.learning_rate, weight_decay=args.weight_decay,
+                    SGD=args.SGD, n_epochs=args.n_epochs, batch_size=args.batch_size,
+                    min_train_masks=args.min_train_masks,
+                    nimg_per_epoch=args.nimg_per_epoch,
+                    nimg_test_per_epoch=args.nimg_test_per_epoch,
+                    save_path=os.path.realpath(args.dir), save_every=args.save_every,
+                    model_name=args.model_name_out)[0]
+                model.pretrained_model = cpmodel_path
+                logger.info(">>>> model trained and saved to %s" % cpmodel_path)
+
+            # train size model
+            if args.train_size:
+                sz_model = models.SizeModel(cp_model=model, device=device)
+                # data has already been normalized and reshaped
+                sz_model.params = train.train_size(
+                    model.net, model.pretrained_model, images, labels,
+                    train_files=image_names, test_data=test_images,
+                    test_labels=test_labels, test_files=image_names_test,
+                    train_probs=train_probs, test_probs=test_probs,
+                    load_files=load_files, channels=channels,
+                    min_train_masks=args.min_train_masks,
+                    channel_axis=args.channel_axis, rgb=(nchan == 3),
+                    nimg_per_epoch=args.nimg_per_epoch, normalize=normalize,
+                    nimg_test_per_epoch=args.nimg_test_per_epoch,
+                    batch_size=args.batch_size)
+                if test_images is not None:
+                    test_masks = [lbl[0] for lbl in test_labels
+                                 ] if test_labels is not None else test_labels
+                    predicted_diams, diams_style = sz_model.eval(
+                        test_images, channels=channels)
+                    ccs = np.corrcoef(
+                        diams_style,
+                        np.array([utils.diameters(lbl)[0] for lbl in test_masks]))[0, 1]
+                    cc = np.corrcoef(
+                        predicted_diams,
+                        np.array([utils.diameters(lbl)[0] for lbl in test_masks]))[0, 1]
+                    logger.info(
+                        "style test correlation: %0.4f; final test correlation: %0.4f" %
+                        (ccs, cc))
+                    np.save(
+                        os.path.join(
+                            args.test_dir,
+                            "%s_predicted_diams.npy" % os.path.split(cpmodel_path)[1]),
+                        {
+                            "predicted_diams": predicted_diams,
+                            "diams_style": diams_style
+                        })
+
+
+if __name__ == "__main__":
+    main()

cellpose/cli.py

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