color

app.py

@@ -100,7 +100,7 @@ CLS_CLASS_NAMES = ['Nucleus', 'Endoplasmic Reticulum', 'Giantin', 'GPP130', 'Lys
 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", "Magma", "Plasma"]
+COLOR_MAPS = ["Grayscale", "Green (GFP)", "Red (RFP)", "Blue (DAPI)", "Magenta", "Cyan", "Yellow", "Fire", "Viridis", "Inferno", "Magma", "Plasma", "Red Hot", "Cyan Hot", "Magenta Hot"]
 
 # --- Helper Functions ---
 def sanitize_prompt_for_filename(prompt):
@@ -120,6 +120,22 @@ def generate_colorbar_preview(color_name):
     
     gradient = np.linspace(0, 1, 256).reshape(1, 256)
     rgb = np.zeros((1, 256, 3))
+
+    if "Hot" in color_name:
+        low_half = np.clip(gradient * 2, 0, 1)
+        high_half = np.clip((gradient - 0.5) * 2, 0, 1)
+        if color_name == "Red Hot":
+            rgb[..., 0] = low_half
+            rgb[..., 1] = high_half
+            rgb[..., 2] = high_half
+        elif color_name == "Cyan Hot":
+            rgb[..., 0] = high_half
+            rgb[..., 1] = low_half
+            rgb[..., 2] = low_half
+        elif color_name == "Magenta Hot":
+            rgb[..., 0] = low_half
+            rgb[..., 1] = high_half
+            rgb[..., 2] = low_half
     
     if color_name == "Green (GFP)": rgb[..., 1] = gradient
     elif color_name == "Red (RFP)": rgb[..., 0] = gradient
@@ -155,6 +171,22 @@ def apply_pseudocolor(image_np, color_name="Grayscale"):
     
     h, w = norm_img.shape
     rgb = np.zeros((h, w, 3), dtype=np.float32)
+
+    if "Hot" in color_name:
+        low_half = np.clip(norm_img * 2, 0, 1)
+        high_half = np.clip((norm_img - 0.5) * 2, 0, 1)
+        if color_name == "Red Hot":
+            rgb[..., 0] = low_half
+            rgb[..., 1] = high_half
+            rgb[..., 2] = high_half
+        elif color_name == "Cyan Hot":
+            rgb[..., 0] = high_half
+            rgb[..., 1] = low_half
+            rgb[..., 2] = low_half
+        elif color_name == "Magenta Hot":
+            rgb[..., 0] = low_half
+            rgb[..., 1] = high_half
+            rgb[..., 2] = low_half
     
     if color_name == "Green (GFP)": rgb[..., 1] = norm_img
     elif color_name == "Red (RFP)": rgb[..., 0] = norm_img
@@ -175,6 +207,23 @@ def apply_pseudocolor(image_np, color_name="Grayscale"):
     
     return Image.fromarray((rgb * 255).astype(np.uint8))
 
+def update_sr_settings(model_name):
+    """
+    - Microtubules -> Cyan Hot
+    - F-actin      -> Red Hot
+    - CCPs         -> Green (GFP)
+    - ER           -> Magenta Hot
+    """
+    if model_name == "Checkpoint ER": 
+        return "ER of COS-7", "Magenta Hot"
+    if model_name == "Checkpoint Microtubules": 
+        return "Microtubules of COS-7", "Cyan Hot"
+    if model_name == "Checkpoint CCPs": 
+        return "CCPs of COS-7", "Green (GFP)"
+    elif model_name == "Checkpoint F-actin": 
+        return "F-actin of COS-7", "Red Hot"  
+    return "", "Grayscale"
+
 def save_temp_tiff(image_np, prefix="output"):
     tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".tif", prefix=f"{prefix}_")
     if image_np.dtype == np.float16: save_data = image_np.astype(np.float32)
@@ -598,7 +647,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
             
             sr_gal = gr.Gallery(value=sr_examples, label="Examples", columns=6, height="auto")
             
-            sr_model.change(update_sr_prompt, sr_model, sr_prompt)
+            # sr_model.change(update_sr_prompt, sr_model, sr_prompt)
+            sr_model.change(update_sr_settings, inputs=sr_model, outputs=[sr_prompt, sr_color])
             # Run returns both input/output states and displays
             sr_btn.click(run_super_resolution, [sr_file, sr_model, sr_prompt, sr_steps, sr_seed, sr_color], [sr_in_disp, sr_out_disp, sr_dl, sr_input_state, sr_raw_state, sr_colorbar])
             # Change color updates both displays