slide bar

app.py

@@ -199,12 +199,106 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
     color_idx = mask % num_colors
     return palette_arr[color_idx]
 
+
+def render_seg_overlay(img_np, inst_mask, overlay_alpha):
+    """Render segmentation overlay from cached image/mask."""
+    if img_np is None or inst_mask is None:
+        return None
+
+    overlay = img_np.copy()
+    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))
+
+    for inst_id in np.unique(inst_mask):
+        if inst_id == 0:
+            continue
+        binary_mask = (inst_mask == inst_id).astype(np.uint8)
+        color = get_well_spaced_color(inst_id)
+        overlay[binary_mask == 1] = (1 - alpha) * overlay[binary_mask == 1] + alpha * color
+
+        contours = measure.find_contours(binary_mask, 0.5)
+        for contour in contours:
+            contour = contour.astype(np.int32)
+            valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
+            valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
+            overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]
+
+    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
+    return Image.fromarray(overlay)
+
+
+def render_count_overlay(img_np, density_normalized, overlay_alpha):
+    """Render counting heatmap overlay from cached image/density."""
+    if img_np is None or density_normalized is None:
+        return None
+
+    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))
+    cmap = cm.get_cmap("jet")
+    density_colored = cmap(density_normalized)[:, :, :3]
+
+    overlay = img_np.copy()
+    threshold = 0.01
+    significant_mask = density_normalized > threshold
+    overlay[significant_mask] = (1 - alpha) * overlay[significant_mask] + alpha * density_colored[significant_mask]
+    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
+    return Image.fromarray(overlay)
+
+
+def update_seg_overlay_alpha(overlay_alpha, seg_vis_cache):
+    """Live update segmentation visualization without rerunning inference."""
+    if not seg_vis_cache:
+        return None
+    return render_seg_overlay(seg_vis_cache.get("img_np"), seg_vis_cache.get("inst_mask"), overlay_alpha)
+
+
+def update_count_overlay_alpha(overlay_alpha, count_vis_cache):
+    """Live update counting visualization without rerunning inference."""
+    if not count_vis_cache:
+        return None
+    return render_count_overlay(count_vis_cache.get("img_np"), count_vis_cache.get("density_normalized"), overlay_alpha)
+
+
+def update_tracking_overlay_alpha(overlay_alpha, track_vis_cache):
+    """Regenerate tracking visualization at new opacity using cached outputs."""
+    if not track_vis_cache:
+        return None
+
+    tif_dir = track_vis_cache.get("tif_dir")
+    output_dir = track_vis_cache.get("output_dir")
+    valid_tif_files = track_vis_cache.get("valid_tif_files")
+    if not tif_dir or not output_dir or not valid_tif_files:
+        return None
+
+    try:
+        return create_tracking_visualization(
+            tif_dir=tif_dir,
+            output_dir=output_dir,
+            valid_tif_files=valid_tif_files,
+            overlay_alpha=overlay_alpha
+        )
+    except Exception as e:
+        print(f"⚠️ Failed to update tracking opacity: {e}")
+        return None
+
+
+def cleanup_tracking_cache(track_vis_cache):
+    """Delete cached tracking temp directories from the previous run."""
+    if not track_vis_cache:
+        return
+    for key in ["input_temp_dir", "output_dir"]:
+        path = track_vis_cache.get(key)
+        if path and os.path.isdir(path):
+            try:
+                shutil.rmtree(path)
+            except Exception:
+                pass
+
+
 # @spaces.GPU
 def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
     """Segmentation handler - supports bounding box, returns colorized overlay and original mask path"""
     if annot_value is None or len(annot_value) < 1:
         print("❌ No annotation input")
-        return None, None
+        return None, None, {}
 
     img_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
@@ -223,7 +317,7 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
         print("📏 mask shape:", mask.shape, "dtype:", mask.dtype, "unique:", np.unique(mask))
     except Exception as e:
         print(f"❌ Inference failed: {str(e)}")
-        return None, None
+        return None, None, {}
 
     temp_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".tif")
     mask_img = Image.fromarray(mask.astype(np.uint16))
@@ -235,7 +329,7 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
         print("📷 Image mode:", img.mode, "size:", img.size)
     except Exception as e:
         print(f"❌ Failed to open image: {e}")
-        return None, None
+        return None, None, {}
 
     try:
         img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
@@ -244,7 +338,7 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
             img_np = img_np / 255.0
     except Exception as e:
         print(f"❌ Error in image conversion/resizing: {e}")
-        return None, None
+        return None, None, {}
 
     mask_np = np.array(mask)
     inst_mask = mask_np.astype(np.int32)
@@ -254,35 +348,18 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
 
     if num_instances == 0:
         print("⚠️ No instance found, returning dummy red image")
-        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None
-
-    overlay = img_np.copy()
-    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))
-
-    for inst_id in np.unique(inst_mask):
-        if inst_id == 0:
-            continue
-        binary_mask = (inst_mask == inst_id).astype(np.uint8)
-        color = get_well_spaced_color(inst_id)
-        overlay[binary_mask == 1] = (1 - alpha) * overlay[binary_mask == 1] + alpha * color
-
-        contours = measure.find_contours(binary_mask, 0.5)
-        for contour in contours:
-            contour = contour.astype(np.int32)
-            valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
-            valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
-            overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]  # 黄色轮廓
-
-    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
+        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None, {}
 
-    return Image.fromarray(overlay), temp_mask_file.name
+    overlay_img = render_seg_overlay(img_np, inst_mask, overlay_alpha)
+    seg_vis_cache = {"img_np": img_np, "inst_mask": inst_mask}
+    return overlay_img, temp_mask_file.name, seg_vis_cache
 
 
 # @spaces.GPU
 def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
     """Counting handler - supports bounding box, returns only density map"""
     if annot_value is None or len(annot_value) < 1:
-        return None, "⚠️ Please provide an image."
+        return None, None, "⚠️ Please provide an image.", {}
     
     image_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
@@ -307,7 +384,7 @@ def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
         )
         
         if 'error' in result:
-            return None, f"❌ Counting failed: {result['error']}"
+            return None, None, f"❌ Counting failed: {result['error']}", {}
         
         count = result['count']
         density_map = result['density_map']
@@ -321,7 +398,7 @@ def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
             print("📷 Image mode:", img.mode, "size:", img.size)
         except Exception as e:
             print(f"❌ Failed to open image: {e}")
-            return None, None
+            return None, None, f"❌ Failed to open image: {str(e)}", {}
 
         try:
             img_rgb = img.convert("RGB").resize(density_map.shape[::-1], resample=Image.BILINEAR)
@@ -331,31 +408,14 @@ def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
                 img_np = img_np / 255.0
         except Exception as e:
             print(f"❌ Error in image conversion/resizing: {e}")
-            return None, None
+            return None, None, f"❌ Error in image conversion/resizing: {str(e)}", {}
 
         
         density_normalized = density_map.copy()
         if density_normalized.max() > 0:
             density_normalized = (density_normalized - density_normalized.min()) / (density_normalized.max() - density_normalized.min())
         
-        cmap = cm.get_cmap("jet")
-        alpha = float(np.clip(overlay_alpha, 0.0, 1.0))
-        density_colored = cmap(density_normalized)[:, :, :3]  # RGB only, ignore alpha
-        
-        overlay = img_np.copy()
-        
-        threshold = 0.01  # Only overlay where density > 1% of max
-        significant_mask = density_normalized > threshold
-        
-        overlay[significant_mask] = (1 - alpha) * overlay[significant_mask] + alpha * density_colored[significant_mask]
-        
-        # Clip and convert to uint8
-        overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
-        
-        
-
-
-
+        overlay_img = render_count_overlay(img_np, density_normalized, overlay_alpha)
         result_text = f"✅ Detected {round(count)} objects"
         if use_box_choice == "Yes" and box_array:
             result_text += f"\n📦 Using bounding box: {box_array}"
@@ -363,14 +423,15 @@ def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
 
         print(f"✅ Counting done - Count: {count:.1f}")
 
-        return Image.fromarray(overlay), temp_density_file.name, result_text
+        count_vis_cache = {"img_np": img_np, "density_normalized": density_normalized}
+        return overlay_img, temp_density_file.name, result_text, count_vis_cache
         
         
     except Exception as e:
         print(f"❌ Counting error: {e}")
         import traceback
         traceback.print_exc()
-        return None, f"❌ Counting failed: {str(e)}"
+        return None, None, f"❌ Counting failed: {str(e)}", {}
 
 
 def find_tif_dir(root_dir):
@@ -664,7 +725,7 @@ def create_tracking_visualization(tif_dir, output_dir, valid_tif_files, overlay_
             return valid_tif_files[0]
 
 # @spaces.GPU
-def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay_alpha):
+def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay_alpha, prev_track_vis_cache):
     """
     Tracking handler - processes a ZIP of TIF frames, supports bounding box, returns visualization and results ZIP
     
@@ -678,8 +739,9 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
         Uploaded ZIP file containing TIF sequence
     """
     if zip_file_obj is None:
-        return None, "⚠️  Please upload a ZIP file containing video frames (.zip)", None, None
+        return None, "⚠️  Please upload a ZIP file containing video frames (.zip)", None, None, {}
     
+    cleanup_tracking_cache(prev_track_vis_cache)
     temp_dir = None
     output_temp_dir = None
     
@@ -727,7 +789,7 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
         tif_dir = find_valid_tif_dir(temp_dir)
         
         if tif_dir is None:
-            return None, "❌ Did not find valid TIF directory", None, None
+            return None, "❌ Did not find valid TIF directory", None, None, {}
         
         # Validate TIFF files
         tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 
@@ -736,7 +798,7 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
                           if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]
         
         if len(valid_tif_files) == 0:
-            return None, "❌ Did not find valid TIF files", None, None
+            return None, "❌ Did not find valid TIF files", None, None, {}
 
         print(f"📈 Using {len(valid_tif_files)} TIF files")
 
@@ -757,7 +819,7 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
         )
         
         if 'error' in result:
-            return None, f"❌ Tracking failed: {result['error']}", None, None
+            return None, f"❌ Tracking failed: {result['error']}", None, None, {}
         
         # Create visualization video of tracked objects
         print("\n🎬 Creating tracking visualization...")
@@ -805,18 +867,17 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
 
         print(f"\n✅ Tracking completed")
 
-        # Clean up input temp directory (keep output temp for download)
-        if temp_dir:
-            try:
-                shutil.rmtree(temp_dir)
-                print(f"🗑️  Cleared input temp directory")
-            except:
-                pass
-        
-        return results_zip, result_text, gr.update(visible=True), tracking_video
+        track_vis_cache = {
+            "tif_dir": tif_dir,
+            "valid_tif_files": valid_tif_files,
+            "output_dir": output_temp_dir,
+            "input_temp_dir": temp_dir,
+        }
+
+        return results_zip, result_text, gr.update(visible=True), tracking_video, track_vis_cache
 
     except zipfile.BadZipFile:
-        return None, "❌ Not a valid ZIP file", None, None
+        return None, "❌ Not a valid ZIP file", None, None, {}
     except Exception as e:
         import traceback
         traceback.print_exc()
@@ -828,7 +889,7 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
                     shutil.rmtree(d)
                 except:
                     pass
-        return None, f"❌ Tracking failed: {str(e)}", None, None
+        return None, f"❌ Tracking failed: {str(e)}", None, None, {}
 
 
 
@@ -890,6 +951,9 @@ with gr.Blocks(
     # 全局状态
     current_query_id = gr.State(str(uuid.uuid4()))
     user_uploaded_examples = gr.State(example_images_seg.copy())
+    seg_vis_state = gr.State({})
+    count_vis_state = gr.State({})
+    track_vis_state = gr.State({})
     
     with gr.Tabs():
         # ===== Tab 1: Segmentation =====
@@ -992,14 +1056,19 @@ with gr.Blocks(
             run_seg_btn.click(
                 fn=segment_with_choice,
                 inputs=[use_box_radio, annotator, seg_alpha_slider],
-                outputs=[seg_output, download_mask_btn]
+                outputs=[seg_output, download_mask_btn, seg_vis_state]
+            )
+            seg_alpha_slider.input(
+                fn=update_seg_overlay_alpha,
+                inputs=[seg_alpha_slider, seg_vis_state],
+                outputs=seg_output
             )
 
             # click event for clear button
             clear_btn.click(
-                fn=lambda: None,
+                fn=lambda: (None, {}),
                 inputs=None,
-                outputs=annotator
+                outputs=[annotator, seg_vis_state]
             )
             
             # init Gallery with example images
@@ -1220,14 +1289,19 @@ with gr.Blocks(
             count_btn.click(
                 fn=count_cells_handler,
                 inputs=[count_use_box_radio, count_annotator, count_alpha_slider],
-                outputs=[count_output, download_density_btn, count_status]
+                outputs=[count_output, download_density_btn, count_status, count_vis_state]
+            )
+            count_alpha_slider.input(
+                fn=update_count_overlay_alpha,
+                inputs=[count_alpha_slider, count_vis_state],
+                outputs=count_output
             )
 
             # Clear selection
             clear_btn.click(
-                fn=lambda: None,
+                fn=lambda: (None, {}),
                 inputs=None,
-                outputs=count_annotator
+                outputs=[count_annotator, count_vis_state]
             )
 
             # Submit feedback
@@ -1548,15 +1622,20 @@ with gr.Blocks(
             # Run tracking
             track_btn.click(
                 fn=track_video_handler,
-                inputs=[track_use_box_radio, track_first_frame_annotator, track_zip_upload, track_alpha_slider],
-                outputs=[track_download, track_output, track_download, track_first_frame_preview]
+                inputs=[track_use_box_radio, track_first_frame_annotator, track_zip_upload, track_alpha_slider, track_vis_state],
+                outputs=[track_download, track_output, track_download, track_first_frame_preview, track_vis_state]
+            )
+            track_alpha_slider.change(
+                fn=update_tracking_overlay_alpha,
+                inputs=[track_alpha_slider, track_vis_state],
+                outputs=track_first_frame_preview
             )
 
             # Clear selection
             clear_btn.click(
-                fn=lambda: None,
+                fn=lambda: (None, {}),
                 inputs=None,
-                outputs=track_first_frame_annotator
+                outputs=[track_first_frame_annotator, track_vis_state]
             )
 
             # Submit feedback