Pseudo-color

app.py

@@ -1,7 +1,7 @@
 import torch
 import numpy as np
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageOps
 import os
 import json
 import glob
@@ -17,7 +17,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 import spaces
 from collections import OrderedDict
-
+import tempfile
+import zipfile
+import matplotlib.cm as cm
 
 # --- Imports from both scripts ---
 from diffusers import DDPMScheduler, DDIMScheduler
@@ -54,17 +56,14 @@ LOGO_PATH = "utils/logo2_transparent.png"
 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 = "FluoGen-Group/FluoGen-demo-test-ckpts"
-MODELS_ROOT_DIR = snapshot_download(repo_id=REPO_ID, token=hf_token) #"models_collection"
+MODELS_ROOT_DIR = snapshot_download(repo_id=REPO_ID, token=hf_token) 
 
-# --- Tab 1: Mask-to-Image Config (Formerly Segmentation-to-Image) ---
+# --- Tab 1: Mask-to-Image Config ---
 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_PRETRAINED_MODEL_PATH = f"{MODELS_ROOT_DIR}/stable-diffusion-v1-5"
 T2I_UNET_PATH = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"
@@ -96,22 +95,16 @@ SEG_MODELS = {
 }
 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_CLASS_NAMES = ['Nucleus', 'Endoplasmic Reticulum', 'Giantin', 'GPP130', 'Lysosomes', 'Mitochondria', 'Nucleolus', 'Actin', 'Endosomes', 'Microtubules']
 CLS_EXAMPLE_IMG_DIR = "example_images_cls"
 
+# --- Constants for Visualization ---
+COLOR_MAPS = ["Grayscale", "Green (GFP)", "Red (RFP)", "Blue (DAPI)", "Magenta", "Cyan", "Yellow", "Fire", "Viridis", "Inferno"]
 
 # --- Helper Functions ---
 def sanitize_prompt_for_filename(prompt):
@@ -123,77 +116,104 @@ def min_max_norm(x):
     if max_val - min_val < 1e-8: return np.zeros_like(x)
     return (x - min_val) / (max_val - min_val)
 
+def apply_pseudocolor(image_np, color_name="Grayscale"):
+    """
+    Applies a pseudocolor to a single channel numpy image.
+    image_np: Single channel numpy array (any bit depth).
+    Returns: PIL Image in RGB.
+    """
+    # Normalize to 0-1 for processing
+    norm_img = min_max_norm(np.squeeze(image_np))
+    
+    if color_name == "Grayscale":
+        # Just convert to uint8 L
+        return Image.fromarray((norm_img * 255).astype(np.uint8)).convert("RGB")
+    
+    # Create RGB canvas
+    h, w = norm_img.shape
+    rgb = np.zeros((h, w, 3), dtype=np.float32)
+    
+    if color_name == "Green (GFP)":
+        rgb[..., 1] = norm_img
+    elif color_name == "Red (RFP)":
+        rgb[..., 0] = norm_img
+    elif color_name == "Blue (DAPI)":
+        rgb[..., 2] = norm_img
+    elif color_name == "Magenta":
+        rgb[..., 0] = norm_img
+        rgb[..., 2] = norm_img
+    elif color_name == "Cyan":
+        rgb[..., 1] = norm_img
+        rgb[..., 2] = norm_img
+    elif color_name == "Yellow":
+        rgb[..., 0] = norm_img
+        rgb[..., 1] = norm_img
+    elif color_name in ["Fire", "Viridis", "Inferno"]:
+        # Use matplotlib colormaps
+        cmap_map = {"Fire": "inferno", "Viridis": "viridis", "Inferno": "inferno"} # Fire looks like inferno usually
+        if color_name == "Fire": cmap = cm.get_cmap("gnuplot2") # Better fire approximation
+        else: cmap = cm.get_cmap(cmap_map[color_name])
+        
+        colored = cmap(norm_img) # Returns RGBA 0-1
+        rgb = colored[..., :3] # Drop Alpha
+    
+    return Image.fromarray((rgb * 255).astype(np.uint8))
+
+def save_temp_tiff(image_np, prefix="output"):
+    """Saves numpy array to a temp TIFF file and returns path."""
+    tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".tif", prefix=f"{prefix}_")
+    # Ensure compatible type for Tiff (float32 or uint16 preferred for science)
+    if image_np.dtype == np.float16:
+        save_data = image_np.astype(np.float32)
+    else:
+        save_data = image_np
+    tifffile.imwrite(tfile.name, save_data)
+    return tfile.name
+
 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
 
 def update_sr_prompt(model_name):
-    if model_name == "Checkpoint ER":
-        return "ER of COS-7"
-    if model_name == "Checkpoint Microtubules":
-        return "Microtubules of COS-7"
-    if model_name == "Checkpoint CCPs":
-        return "CCPs of COS-7"
-    elif model_name == "Checkpoint F-actin":
-        return "F-actin of COS-7"
-    return "" # 或者返回一个默认值
+    if model_name == "Checkpoint ER": return "ER of COS-7"
+    if model_name == "Checkpoint Microtubules": return "Microtubules of COS-7"
+    if model_name == "Checkpoint CCPs": return "CCPs of COS-7"
+    elif model_name == "Checkpoint F-actin": return "F-actin of COS-7"
+    return ""
 
 PROMPT_TO_MODEL_MAP = {}
 current_t2i_unet_path = None
 def load_all_prompts():
     global PROMPT_TO_MODEL_MAP
     categories = [
-        {
-            "file": "prompts/basic_prompts.json", 
-            "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"
-        },
-        {
-            "file": "prompts/others_prompts.json", 
-            "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/FULL-checkpoint-275000"
-        },
-        {
-            "file": "prompts/hpa_prompts.json", 
-            "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/HPA-checkpoint-40000"
-        }
+        {"file": "prompts/basic_prompts.json", "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"},
+        {"file": "prompts/others_prompts.json", "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/FULL-checkpoint-275000"},
+        {"file": "prompts/hpa_prompts.json", "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/HPA-checkpoint-40000"}
     ]
-    
     combined_prompts = [] 
     for cat in categories:
-        file_path = cat["file"]
-        model_path = cat["model"]
         try:
-            if os.path.exists(file_path):
-                with open(file_path, "r", encoding="utf-8") as f:
+            if os.path.exists(cat["file"]):
+                with open(cat["file"], "r", encoding="utf-8") as f:
                     data = json.load(f)
                     if isinstance(data, list):
                         combined_prompts.extend(data)
-                        for p in data:
-                            PROMPT_TO_MODEL_MAP[p] = model_path
-                        print(f"✓ Loaded {len(data)} prompts from {file_path}")
-        except Exception as e:
-            print(f"✗ Error loading {file_path}: {e}")
-
-    if not combined_prompts:
-        return ["F-actin of COS-7", "ER of COS-7"]
+                        for p in data: PROMPT_TO_MODEL_MAP[p] = cat["model"]
+                        print(f"✓ Loaded {len(data)} prompts from {cat['file']}")
+        except Exception as e: print(f"✗ Error loading {cat['file']}: {e}")
+    if not combined_prompts: return ["F-actin of COS-7", "ER of COS-7"]
     return combined_prompts
 T2I_PROMPTS = load_all_prompts()
 
@@ -216,7 +236,7 @@ except Exception as 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
+    default_controlnet_path = M2I_CONTROLNET_PATH
     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")
@@ -254,19 +274,8 @@ def swap_t2i_unet(pipe, target_unet_path):
             raise gr.Error(f"Failed to load UNet from {target_unet_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))0
-#         if not os.path.exists(example_filepath):
-#             generated_image.save(example_filepath); print(f"✓ New T2I example saved: {example_filepath}")
-#     return generated_image
 @spaces.GPU(duration=120)
-def generate_t2i(prompt, num_inference_steps):
+def generate_t2i(prompt, num_inference_steps, colormap_choice):
     global t2i_pipe
     if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
     target_model_path = PROMPT_TO_MODEL_MAP.get(prompt)
@@ -274,31 +283,31 @@ def generate_t2i(prompt, num_inference_steps):
         target_model_path = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/FULL-checkpoint-275000"
         print(f"ℹ️ Prompt '{prompt}' not found in predefined list. Using Foundation (Full) model.")
     t2i_pipe = swap_t2i_unet(t2i_pipe, target_model_path)
+    
     print(f"\n🚀 Task started... | Prompt: '{prompt}' | Model: {current_t2i_unet_path}")
     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)
+    
+    # 1. Save Raw Data
+    raw_file_path = save_temp_tiff(image_np, prefix="t2i_raw")
+    
+    # 2. Apply Pseudocolor for Display
+    display_image = apply_pseudocolor(image_np, colormap_choice)
+    
     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
+            display_image.save(example_filepath)
+    
+    # Return Display Image AND Path to Raw File
+    return display_image, raw_file_path
 
 @spaces.GPU(duration=120)
-def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed):
+def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed, colormap_choice):
     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)
@@ -314,30 +323,44 @@ def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, se
     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 = []
+    generated_display_images = []
+    generated_raw_files = []
+    
+    # Create a temp dir for the zip file
+    temp_dir = tempfile.mkdtemp()
+    
     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
+        # Save individual raw file
+        raw_name = f"m2i_sample_{i+1}.tif"
+        raw_path = os.path.join(temp_dir, raw_name)
+        
+        # Ensure correct type saving
+        save_data = output_np.astype(np.float32) if output_np.dtype == np.float16 else output_np
+        tifffile.imwrite(raw_path, save_data)
+        generated_raw_files.append(raw_path)
+        
+        # Create display image
+        pil_image = apply_pseudocolor(output_np, colormap_choice)
+        generated_display_images.append(pil_image)
+        print(f"✓ Generated image {i+1}/{int(num_images)}")
+    
+    # Create ZIP file
+    zip_filename = os.path.join(temp_dir, "raw_output_images.zip")
+    with zipfile.ZipFile(zip_filename, 'w') as zipf:
+        for file in generated_raw_files:
+            zipf.write(file, os.path.basename(file))
+            
+    return input_display_image, generated_display_images, zip_filename
 
 @spaces.GPU(duration=120)
-def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps, seed):
+def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps, seed, colormap_choice):
     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.")
@@ -360,25 +383,19 @@ def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps,
     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
+    
+    # Save Raw
+    raw_file_path = save_temp_tiff(output_np, prefix="sr_raw")
+    # Display Color
+    output_display = apply_pseudocolor(output_np, colormap_choice)
         
-    return input_display_image, numpy_to_pil(output_np)
+    return input_display_image, output_display, raw_file_path
 
 @spaces.GPU(duration=120)
-def run_denoising(noisy_image_np, image_type, steps, seed):
+def run_denoising(noisy_image_np, image_type, steps, seed, colormap_choice):
     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}")
@@ -390,27 +407,22 @@ def run_denoising(noisy_image_np, image_type, steps, seed):
     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
+    
+    # Save Raw
+    raw_file_path = save_temp_tiff(output_np, prefix="dn_raw")
+    # Display Color
+    output_display = apply_pseudocolor(output_np, colormap_choice)
         
-    return numpy_to_pil(noisy_image_np, "L"), numpy_to_pil(output_np)
+    return numpy_to_pil(noisy_image_np, "L"), output_display, raw_file_path
 
 @spaces.GPU(duration=120)
 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.")
-        
+    # Segmentation remains mostly same, usually overlay is RGB anyway
+    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 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)
@@ -418,155 +430,72 @@ def run_segmentation(input_image_np, model_name, diameter, flow_threshold, cellp
         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
-        )
+        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")
+    red_mask_layer[mask_output > 0] = [139, 0, 0]
+    blended_image_np = ((0.6 * original_rgb_np + 0.4 * red_mask_layer).astype(np.uint8))
     
     return numpy_to_pil(input_image_np, "L"), numpy_to_pil(blended_image_np, "RGB")
 
 @spaces.GPU(duration=120)
 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.")
-
+    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.")
-        
+    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.fc = nn.Linear(model.fc.in_features, len(CLS_CLASS_NAMES))
         model.load_state_dict(torch.load(model_path, map_location=DEVICE))
-        model.to(DEVICE)
-        model.eval()
+        model.to(DEVICE).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
-    ])
+    transform_test = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
     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)
+for d in [M2I_EXAMPLE_IMG_DIR, T2I_EXAMPLE_IMG_DIR, SR_EXAMPLE_IMG_DIR, DN_EXAMPLE_IMG_DIR, SEG_EXAMPLE_IMG_DIR, CLS_EXAMPLE_IMG_DIR]:
+    os.makedirs(d, 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))
+        t2i_gallery_examples.append((os.path.join(T2I_EXAMPLE_IMG_DIR, filename), filename_to_prompt_map[filename]))
 
 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')):
+        if f.lower().endswith(('.tif', '.tiff', '.png', '.jpg')):
             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}")
+                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)
+                examples.append((numpy_to_pil(img_np, "L"), filepath))
+            except: pass
     return examples
 
 m2i_gallery_examples = load_image_examples(M2I_EXAMPLE_IMG_DIR)
@@ -576,11 +505,8 @@ 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']
+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():
@@ -590,175 +516,112 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Tabs():
         # --- TAB 1: 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_prompt_input = gr.Dropdown(
-                        choices=T2I_PROMPTS,
-                        value=T2I_PROMPTS[0],
-                        label="Search or Type a Prompt",
-                        filterable=True,
-                        allow_custom_value=True 
-                    )
+                    t2i_prompt_input = gr.Dropdown(choices=T2I_PROMPTS, value=T2I_PROMPTS[0], label="Search or Type a Prompt", filterable=True, allow_custom_value=True)
                     t2i_steps_slider = gr.Slider(minimum=10, maximum=200, step=1, value=50, label="Inference Steps")
+                    # New: Color Selector
+                    t2i_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
                     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")
+                    t2i_generated_output = gr.Image(label="Generated Image (Pseudocolor)", type="pil", interactive=False)
+                    # New: Raw Download Button
+                    t2i_raw_download = gr.DownloadButton(label="Download Raw Output (.tif)", visible=True)
+            t2i_gallery = gr.Gallery(value=t2i_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
 
         # --- TAB 2: 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_prompt_input = gr.Textbox(
-                        label="Prompt", 
-                        value="F-actin of COS-7", # 初始值根据你的默认选择设定
-                        interactive=False 
-                    )
+                    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", value="F-actin of COS-7", interactive=False)
                     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_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
                     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_input_display = gr.Image(label="Input (Avg)", 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")
+                    sr_raw_download = gr.DownloadButton(label="Download Raw Output (.tif)", visible=True)
+            sr_gallery = gr.Gallery(value=sr_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
 
         # --- TAB 3: 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_image_type_selector = gr.Dropdown(choices=list(DN_PROMPT_RULES.keys()), value='MICE', label="Select Image Type")
                     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_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
                     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_original_display = gr.Image(label="Original", 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")
+                    dn_raw_download = gr.DownloadButton(label="Download Raw Output (.tif)", visible=True)
+            dn_gallery = gr.Gallery(value=dn_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
 
         # --- TAB 4: 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_cell_type_input = gr.Textbox(label="Cell Type", placeholder="e.g., CoNSS, HeLa")
+                    m2i_num_images_slider = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of Images")
                     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")
+                    m2i_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
+                    m2i_generate_button = gr.Button("Generate Samples", variant="primary")
                 with gr.Column(scale=2):
                     m2i_output_gallery = gr.Gallery(label="Generated Samples", columns=5, object_fit="contain", height="auto")
+                    m2i_raw_download = gr.DownloadButton(label="Download All Raw Samples (.zip)", visible=True)
                     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")
+            m2i_gallery = gr.Gallery(value=m2i_gallery_examples, label="Examples", 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_input_image = gr.Image(type="numpy", label="Upload Image", image_mode="L")
+                    seg_model_selector = gr.Dropdown(choices=list(SEG_MODELS.keys()), value=list(SEG_MODELS.keys())[0], label="Model")
+                    seg_diameter_input = gr.Number(label="Cell Diameter (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")
+                        seg_original_display = gr.Image(label="Original", type="pil", interactive=False)
+                        seg_output_image = gr.Image(label="Segmented Overlay", type="pil", interactive=False)
+            seg_gallery = gr.Gallery(value=seg_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
             
-        # --- NEW TAB 6: Classification ---
+        # --- 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_input_image = gr.Image(type="numpy", label="Upload Image", image_mode="L")
+                    cls_model_selector = gr.Dropdown(choices=list(CLS_MODEL_PATHS.keys()), value=list(CLS_MODEL_PATHS.keys())[0], label="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")
+                    cls_output_label = gr.Label(label="Results", num_top_classes=len(CLS_CLASS_NAMES))
+            cls_gallery = gr.Gallery(value=cls_gallery_examples, label="Examples", 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_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, m2i_color_input], outputs=[m2i_input_display, m2i_output_gallery, m2i_raw_download])
     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_generate_button.click(fn=generate_t2i, inputs=[t2i_prompt_input, t2i_steps_slider, t2i_color_input], outputs=[t2i_generated_output, t2i_raw_download])
     t2i_gallery.select(fn=select_example_prompt, outputs=t2i_prompt_input)
     
     sr_model_selector.change(fn=update_sr_prompt, inputs=sr_model_selector, outputs=sr_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_generate_button.click(fn=run_super_resolution, inputs=[sr_input_file, sr_model_selector, sr_prompt_input, sr_steps_slider, sr_seed_input, sr_color_input], outputs=[sr_input_display, sr_output_image, sr_raw_download])
     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_generate_button.click(fn=run_denoising, inputs=[dn_input_image, dn_image_type_selector, dn_steps_slider, dn_seed_input, dn_color_input], outputs=[dn_original_display, dn_output_image, dn_raw_download])
     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])
@@ -771,5 +634,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 # --- 4. Launch Application ---
 if __name__ == "__main__":
     print("Interface built. Launching server...")
-    demo.launch()
-
+    demo.launch()