updated UI

.gitignore

@@ -5,7 +5,9 @@ dataset/
 
 *_pycache*/
 
+S2FApp/.s2f_env/
 test/
 S2FApp/ckp/spheroid/ckp_spheroid_FN.pth
 S2FApp/ckp/single_cell/ckp_singlecell_GN.pth
 S2FApp/ckp/spheroid/ckp_spheroid_GN.pth
+S2FApp/*.docx

S2FApp/.gitignore

@@ -1,6 +1,7 @@
 __pycache__
 *.py[cod]
-.venv
 venv
+.venv
 .DS_Store
 ckp/*.pth
+*.docx

S2FApp/.streamlit/config.toml

@@ -1,3 +1,13 @@
 [server]
 # Required for file uploads on Hugging Face Spaces (iframe blocks XSRF cookies)
 enableXsrfProtection = false
+
+[theme]
+primaryColor = "#0d9488"
+backgroundColor = "#ffffff"
+secondaryBackgroundColor = "#f8fafc"
+textColor = "#1e293b"
+font = "sans serif"
+
+[logger]
+level = "error"

S2FApp/Dockerfile

@@ -13,15 +13,17 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && rm -rf /var/lib/apt/lists/*
 
 # Copy requirements first for better caching
-COPY requirements.txt .
+COPY requirements.txt requirements-docker.txt ./
 
 # Install Python dependencies - CPU-only PyTorch to fit Space memory limits (avoids OOM)
-RUN pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
-    pip install --no-cache-dir numpy opencv-python scipy scikit-image streamlit matplotlib Pillow plotly huggingface_hub reportlab streamlit-drawable-canvas-fix psutil
+# PyTorch 2.2 + torchvision 0.17 (CPU) - match requirements.txt torch>=2.0
+RUN pip install --no-cache-dir torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir -r requirements-docker.txt
 
 # Copy app code (chown for HF Spaces permissions)
 COPY --chown=user:user app.py predictor.py download_ckp.py ./
 COPY --chown=user:user .streamlit/ .streamlit/
+COPY --chown=user:user static/ static/
 COPY --chown=user:user models/ models/
 COPY --chown=user:user ui/ ui/
 COPY --chown=user:user utils/ utils/

S2FApp/app.py

@@ -22,6 +22,7 @@ from config.constants import (
     MODEL_TYPE_LABELS,
     SAMPLE_EXTENSIONS,
     SAMPLE_THUMBNAIL_LIMIT,
+    THEMES,
 )
 from utils.paths import get_ckp_base, get_ckp_folder, get_sample_folder, list_files_in_folder, model_subfolder
 from utils.segmentation import estimate_cell_mask
@@ -40,7 +41,7 @@ from ui.components import (
 
 CITATION = (
     "Lautaro Baro, Kaveh Shahhosseini, Amparo Andrés-Bordería, Claudio Angione, and Maria Angeles Juanes. "
-    "**\"Shape-to-force (S2F): Predicting Cell Traction Forces from LabelFree Imaging\"**, 2026."
+    "<b>\"Shape-to-force (S2F): Predicting Cell Traction Forces from LabelFree Imaging\"</b>, 2026."
 )
 
 # Measure tool dialog: defined early so it exists before render_result_display uses it
@@ -58,6 +59,7 @@ if HAS_DRAWABLE_CANVAS and ST_DIALOG:
             max_percentile=st.session_state.get("measure_max_percentile", 100),
             clip_min=st.session_state.get("measure_clip_min", 0),
             clip_max=st.session_state.get("measure_clip_max", 1),
+            clip_bounds=st.session_state.get("measure_clip_bounds", False),
         )
         bf_img = st.session_state.get("measure_bf_img")
         original_vals = st.session_state.get("measure_original_vals")
@@ -82,7 +84,8 @@ def _get_measure_dialog_fn():
 
 def _populate_measure_session_state(heatmap, img, pixel_sum, force, key_img, colormap_name,
                                     display_mode, auto_cell_boundary, cell_mask=None,
-                                    min_percentile=0, max_percentile=100, clip_min=0, clip_max=1):
+                                    min_percentile=0, max_percentile=100, clip_min=0, clip_max=1,
+                                    clip_bounds=False):
     """Populate session state for the measure tool. If cell_mask is None and auto_cell_boundary, computes it."""
     if cell_mask is None and auto_cell_boundary:
         cell_mask = estimate_cell_mask(heatmap)
@@ -92,6 +95,7 @@ def _populate_measure_session_state(heatmap, img, pixel_sum, force, key_img, col
     st.session_state["measure_max_percentile"] = max_percentile
     st.session_state["measure_clip_min"] = clip_min
     st.session_state["measure_clip_max"] = clip_max
+    st.session_state["measure_clip_bounds"] = clip_bounds
     st.session_state["measure_bf_img"] = img.copy()
     st.session_state["measure_input_filename"] = key_img or "image"
     st.session_state["measure_original_vals"] = build_original_vals(heatmap, pixel_sum, force)
@@ -101,57 +105,27 @@ def _populate_measure_session_state(heatmap, img, pixel_sum, force, key_img, col
     st.session_state["measure_cell_mask"] = cell_mask if auto_cell_boundary else None
 
 
-st.set_page_config(page_title="Shape2Force (S2F)", page_icon="🦠", layout="centered")
+st.set_page_config(page_title="Shape2Force (S2F)", page_icon="🦠", layout="wide")
 
-# Theme CSS (inject based on sidebar selection)
-def _inject_theme_css(theme):
-    if theme == "Dark":
-        st.markdown("""
-        <style>
-        .stApp { background-color: #0e1117 !important; }
-        .stApp header { background-color: #0e1117 !important; }
-        section[data-testid="stSidebar"] { background-color: #1a1a2e !important; }
-        section[data-testid="stSidebar"] .stMarkdown { color: #fafafa !important; }
-        section[data-testid="stSidebar"] [data-testid="stWidgetLabel"] { color: #e2e8f0 !important; }
-        h1, h2, h3 { color: #fafafa !important; }
-        p { color: #e2e8f0 !important; }
-        .stCaption { color: #94a3b8 !important; }
-        </style>
-        """, unsafe_allow_html=True)
+st.markdown('<link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap" rel="stylesheet">', unsafe_allow_html=True)
 
+_css_path = os.path.join(S2F_ROOT, "static", "s2f_styles.css")
+if os.path.exists(_css_path):
+    with open(_css_path, "r") as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
 
 st.markdown("""
-<style>
-section[data-testid="stSidebar"] { width: 380px !important; }
-@media (max-width: 768px) {
-  section[data-testid="stSidebar"] { width: 100% !important; max-width: 100% !important; }
-}
-section[data-testid="stSidebar"] h2 {
-  font-size: 1.25rem !important;
-  font-weight: 600 !important;
-}
-section[data-testid="stSidebar"] [data-testid="stWidgetLabel"],
-section[data-testid="stSidebar"] [data-testid="stWidgetLabel"] p {
-  font-size: 0.95rem !important;
-  font-weight: 500 !important;
-}
-div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) > div {
-    flex: 1 1 0 !important; min-width: 0 !important;
-}
-div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) button {
-    width: 100% !important; min-width: 100px !important; white-space: nowrap !important;
-}
-div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) > div:nth-child(1) button {
-    background-color: #0d9488 !important; color: white !important; border-color: #0d9488 !important;
-}
-div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) > div:nth-child(1) button:hover {
-    background-color: #0f766e !important; border-color: #0f766e !important; color: white !important;
-}
-</style>
+<div class="s2f-header">
+    <h1>🦠 Shape2Force (S2F)</h1>
+    <p>Predict traction force maps from bright-field microscopy images of cells or spheroids</p>
+</div>
 """, unsafe_allow_html=True)
 
-st.title("🦠 Shape2Force (S2F)")
-st.caption("Predict traction force maps from bright-field microscopy images of cells or spheroids")
+st.markdown(f"""
+<div class="footer-citation">
+    <span>If you find this software useful, please cite: {CITATION}</span>
+</div>
+""", unsafe_allow_html=True)
 
 # Folders
 ckp_base = get_ckp_base(S2F_ROOT)
@@ -173,9 +147,67 @@ def get_cached_sample_thumbnails(model_type, sample_folder, sample_files):
     return cache[cache_key]
 
 
+def _render_sample_selector(model_type, batch_mode):
+    """
+    Render sample image selector (Example mode). Returns (img, imgs_batch, selected_sample, selected_samples).
+    For single mode: img is set, imgs_batch=[]. For batch: img=None, imgs_batch=list of (img, key).
+    """
+    sample_folder = get_sample_folder(S2F_ROOT, model_type)
+    sample_files = list_files_in_folder(sample_folder, SAMPLE_EXTENSIONS)
+    sample_subfolder_name = model_subfolder(model_type)
+    img = None
+    imgs_batch = []
+    selected_sample = None
+    selected_samples = []
+
+    if not sample_files:
+        st.info(f"No example images in samples/{sample_subfolder_name}/. Add images or use Upload.")
+        return img, imgs_batch, selected_sample, selected_samples
+
+    if batch_mode:
+        selected_samples = st.multiselect(
+            f"Select example images (max {BATCH_MAX_IMAGES})",
+            sample_files,
+            default=None,
+            max_selections=BATCH_MAX_IMAGES,
+            key=f"sample_batch_{model_type}",
+        )
+        if selected_samples:
+            for fname in selected_samples[:BATCH_MAX_IMAGES]:
+                sample_path = os.path.join(sample_folder, fname)
+                loaded = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)
+                if loaded is not None:
+                    imgs_batch.append((loaded, fname))
+    else:
+        selected_sample = st.selectbox(
+            f"Select example image (from `samples/{sample_subfolder_name}/`)",
+            sample_files,
+            format_func=lambda x: x,
+            key=f"sample_{model_type}",
+        )
+        if selected_sample:
+            sample_path = os.path.join(sample_folder, selected_sample)
+            img = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)
+
+    thumbnails = get_cached_sample_thumbnails(model_type, sample_folder, sample_files)
+    n_cols = min(5, len(thumbnails))
+    cols = st.columns(n_cols)
+    for i, (fname, sample_img) in enumerate(thumbnails):
+        if sample_img is not None:
+            with cols[i % n_cols]:
+                st.image(sample_img, caption=fname, width=120)
+    return img, imgs_batch, selected_sample, selected_samples
+
+
 # Sidebar
 with st.sidebar:
-    st.header("Settings")
+    st.markdown("""
+    <div class="sidebar-brand">
+        <span class="brand-text">Shape2Force</span>
+    </div>
+    """, unsafe_allow_html=True)
+
+    st.markdown('<div class="sidebar-section"><span class="section-title">Model</span></div>', unsafe_allow_html=True)
 
     model_type = st.radio(
         "Model type",
@@ -193,6 +225,7 @@ with st.sidebar:
         checkpoint = st.selectbox(
             "Checkpoint",
             ckp_files,
+            key=f"checkpoint_{model_type}",
             help=f"Select a .pth file from ckp/{ckp_subfolder_name}/",
         )
     else:
@@ -232,11 +265,7 @@ with st.sidebar:
         except FileNotFoundError:
             st.error("config/substrate_settings.json not found")
 
-    auto_cell_boundary = st.toggle(
-        "Auto boundary",
-        value=False,
-        help="When on: estimate cell region from force map and use it for metrics (red contour). When off: use entire map.",
-    )
+    st.markdown('<div class="sidebar-section"><span class="section-title">Analysis</span></div>', unsafe_allow_html=True)
 
     batch_mode = st.toggle(
         "Batch mode",
@@ -244,45 +273,73 @@ with st.sidebar:
         help=f"Process up to {BATCH_MAX_IMAGES} images at once. Upload multiple files or select multiple examples.",
     )
 
-    display_mode = st.radio(
-        "Heatmap display",
-        ["Default", "Percentile", "Range"],
-        horizontal=True,
-        help="Default: full 0–1 range. Percentile: map a percentile range to improve contrast when few bright pixels dominate. Range: show only values in [min, max]; others hidden (black).",
+    auto_cell_boundary = st.toggle(
+        "Auto boundary",
+        value=False,
+        help="When on: estimate cell region from force map and use it for metrics (red contour). When off: use entire map.",
     )
-    min_percentile, max_percentile = 0, 100
-    clip_min, clip_max = 0.0, 1.0
-    if display_mode == "Percentile":
-        col_pmin, col_pmax = st.columns(2)
-        with col_pmin:
-            min_percentile = st.slider("Min percentile", 0, 100, 2, 1, help="Values below this percentile → black")
-        with col_pmax:
-            max_percentile = st.slider("Max percentile", 0, 100, 99, 1, help="Values above this percentile → white")
-        if min_percentile >= max_percentile:
-            st.warning("Min percentile must be less than max. Using min=0, max=100.")
-            min_percentile, max_percentile = 0, 100
-    elif display_mode == "Range":
-        col_cmin, col_cmax = st.columns(2)
-        with col_cmin:
-            clip_min = st.number_input("Min", value=0.0, min_value=0.0, max_value=1.0, step=0.01, format="%.3f",
-                                       help="Values below this range → hidden (black)")
-        with col_cmax:
-            clip_max = st.number_input("Max", value=1.0, min_value=0.0, max_value=1.0, step=0.01, format="%.3f",
-                                       help="Values above this range → hidden (black)")
-        if clip_min >= clip_max:
-            st.warning("Min must be less than max. Using min=0, max=1.")
-            clip_min, clip_max = 0.0, 1.0
-    colormap_name = st.selectbox(
-        "Heatmap colormap",
-        list(COLORMAPS.keys()),
-        help="Color scheme for the force map. Viridis is often preferred for accessibility.",
+
+    clip_min, clip_max = st.slider(
+        "Force Range",
+        min_value=0.0,
+        max_value=1.0,
+        value=(0.0, 1.0),
+        step=0.01,
+        format="%.2f",
+        help="Min–max range for force values. Values outside are set to 0; inside are rescaled so max shows as red.",
     )
+    if clip_min >= clip_max:
+        clip_min, clip_max = 0.0, 1.0
+    display_mode = "Range" if (clip_min != 0.0 or clip_max != 1.0) else "Default"
+    clip_bounds = False if display_mode == "Range" else True
+    min_percentile, max_percentile = 0, 100
+
+    st.markdown('<div class="sidebar-section"><span class="section-title">Display</span></div>', unsafe_allow_html=True)
+
+    cm_col_lbl, cm_col_sb = st.columns([1, 2])
+    with cm_col_lbl:
+        st.markdown('<p class="selectbox-label">Colormap</p>', unsafe_allow_html=True)
+    with cm_col_sb:
+        colormap_name = st.selectbox(
+            "Colormap",
+            list(COLORMAPS.keys()),
+            key="s2f_colormap",
+            label_visibility="collapsed",
+            help="Color scheme for the force map. Viridis is often preferred for accessibility.",
+        )
+
+    th_col_lbl, th_col_sb = st.columns([1, 2])
+    with th_col_lbl:
+        st.markdown('<p class="selectbox-label">Theme</p>', unsafe_allow_html=True)
+    with th_col_sb:
+        theme_name = st.selectbox(
+            "Theme",
+            list(THEMES.keys()),
+            index=0,
+            key="s2f_theme",
+            label_visibility="collapsed",
+            help="App accent color theme.",
+        )
+
 
-    theme = st.radio("Theme", ["Light", "Dark"], horizontal=True, key="theme_selector")
-    _inject_theme_css(theme)
+# Inject theme CSS (main area so it applies globally)
+primary, primary_dark, primary_darker, primary_rgb = THEMES[theme_name]
+st.markdown(
+    f"""
+    <style>
+    :root {{
+        --s2f-primary: {primary};
+        --s2f-primary-dark: {primary_dark};
+        --s2f-primary-darker: {primary_darker};
+        --s2f-primary-rgb: {primary_rgb};
+    }}
+    </style>
+    """,
+    unsafe_allow_html=True,
+)
 
 # Main area: image input
-img_source = st.radio("Image source", ["Upload", "Example"], horizontal=True, label_visibility="collapsed")
+img_source = st.radio("Image source", ["Upload", "Example"], horizontal=True, label_visibility="collapsed", key="s2f_img_source")
 img = None
 imgs_batch = []  # list of (img, key_img) for batch mode
 uploaded = None
@@ -309,32 +366,7 @@ if batch_mode:
                     imgs_batch.append((decoded, u.name))
                 u.seek(0)
     else:
-        sample_folder = get_sample_folder(S2F_ROOT, model_type)
-        sample_files = list_files_in_folder(sample_folder, SAMPLE_EXTENSIONS)
-        sample_subfolder_name = model_subfolder(model_type)
-        if sample_files:
-            selected_samples = st.multiselect(
-                f"Select example images (max {BATCH_MAX_IMAGES})",
-                sample_files,
-                default=None,
-                max_selections=BATCH_MAX_IMAGES,
-                key=f"sample_batch_{model_type}",
-            )
-            if selected_samples:
-                for fname in selected_samples[:BATCH_MAX_IMAGES]:
-                    sample_path = os.path.join(sample_folder, fname)
-                    loaded = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)
-                    if loaded is not None:
-                        imgs_batch.append((loaded, fname))
-            thumbnails = get_cached_sample_thumbnails(model_type, sample_folder, sample_files)
-            n_cols = min(5, len(thumbnails))
-            cols = st.columns(n_cols)
-            for i, (fname, sample_img) in enumerate(thumbnails):
-                if sample_img is not None:
-                    with cols[i % n_cols]:
-                        st.image(sample_img, caption=fname, width=120)
-        else:
-            st.info(f"No example images in samples/{sample_subfolder_name}/. Add images or use Upload.")
+        img, imgs_batch, selected_sample, selected_samples = _render_sample_selector(model_type, batch_mode=True)
 else:
     # Single image mode
     if img_source == "Upload":
@@ -349,37 +381,20 @@ else:
             img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
             uploaded.seek(0)
     else:
-        sample_folder = get_sample_folder(S2F_ROOT, model_type)
-        sample_files = list_files_in_folder(sample_folder, SAMPLE_EXTENSIONS)
-        sample_subfolder_name = model_subfolder(model_type)
-        if sample_files:
-            selected_sample = st.selectbox(
-                f"Select example image (from `samples/{sample_subfolder_name}/`)",
-                sample_files,
-                format_func=lambda x: x,
-                key=f"sample_{model_type}",
-            )
-            if selected_sample:
-                sample_path = os.path.join(sample_folder, selected_sample)
-                img = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)
-            thumbnails = get_cached_sample_thumbnails(model_type, sample_folder, sample_files)
-            n_cols = min(5, len(thumbnails))
-            cols = st.columns(n_cols)
-            for i, (fname, sample_img) in enumerate(thumbnails):
-                if sample_img is not None:
-                    with cols[i % n_cols]:
-                        st.image(sample_img, caption=fname, width=120)
-        else:
-            st.info(f"No example images in samples/{sample_subfolder_name}/. Add images or use Upload.")
-
-col_btn, col_model, col_path = st.columns([1, 1, 3])
+        img, imgs_batch, selected_sample, selected_samples = _render_sample_selector(model_type, batch_mode=False)
+
+st.markdown("")
+col_btn, col_info = st.columns([1, 3])
 with col_btn:
-    run = st.button("Run prediction", type="primary")
-with col_model:
-    st.markdown(f"<span style='display: inline-flex; align-items: center; height: 38px;'>{MODEL_TYPE_LABELS[model_type]}</span>", unsafe_allow_html=True)
-with col_path:
+    run = st.button("Run prediction", type="primary", use_container_width=True)
+with col_info:
     ckp_path = f"ckp/{ckp_subfolder_name}/{checkpoint}" if checkpoint else f"ckp/{ckp_subfolder_name}/"
-    st.markdown(f"<span style='display: inline-flex; align-items: center; height: 38px;'>Checkpoint: <code>{ckp_path}</code></span>", unsafe_allow_html=True)
+    st.markdown(f"""
+    <div class="run-info">
+        <span class="run-info-tag">{MODEL_TYPE_LABELS[model_type]}</span>
+        <code>{ckp_path}</code>
+    </div>
+    """, unsafe_allow_html=True)
 
 has_image = img is not None
 has_batch = len(imgs_batch) > 0
@@ -400,18 +415,50 @@ just_ran = run and checkpoint and has_image and not batch_mode
 just_ran_batch = run and checkpoint and has_batch and batch_mode
 
 
-def get_or_create_predictor(model_type, checkpoint, ckp_folder):
-    """Cache predictor in session state. Invalidate when model/checkpoint changes."""
-    cache_key = (model_type, checkpoint)
-    if "predictor" not in st.session_state or st.session_state.get("predictor_key") != cache_key:
-        from predictor import S2FPredictor
-        st.session_state["predictor"] = S2FPredictor(
-            model_type=model_type,
-            checkpoint_path=checkpoint,
-            ckp_folder=ckp_folder,
-        )
-        st.session_state["predictor_key"] = cache_key
-    return st.session_state["predictor"]
+@st.cache_resource
+def _load_predictor(model_type, checkpoint, ckp_folder):
+    """Load and cache predictor. Invalidated when model_type or checkpoint changes."""
+    from predictor import S2FPredictor
+    return S2FPredictor(
+        model_type=model_type,
+        checkpoint_path=checkpoint,
+        ckp_folder=ckp_folder,
+    )
+
+
+def _prepare_and_render_cached_result(r, key_img, colormap_name, display_mode, auto_cell_boundary,
+                                     min_percentile, max_percentile, clip_min, clip_max, clip_bounds,
+                                     download_key_suffix="", check_measure_dialog=False):
+    """Prepare display from cached result and render. Used by both just_ran and has_cached paths."""
+    img, heatmap, force, pixel_sum = r["img"], r["heatmap"], r["force"], r["pixel_sum"]
+    display_heatmap = apply_display_scale(
+        heatmap, display_mode,
+        min_percentile=min_percentile,
+        max_percentile=max_percentile,
+        clip_min=clip_min,
+        clip_max=clip_max,
+        clip_bounds=clip_bounds,
+    )
+    cell_mask = estimate_cell_mask(heatmap) if auto_cell_boundary else None
+    _populate_measure_session_state(
+        heatmap, img, pixel_sum, force, key_img, colormap_name,
+        display_mode, auto_cell_boundary, cell_mask=cell_mask,
+        min_percentile=min_percentile, max_percentile=max_percentile,
+        clip_min=clip_min, clip_max=clip_max, clip_bounds=clip_bounds,
+    )
+    if check_measure_dialog and st.session_state.pop("open_measure_dialog", False):
+        measure_region_dialog()
+    st.success("Prediction complete!")
+    render_result_display(
+        img, heatmap, display_heatmap, pixel_sum, force, key_img,
+        download_key_suffix=download_key_suffix,
+        colormap_name=colormap_name,
+        display_mode=display_mode,
+        measure_region_dialog=_get_measure_dialog_fn(),
+        auto_cell_boundary=auto_cell_boundary,
+        cell_mask=cell_mask,
+        clip_min=clip_min, clip_max=clip_max, clip_bounds=clip_bounds,
+    )
 
 
 if just_ran_batch:
@@ -419,27 +466,24 @@ if just_ran_batch:
     st.session_state["batch_results"] = None
     with st.spinner("Loading model and predicting..."):
         try:
-            predictor = get_or_create_predictor(model_type, checkpoint, ckp_folder)
+            predictor = _load_predictor(model_type, checkpoint, ckp_folder)
             sub_val = substrate_val if model_type == "single_cell" and not use_manual else DEFAULT_SUBSTRATE
-            batch_results = []
-            progress_bar = st.progress(0, text="Processing images...")
-            for idx, (img_b, key_b) in enumerate(imgs_batch):
-                progress_bar.progress((idx + 1) / len(imgs_batch), text=f"Processing {key_b}...")
-                heatmap, force, pixel_sum = predictor.predict(
-                    image_array=img_b,
-                    substrate=sub_val,
-                    substrate_config=substrate_config if model_type == "single_cell" else None,
-                )
-                cell_mask = estimate_cell_mask(heatmap) if auto_cell_boundary else None
-                batch_results.append({
+            pred_results = predictor.predict_batch(
+                imgs_batch,
+                substrate=sub_val,
+                substrate_config=substrate_config if model_type == "single_cell" else None,
+            )
+            batch_results = [
+                {
                     "img": img_b.copy(),
                     "heatmap": heatmap.copy(),
                     "force": force,
                     "pixel_sum": pixel_sum,
                     "key_img": key_b,
-                    "cell_mask": cell_mask,
-                })
-            progress_bar.empty()
+                    "cell_mask": estimate_cell_mask(heatmap) if auto_cell_boundary else None,
+                }
+                for (img_b, key_b), (heatmap, force, pixel_sum) in zip(imgs_batch, pred_results)
+            ]
             st.session_state["batch_results"] = batch_results
             st.success(f"Prediction complete for {len(batch_results)} image(s)!")
             render_batch_results(
@@ -451,6 +495,7 @@ if just_ran_batch:
                 clip_min=clip_min,
                 clip_max=clip_max,
                 auto_cell_boundary=auto_cell_boundary,
+                clip_bounds=clip_bounds,
             )
         except Exception as e:
             st.error(f"Prediction failed: {e}")
@@ -467,88 +512,44 @@ elif batch_mode and st.session_state.get("batch_results"):
         clip_min=clip_min,
         clip_max=clip_max,
         auto_cell_boundary=auto_cell_boundary,
+        clip_bounds=clip_bounds,
     )
 
 elif just_ran:
     st.session_state["prediction_result"] = None
     with st.spinner("Loading model and predicting..."):
         try:
-            predictor = get_or_create_predictor(model_type, checkpoint, ckp_folder)
+            predictor = _load_predictor(model_type, checkpoint, ckp_folder)
             sub_val = substrate_val if model_type == "single_cell" and not use_manual else DEFAULT_SUBSTRATE
             heatmap, force, pixel_sum = predictor.predict(
                 image_array=img,
                 substrate=sub_val,
                 substrate_config=substrate_config if model_type == "single_cell" else None,
             )
-
-            st.success("Prediction complete!")
-
-            display_heatmap = apply_display_scale(
-                heatmap, display_mode,
-                min_percentile=min_percentile,
-                max_percentile=max_percentile,
-                clip_min=clip_min,
-                clip_max=clip_max,
-            )
-
             cache_key = (model_type, checkpoint, key_img)
-            st.session_state["prediction_result"] = {
+            r = {
                 "img": img.copy(),
                 "heatmap": heatmap.copy(),
                 "force": force,
                 "pixel_sum": pixel_sum,
                 "cache_key": cache_key,
             }
-            cell_mask = estimate_cell_mask(heatmap) if auto_cell_boundary else None
-            _populate_measure_session_state(
-                heatmap, img, pixel_sum, force, key_img, colormap_name,
-                display_mode, auto_cell_boundary, cell_mask=cell_mask,
-                min_percentile=min_percentile, max_percentile=max_percentile,
-                clip_min=clip_min, clip_max=clip_max,
+            st.session_state["prediction_result"] = r
+            _prepare_and_render_cached_result(
+                r, key_img, colormap_name, display_mode, auto_cell_boundary,
+                min_percentile, max_percentile, clip_min, clip_max, clip_bounds,
+                download_key_suffix="", check_measure_dialog=False,
             )
-            render_result_display(
-                img, heatmap, display_heatmap, pixel_sum, force, key_img,
-                colormap_name=colormap_name,
-                display_mode=display_mode,
-                measure_region_dialog=_get_measure_dialog_fn(),
-                auto_cell_boundary=auto_cell_boundary,
-                cell_mask=cell_mask,
-            )
-
         except Exception as e:
             st.error(f"Prediction failed: {e}")
             st.code(traceback.format_exc())
 
 elif has_cached:
     r = st.session_state["prediction_result"]
-    img, heatmap, force, pixel_sum = r["img"], r["heatmap"], r["force"], r["pixel_sum"]
-    display_heatmap = apply_display_scale(
-        heatmap, display_mode,
-        min_percentile=min_percentile,
-        max_percentile=max_percentile,
-        clip_min=clip_min,
-        clip_max=clip_max,
-    )
-    cell_mask = estimate_cell_mask(heatmap) if auto_cell_boundary else None
-    _populate_measure_session_state(
-        heatmap, img, pixel_sum, force, key_img, colormap_name,
-        display_mode, auto_cell_boundary, cell_mask=cell_mask,
-        min_percentile=min_percentile, max_percentile=max_percentile,
-        clip_min=clip_min, clip_max=clip_max,
-    )
-
-    if st.session_state.pop("open_measure_dialog", False):
-        measure_region_dialog()
-
-    st.success("Prediction complete!")
-    render_result_display(
-        img, heatmap, display_heatmap, pixel_sum, force, key_img,
-        download_key_suffix="_cached",
-        colormap_name=colormap_name,
-        display_mode=display_mode,
-        measure_region_dialog=_get_measure_dialog_fn(),
-        auto_cell_boundary=auto_cell_boundary,
-        cell_mask=cell_mask,
+    _prepare_and_render_cached_result(
+        r, key_img, colormap_name, display_mode, auto_cell_boundary,
+        min_percentile, max_percentile, clip_min, clip_max, clip_bounds,
+        download_key_suffix="_cached", check_measure_dialog=True,
     )
 
 elif run and not checkpoint:
@@ -558,6 +559,5 @@ elif run and not has_image and not has_batch:
 elif run and batch_mode and not has_batch:
     st.warning(f"Please upload or select 1–{BATCH_MAX_IMAGES} images for batch processing.")
 
-st.sidebar.divider()
+st.sidebar.markdown('<div class="sidebar-section"><span class="section-title"></span></div>', unsafe_allow_html=True)
 render_system_status()
-st.sidebar.caption("<br>If you find this software useful, please cite:<br>" + CITATION, unsafe_allow_html=True)

S2FApp/config/constants.py

@@ -25,6 +25,16 @@ TOOL_LABELS = {"polygon": "Polygon", "rect": "Rectangle", "circle": "Circle"}
 # File extensions
 SAMPLE_EXTENSIONS = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
 
+# UI themes: primary, primary-dark, primary-darker, rgb (for rgba)
+THEMES = {
+    "Teal": ("#0d9488", "#0f766e", "#115e59", "13, 148, 136"),
+    "Blue": ("#2563eb", "#1d4ed8", "#1e40af", "37, 99, 235"),
+    "Indigo": ("#6366f1", "#4f46e5", "#4338ca", "99, 102, 241"),
+    "Purple": ("#7c3aed", "#6d28d9", "#5b21b6", "124, 58, 237"),
+    "Amber": ("#f59e0b", "#d97706", "#b45309", "245, 158, 11"),
+    "Emerald": ("#10b981", "#059669", "#047857", "16, 185, 129"),
+}
+
 # Colormaps (OpenCV)
 COLORMAPS = {
     "Jet": cv2.COLORMAP_JET,

S2FApp/predictor.py

@@ -177,3 +177,50 @@ class S2FPredictor:
         pixel_sum = float(np.sum(heatmap))
 
         return heatmap, force, pixel_sum
+
+    def predict_batch(self, images, substrate=None, substrate_config=None):
+        """
+        Run prediction on a batch of images (single forward pass).
+
+        Args:
+            images: List of (img_array, key) or list of img arrays. img_array: (H, W) or (H, W, C).
+            substrate: Substrate name for single-cell mode (same for all images).
+            substrate_config: Optional dict with 'pixelsize' and 'young' (same for all).
+
+        Returns:
+            List of (heatmap, force, pixel_sum) tuples.
+        """
+        imgs = []
+        for item in images:
+            img = item[0] if isinstance(item, tuple) else item
+            img = np.asarray(img, dtype=np.float32)
+            if img.ndim == 3:
+                img = img[:, :, 0] if img.shape[-1] >= 1 else img
+            if img.max() > 1.0:
+                img = img / 255.0
+            img = cv2.resize(img, (MODEL_INPUT_SIZE, MODEL_INPUT_SIZE))
+            imgs.append(img)
+        x = torch.from_numpy(np.stack(imgs)).float().unsqueeze(1).to(self.device)  # [B, 1, H, W]
+
+        if self.model_type == "single_cell" and self.norm_params is not None:
+            sub = substrate if substrate is not None else DEFAULT_SUBSTRATE
+            settings_ch = create_settings_channels_single(
+                sub, self.device, x.shape[2], x.shape[3],
+                config_path=self.config_path, substrate_config=substrate_config
+            )
+            settings_batch = settings_ch.expand(x.shape[0], -1, -1, -1)
+            x = torch.cat([x, settings_batch], dim=1)  # [B, 3, H, W]
+
+        with torch.no_grad():
+            pred = self.generator(x)
+
+        if self._use_tanh_output:
+            pred = (pred + 1.0) / 2.0
+
+        results = []
+        for i in range(pred.shape[0]):
+            heatmap = pred[i, 0].cpu().numpy()
+            force = sum_force_map(pred[i : i + 1]).item()
+            pixel_sum = float(np.sum(heatmap))
+            results.append((heatmap, force, pixel_sum))
+        return results

S2FApp/requirements-docker.txt

@@ -0,0 +1,13 @@
+# Shape2Force App - deps for Docker (torch/torchvision installed separately from CPU wheel)
+numpy>=1.20.0
+opencv-python>=4.5.0
+scipy>=1.7.0
+scikit-image>=0.19.0
+streamlit>=1.28.0
+streamlit-drawable-canvas-fix>=0.9.8
+matplotlib>=3.5.0
+Pillow>=9.0.0
+plotly>=5.14.0
+huggingface_hub>=0.20.0
+reportlab>=4.0.0
+psutil>=5.9.0

S2FApp/static/s2f_styles.css

@@ -0,0 +1,471 @@
+/* === Theme variables (overridden by theme selector) === */
+:root {
+    --s2f-primary: #0d9488;
+    --s2f-primary-dark: #0f766e;
+    --s2f-primary-darker: #115e59;
+    --s2f-primary-rgb: 13, 148, 136;
+}
+
+/* === Typography === */
+html, body, .stApp {
+    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
+}
+
+/* === Header banner === */
+.s2f-header {
+    background: linear-gradient(135deg, var(--s2f-primary) 0%, var(--s2f-primary-dark) 40%, var(--s2f-primary-darker) 100%);
+    padding: 1.1rem 1.5rem 1rem;
+    border-radius: 12px;
+    margin-bottom: 1.2rem;
+    color: white;
+    position: relative;
+    overflow: hidden;
+    box-shadow: 0 4px 20px rgba(var(--s2f-primary-rgb), 0.25);
+}
+.s2f-header::before {
+    content: '';
+    position: absolute;
+    top: -50%;
+    right: -15%;
+    width: 300px;
+    height: 300px;
+    background: radial-gradient(circle, rgba(255,255,255,0.08) 0%, transparent 70%);
+    border-radius: 50%;
+}
+.s2f-header::after {
+    content: '';
+    position: absolute;
+    bottom: -30%;
+    left: 10%;
+    width: 200px;
+    height: 200px;
+    background: radial-gradient(circle, rgba(255,255,255,0.05) 0%, transparent 70%);
+    border-radius: 50%;
+}
+.s2f-header h1 {
+    font-size: 1.85rem !important;
+    font-weight: 700 !important;
+    margin: 0 0 0.35rem !important;
+    color: white !important;
+    letter-spacing: -0.02em;
+    position: relative;
+    z-index: 1;
+}
+.s2f-header p {
+    font-size: 0.95rem !important;
+    color: rgba(255,255,255,0.85) !important;
+    margin: 0 !important;
+    font-weight: 400;
+    position: relative;
+    z-index: 1;
+}
+
+/* === Sidebar === */
+section[data-testid="stSidebar"] {
+    width: 360px !important;
+    background: linear-gradient(180deg, #f8fafc 0%, #f1f5f9 100%) !important;
+    border-right: 1px solid #e2e8f0 !important;
+}
+@media (max-width: 768px) {
+    section[data-testid="stSidebar"] { width: 100% !important; max-width: 100% !important; }
+}
+section[data-testid="stSidebar"] [data-testid="stWidgetLabel"],
+section[data-testid="stSidebar"] [data-testid="stWidgetLabel"] p {
+    font-size: 0.9rem !important;
+    font-weight: 500 !important;
+    color: #334155 !important;
+}
+.sidebar-section {
+    display: flex;
+    align-items: center;
+    gap: 8px;
+    padding: 0.5rem 0 0.25rem;
+    margin-top: 0.6rem;
+    border-bottom: 2px solid #94a3b8;
+    margin-bottom: 0.75rem;
+}
+.sidebar-section .section-title {
+    font-size: 0.78rem;
+    font-weight: 700;
+    color: var(--s2f-primary-dark);
+    text-transform: uppercase;
+    letter-spacing: 0.06em;
+}
+.sidebar-brand {
+    display: flex;
+    align-items: center;
+    gap: 10px;
+    padding-bottom: 0.8rem;
+    margin-bottom: 0.5rem;
+    border-bottom: 1px solid #e2e8f0;
+}
+.sidebar-brand .brand-icon {
+    font-size: 1.6rem;
+    line-height: 1;
+}
+.sidebar-brand .brand-text {
+    font-size: 1.1rem;
+    font-weight: 700;
+    color: var(--s2f-primary-dark);
+    letter-spacing: -0.01em;
+}
+
+/* === Metric cards === */
+[data-testid="stMetric"] {
+    background: linear-gradient(145deg, #ffffff 0%, #f8fafc 100%);
+    border: 1px solid #e2e8f0;
+    border-radius: 12px;
+    padding: 0.85rem 1rem !important;
+    box-shadow: 0 1px 4px rgba(0,0,0,0.06);
+    transition: box-shadow 0.2s ease, transform 0.2s ease;
+}
+[data-testid="stMetric"]:hover {
+    box-shadow: 0 4px 14px rgba(0,0,0,0.1);
+    transform: translateY(-1px);
+}
+[data-testid="stMetric"] label {
+    font-size: 0.75rem !important;
+    font-weight: 600 !important;
+    color: #64748b !important;
+    text-transform: uppercase;
+    letter-spacing: 0.04em;
+}
+[data-testid="stMetric"] [data-testid="stMetricValue"] {
+    font-size: 1.25rem !important;
+    font-weight: 700 !important;
+    color: var(--s2f-primary-dark) !important;
+}
+
+/* === Buttons === */
+.stButton > button[kind="primary"], button[kind="primary"] {
+    background: linear-gradient(135deg, var(--s2f-primary), var(--s2f-primary-dark)) !important;
+    border: none !important;
+    border-radius: 10px !important;
+    font-weight: 600 !important;
+    letter-spacing: 0.02em !important;
+    box-shadow: 0 2px 10px rgba(var(--s2f-primary-rgb), 0.3) !important;
+    transition: all 0.2s ease !important;
+}
+.stButton > button[kind="primary"]:hover, button[kind="primary"]:hover {
+    background: linear-gradient(135deg, var(--s2f-primary-dark), var(--s2f-primary-darker)) !important;
+    box-shadow: 0 4px 18px rgba(var(--s2f-primary-rgb), 0.4) !important;
+    transform: translateY(-1px) !important;
+}
+.stButton > button:not([kind="primary"]) {
+    border-radius: 10px !important;
+    font-weight: 500 !important;
+    border: 1px solid #cbd5e1 !important;
+    transition: all 0.2s ease !important;
+}
+.stButton > button:not([kind="primary"]):hover {
+    border-color: var(--s2f-primary) !important;
+    color: var(--s2f-primary) !important;
+    background: rgba(var(--s2f-primary-rgb), 0.04) !important;
+}
+[data-testid="stDownloadButton"] button {
+    border-radius: 10px !important;
+    font-weight: 500 !important;
+    transition: all 0.2s ease !important;
+}
+[data-testid="stDownloadButton"] button:hover {
+    border-color: var(--s2f-primary) !important;
+    color: var(--s2f-primary) !important;
+}
+
+/* === Action buttons row (measure + downloads) === */
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) {
+    background: #f8fafc;
+    border: 1px solid #e2e8f0;
+    border-radius: 14px;
+    padding: 0.65rem 0.6rem !important;
+    margin-top: 0.5rem;
+    gap: 0.5rem !important;
+}
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) > div {
+    flex: 1 1 0 !important; min-width: 0 !important;
+}
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) button {
+    width: 100% !important;
+    min-width: 0 !important;
+    white-space: nowrap !important;
+    border-radius: 10px !important;
+    font-size: 0.82rem !important;
+    font-weight: 600 !important;
+    padding: 0.55rem 0.8rem !important;
+    letter-spacing: 0.01em !important;
+}
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) > div:nth-child(1) button {
+    background: linear-gradient(135deg, var(--s2f-primary), var(--s2f-primary-dark)) !important;
+    color: white !important;
+    border-color: transparent !important;
+    box-shadow: 0 2px 8px rgba(var(--s2f-primary-rgb), 0.25) !important;
+}
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) > div:nth-child(1) button:hover {
+    background: linear-gradient(135deg, var(--s2f-primary-dark), var(--s2f-primary-darker)) !important;
+    box-shadow: 0 4px 14px rgba(var(--s2f-primary-rgb), 0.35) !important;
+    transform: translateY(-1px) !important;
+}
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) [data-testid="stDownloadButton"] button {
+    background: white !important;
+    border: 1px solid #e2e8f0 !important;
+    color: #334155 !important;
+    box-shadow: 0 1px 2px rgba(0,0,0,0.04) !important;
+}
+div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([data-testid="stButton"]) [data-testid="stDownloadButton"] button:hover {
+    background: rgba(var(--s2f-primary-rgb), 0.06) !important;
+    border-color: var(--s2f-primary) !important;
+    color: var(--s2f-primary-dark) !important;
+    box-shadow: 0 2px 6px rgba(var(--s2f-primary-rgb), 0.12) !important;
+    transform: translateY(-1px) !important;
+}
+
+/* === Expanders === */
+.stExpander {
+    border: 1px solid #e2e8f0 !important;
+    border-radius: 12px !important;
+    overflow: hidden;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.04);
+}
+
+/* === File uploader === */
+[data-testid="stFileUploader"] section {
+    border: 2px dashed #cbd5e1 !important;
+    border-radius: 12px !important;
+    transition: border-color 0.2s ease;
+}
+[data-testid="stFileUploader"] section:hover {
+    border-color: var(--s2f-primary) !important;
+}
+
+/* === Result labels === */
+.result-label {
+    display: flex;
+    align-items: center;
+    gap: 8px;
+    font-size: 0.92rem;
+    font-weight: 600;
+    color: #334155;
+    padding: 0.4rem 0;
+}
+.result-badge {
+    font-size: 0.68rem;
+    font-weight: 700;
+    padding: 2px 8px;
+    border-radius: 4px;
+    letter-spacing: 0.06em;
+}
+.result-badge.input {
+    background: #e2e8f0;
+    color: #475569;
+}
+.result-badge.output {
+    background: rgba(var(--s2f-primary-rgb), 0.15);
+    color: var(--s2f-primary-dark);
+}
+
+/* === Scale visualization === */
+.scale-viz {
+    margin: 0.3rem 0 0.5rem;
+    font-size: 0.78rem;
+    color: #64748b;
+}
+.sv-track {
+    display: flex;
+    align-items: center;
+    gap: 6px;
+}
+.sv-end {
+    font-weight: 600;
+    font-size: 0.72rem;
+    color: #94a3b8;
+    min-width: 14px;
+    text-align: center;
+}
+.sv-bar {
+    flex: 1;
+    height: 10px;
+    background: #e2e8f0;
+    border-radius: 5px;
+    position: relative;
+    overflow: visible;
+}
+.sv-active {
+    position: absolute;
+    top: 0;
+    height: 100%;
+    background: linear-gradient(90deg, var(--s2f-primary), var(--s2f-primary-dark));
+    border-radius: 5px;
+    box-shadow: 0 1px 4px rgba(var(--s2f-primary-rgb), 0.3);
+}
+.sv-lbl {
+    position: absolute;
+    top: 14px;
+    font-size: 0.7rem;
+    font-weight: 700;
+    color: var(--s2f-primary-dark);
+    white-space: nowrap;
+}
+.sv-lbl-l { left: 0; }
+.sv-lbl-r { right: 0; }
+.sv-note {
+    display: flex;
+    align-items: center;
+    gap: 4px;
+    margin-top: 10px;
+    font-size: 0.75rem;
+    color: #64748b;
+}
+.sv-pill {
+    background: rgba(var(--s2f-primary-rgb), 0.15);
+    color: var(--s2f-primary-dark);
+    font-weight: 700;
+    padding: 1px 6px;
+    border-radius: 4px;
+    font-size: 0.72rem;
+}
+
+/* === Run prediction info bar === */
+.run-info {
+    display: flex;
+    align-items: center;
+    gap: 10px;
+    height: 42px;
+    font-size: 0.85rem;
+    color: #64748b;
+}
+.run-info-tag {
+    background: rgba(var(--s2f-primary-rgb), 0.1);
+    color: var(--s2f-primary-dark);
+    font-weight: 600;
+    font-size: 0.78rem;
+    padding: 3px 10px;
+    border-radius: 6px;
+}
+.run-info code {
+    background: #f1f5f9;
+    padding: 2px 8px;
+    border-radius: 4px;
+    font-size: 0.8rem;
+    color: #475569;
+}
+
+/* === Messages === */
+.stSuccess {
+    background: linear-gradient(135deg, #ecfdf5 0%, #d1fae5 100%) !important;
+    border-left: 4px solid #10b981 !important;
+    border-radius: 8px !important;
+}
+.stWarning { border-radius: 8px !important; }
+.stInfo { border-radius: 8px !important; }
+
+/* === Selectbox, toggle === */
+[data-testid="stSelectbox"] > div > div { border-radius: 8px !important; }
+.selectbox-label {
+    margin: 0;
+    padding-top: 0.4rem;
+    font-size: 0.9rem;
+    font-weight: 500;
+    color: #334155;
+    line-height: 1.2;
+}
+
+/* === Dataframe === */
+[data-testid="stDataFrame"] {
+    border-radius: 10px !important;
+    overflow: hidden;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.06);
+}
+
+/* === Batch colorbar (minimal, no box) === */
+.colorbar-table-header {
+    width: 100%;
+    margin-bottom: 0.5rem;
+    padding: 0;
+    background: transparent;
+    border: none;
+    box-shadow: none;
+}
+.colorbar-ticks {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    margin-bottom: 5px;
+    padding: 0 1px;
+    font-size: 0.65rem;
+    font-weight: 600;
+    color: #64748b;
+    letter-spacing: 0.03em;
+}
+.colorbar-ticks .cb-tick {
+    font-variant-numeric: tabular-nums;
+}
+.colorbar-bar {
+    width: 100%;
+    height: 6px;
+    background-size: 100% 100%;
+    background-repeat: no-repeat;
+    background-position: center;
+    border-radius: 3px;
+    box-shadow: inset 0 1px 1px rgba(0,0,0,0.05);
+}
+/* === Divider === */
+hr { border-color: #cbd5e1 !important; opacity: 0.7; }
+
+/* === Plotly chart === */
+.stPlotlyChart { border-radius: 12px; overflow: hidden; }
+
+/* === System status === */
+.system-status {
+    font-size: 0.78rem;
+    margin-top: 0.5rem;
+    padding: 8px 12px;
+    border-radius: 8px;
+    border: 1px solid rgba(148, 163, 184, 0.25);
+    background: rgba(148, 163, 184, 0.08);
+    color: inherit;
+    display: flex;
+    align-items: center;
+    gap: 6px;
+}
+.system-status .status-dot {
+    width: 6px;
+    height: 6px;
+    border-radius: 50%;
+    background: #10b981;
+    display: inline-block;
+    flex-shrink: 0;
+}
+
+/* === Footer citation === */
+.footer-citation {
+    position: fixed;
+    bottom: 0;
+    left: 360px;
+    right: 0;
+    z-index: 999;
+    padding: 0.45rem 1rem;
+    background: #f1f5f9;
+    border-top: 1px solid #e2e8f0;
+    font-size: 0.7rem;
+    color: #64748b;
+    text-align: center;
+    line-height: 1.4;
+}
+.block-container {
+    padding-bottom: 2.5rem !important;
+    max-width: 1050px !important;
+}
+section[data-testid="stSidebar"] > div:first-child {
+    padding-top: 1rem !important;
+}
+
+/* === Responsive === */
+@media (max-width: 768px) {
+    .s2f-header {
+        padding: 1.5rem;
+        border-radius: 12px;
+    }
+    .s2f-header h1 {
+        font-size: 1.4rem !important;
+    }
+}

S2FApp/ui/components.py

@@ -1,837 +1,29 @@
-"""UI components for S2F App."""
-import csv
-import html
-import io
-import os
-import zipfile
-
-import cv2
-import numpy as np
+"""UI components for S2F App. Re-exports from submodules for backward compatibility."""
 import streamlit as st
-from PIL import Image
-import plotly.graph_objects as go
-from plotly.subplots import make_subplots
-
-from config.constants import (
-    CANVAS_SIZE,
-    COLORMAPS,
-    DRAW_TOOLS,
-    TOOL_LABELS,
-)
-from utils.display import apply_display_scale, cv_colormap_to_plotly_colorscale
-from utils.report import (
-    heatmap_to_rgb,
-    heatmap_to_rgb_with_contour,
-    heatmap_to_png_bytes,
-    create_pdf_report,
-    create_measure_pdf_report,
-)
-from utils.segmentation import estimate_cell_mask
 
-try:
-    from streamlit_drawable_canvas import st_canvas
-    HAS_DRAWABLE_CANVAS = True
-except (ImportError, AttributeError):
-    HAS_DRAWABLE_CANVAS = False
-
-try:
-    import psutil
-    HAS_PSUTIL = True
-except ImportError:
-    HAS_PSUTIL = False
-
-# Resolve st.dialog early to fix ordering bug (used in _render_result_display)
+# Resolve st.dialog early to fix ordering bug (used in measure dialog)
 ST_DIALOG = getattr(st, "dialog", None) or getattr(st, "experimental_dialog", None)
 
+from ui.system_status import render_system_status
+from ui.result_display import render_batch_results, render_result_display
+from ui.measure_tool import (
+    build_original_vals,
+    build_cell_vals,
+    render_region_canvas,
+    parse_canvas_shapes_to_masks,
+    compute_region_metrics,
+    HAS_DRAWABLE_CANVAS,
+)
 
-def _get_container_memory():
-    """
-    Read memory from cgroups when running in a container (Docker, HF Spaces).
-    psutil reports host memory in containers, which can be misleading (e.g. 128 GB vs 16 GB limit).
-    Returns (used_bytes, total_bytes) or None to fall back to psutil.
-    """
-    try:
-        # cgroup v2 (modern Docker, HF Spaces)
-        for base in ("/sys/fs/cgroup", "/sys/fs/cgroup/self"):
-            try:
-                with open(f"{base}/memory.max", "r") as f:
-                    max_val = f.read().strip()
-                if max_val == "max":
-                    return None  # No limit, use psutil
-                total = int(max_val)
-                with open(f"{base}/memory.current", "r") as f:
-                    used = int(f.read().strip())
-                return (used, total)
-            except (FileNotFoundError, ValueError):
-                continue
-        # cgroup v1
-        try:
-            with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as f:
-                total = int(f.read().strip())
-            with open("/sys/fs/cgroup/memory/memory.usage_in_bytes", "r") as f:
-                used = int(f.read().strip())
-            if total > 2**50:  # Often 9223372036854771712 when unlimited
-                return None
-            return (used, total)
-        except (FileNotFoundError, ValueError):
-            pass
-    except Exception:
-        pass
-    return None
-
-
-def render_system_status():
-    """Render a small live CPU/memory status panel in the sidebar."""
-    if not HAS_PSUTIL:
-        return
-    try:
-        cpu = psutil.cpu_percent(interval=0.1)
-        # Prefer cgroup memory in containers (Docker, HF Spaces); psutil shows host memory
-        container_mem = _get_container_memory()
-        if container_mem is not None:
-            used_bytes, total_bytes = container_mem
-            mem_used_gb = used_bytes / (1024**3)
-            mem_total_gb = total_bytes / (1024**3)
-            mem_pct = 100 * used_bytes / total_bytes if total_bytes > 0 else 0
-        else:
-            mem = psutil.virtual_memory()
-            mem_used_gb = mem.used / (1024**3)
-            mem_total_gb = mem.total / (1024**3)
-            mem_pct = mem.percent
-        st.sidebar.markdown(
-            f"""
-            <div style="
-                font-size: 0.8rem; margin-top: 0.5rem; padding: 6px 10px;
-                border-radius: 6px;
-                border: 1px solid rgba(148, 163, 184, 0.3);
-                background: rgba(148, 163, 184, 0.1);
-                color: inherit;
-            ">
-                <strong>System</strong> CPU {cpu:.0f}% · Mem {mem_pct:.0f}% ({mem_used_gb:.1f}/{mem_total_gb:.1f} GB)
-            </div>
-            """,
-            unsafe_allow_html=True,
-        )
-    except Exception:
-        pass
-
-
-def render_batch_results(batch_results, colormap_name="Jet", display_mode="Default",
-                        min_percentile=0, max_percentile=100, clip_min=0, clip_max=1,
-                        auto_cell_boundary=False):
-    """
-    Render batch prediction results: summary table, bright-field row, heatmap row, and bulk download.
-    batch_results: list of dicts with img, heatmap, force, pixel_sum, key_img, cell_mask.
-    cell_mask is computed on-the-fly when auto_cell_boundary is True and not stored.
-    """
-    if not batch_results:
-        return
-    st.markdown("### Batch results")
-    # Resolve cell_mask for each result (compute if needed when auto_cell_boundary toggled on)
-    for r in batch_results:
-        if auto_cell_boundary and (r.get("cell_mask") is None or not np.any(r.get("cell_mask", 0) > 0)):
-            r["_cell_mask"] = estimate_cell_mask(r["heatmap"])
-        else:
-            r["_cell_mask"] = r.get("cell_mask") if auto_cell_boundary else None
-    # Build table rows - consistent column names for both modes
-    headers = ["Image", "Force", "Sum", "Max", "Mean"]
-    rows = []
-    csv_rows = [["image"] + headers[1:]]
-    for r in batch_results:
-        heatmap = r["heatmap"]
-        cell_mask = r.get("_cell_mask")
-        key = r["key_img"] or "image"
-        if auto_cell_boundary and cell_mask is not None and np.any(cell_mask > 0):
-            vals = heatmap[cell_mask > 0]
-            cell_pixel_sum = float(np.sum(vals))
-            cell_force = cell_pixel_sum * (r["force"] / r["pixel_sum"]) if r["pixel_sum"] > 0 else cell_pixel_sum
-            cell_mean = cell_pixel_sum / np.sum(cell_mask) if np.sum(cell_mask) > 0 else 0
-            row = [key, f"{cell_force:.2f}", f"{cell_pixel_sum:.2f}",
-                   f"{np.max(heatmap):.4f}", f"{cell_mean:.4f}"]
-        else:
-            row = [key, f"{r['force']:.2f}", f"{r['pixel_sum']:.2f}",
-                   f"{np.max(heatmap):.4f}", f"{np.mean(heatmap):.4f}"]
-        rows.append(row)
-        csv_rows.append([os.path.splitext(key)[0]] + row[1:])
-    # Bright-field row
-    st.markdown("**Input: Bright-field images**")
-    n_cols = min(5, len(batch_results))
-    bf_cols = st.columns(n_cols)
-    for i, r in enumerate(batch_results):
-        img = r["img"]
-        if img.ndim == 2:
-            img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
-        else:
-            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        with bf_cols[i % n_cols]:
-            st.image(img_rgb, caption=r["key_img"], use_container_width=True)
-    # Heatmap row
-    st.markdown("**Output: Predicted force maps**")
-    hm_cols = st.columns(n_cols)
-    for i, r in enumerate(batch_results):
-        display_heatmap = apply_display_scale(
-            r["heatmap"], display_mode,
-            min_percentile=min_percentile, max_percentile=max_percentile,
-            clip_min=clip_min, clip_max=clip_max,
-        )
-        hm_rgb = heatmap_to_rgb_with_contour(
-            display_heatmap, colormap_name,
-            r.get("_cell_mask") if auto_cell_boundary else None,
-        )
-        with hm_cols[i % n_cols]:
-            st.image(hm_rgb, caption=r["key_img"], use_container_width=True)
-    # Table
-    st.dataframe(
-        {h: [r[i] for r in rows] for i, h in enumerate(headers)},
-        use_container_width=True,
-        hide_index=True,
-    )
-    # Histograms in accordion (one per row for visibility)
-    with st.expander("Force distribution (histograms)", expanded=False):
-        for i, r in enumerate(batch_results):
-            heatmap = r["heatmap"]
-            cell_mask = r.get("_cell_mask")
-            vals = heatmap[cell_mask > 0] if (cell_mask is not None and np.any(cell_mask > 0) and auto_cell_boundary) else heatmap.flatten()
-            vals = vals[vals > 0] if np.any(vals > 0) else vals
-            st.markdown(f"**{r['key_img']}**")
-            if len(vals) > 0:
-                fig = go.Figure(data=[go.Histogram(x=vals, nbinsx=50, marker_color="#0d9488")])
-                fig.update_layout(
-                    height=220, margin=dict(l=40, r=20, t=10, b=40),
-                    xaxis_title="Force value", yaxis_title="Count",
-                    showlegend=False,
-                )
-                st.plotly_chart(fig, use_container_width=True, config={"displayModeBar": False})
-            else:
-                st.caption("No data")
-            if i < len(batch_results) - 1:
-                st.divider()
-    # Bulk downloads: CSV and heatmaps (zip)
-    buf_csv = io.StringIO()
-    csv.writer(buf_csv).writerows(csv_rows)
-    zip_buf = io.BytesIO()
-    with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf:
-        for r in batch_results:
-            display_heatmap = apply_display_scale(
-                r["heatmap"], display_mode,
-                min_percentile=min_percentile, max_percentile=max_percentile,
-                clip_min=clip_min, clip_max=clip_max,
-            )
-            hm_bytes = heatmap_to_png_bytes(
-                display_heatmap, colormap_name,
-                r.get("_cell_mask") if auto_cell_boundary else None,
-            )
-            base = os.path.splitext(r["key_img"] or "image")[0]
-            zf.writestr(f"{base}_heatmap.png", hm_bytes.getvalue())
-    zip_buf.seek(0)
-    dl_col1, dl_col2 = st.columns(2)
-    with dl_col1:
-        st.download_button(
-            "Download all as CSV",
-            data=buf_csv.getvalue(),
-            file_name="s2f_batch_results.csv",
-            mime="text/csv",
-            key="download_batch_csv",
-            icon=":material/download:",
-        )
-    with dl_col2:
-        st.download_button(
-            "Download all heatmaps",
-            data=zip_buf.getvalue(),
-            file_name="s2f_batch_heatmaps.zip",
-            mime="application/zip",
-            key="download_batch_heatmaps",
-            icon=":material/image:",
-        )
-
-
-# Distinct colors for each region (RGB - heatmap_rgb is RGB)
-_REGION_COLORS = [
-    (255, 102, 0),   # orange
-    (255, 165, 0),   # orange-red
-    (255, 255, 0),   # yellow
-    (255, 0, 255),   # magenta
-    (0, 255, 127),   # spring green
-    (0, 128, 255),   # blue
-    (203, 192, 255), # lavender
-    (255, 215, 0),   # gold
+__all__ = [
+    "ST_DIALOG",
+    "HAS_DRAWABLE_CANVAS",
+    "render_system_status",
+    "render_batch_results",
+    "render_result_display",
+    "build_original_vals",
+    "build_cell_vals",
+    "render_region_canvas",
+    "parse_canvas_shapes_to_masks",
+    "compute_region_metrics",
 ]
-
-
-def _draw_region_overlay(annotated, mask, color, fill_alpha=0.3, stroke_width=2):
-    """Draw single region overlay on annotated heatmap (fill + alpha blend + contour). Modifies annotated in place."""
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    overlay = annotated.copy()
-    cv2.fillPoly(overlay, contours, color)
-    mask_3d = np.stack([mask] * 3, axis=-1).astype(bool)
-    annotated[mask_3d] = (
-        (1 - fill_alpha) * annotated[mask_3d].astype(np.float32)
-        + fill_alpha * overlay[mask_3d].astype(np.float32)
-    ).astype(np.uint8)
-    cv2.drawContours(annotated, contours, -1, color, stroke_width)
-
-
-def make_annotated_heatmap(heatmap_rgb, mask, fill_alpha=0.3, stroke_color=(255, 102, 0), stroke_width=2):
-    """Composite heatmap with drawn region overlay."""
-    annotated = heatmap_rgb.copy()
-    _draw_region_overlay(annotated, mask, stroke_color, fill_alpha, stroke_width)
-    return annotated
-
-
-def make_annotated_heatmap_multi_regions(heatmap_rgb, masks, labels, cell_mask=None, fill_alpha=0.3):
-    """Draw each region separately with distinct color and label (R1, R2, ...). No merging."""
-    annotated = heatmap_rgb.copy()
-    if cell_mask is not None and np.any(cell_mask > 0):
-        contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cv2.drawContours(annotated, contours, -1, (255, 0, 0), 2)
-    for i, mask in enumerate(masks):
-        color = _REGION_COLORS[i % len(_REGION_COLORS)]
-        _draw_region_overlay(annotated, mask, color, fill_alpha, stroke_width=2)
-        # Label at centroid
-        M = cv2.moments(mask)
-        if M["m00"] > 0:
-            cx = int(M["m10"] / M["m00"])
-            cy = int(M["m01"] / M["m00"])
-            label = labels[i] if i < len(labels) else f"R{i + 1}"
-            cv2.putText(
-                annotated, label, (cx - 12, cy + 5),
-                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA
-            )
-            cv2.putText(
-                annotated, label, (cx - 12, cy + 5),
-                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1, cv2.LINE_AA
-            )
-    return annotated
-
-
-def _obj_to_pts(obj, scale_x, scale_y, heatmap_w, heatmap_h):
-    """Convert a single canvas object to polygon points in heatmap coords. Returns None if invalid."""
-    obj_type = obj.get("type", "")
-    pts = []
-    if obj_type == "rect":
-        left = obj.get("left", 0)
-        top = obj.get("top", 0)
-        w = obj.get("width", 0)
-        h = obj.get("height", 0)
-        pts = np.array([
-            [left, top], [left + w, top], [left + w, top + h], [left, top + h]
-        ], dtype=np.float32)
-    elif obj_type == "circle" or obj_type == "ellipse":
-        left = obj.get("left", 0)
-        top = obj.get("top", 0)
-        width = obj.get("width", 0)
-        height = obj.get("height", 0)
-        radius = obj.get("radius", 0)
-        angle_deg = obj.get("angle", 0)
-        if radius > 0:
-            rx = ry = radius
-            angle_rad = np.deg2rad(angle_deg)
-            cx = left + radius * np.cos(angle_rad)
-            cy = top + radius * np.sin(angle_rad)
-        else:
-            rx = width / 2 if width > 0 else 0
-            ry = height / 2 if height > 0 else 0
-            if rx <= 0 or ry <= 0:
-                return None
-            cx = left + rx
-            cy = top + ry
-        if rx <= 0 or ry <= 0:
-            return None
-        n = 32
-        angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
-        pts = np.column_stack([cx + rx * np.cos(angles), cy + ry * np.sin(angles)]).astype(np.float32)
-    elif obj_type == "path":
-        path = obj.get("path", [])
-        for cmd in path:
-            if isinstance(cmd, (list, tuple)) and len(cmd) >= 3:
-                if cmd[0] in ("M", "L"):
-                    pts.append([float(cmd[1]), float(cmd[2])])
-                elif cmd[0] == "Q" and len(cmd) >= 5:
-                    pts.append([float(cmd[3]), float(cmd[4])])
-                elif cmd[0] == "C" and len(cmd) >= 7:
-                    pts.append([float(cmd[5]), float(cmd[6])])
-        if len(pts) < 3:
-            return None
-        pts = np.array(pts, dtype=np.float32)
-    else:
-        return None
-    pts[:, 0] *= scale_x
-    pts[:, 1] *= scale_y
-    pts = np.clip(pts, 0, [heatmap_w - 1, heatmap_h - 1]).astype(np.int32)
-    return pts
-
-
-def parse_canvas_shapes_to_masks(json_data, canvas_h, canvas_w, heatmap_h, heatmap_w):
-    """Parse drawn shapes and return a list of individual masks (one per shape)."""
-    if not json_data or "objects" not in json_data or not json_data["objects"]:
-        return []
-    scale_x = heatmap_w / canvas_w
-    scale_y = heatmap_h / canvas_h
-    masks = []
-    for obj in json_data["objects"]:
-        pts = _obj_to_pts(obj, scale_x, scale_y, heatmap_w, heatmap_h)
-        if pts is None:
-            continue
-        mask = np.zeros((heatmap_h, heatmap_w), dtype=np.uint8)
-        cv2.fillPoly(mask, [pts], 1)
-        masks.append(mask)
-    return masks
-
-
-def build_original_vals(raw_heatmap, pixel_sum, force):
-    """Build original_vals dict for measure tool (full map)."""
-    return {
-        "pixel_sum": pixel_sum,
-        "force": force,
-        "max": float(np.max(raw_heatmap)),
-        "mean": float(np.mean(raw_heatmap)),
-    }
-
-
-def build_cell_vals(raw_heatmap, cell_mask, pixel_sum, force):
-    """Build cell_vals dict for measure tool (estimated cell area). Returns None if invalid."""
-    cell_pixel_sum, cell_force, cell_mean = _compute_cell_metrics(raw_heatmap, cell_mask, pixel_sum, force)
-    if cell_pixel_sum is None:
-        return None
-    region_values = raw_heatmap * cell_mask
-    region_nonzero = region_values[cell_mask > 0]
-    cell_max = float(np.max(region_nonzero)) if len(region_nonzero) > 0 else 0
-    return {
-        "pixel_sum": cell_pixel_sum,
-        "force": cell_force,
-        "max": cell_max,
-        "mean": cell_mean,
-    }
-
-
-def compute_region_metrics(raw_heatmap, mask, original_vals=None):
-    """Compute region metrics from mask."""
-    area_px = int(np.sum(mask))
-    region_values = raw_heatmap * mask
-    region_nonzero = region_values[mask > 0]
-    force_sum = float(np.sum(region_values))
-    density = force_sum / area_px if area_px > 0 else 0
-    region_max = float(np.max(region_nonzero)) if len(region_nonzero) > 0 else 0
-    region_mean = float(np.mean(region_nonzero)) if len(region_nonzero) > 0 else 0
-    region_force_scaled = (
-        force_sum * (original_vals["force"] / original_vals["pixel_sum"])
-        if original_vals and original_vals.get("pixel_sum", 0) > 0
-        else force_sum
-    )
-    return {
-        "area_px": area_px,
-        "force_sum": force_sum,
-        "density": density,
-        "max": region_max,
-        "mean": region_mean,
-        "force_scaled": region_force_scaled,
-    }
-
-
-def render_region_metrics_and_downloads(metrics_list, masks, heatmap_rgb, input_filename, key_suffix, has_original_vals,
-                                        first_region_label=None, bf_img=None, cell_mask=None, colormap_name="Jet"):
-    """
-    Render per-shape metrics table and download buttons.
-    first_region_label: custom label for first row (e.g. 'Auto boundary').
-    masks: list of region masks (user-drawn only; used for labeled heatmap with R1, R2...).
-    """
-    base_name = os.path.splitext(input_filename or "image")[0]
-    st.markdown("**Regions (each selection = one row)**")
-    if has_original_vals:
-        headers = ["Region", "Area", "F.sum", "Force", "Max", "Mean"]
-        csv_rows = [["image", "region"] + headers[1:]]
-    else:
-        headers = ["Region", "Area (px²)", "Force sum", "Mean"]
-        csv_rows = [["image", "region", "Area", "Force sum", "Mean"]]
-    table_rows = [headers]
-    for i, metrics in enumerate(metrics_list, 1):
-        region_label = first_region_label if (i == 1 and first_region_label) else f"Region {i - (1 if first_region_label else 0)}"
-        if has_original_vals:
-            row = [region_label, str(metrics["area_px"]), f"{metrics['force_sum']:.3f}", f"{metrics['force_scaled']:.1f}",
-                   f"{metrics['max']:.3f}", f"{metrics['mean']:.4f}"]
-            csv_rows.append([base_name, region_label, metrics["area_px"], f"{metrics['force_sum']:.3f}",
-                            f"{metrics['force_scaled']:.1f}", f"{metrics['max']:.3f}", f"{metrics['mean']:.4f}"])
-        else:
-            row = [region_label, str(metrics["area_px"]), f"{metrics['force_sum']:.4f}", f"{metrics['mean']:.6f}"]
-            csv_rows.append([base_name, region_label, metrics["area_px"], f"{metrics['force_sum']:.4f}",
-                            f"{metrics['mean']:.6f}"])
-        table_rows.append(row)
-    # Render as HTML table to avoid Streamlit's default row/column indices
-    header = table_rows[0]
-    body = table_rows[1:]
-    th_cells = "".join(
-        f'<th style="border: 1px solid #ddd; padding: 8px; text-align: left;">{html.escape(str(h))}</th>'
-        for h in header
-    )
-    rows_html = [
-        "<tr>"
-        + "".join(
-            f'<td style="border: 1px solid #ddd; padding: 8px;">{html.escape(str(c))}</td>'
-            for c in row
-        )
-        + "</tr>"
-        for row in body
-    ]
-    table_html = (
-        f'<table style="border-collapse: collapse; width: 100%;">'
-        f"<thead><tr>{th_cells}</tr></thead>"
-        f"<tbody>{''.join(rows_html)}</tbody></table>"
-    )
-    st.markdown(table_html, unsafe_allow_html=True)
-    buf_csv = io.StringIO()
-    csv.writer(buf_csv).writerows(csv_rows)
-    # Annotated heatmap: each region separate with R1, R2 labels (no merging)
-    # heatmap_rgb already has cell contour if applicable
-    region_labels = [f"R{i + 1}" for i in range(len(masks))]
-    heatmap_labeled = make_annotated_heatmap_multi_regions(heatmap_rgb.copy(), masks, region_labels, cell_mask=None)
-    buf_img = io.BytesIO()
-    Image.fromarray(heatmap_labeled).save(buf_img, format="PNG")
-    buf_img.seek(0)
-    # PDF report (requires bf_img)
-    pdf_bytes = None
-    if bf_img is not None:
-        pdf_bytes = create_measure_pdf_report(bf_img, heatmap_labeled, table_rows, base_name)
-    n_cols = 3 if pdf_bytes is not None else 2
-    dl_cols = st.columns(n_cols)
-    with dl_cols[0]:
-        st.download_button("Download all regions", data=buf_csv.getvalue(),
-            file_name=f"{base_name}_all_regions.csv", mime="text/csv",
-            key=f"download_all_regions_{key_suffix}", icon=":material/download:")
-    with dl_cols[1]:
-        st.download_button("Download heatmap", data=buf_img.getvalue(),
-            file_name=f"{base_name}_annotated_heatmap.png", mime="image/png",
-            key=f"download_annotated_{key_suffix}", icon=":material/image:")
-    if pdf_bytes is not None:
-        with dl_cols[2]:
-            st.download_button("Download report", data=pdf_bytes,
-                file_name=f"{base_name}_measure_report.pdf", mime="application/pdf",
-                key=f"download_measure_pdf_{key_suffix}", icon=":material/picture_as_pdf:")
-
-
-def _draw_contour_on_image(img_rgb, mask, stroke_color=(255, 0, 0), stroke_width=2):
-    """Draw contour from mask on RGB image. Resizes mask to match img if needed."""
-    h, w = img_rgb.shape[:2]
-    if mask.shape[:2] != (h, w):
-        mask = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if contours:
-        cv2.drawContours(img_rgb, contours, -1, stroke_color, stroke_width)
-    return img_rgb
-
-
-def render_region_canvas(display_heatmap, raw_heatmap=None, bf_img=None, original_vals=None, cell_vals=None,
-                         cell_mask=None, key_suffix="", input_filename=None, colormap_name="Jet"):
-    """Render drawable canvas and region metrics. When cell_vals: show cell area (replaces Full map). Else: show Full map."""
-    raw_heatmap = raw_heatmap if raw_heatmap is not None else display_heatmap
-    h, w = display_heatmap.shape
-    heatmap_rgb = heatmap_to_rgb_with_contour(display_heatmap, colormap_name, cell_mask)
-    pil_bg = Image.fromarray(heatmap_rgb).resize((CANVAS_SIZE, CANVAS_SIZE), Image.Resampling.LANCZOS)
-
-    st.markdown("""
-    <style>
-        [data-testid="stDialog"] [data-testid="stSelectbox"], [data-testid="stExpander"] [data-testid="stSelectbox"],
-        [data-testid="stDialog"] [data-testid="stSelectbox"] > div, [data-testid="stExpander"] [data-testid="stSelectbox"] > div {
-            width: 100% !important; max-width: 100% !important;
-        }
-        [data-testid="stDialog"] [data-testid="stMetric"] label, [data-testid="stDialog"] [data-testid="stMetric"] [data-testid="stMetricValue"],
-        [data-testid="stExpander"] [data-testid="stMetric"] label, [data-testid="stExpander"] [data-testid="stMetric"] [data-testid="stMetricValue"] {
-            font-size: 0.95rem !important;
-        }
-        [data-testid="stDialog"] img, [data-testid="stExpander"] img { border-radius: 0 !important; }
-    </style>
-    """, unsafe_allow_html=True)
-
-    if bf_img is not None:
-        bf_resized = cv2.resize(bf_img, (CANVAS_SIZE, CANVAS_SIZE))
-        bf_rgb = cv2.cvtColor(bf_resized, cv2.COLOR_GRAY2RGB) if bf_img.ndim == 2 else cv2.cvtColor(bf_resized, cv2.COLOR_BGR2RGB)
-        left_col, right_col = st.columns(2, gap=None)
-        with left_col:
-            draw_mode = st.selectbox("Tool", DRAW_TOOLS, format_func=lambda x: TOOL_LABELS[x], key=f"draw_mode_region_{key_suffix}")
-            st.caption("Left-click add, right-click close.  \nForce map (draw region)")
-            canvas_result = st_canvas(
-                fill_color="rgba(255, 165, 0, 0.3)", stroke_width=2, stroke_color="#ff6600",
-                background_image=pil_bg, drawing_mode=draw_mode, update_streamlit=True,
-                height=CANVAS_SIZE, width=CANVAS_SIZE, display_toolbar=True,
-                key=f"region_measure_canvas_{key_suffix}",
-            )
-        with right_col:
-            vals = cell_vals if cell_vals else original_vals
-            if vals:
-                label = "Cell area" if cell_vals else "Full map"
-                st.markdown(f'<p style="font-weight: 400; color: #334155; font-size: 0.95rem; margin: 0 20px 4px 4px;">{label}</p>', unsafe_allow_html=True)
-                st.markdown(f"""
-                <div style="width: 100%; box-sizing: border-box; border: 1px solid #e2e8f0; border-radius: 10px;
-                    padding: 10px 12px; margin: 0 10px 20px 10px; background: linear-gradient(145deg, #f8fafc 0%, #f1f5f9 100%);
-                    box-shadow: 0 1px 3px rgba(0,0,0,0.06);">
-                    <div style="display: flex; flex-wrap: wrap; gap: 5px; font-size: 0.9rem;">
-                        <span><strong>Sum:</strong> {vals['pixel_sum']:.1f}</span>
-                        <span><strong>Force:</strong> {vals['force']:.1f}</span>
-                        <span><strong>Max:</strong> {vals['max']:.3f}</span>
-                        <span><strong>Mean:</strong> {vals['mean']:.3f}</span>
-                    </div>
-                </div>
-                """, unsafe_allow_html=True)
-            st.caption("Bright-field")
-            bf_display = bf_rgb.copy()
-            if cell_mask is not None and np.any(cell_mask > 0):
-                bf_display = _draw_contour_on_image(bf_display, cell_mask, stroke_color=(255, 0, 0), stroke_width=2)
-            st.image(bf_display, width=CANVAS_SIZE)
-    else:
-        st.markdown("**Draw a region** on the heatmap.")
-        draw_mode = st.selectbox("Drawing tool", DRAW_TOOLS,
-            format_func=lambda x: "Polygon (free shape)" if x == "polygon" else TOOL_LABELS[x],
-            key=f"draw_mode_region_{key_suffix}")
-        st.caption("Polygon: left-click to add points, right-click to close.")
-        canvas_result = st_canvas(
-            fill_color="rgba(255, 165, 0, 0.3)", stroke_width=2, stroke_color="#ff6600",
-            background_image=pil_bg, drawing_mode=draw_mode, update_streamlit=True,
-            height=CANVAS_SIZE, width=CANVAS_SIZE, display_toolbar=True,
-            key=f"region_measure_canvas_{key_suffix}",
-        )
-
-    if canvas_result.json_data:
-        masks = parse_canvas_shapes_to_masks(canvas_result.json_data, CANVAS_SIZE, CANVAS_SIZE, h, w)
-        if masks:
-            metrics_list = [compute_region_metrics(raw_heatmap, m, original_vals) for m in masks]
-            if cell_mask is not None and np.any(cell_mask > 0):
-                cell_metrics = compute_region_metrics(raw_heatmap, cell_mask, original_vals)
-                metrics_list = [cell_metrics] + metrics_list
-            render_region_metrics_and_downloads(
-                metrics_list, masks, heatmap_rgb, input_filename, key_suffix, original_vals is not None,
-                first_region_label="Auto boundary" if (cell_mask is not None and np.any(cell_mask > 0)) else None,
-                bf_img=bf_img, cell_mask=cell_mask, colormap_name=colormap_name,
-            )
-
-
-def _compute_cell_metrics(raw_heatmap, cell_mask, pixel_sum, force):
-    """Compute metrics over estimated cell area only."""
-    area_px = int(np.sum(cell_mask))
-    if area_px == 0:
-        return None, None, None
-    region_values = raw_heatmap * cell_mask
-    cell_pixel_sum = float(np.sum(region_values))
-    cell_force = cell_pixel_sum * (force / pixel_sum) if pixel_sum > 0 else cell_pixel_sum
-    cell_mean = cell_pixel_sum / area_px
-    return cell_pixel_sum, cell_force, cell_mean
-
-
-def _add_cell_contour_to_fig(fig_pl, cell_mask, row=1, col=2):
-    """Add red contour overlay to Plotly heatmap subplot."""
-    if cell_mask is None or not np.any(cell_mask > 0):
-        return
-    contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if not contours:
-        return
-    # Use largest contour
-    cnt = max(contours, key=cv2.contourArea)
-    pts = cnt.squeeze()
-    if pts.ndim == 1:
-        pts = pts.reshape(1, 2)
-    x, y = pts[:, 0].tolist(), pts[:, 1].tolist()
-    if x[0] != x[-1] or y[0] != y[-1]:
-        x.append(x[0])
-        y.append(y[0])
-    fig_pl.add_trace(
-        go.Scatter(x=x, y=y, mode="lines", line=dict(color="red", width=2), showlegend=False),
-        row=row, col=col
-    )
-
-
-def render_result_display(img, raw_heatmap, display_heatmap, pixel_sum, force, key_img, download_key_suffix="",
-                         colormap_name="Jet", display_mode="Default", measure_region_dialog=None, auto_cell_boundary=True,
-                         cell_mask=None):
-    """
-    Render prediction result: plot, metrics, expander, and download/measure buttons.
-    measure_region_dialog: callable to open measure dialog (when ST_DIALOG available).
-    auto_cell_boundary: when True, use estimated cell area for metrics; when False, use entire map.
-    cell_mask: optional precomputed cell mask; if None and auto_cell_boundary, will be computed.
-    """
-    if cell_mask is None and auto_cell_boundary:
-        cell_mask = estimate_cell_mask(raw_heatmap)
-    elif not auto_cell_boundary:
-        cell_mask = None
-    cell_pixel_sum, cell_force, cell_mean = _compute_cell_metrics(raw_heatmap, cell_mask, pixel_sum, force) if cell_mask is not None else (None, None, None)
-    use_cell_metrics = auto_cell_boundary and cell_pixel_sum is not None and cell_force is not None and cell_mean is not None
-
-    base_name = os.path.splitext(key_img or "image")[0]
-    if use_cell_metrics:
-        main_csv_rows = [
-            ["image", "Cell sum", "Cell force (scaled)", "Heatmap max", "Heatmap mean"],
-            [base_name, f"{cell_pixel_sum:.2f}", f"{cell_force:.2f}",
-             f"{np.max(raw_heatmap):.4f}", f"{cell_mean:.4f}"],
-        ]
-    else:
-        main_csv_rows = [
-            ["image", "Sum of all pixels", "Cell force (scaled)", "Heatmap max", "Heatmap mean"],
-            [base_name, f"{pixel_sum:.2f}", f"{force:.2f}",
-             f"{np.max(raw_heatmap):.4f}", f"{np.mean(raw_heatmap):.4f}"],
-        ]
-    buf_main_csv = io.StringIO()
-    csv.writer(buf_main_csv).writerows(main_csv_rows)
-
-    buf_hm = heatmap_to_png_bytes(display_heatmap, colormap_name, cell_mask=cell_mask)
-
-    tit1, tit2 = st.columns(2)
-    with tit1:
-        st.markdown('<p style="font-size: 1.1rem; color: black; font-weight: 600;">Input: Bright-field image</p>', unsafe_allow_html=True)
-    with tit2:
-        st.markdown('<p style="font-size: 1.1rem; color: black; font-weight: 600;">Output: Predicted traction force map</p>', unsafe_allow_html=True)
-    fig_pl = make_subplots(rows=1, cols=2)
-    fig_pl.add_trace(go.Heatmap(z=img, colorscale="gray", showscale=False), row=1, col=1)
-    plotly_colorscale = cv_colormap_to_plotly_colorscale(colormap_name)
-    zmin, zmax = 0.0, 1.0
-    fig_pl.add_trace(go.Heatmap(z=display_heatmap, colorscale=plotly_colorscale, zmin=zmin, zmax=zmax, showscale=True,
-        colorbar=dict(len=0.4, thickness=12)), row=1, col=2)
-    _add_cell_contour_to_fig(fig_pl, cell_mask, row=1, col=2)
-    fig_pl.update_layout(
-        height=400,
-        margin=dict(l=10, r=10, t=10, b=10),
-        xaxis=dict(scaleanchor="y", scaleratio=1),
-        xaxis2=dict(scaleanchor="y2", scaleratio=1),
-    )
-    fig_pl.update_xaxes(showticklabels=False)
-    fig_pl.update_yaxes(showticklabels=False, autorange="reversed")
-    st.plotly_chart(fig_pl, use_container_width=True, config={"displayModeBar": True, "responsive": True})
-
-    col1, col2, col3, col4 = st.columns(4)
-    if use_cell_metrics:
-        with col1:
-            st.metric("Cell sum", f"{cell_pixel_sum:.2f}", help="Sum over estimated cell area (background excluded)")
-        with col2:
-            st.metric("Cell force (scaled)", f"{cell_force:.2f}", help="Total traction force in physical units")
-        with col3:
-            st.metric("Heatmap max", f"{np.max(raw_heatmap):.4f}", help="Peak force intensity in the map")
-        with col4:
-            st.metric("Heatmap mean", f"{cell_mean:.4f}", help="Mean force over estimated cell area")
-    else:
-        with col1:
-            st.metric("Sum of all pixels", f"{pixel_sum:.2f}", help="Raw sum of all pixel values in the force map")
-        with col2:
-            st.metric("Cell force (scaled)", f"{force:.2f}", help="Total traction force in physical units")
-        with col3:
-            st.metric("Heatmap max", f"{np.max(raw_heatmap):.4f}", help="Peak force intensity in the map")
-        with col4:
-            st.metric("Heatmap mean", f"{np.mean(raw_heatmap):.4f}", help="Average force intensity (full FOV)")
-
-    # Statistics panel (mean, std, percentiles, histogram)
-    with st.expander("Statistics"):
-        vals = raw_heatmap[cell_mask > 0] if (cell_mask is not None and np.any(cell_mask > 0) and use_cell_metrics) else raw_heatmap.flatten()
-        if len(vals) > 0:
-            st.markdown("**Summary**")
-            stat_col1, stat_col2, stat_col3 = st.columns(3)
-            with stat_col1:
-                st.metric("Mean", f"{float(np.mean(vals)):.4f}")
-                st.metric("Std", f"{float(np.std(vals)):.4f}")
-            with stat_col2:
-                p25, p50, p75 = float(np.percentile(vals, 25)), float(np.percentile(vals, 50)), float(np.percentile(vals, 75))
-                st.metric("P25", f"{p25:.4f}")
-                st.metric("P50 (median)", f"{p50:.4f}")
-                st.metric("P75", f"{p75:.4f}")
-            with stat_col3:
-                p90 = float(np.percentile(vals, 90))
-                st.metric("P90", f"{p90:.4f}")
-            st.markdown("**Histogram**")
-            hist_fig = go.Figure(data=[go.Histogram(x=vals, nbinsx=50, marker_color="#0d9488")])
-            hist_fig.update_layout(
-                height=220, margin=dict(l=40, r=20, t=20, b=40),
-                xaxis_title="Force value", yaxis_title="Count",
-                showlegend=False,
-            )
-            st.plotly_chart(hist_fig, use_container_width=True, config={"displayModeBar": False})
-        else:
-            st.caption("No nonzero values to compute statistics.")
-
-    with st.expander("How to read the results"):
-        if use_cell_metrics:
-            st.markdown("""
-**Input (left):** Bright-field microscopy image of a cell or spheroid on a substrate.  
-This is the raw image you provided—it shows cell shape but not forces.
-
-**Output (right):** Predicted traction force map.  
-- **Color** indicates force magnitude: blue = low, red = high  
-- **Brighter/warmer colors** = stronger forces exerted by the cell on the substrate  
-- **Red border = estimated cell area** (background excluded from metrics)
-- Values are normalized to [0, 1] for visualization
-
-**Metrics (auto cell boundary on):**
-- **Cell sum:** Sum over estimated cell area (background excluded)
-- **Cell force (scaled):** Total traction force in physical units
-- **Heatmap max:** Peak force intensity in the map
-- **Heatmap mean:** Mean force over the estimated cell area
-            """)
-        else:
-            st.markdown("""
-**Input (left):** Bright-field microscopy image of a cell or spheroid on a substrate.  
-This is the raw image you provided—it shows cell shape but not forces.
-
-**Output (right):** Predicted traction force map.  
-- **Color** indicates force magnitude: blue = low, red = high  
-- **Brighter/warmer colors** = stronger forces exerted by the cell on the substrate  
-- Values are normalized to [0, 1] for visualization
-
-**Metrics (auto cell boundary off):**
-- **Sum of all pixels:** Raw sum over entire map
-- **Cell force (scaled):** Total traction force in physical units
-- **Heatmap max/mean:** Peak and average force intensity (full field of view)
-            """)
-
-    original_vals = build_original_vals(raw_heatmap, pixel_sum, force)
-
-    pdf_bytes = create_pdf_report(
-        img, display_heatmap, raw_heatmap, pixel_sum, force, base_name, colormap_name,
-        cell_mask=cell_mask, cell_pixel_sum=cell_pixel_sum, cell_force=cell_force, cell_mean=cell_mean
-    )
-
-    btn_col1, btn_col2, btn_col3, btn_col4 = st.columns(4)
-    with btn_col1:
-        if HAS_DRAWABLE_CANVAS and measure_region_dialog is not None:
-            if st.button("Measure tool", key="open_measure", icon=":material/straighten:"):
-                st.session_state["open_measure_dialog"] = True
-                st.rerun()
-        elif HAS_DRAWABLE_CANVAS:
-            with st.expander("Measure tool"):
-                expander_cell_vals = build_cell_vals(raw_heatmap, cell_mask, pixel_sum, force) if (auto_cell_boundary and cell_mask is not None) else None
-                expander_cell_mask = cell_mask if auto_cell_boundary else None
-                render_region_canvas(
-                    display_heatmap,
-                    raw_heatmap=raw_heatmap,
-                    bf_img=img,
-                    original_vals=original_vals,
-                    cell_vals=expander_cell_vals,
-                    cell_mask=expander_cell_mask,
-                    key_suffix="expander",
-                    input_filename=key_img,
-                    colormap_name=colormap_name,
-                )
-        else:
-            st.caption("Install `streamlit-drawable-canvas-fix` for region measurement: `pip install streamlit-drawable-canvas-fix`")
-    with btn_col2:
-        st.download_button(
-            "Download heatmap",
-            width="stretch",
-            data=buf_hm.getvalue(),
-            file_name="s2f_heatmap.png",
-            mime="image/png",
-            key=f"download_heatmap{download_key_suffix}",
-            icon=":material/download:",
-        )
-    with btn_col3:
-        st.download_button(
-            "Download values",
-            width="stretch",
-            data=buf_main_csv.getvalue(),
-            file_name=f"{base_name}_main_values.csv",
-            mime="text/csv",
-            key=f"download_main_values{download_key_suffix}",
-            icon=":material/download:",
-        )
-    with btn_col4:
-        st.download_button(
-            "Download report",
-            width="stretch",
-            data=pdf_bytes,
-            file_name=f"{base_name}_report.pdf",
-            mime="application/pdf",
-            key=f"download_pdf{download_key_suffix}",
-            icon=":material/picture_as_pdf:",
-        )

S2FApp/ui/heatmaps.py

@@ -0,0 +1,119 @@
+"""Heatmap visualization utilities (colorbar, overlays, Plotly)."""
+import base64
+
+import cv2
+import numpy as np
+import streamlit as st
+import plotly.graph_objects as go
+
+from config.constants import COLORMAPS
+
+
+def _colormap_gradient_base64(colormap_name, width=512):
+    """Generate a horizontal gradient bar as base64 PNG for the given colormap."""
+    cv2_cmap = COLORMAPS.get(colormap_name, cv2.COLORMAP_JET)
+    gradient = np.linspace(0, 255, width, dtype=np.uint8).reshape(1, -1)
+    rgb = cv2.cvtColor(cv2.applyColorMap(gradient, cv2_cmap), cv2.COLOR_BGR2RGB)
+    bar = np.repeat(rgb, 6, axis=0)
+    _, buf = cv2.imencode(".png", cv2.cvtColor(bar, cv2.COLOR_RGB2BGR))
+    return base64.b64encode(buf.tobytes()).decode("utf-8")
+
+# Distinct colors for each region (RGB - heatmap_rgb is RGB)
+_REGION_COLORS = [
+    (0, 188, 212),   # cyan (matches drawing tool)
+    (0, 230, 118),   # green
+    (255, 235, 59),  # yellow
+    (171, 71, 188),  # purple
+    (0, 150, 255),   # blue
+    (255, 167, 38),  # amber
+    (124, 179, 66),  # light green
+    (233, 30, 99),   # pink
+]
+
+
+def _draw_region_overlay(annotated, mask, color, fill_alpha=0.3, stroke_width=2):
+    """Draw single region overlay on annotated heatmap (fill + alpha blend + contour). Modifies annotated in place."""
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    overlay = annotated.copy()
+    cv2.fillPoly(overlay, contours, color)
+    mask_3d = np.stack([mask] * 3, axis=-1).astype(bool)
+    annotated[mask_3d] = (
+        (1 - fill_alpha) * annotated[mask_3d].astype(np.float32)
+        + fill_alpha * overlay[mask_3d].astype(np.float32)
+    ).astype(np.uint8)
+    cv2.drawContours(annotated, contours, -1, color, stroke_width)
+
+
+def render_horizontal_colorbar(colormap_name, clip_min=0, clip_max=1, is_rescale=False):
+    """Render a compact horizontal colorbar for batch mode, anchored above the table."""
+    ticks = [0, 0.25, 0.5, 0.75, 1]
+    if is_rescale:
+        rng = clip_max - clip_min
+        labels = [f"{clip_min + t * rng:.2f}" for t in ticks]
+    else:
+        labels = [f"{t:.2f}" for t in ticks]
+
+    data_url = _colormap_gradient_base64(colormap_name)
+    labels_html = "".join(f'<span class="cb-tick">{l}</span>' for l in labels)
+    html = f"""
+    <div class="colorbar-table-header">
+        <div class="colorbar-ticks">{labels_html}</div>
+        <div class="colorbar-bar" style="background-image: url(data:image/png;base64,{data_url});"></div>
+    </div>
+    """
+    st.markdown(html, unsafe_allow_html=True)
+
+
+def make_annotated_heatmap(heatmap_rgb, mask, fill_alpha=0.3, stroke_color=(0, 188, 212), stroke_width=2):
+    """Composite heatmap with drawn region overlay."""
+    annotated = heatmap_rgb.copy()
+    _draw_region_overlay(annotated, mask, stroke_color, fill_alpha, stroke_width)
+    return annotated
+
+
+def make_annotated_heatmap_multi_regions(heatmap_rgb, masks, labels, cell_mask=None, fill_alpha=0.3):
+    """Draw each region separately with distinct color and label (R1, R2, ...). No merging."""
+    annotated = heatmap_rgb.copy()
+    if cell_mask is not None and np.any(cell_mask > 0):
+        contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(annotated, contours, -1, (255, 0, 0), 5)
+    for i, mask in enumerate(masks):
+        color = _REGION_COLORS[i % len(_REGION_COLORS)]
+        _draw_region_overlay(annotated, mask, color, fill_alpha, stroke_width=2)
+        # Label at centroid
+        M = cv2.moments(mask)
+        if M["m00"] > 0:
+            cx = int(M["m10"] / M["m00"])
+            cy = int(M["m01"] / M["m00"])
+            label = labels[i] if i < len(labels) else f"R{i + 1}"
+            cv2.putText(
+                annotated, label, (cx - 12, cy + 5),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA
+            )
+            cv2.putText(
+                annotated, label, (cx - 12, cy + 5),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1, cv2.LINE_AA
+            )
+    return annotated
+
+
+def add_cell_contour_to_fig(fig_pl, cell_mask, row=1, col=2):
+    """Add red contour overlay to Plotly heatmap subplot."""
+    if cell_mask is None or not np.any(cell_mask > 0):
+        return
+    contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return
+    # Use largest contour
+    cnt = max(contours, key=cv2.contourArea)
+    pts = cnt.squeeze()
+    if pts.ndim == 1:
+        pts = pts.reshape(1, 2)
+    x, y = pts[:, 0].tolist(), pts[:, 1].tolist()
+    if x[0] != x[-1] or y[0] != y[-1]:
+        x.append(x[0])
+        y.append(y[0])
+    fig_pl.add_trace(
+        go.Scatter(x=x, y=y, mode="lines", line=dict(color="red", width=4), showlegend=False),
+        row=row, col=col
+    )

S2FApp/ui/measure_tool.py

@@ -0,0 +1,333 @@
+"""Measure tool: drawable canvas, region metrics, and downloads."""
+import csv
+import html
+import io
+import os
+
+import cv2
+import numpy as np
+import streamlit as st
+from PIL import Image
+
+from config.constants import CANVAS_SIZE, DRAW_TOOLS, TOOL_LABELS
+from utils.report import heatmap_to_rgb_with_contour, create_measure_pdf_report
+from ui.heatmaps import make_annotated_heatmap_multi_regions
+
+try:
+    from streamlit_drawable_canvas import st_canvas
+    HAS_DRAWABLE_CANVAS = True
+except (ImportError, AttributeError):
+    HAS_DRAWABLE_CANVAS = False
+
+
+def _obj_to_pts(obj, scale_x, scale_y, heatmap_w, heatmap_h):
+    """Convert a single canvas object to polygon points in heatmap coords. Returns None if invalid."""
+    obj_type = obj.get("type", "")
+    pts = []
+    if obj_type == "rect":
+        left = obj.get("left", 0)
+        top = obj.get("top", 0)
+        w = obj.get("width", 0)
+        h = obj.get("height", 0)
+        pts = np.array([
+            [left, top], [left + w, top], [left + w, top + h], [left, top + h]
+        ], dtype=np.float32)
+    elif obj_type == "circle" or obj_type == "ellipse":
+        left = obj.get("left", 0)
+        top = obj.get("top", 0)
+        width = obj.get("width", 0)
+        height = obj.get("height", 0)
+        radius = obj.get("radius", 0)
+        angle_deg = obj.get("angle", 0)
+        if radius > 0:
+            rx = ry = radius
+            angle_rad = np.deg2rad(angle_deg)
+            cx = left + radius * np.cos(angle_rad)
+            cy = top + radius * np.sin(angle_rad)
+        else:
+            rx = width / 2 if width > 0 else 0
+            ry = height / 2 if height > 0 else 0
+            if rx <= 0 or ry <= 0:
+                return None
+            cx = left + rx
+            cy = top + ry
+        if rx <= 0 or ry <= 0:
+            return None
+        n = 32
+        angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
+        pts = np.column_stack([cx + rx * np.cos(angles), cy + ry * np.sin(angles)]).astype(np.float32)
+    elif obj_type == "path":
+        path = obj.get("path", [])
+        for cmd in path:
+            if isinstance(cmd, (list, tuple)) and len(cmd) >= 3:
+                if cmd[0] in ("M", "L"):
+                    pts.append([float(cmd[1]), float(cmd[2])])
+                elif cmd[0] == "Q" and len(cmd) >= 5:
+                    pts.append([float(cmd[3]), float(cmd[4])])
+                elif cmd[0] == "C" and len(cmd) >= 7:
+                    pts.append([float(cmd[5]), float(cmd[6])])
+        if len(pts) < 3:
+            return None
+        pts = np.array(pts, dtype=np.float32)
+    else:
+        return None
+    pts[:, 0] *= scale_x
+    pts[:, 1] *= scale_y
+    pts = np.clip(pts, 0, [heatmap_w - 1, heatmap_h - 1]).astype(np.int32)
+    return pts
+
+
+def parse_canvas_shapes_to_masks(json_data, canvas_h, canvas_w, heatmap_h, heatmap_w):
+    """Parse drawn shapes and return a list of individual masks (one per shape)."""
+    if not json_data or "objects" not in json_data or not json_data["objects"]:
+        return []
+    scale_x = heatmap_w / canvas_w
+    scale_y = heatmap_h / canvas_h
+    masks = []
+    for obj in json_data["objects"]:
+        pts = _obj_to_pts(obj, scale_x, scale_y, heatmap_w, heatmap_h)
+        if pts is None:
+            continue
+        mask = np.zeros((heatmap_h, heatmap_w), dtype=np.uint8)
+        cv2.fillPoly(mask, [pts], 1)
+        masks.append(mask)
+    return masks
+
+
+def build_original_vals(raw_heatmap, pixel_sum, force):
+    """Build original_vals dict for measure tool (full map)."""
+    return {
+        "pixel_sum": pixel_sum,
+        "force": force,
+        "max": float(np.max(raw_heatmap)),
+        "mean": float(np.mean(raw_heatmap)),
+    }
+
+
+def _compute_cell_metrics(raw_heatmap, cell_mask, pixel_sum, force):
+    """Compute metrics over estimated cell area only."""
+    area_px = int(np.sum(cell_mask))
+    if area_px == 0:
+        return None, None, None
+    region_values = raw_heatmap * cell_mask
+    cell_pixel_sum = float(np.sum(region_values))
+    cell_force = cell_pixel_sum * (force / pixel_sum) if pixel_sum > 0 else cell_pixel_sum
+    cell_mean = cell_pixel_sum / area_px
+    return cell_pixel_sum, cell_force, cell_mean
+
+
+def build_cell_vals(raw_heatmap, cell_mask, pixel_sum, force):
+    """Build cell_vals dict for measure tool (estimated cell area). Returns None if invalid."""
+    cell_pixel_sum, cell_force, cell_mean = _compute_cell_metrics(raw_heatmap, cell_mask, pixel_sum, force)
+    if cell_pixel_sum is None:
+        return None
+    region_values = raw_heatmap * cell_mask
+    region_nonzero = region_values[cell_mask > 0]
+    cell_max = float(np.max(region_nonzero)) if len(region_nonzero) > 0 else 0
+    return {
+        "pixel_sum": cell_pixel_sum,
+        "force": cell_force,
+        "max": cell_max,
+        "mean": cell_mean,
+    }
+
+
+def compute_region_metrics(raw_heatmap, mask, original_vals=None):
+    """Compute region metrics from mask."""
+    area_px = int(np.sum(mask))
+    region_values = raw_heatmap * mask
+    region_nonzero = region_values[mask > 0]
+    force_sum = float(np.sum(region_values))
+    density = force_sum / area_px if area_px > 0 else 0
+    region_max = float(np.max(region_nonzero)) if len(region_nonzero) > 0 else 0
+    region_mean = float(np.mean(region_nonzero)) if len(region_nonzero) > 0 else 0
+    region_force_scaled = (
+        force_sum * (original_vals["force"] / original_vals["pixel_sum"])
+        if original_vals and original_vals.get("pixel_sum", 0) > 0
+        else force_sum
+    )
+    return {
+        "area_px": area_px,
+        "force_sum": force_sum,
+        "density": density,
+        "max": region_max,
+        "mean": region_mean,
+        "force_scaled": region_force_scaled,
+    }
+
+
+def _draw_contour_on_image(img_rgb, mask, stroke_color=(255, 0, 0), stroke_width=5):
+    """Draw contour from mask on RGB image. Resizes mask to match img if needed."""
+    h, w = img_rgb.shape[:2]
+    if mask.shape[:2] != (h, w):
+        mask = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if contours:
+        cv2.drawContours(img_rgb, contours, -1, stroke_color, stroke_width)
+    return img_rgb
+
+
+def render_region_metrics_and_downloads(metrics_list, masks, heatmap_rgb, input_filename, key_suffix, has_original_vals,
+                                        first_region_label=None, bf_img=None, cell_mask=None, colormap_name="Jet"):
+    """
+    Render per-shape metrics table and download buttons.
+    first_region_label: custom label for first row (e.g. 'Auto boundary').
+    masks: list of region masks (user-drawn only; used for labeled heatmap with R1, R2...).
+    """
+    base_name = os.path.splitext(input_filename or "image")[0]
+    st.markdown("**Regions (each selection = one row)**")
+    if has_original_vals:
+        headers = ["Region", "Area", "F.sum", "Force", "Max", "Mean"]
+        csv_rows = [["image", "region"] + headers[1:]]
+    else:
+        headers = ["Region", "Area (px²)", "Force sum", "Mean"]
+        csv_rows = [["image", "region", "Area", "Force sum", "Mean"]]
+    table_rows = [headers]
+    for i, metrics in enumerate(metrics_list, 1):
+        region_label = first_region_label if (i == 1 and first_region_label) else f"Region {i - (1 if first_region_label else 0)}"
+        if has_original_vals:
+            row = [region_label, str(metrics["area_px"]), f"{metrics['force_sum']:.3f}", f"{metrics['force_scaled']:.1f}",
+                   f"{metrics['max']:.3f}", f"{metrics['mean']:.4f}"]
+            csv_rows.append([base_name, region_label, metrics["area_px"], f"{metrics['force_sum']:.3f}",
+                            f"{metrics['force_scaled']:.1f}", f"{metrics['max']:.3f}", f"{metrics['mean']:.4f}"])
+        else:
+            row = [region_label, str(metrics["area_px"]), f"{metrics['force_sum']:.4f}", f"{metrics['mean']:.6f}"]
+            csv_rows.append([base_name, region_label, metrics["area_px"], f"{metrics['force_sum']:.4f}",
+                            f"{metrics['mean']:.6f}"])
+        table_rows.append(row)
+    # Render as HTML table to avoid Streamlit's default row/column indices
+    header = table_rows[0]
+    body = table_rows[1:]
+    th_cells = "".join(
+        f'<th style="border: 1px solid #ddd; padding: 8px; text-align: left;">{html.escape(str(h))}</th>'
+        for h in header
+    )
+    rows_html = [
+        "<tr>"
+        + "".join(
+            f'<td style="border: 1px solid #ddd; padding: 8px;">{html.escape(str(c))}</td>'
+            for c in row
+        )
+        + "</tr>"
+        for row in body
+    ]
+    table_html = (
+        f'<table style="border-collapse: collapse; width: 100%;">'
+        f"<thead><tr>{th_cells}</tr></thead>"
+        f"<tbody>{''.join(rows_html)}</tbody></table>"
+    )
+    st.markdown(table_html, unsafe_allow_html=True)
+    buf_csv = io.StringIO()
+    csv.writer(buf_csv).writerows(csv_rows)
+    # Annotated heatmap: each region separate with R1, R2 labels (no merging)
+    region_labels = [f"R{i + 1}" for i in range(len(masks))]
+    heatmap_labeled = make_annotated_heatmap_multi_regions(heatmap_rgb.copy(), masks, region_labels, cell_mask=None)
+    buf_img = io.BytesIO()
+    Image.fromarray(heatmap_labeled).save(buf_img, format="PNG")
+    buf_img.seek(0)
+    # PDF report (requires bf_img)
+    pdf_bytes = None
+    if bf_img is not None:
+        pdf_bytes = create_measure_pdf_report(bf_img, heatmap_labeled, table_rows, base_name)
+    n_cols = 3 if pdf_bytes is not None else 2
+    dl_cols = st.columns(n_cols)
+    with dl_cols[0]:
+        st.download_button("Download all regions", data=buf_csv.getvalue(),
+            file_name=f"{base_name}_all_regions.csv", mime="text/csv",
+            key=f"download_all_regions_{key_suffix}", icon=":material/download:")
+    with dl_cols[1]:
+        st.download_button("Download heatmap", data=buf_img.getvalue(),
+            file_name=f"{base_name}_annotated_heatmap.png", mime="image/png",
+            key=f"download_annotated_{key_suffix}", icon=":material/image:")
+    if pdf_bytes is not None:
+        with dl_cols[2]:
+            st.download_button("Download report", data=pdf_bytes,
+                file_name=f"{base_name}_measure_report.pdf", mime="application/pdf",
+                key=f"download_measure_pdf_{key_suffix}", icon=":material/picture_as_pdf:")
+
+
+def render_region_canvas(display_heatmap, raw_heatmap=None, bf_img=None, original_vals=None, cell_vals=None,
+                         cell_mask=None, key_suffix="", input_filename=None, colormap_name="Jet"):
+    """Render drawable canvas and region metrics. When cell_vals: show cell area (replaces Full map). Else: show Full map."""
+    if not HAS_DRAWABLE_CANVAS:
+        st.caption("Install `streamlit-drawable-canvas-fix` for region measurement: `pip install streamlit-drawable-canvas-fix`")
+        return
+    raw_heatmap = raw_heatmap if raw_heatmap is not None else display_heatmap
+    h, w = display_heatmap.shape
+    heatmap_rgb = heatmap_to_rgb_with_contour(display_heatmap, colormap_name, cell_mask)
+    pil_bg = Image.fromarray(heatmap_rgb).resize((CANVAS_SIZE, CANVAS_SIZE), Image.Resampling.LANCZOS)
+
+    st.markdown("""
+    <style>
+        [data-testid="stDialog"] [data-testid="stSelectbox"], [data-testid="stExpander"] [data-testid="stSelectbox"],
+        [data-testid="stDialog"] [data-testid="stSelectbox"] > div, [data-testid="stExpander"] [data-testid="stSelectbox"] > div {
+            width: 100% !important; max-width: 100% !important;
+        }
+        [data-testid="stDialog"] [data-testid="stMetric"] label, [data-testid="stDialog"] [data-testid="stMetric"] [data-testid="stMetricValue"],
+        [data-testid="stExpander"] [data-testid="stMetric"] label, [data-testid="stExpander"] [data-testid="stMetric"] [data-testid="stMetricValue"] {
+            font-size: 0.95rem !important;
+        }
+        [data-testid="stDialog"] img, [data-testid="stExpander"] img { border-radius: 0 !important; }
+    </style>
+    """, unsafe_allow_html=True)
+
+    if bf_img is not None:
+        bf_resized = cv2.resize(bf_img, (CANVAS_SIZE, CANVAS_SIZE))
+        bf_rgb = cv2.cvtColor(bf_resized, cv2.COLOR_GRAY2RGB) if bf_img.ndim == 2 else cv2.cvtColor(bf_resized, cv2.COLOR_BGR2RGB)
+        left_col, right_col = st.columns(2, gap=None)
+        with left_col:
+            draw_mode = st.selectbox("Tool", DRAW_TOOLS, format_func=lambda x: TOOL_LABELS[x], key=f"draw_mode_region_{key_suffix}")
+            st.caption("Left-click add, right-click close.  \nForce map (draw region)")
+            canvas_result = st_canvas(
+                fill_color="rgba(0, 188, 212, 0.25)", stroke_width=2, stroke_color="#00bcd4",
+                background_image=pil_bg, drawing_mode=draw_mode, update_streamlit=True,
+                height=CANVAS_SIZE, width=CANVAS_SIZE, display_toolbar=True,
+                key=f"region_measure_canvas_{key_suffix}",
+            )
+        with right_col:
+            vals = cell_vals if cell_vals else original_vals
+            if vals:
+                label = "Cell area" if cell_vals else "Full map"
+                st.markdown(f'<p style="font-weight: 400; color: #334155; font-size: 0.95rem; margin: 0 20px 4px 4px;">{label}</p>', unsafe_allow_html=True)
+                st.markdown(f"""
+                <div style="width: 100%; box-sizing: border-box; border: 1px solid #e2e8f0; border-radius: 10px;
+                    padding: 10px 12px; margin: 0 10px 20px 10px; background: linear-gradient(145deg, #f8fafc 0%, #f1f5f9 100%);
+                    box-shadow: 0 1px 3px rgba(0,0,0,0.06);">
+                    <div style="display: flex; flex-wrap: wrap; gap: 5px; font-size: 0.9rem;">
+                        <span><strong>Sum:</strong> {vals['pixel_sum']:.1f}</span>
+                        <span><strong>Force:</strong> {vals['force']:.1f}</span>
+                        <span><strong>Max:</strong> {vals['max']:.3f}</span>
+                        <span><strong>Mean:</strong> {vals['mean']:.3f}</span>
+                    </div>
+                </div>
+                """, unsafe_allow_html=True)
+            st.caption("Bright-field")
+            bf_display = bf_rgb.copy()
+            if cell_mask is not None and np.any(cell_mask > 0):
+                bf_display = _draw_contour_on_image(bf_display, cell_mask, stroke_color=(255, 0, 0), stroke_width=5)
+            st.image(bf_display, width=CANVAS_SIZE)
+    else:
+        st.markdown("**Draw a region** on the heatmap.")
+        draw_mode = st.selectbox("Drawing tool", DRAW_TOOLS,
+            format_func=lambda x: "Polygon (free shape)" if x == "polygon" else TOOL_LABELS[x],
+            key=f"draw_mode_region_{key_suffix}")
+        st.caption("Polygon: left-click to add points, right-click to close.")
+        canvas_result = st_canvas(
+            fill_color="rgba(0, 188, 212, 0.25)", stroke_width=2, stroke_color="#00bcd4",
+            background_image=pil_bg, drawing_mode=draw_mode, update_streamlit=True,
+            height=CANVAS_SIZE, width=CANVAS_SIZE, display_toolbar=True,
+            key=f"region_measure_canvas_{key_suffix}",
+        )
+
+    if canvas_result.json_data:
+        masks = parse_canvas_shapes_to_masks(canvas_result.json_data, CANVAS_SIZE, CANVAS_SIZE, h, w)
+        if masks:
+            metrics_list = [compute_region_metrics(raw_heatmap, m, original_vals) for m in masks]
+            if cell_mask is not None and np.any(cell_mask > 0):
+                cell_metrics = compute_region_metrics(raw_heatmap, cell_mask, original_vals)
+                metrics_list = [cell_metrics] + metrics_list
+            render_region_metrics_and_downloads(
+                metrics_list, masks, heatmap_rgb, input_filename, key_suffix, original_vals is not None,
+                first_region_label="Auto boundary" if (cell_mask is not None and np.any(cell_mask > 0)) else None,
+                bf_img=bf_img, cell_mask=cell_mask, colormap_name=colormap_name,
+            )

S2FApp/ui/result_display.py

@@ -0,0 +1,358 @@
+"""Result display: single and batch prediction views."""
+import csv
+import io
+import os
+import zipfile
+
+import cv2
+import numpy as np
+import streamlit as st
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+from utils.display import apply_display_scale, cv_colormap_to_plotly_colorscale
+from utils.report import heatmap_to_rgb_with_contour, heatmap_to_png_bytes, create_pdf_report
+from utils.segmentation import estimate_cell_mask
+from ui.heatmaps import render_horizontal_colorbar, add_cell_contour_to_fig
+from ui.measure_tool import (
+    build_original_vals,
+    build_cell_vals,
+    render_region_canvas,
+    _compute_cell_metrics,
+    HAS_DRAWABLE_CANVAS,
+)
+
+
+def render_batch_results(batch_results, colormap_name="Jet", display_mode="Default",
+                        min_percentile=0, max_percentile=100, clip_min=0, clip_max=1,
+                        auto_cell_boundary=False, clip_bounds=False):
+    """
+    Render batch prediction results: summary table, bright-field row, heatmap row, and bulk download.
+    batch_results: list of dicts with img, heatmap, force, pixel_sum, key_img, cell_mask.
+    cell_mask is computed on-the-fly when auto_cell_boundary is True and not stored.
+    """
+    if not batch_results:
+        return
+
+    # Resolve cell_mask and precompute display_heatmap for each result
+    for r in batch_results:
+        if auto_cell_boundary and (r.get("cell_mask") is None or not np.any(r.get("cell_mask", 0) > 0)):
+            r["_cell_mask"] = estimate_cell_mask(r["heatmap"])
+        else:
+            r["_cell_mask"] = r.get("cell_mask") if auto_cell_boundary else None
+        r["_display_heatmap"] = apply_display_scale(
+            r["heatmap"], display_mode,
+            min_percentile=min_percentile, max_percentile=max_percentile,
+            clip_min=clip_min, clip_max=clip_max, clip_bounds=clip_bounds,
+        )
+    # Build table rows - consistent column names for both modes
+    headers = ["Image", "Force", "Sum", "Max", "Mean"]
+    rows = []
+    csv_rows = [["image"] + headers[1:]]
+    for r in batch_results:
+        heatmap = r["heatmap"]
+        cell_mask = r.get("_cell_mask")
+        key = r["key_img"] or "image"
+        if auto_cell_boundary and cell_mask is not None and np.any(cell_mask > 0):
+            vals = heatmap[cell_mask > 0]
+            cell_pixel_sum = float(np.sum(vals))
+            cell_force = cell_pixel_sum * (r["force"] / r["pixel_sum"]) if r["pixel_sum"] > 0 else cell_pixel_sum
+            cell_mean = cell_pixel_sum / np.sum(cell_mask) if np.sum(cell_mask) > 0 else 0
+            row = [key, f"{cell_force:.2f}", f"{cell_pixel_sum:.2f}",
+                   f"{np.max(heatmap):.4f}", f"{cell_mean:.4f}"]
+        else:
+            row = [key, f"{r['force']:.2f}", f"{r['pixel_sum']:.2f}",
+                   f"{np.max(heatmap):.4f}", f"{np.mean(heatmap):.4f}"]
+        rows.append(row)
+        csv_rows.append([os.path.splitext(key)[0]] + row[1:])
+    st.markdown('<div class="result-label"><span class="result-badge input">INPUT</span> Bright-field images</div>', unsafe_allow_html=True)
+    n_cols = min(5, len(batch_results))
+    bf_cols = st.columns(n_cols)
+    for i, r in enumerate(batch_results):
+        img = r["img"]
+        if img.ndim == 2:
+            img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+        else:
+            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        with bf_cols[i % n_cols]:
+            st.image(img_rgb, caption=r["key_img"], use_container_width=True)
+    is_rescale_b = display_mode == "Range" and clip_max > clip_min and not (clip_min == 0 and clip_max == 1)
+    st.markdown('<div class="result-label"><span class="result-badge output">OUTPUT</span> Predicted force maps</div>', unsafe_allow_html=True)
+    hm_cols = st.columns(n_cols)
+    for i, r in enumerate(batch_results):
+        hm_rgb = heatmap_to_rgb_with_contour(
+            r["_display_heatmap"], colormap_name,
+            r.get("_cell_mask") if auto_cell_boundary else None,
+        )
+        with hm_cols[i % n_cols]:
+            st.image(hm_rgb, caption=r["key_img"], use_container_width=True)
+    render_horizontal_colorbar(colormap_name, clip_min, clip_max, is_rescale_b)
+    # Table
+    st.dataframe(
+        {h: [r[i] for r in rows] for i, h in enumerate(headers)},
+        use_container_width=True,
+        hide_index=True,
+    )
+    # Histograms in accordion (one per row for visibility)
+    with st.expander("Force distribution (histograms)", expanded=False):
+        for i, r in enumerate(batch_results):
+            heatmap = r["heatmap"]
+            cell_mask = r.get("_cell_mask")
+            vals = heatmap[cell_mask > 0] if (cell_mask is not None and np.any(cell_mask > 0) and auto_cell_boundary) else heatmap.flatten()
+            vals = vals[vals > 0] if np.any(vals > 0) else vals
+            st.markdown(f"**{r['key_img']}**")
+            if len(vals) > 0:
+                fig = go.Figure(data=[go.Histogram(x=vals, nbinsx=50, marker_color="#0d9488")])
+                fig.update_layout(
+                    height=220, margin=dict(l=40, r=20, t=10, b=40),
+                    xaxis_title="Force value", yaxis_title="Count",
+                    showlegend=False,
+                )
+                st.plotly_chart(fig, use_container_width=True, config={"displayModeBar": False})
+            else:
+                st.caption("No data")
+            if i < len(batch_results) - 1:
+                st.divider()
+    # Bulk downloads: CSV and heatmaps (zip)
+    buf_csv = io.StringIO()
+    csv.writer(buf_csv).writerows(csv_rows)
+    zip_buf = io.BytesIO()
+    with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf:
+        for r in batch_results:
+            hm_bytes = heatmap_to_png_bytes(
+                r["_display_heatmap"], colormap_name,
+                r.get("_cell_mask") if auto_cell_boundary else None,
+            )
+            base = os.path.splitext(r["key_img"] or "image")[0]
+            zf.writestr(f"{base}_heatmap.png", hm_bytes.getvalue())
+    zip_buf.seek(0)
+    dl_col1, dl_col2 = st.columns(2)
+    with dl_col1:
+        st.download_button(
+            "Download all as CSV",
+            data=buf_csv.getvalue(),
+            file_name="s2f_batch_results.csv",
+            mime="text/csv",
+            key="download_batch_csv",
+            icon=":material/download:",
+        )
+    with dl_col2:
+        st.download_button(
+            "Download all heatmaps",
+            data=zip_buf.getvalue(),
+            file_name="s2f_batch_heatmaps.zip",
+            mime="application/zip",
+            key="download_batch_heatmaps",
+            icon=":material/image:",
+        )
+
+
+def render_result_display(img, raw_heatmap, display_heatmap, pixel_sum, force, key_img, download_key_suffix="",
+                         colormap_name="Jet", display_mode="Default", measure_region_dialog=None, auto_cell_boundary=True,
+                         cell_mask=None, clip_min=0.0, clip_max=1.0, clip_bounds=False):
+    """
+    Render prediction result: plot, metrics, expander, and download/measure buttons.
+    measure_region_dialog: callable to open measure dialog (when ST_DIALOG available).
+    auto_cell_boundary: when True, use estimated cell area for metrics; when False, use entire map.
+    cell_mask: optional precomputed cell mask; if None and auto_cell_boundary, will be computed.
+    """
+    if cell_mask is None and auto_cell_boundary:
+        cell_mask = estimate_cell_mask(raw_heatmap)
+    elif not auto_cell_boundary:
+        cell_mask = None
+    cell_pixel_sum, cell_force, cell_mean = _compute_cell_metrics(raw_heatmap, cell_mask, pixel_sum, force) if cell_mask is not None else (None, None, None)
+    use_cell_metrics = auto_cell_boundary and cell_pixel_sum is not None and cell_force is not None and cell_mean is not None
+
+    base_name = os.path.splitext(key_img or "image")[0]
+    if use_cell_metrics:
+        main_csv_rows = [
+            ["image", "Cell sum", "Cell force (scaled)", "Heatmap max", "Heatmap mean"],
+            [base_name, f"{cell_pixel_sum:.2f}", f"{cell_force:.2f}",
+             f"{np.max(raw_heatmap):.4f}", f"{cell_mean:.4f}"],
+        ]
+    else:
+        main_csv_rows = [
+            ["image", "Sum of all pixels", "Cell force (scaled)", "Heatmap max", "Heatmap mean"],
+            [base_name, f"{pixel_sum:.2f}", f"{force:.2f}",
+             f"{np.max(raw_heatmap):.4f}", f"{np.mean(raw_heatmap):.4f}"],
+        ]
+    buf_main_csv = io.StringIO()
+    csv.writer(buf_main_csv).writerows(main_csv_rows)
+
+    buf_hm = heatmap_to_png_bytes(display_heatmap, colormap_name, cell_mask=cell_mask)
+
+    is_rescale = display_mode == "Range" and clip_max > clip_min and not (clip_min == 0.0 and clip_max == 1.0)
+
+    tit1, tit2 = st.columns(2)
+    with tit1:
+        st.markdown('<div class="result-label"><span class="result-badge input">INPUT</span> Bright-field image</div>', unsafe_allow_html=True)
+    with tit2:
+        st.markdown('<div class="result-label"><span class="result-badge output">OUTPUT</span> Predicted force map</div>', unsafe_allow_html=True)
+    fig_pl = make_subplots(rows=1, cols=2)
+    fig_pl.add_trace(go.Heatmap(z=img, colorscale="gray", showscale=False), row=1, col=1)
+    plotly_colorscale = cv_colormap_to_plotly_colorscale(colormap_name)
+    colorbar_cfg = dict(len=0.4, thickness=12, tickmode="array")
+    tick_positions = [0, 0.25, 0.5, 0.75, 1]
+    if is_rescale:
+        rng = clip_max - clip_min
+        colorbar_cfg["tickvals"] = tick_positions
+        colorbar_cfg["ticktext"] = [f"{clip_min + t * rng:.2f}" for t in tick_positions]
+    else:
+        colorbar_cfg["tickvals"] = tick_positions
+        colorbar_cfg["ticktext"] = [f"{t:.2f}" for t in tick_positions]
+    fig_pl.add_trace(go.Heatmap(z=display_heatmap, colorscale=plotly_colorscale, zmin=0.0, zmax=1.0, showscale=True,
+        colorbar=colorbar_cfg), row=1, col=2)
+    add_cell_contour_to_fig(fig_pl, cell_mask, row=1, col=2)
+    fig_pl.update_layout(
+        height=400,
+        margin=dict(l=10, r=10, t=10, b=10),
+        xaxis=dict(scaleanchor="y", scaleratio=1),
+        xaxis2=dict(scaleanchor="y2", scaleratio=1),
+    )
+    fig_pl.update_xaxes(showticklabels=False, showgrid=False, zeroline=False)
+    fig_pl.update_yaxes(showticklabels=False, autorange="reversed", showgrid=False, zeroline=False)
+    st.plotly_chart(fig_pl, use_container_width=True, config={"displayModeBar": True, "responsive": True})
+
+    col1, col2, col3, col4 = st.columns(4)
+    if use_cell_metrics:
+        with col1:
+            st.metric("Cell sum", f"{cell_pixel_sum:.2f}", help="Sum over estimated cell area (background excluded)")
+        with col2:
+            st.metric("Cell force (scaled)", f"{cell_force:.2f}", help="Total traction force in physical units")
+        with col3:
+            st.metric("Heatmap max", f"{np.max(raw_heatmap):.4f}", help="Peak force intensity in the map")
+        with col4:
+            st.metric("Heatmap mean", f"{cell_mean:.4f}", help="Mean force over estimated cell area")
+    else:
+        with col1:
+            st.metric("Sum of all pixels", f"{pixel_sum:.2f}", help="Raw sum of all pixel values in the force map")
+        with col2:
+            st.metric("Cell force (scaled)", f"{force:.2f}", help="Total traction force in physical units")
+        with col3:
+            st.metric("Heatmap max", f"{np.max(raw_heatmap):.4f}", help="Peak force intensity in the map")
+        with col4:
+            st.metric("Heatmap mean", f"{np.mean(raw_heatmap):.4f}", help="Average force intensity (full FOV)")
+
+    # Statistics panel (mean, std, percentiles, histogram)
+    with st.expander("Statistics"):
+        vals = raw_heatmap[cell_mask > 0] if (cell_mask is not None and np.any(cell_mask > 0) and use_cell_metrics) else raw_heatmap.flatten()
+        if len(vals) > 0:
+            stat_col1, stat_col2, stat_col3 = st.columns(3)
+            p25, p50, p75, p90 = (
+                float(np.percentile(vals, 25)), float(np.percentile(vals, 50)),
+                float(np.percentile(vals, 75)), float(np.percentile(vals, 90)),
+            )
+            with stat_col1:
+                st.metric("Mean", f"{float(np.mean(vals)):.4f}")
+                st.metric("Std", f"{float(np.std(vals)):.4f}")
+            with stat_col2:
+                st.metric("P25", f"{p25:.4f}")
+                st.metric("P50 (median)", f"{p50:.4f}")
+            with stat_col3:
+                st.metric("P75", f"{p75:.4f}")
+                st.metric("P90", f"{p90:.4f}")
+            st.markdown("**Histogram**")
+            hist_fig = go.Figure(data=[go.Histogram(x=vals, nbinsx=50, marker_color="#0d9488")])
+            hist_fig.update_layout(
+                height=220, margin=dict(l=40, r=20, t=20, b=40),
+                xaxis_title="Force value", yaxis_title="Count",
+                showlegend=False,
+            )
+            st.plotly_chart(hist_fig, use_container_width=True, config={"displayModeBar": False})
+        else:
+            st.caption("No nonzero values to compute statistics.")
+
+    with st.expander("How to read the results"):
+        if use_cell_metrics:
+            st.markdown("""
+**Input (left):** Bright-field microscopy image of a cell or spheroid on a substrate.  
+This is the raw image you provided—it shows cell shape but not forces.
+
+**Output (right):** Predicted traction force map.  
+- **Color** indicates force magnitude: blue = low, red = high  
+- **Brighter/warmer colors** = stronger forces exerted by the cell on the substrate  
+- **Red border = estimated cell area** (background excluded from metrics)
+- Values are normalized to [0, 1] for visualization
+
+**Metrics (auto cell boundary on):**
+- **Cell sum:** Sum over estimated cell area (background excluded)
+- **Cell force (scaled):** Total traction force in physical units
+- **Heatmap max:** Peak force intensity in the map
+- **Heatmap mean:** Mean force over the estimated cell area
+            """)
+        else:
+            st.markdown("""
+**Input (left):** Bright-field microscopy image of a cell or spheroid on a substrate.  
+This is the raw image you provided—it shows cell shape but not forces.
+
+**Output (right):** Predicted traction force map.  
+- **Color** indicates force magnitude: blue = low, red = high  
+- **Brighter/warmer colors** = stronger forces exerted by the cell on the substrate  
+- Values are normalized to [0, 1] for visualization
+
+**Metrics (auto cell boundary off):**
+- **Sum of all pixels:** Raw sum over entire map
+- **Cell force (scaled):** Total traction force in physical units
+- **Heatmap max/mean:** Peak and average force intensity (full field of view)
+            """)
+
+    original_vals = build_original_vals(raw_heatmap, pixel_sum, force)
+
+    pdf_bytes = create_pdf_report(
+        img, display_heatmap, raw_heatmap, pixel_sum, force, base_name, colormap_name,
+        cell_mask=cell_mask, cell_pixel_sum=cell_pixel_sum, cell_force=cell_force, cell_mean=cell_mean
+    )
+
+    btn_col1, btn_col2, btn_col3, btn_col4 = st.columns(4)
+    with btn_col1:
+        if HAS_DRAWABLE_CANVAS and measure_region_dialog is not None:
+            if st.button("Measure tool", key="open_measure", icon=":material/straighten:"):
+                st.session_state["open_measure_dialog"] = True
+                st.rerun()
+        elif HAS_DRAWABLE_CANVAS:
+            with st.expander("Measure tool"):
+                expander_cell_vals = build_cell_vals(raw_heatmap, cell_mask, pixel_sum, force) if (auto_cell_boundary and cell_mask is not None) else None
+                expander_cell_mask = cell_mask if auto_cell_boundary else None
+                render_region_canvas(
+                    display_heatmap,
+                    raw_heatmap=raw_heatmap,
+                    bf_img=img,
+                    original_vals=original_vals,
+                    cell_vals=expander_cell_vals,
+                    cell_mask=expander_cell_mask,
+                    key_suffix="expander",
+                    input_filename=key_img,
+                    colormap_name=colormap_name,
+                )
+        else:
+            st.caption("Install `streamlit-drawable-canvas-fix` for region measurement: `pip install streamlit-drawable-canvas-fix`")
+    with btn_col2:
+        st.download_button(
+            "Download heatmap",
+            width="stretch",
+            data=buf_hm.getvalue(),
+            file_name="s2f_heatmap.png",
+            mime="image/png",
+            key=f"download_heatmap{download_key_suffix}",
+            icon=":material/download:",
+        )
+    with btn_col3:
+        st.download_button(
+            "Download values",
+            width="stretch",
+            data=buf_main_csv.getvalue(),
+            file_name=f"{base_name}_main_values.csv",
+            mime="text/csv",
+            key=f"download_main_values{download_key_suffix}",
+            icon=":material/download:",
+        )
+    with btn_col4:
+        st.download_button(
+            "Download report",
+            width="stretch",
+            data=pdf_bytes,
+            file_name=f"{base_name}_report.pdf",
+            mime="application/pdf",
+            key=f"download_pdf{download_key_suffix}",
+            icon=":material/picture_as_pdf:",
+        )

S2FApp/ui/system_status.py

@@ -0,0 +1,74 @@
+"""System status UI component (CPU/memory)."""
+import streamlit as st
+
+try:
+    import psutil
+    HAS_PSUTIL = True
+except ImportError:
+    HAS_PSUTIL = False
+
+
+def _get_container_memory():
+    """
+    Read memory from cgroups when running in a container (Docker, HF Spaces).
+    psutil reports host memory in containers, which can be misleading (e.g. 128 GB vs 16 GB limit).
+    Returns (used_bytes, total_bytes) or None to fall back to psutil.
+    """
+    try:
+        # cgroup v2 (modern Docker, HF Spaces)
+        for base in ("/sys/fs/cgroup", "/sys/fs/cgroup/self"):
+            try:
+                with open(f"{base}/memory.max", "r") as f:
+                    max_val = f.read().strip()
+                if max_val == "max":
+                    return None  # No limit, use psutil
+                total = int(max_val)
+                with open(f"{base}/memory.current", "r") as f:
+                    used = int(f.read().strip())
+                return (used, total)
+            except (FileNotFoundError, ValueError):
+                continue
+        # cgroup v1
+        try:
+            with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as f:
+                total = int(f.read().strip())
+            with open("/sys/fs/cgroup/memory/memory.usage_in_bytes", "r") as f:
+                used = int(f.read().strip())
+            if total > 2**50:  # Often 9223372036854771712 when unlimited
+                return None
+            return (used, total)
+        except (FileNotFoundError, ValueError):
+            pass
+    except Exception:
+        pass
+    return None
+
+
+def render_system_status():
+    """Render CPU/memory status in the sidebar (always visible)."""
+    if not HAS_PSUTIL:
+        return
+    try:
+        cpu = psutil.cpu_percent(interval=0.1)
+        container_mem = _get_container_memory()
+        if container_mem is not None:
+            used_bytes, total_bytes = container_mem
+            mem_used_gb = used_bytes / (1024**3)
+            mem_total_gb = total_bytes / (1024**3)
+            mem_pct = 100 * used_bytes / total_bytes if total_bytes > 0 else 0
+        else:
+            mem = psutil.virtual_memory()
+            mem_used_gb = mem.used / (1024**3)
+            mem_total_gb = mem.total / (1024**3)
+            mem_pct = mem.percent
+        st.sidebar.markdown(
+            f"""
+            <div class="system-status">
+                <span class="status-dot"></span>
+                <span><strong>System</strong>&ensp;CPU {cpu:.0f}%&ensp;·&ensp;Mem {mem_pct:.0f}% ({mem_used_gb:.1f}/{mem_total_gb:.1f} GB)</span>
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
+    except Exception:
+        pass

S2FApp/utils/display.py

@@ -19,12 +19,50 @@ def cv_colormap_to_plotly_colorscale(colormap_name, n_samples=None):
     return scale
 
 
-def apply_display_scale(heatmap, mode, min_percentile=0, max_percentile=100, clip_min=0, clip_max=1):
+def build_range_colorscale(colormap_name, clip_min, clip_max, n_range_samples=32):
     """
-    Apply display scaling (Fiji/ImageJ style). Display only—does not change underlying values.
+    Build a Plotly colorscale for Range mode: normal gradient in [clip_min, clip_max],
+    the "zero" color everywhere else (0 → clip_min and clip_max → 1).
+    """
+    cv2_cmap = COLORMAPS.get(colormap_name, cv2.COLORMAP_JET)
+
+    zero_px = np.array([[0]], dtype=np.uint8)
+    zero_rgb = cv2.applyColorMap(zero_px, cv2_cmap)
+    zero_rgb = cv2.cvtColor(zero_rgb, cv2.COLOR_BGR2RGB)
+    zr, zg, zb = zero_rgb[0, 0]
+    zero_color = f"rgb({zr},{zg},{zb})"
+
+    eps = 0.0005
+    scale = []
+
+    scale.append([0.0, zero_color])
+    if clip_min > eps:
+        scale.append([clip_min - eps, zero_color])
+
+    positions = np.linspace(clip_min, clip_max, n_range_samples)
+    pixel_vals = np.clip((positions * 255).astype(np.uint8), 0, 255).reshape(1, -1)
+    rgb = cv2.applyColorMap(pixel_vals, cv2_cmap)
+    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+    for i, pos in enumerate(positions):
+        r, g, b = rgb[0, i]
+        scale.append([float(pos), f"rgb({r},{g},{b})"])
+
+    if clip_max < 1.0 - eps:
+        scale.append([clip_max + eps, zero_color])
+    scale.append([1.0, zero_color])
+
+    return scale
+
+
+def apply_display_scale(heatmap, mode, min_percentile=0, max_percentile=100,
+                        clip_min=0, clip_max=1, clip_bounds=False):
+    """
+    Apply display scaling. Display only—does not change underlying values.
     - Default: full 0–1 range as-is.
-    - Percentile: map min..max percentiles to 0..1.
-    - Range: show only values in [clip_min, clip_max]; others hidden (black).
+    - Range: keep original values inside [clip_min, clip_max].
+      clip_bounds=False → zero out outside.  clip_bounds=True → clamp to bounds.
+    - Rescale: map [clip_min, clip_max] → [0, 1].
+      clip_bounds=False → zero out outside.  clip_bounds=True → clamp to bounds first.
     """
     if mode == "Default" or mode == "Auto" or mode == "Full":
         return np.clip(heatmap, 0, 1).astype(np.float32)
@@ -36,9 +74,24 @@ def apply_display_scale(heatmap, mode, min_percentile=0, max_percentile=100, cli
             return np.clip(out, 0, 1).astype(np.float32)
         return np.clip(heatmap, 0, 1).astype(np.float32)
     if mode == "Range" or mode == "Filter" or mode == "Threshold":
+        # Range: filter (discard outside) + rescale [clip_min, clip_max] → [0, 1] so max shows as red
+        vmin, vmax = float(clip_min), float(clip_max)
+        if vmax > vmin:
+            h = heatmap.astype(np.float32)
+            if clip_bounds:
+                return np.clip(h, vmin, vmax).astype(np.float32)
+            mask = (h >= vmin) & (h <= vmax)
+            out = np.zeros_like(h)
+            out[mask] = (h[mask] - vmin) / (vmax - vmin)
+            return np.clip(out, 0, 1).astype(np.float32)
+        return np.clip(heatmap, 0, 1).astype(np.float32)
+    if mode == "Rescale":
         vmin, vmax = float(clip_min), float(clip_max)
         if vmax > vmin:
             h = heatmap.astype(np.float32)
+            if clip_bounds:
+                clamped = np.clip(h, vmin, vmax)
+                return ((clamped - vmin) / (vmax - vmin)).astype(np.float32)
             mask = (h >= vmin) & (h <= vmax)
             out = np.zeros_like(h)
             out[mask] = (h[mask] - vmin) / (vmax - vmin)

S2FApp/utils/report.py

@@ -48,27 +48,34 @@ def _pdf_image_layout(page_w_pt, page_h_pt, margin=72, n_images=2):
     }
 
 
-def heatmap_to_rgb(scaled_heatmap, colormap_name="Jet"):
-    """Convert scaled heatmap (float 0-1) to RGB array using the given colormap."""
-    heatmap_uint8 = (np.clip(scaled_heatmap, 0, 1) * 255).astype(np.uint8)
+def heatmap_to_rgb(scaled_heatmap, colormap_name="Jet", zmin=None, zmax=None):
+    """Convert scaled heatmap to RGB array using the given colormap.
+    If zmin, zmax are provided (e.g. for Range mode), map [zmin,zmax] to 0-1 for coloring."""
+    arr = np.asarray(scaled_heatmap, dtype=np.float32)
+    if zmin is not None and zmax is not None and zmax > zmin:
+        arr = np.clip((arr - zmin) / (zmax - zmin), 0, 1)
+    else:
+        arr = np.clip(arr, 0, 1)
+    heatmap_uint8 = (arr * 255).astype(np.uint8)
     cv2_colormap = COLORMAPS.get(colormap_name, cv2.COLORMAP_JET)
     heatmap_rgb = cv2.cvtColor(cv2.applyColorMap(heatmap_uint8, cv2_colormap), cv2.COLOR_BGR2RGB)
     return heatmap_rgb
 
 
-def heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name="Jet", cell_mask=None):
+def heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name="Jet", cell_mask=None, zmin=None, zmax=None):
     """Convert heatmap to RGB, optionally drawing red cell contour. Mask must match heatmap shape."""
-    heatmap_rgb = heatmap_to_rgb(scaled_heatmap, colormap_name)
+    heatmap_rgb = heatmap_to_rgb(scaled_heatmap, colormap_name, zmin=zmin, zmax=zmax)
     if cell_mask is not None and np.any(cell_mask > 0):
         contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         if contours:
-            cv2.drawContours(heatmap_rgb, contours, -1, (255, 0, 0), 2)
+            cv2.drawContours(heatmap_rgb, contours, -1, (255, 0, 0), 5)
     return heatmap_rgb
 
 
-def heatmap_to_png_bytes(scaled_heatmap, colormap_name="Jet", cell_mask=None):
-    """Convert scaled heatmap (float 0-1) to PNG bytes buffer. Optionally draw red cell contour."""
-    heatmap_rgb = heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name, cell_mask)
+def heatmap_to_png_bytes(scaled_heatmap, colormap_name="Jet", cell_mask=None, zmin=None, zmax=None):
+    """Convert scaled heatmap to PNG bytes buffer. Optionally draw red cell contour.
+    If zmin, zmax provided (Range mode), map that range to full colormap."""
+    heatmap_rgb = heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name, cell_mask, zmin=zmin, zmax=zmax)
     buf = io.BytesIO()
     Image.fromarray(heatmap_rgb).save(buf, format="PNG")
     buf.seek(0)
@@ -76,7 +83,7 @@ def heatmap_to_png_bytes(scaled_heatmap, colormap_name="Jet", cell_mask=None):
 
 
 def create_pdf_report(img, display_heatmap, raw_heatmap, pixel_sum, force, base_name, colormap_name="Jet",
-                      cell_mask=None, cell_pixel_sum=None, cell_force=None, cell_mean=None):
+                      cell_mask=None, cell_pixel_sum=None, cell_force=None, cell_mean=None, zmin=None, zmax=None):
     """Create a PDF report with input image, heatmap, and metrics."""
     buf = io.BytesIO()
     c = canvas.Canvas(buf, pagesize=A4)
@@ -108,7 +115,7 @@ def create_pdf_report(img, display_heatmap, raw_heatmap, pixel_sum, force, base_
     bf_label_w = c.stringWidth("Bright-field", "Helvetica", 9)
     c.drawString(bf_x + (img_w - bf_label_w) / 2, img_bottom - 14, "Bright-field")
 
-    heatmap_rgb = heatmap_to_rgb_with_contour(display_heatmap, colormap_name, cell_mask)
+    heatmap_rgb = heatmap_to_rgb_with_contour(display_heatmap, colormap_name, cell_mask, zmin=zmin, zmax=zmax)
     hm_buf = io.BytesIO()
     Image.fromarray(heatmap_rgb).save(hm_buf, format="PNG")
     hm_buf.seek(0)

S2FApp/utils/segmentation.py

@@ -2,7 +2,7 @@
 import numpy as np
 from scipy.ndimage import gaussian_filter
 from skimage.filters import threshold_otsu
-from skimage.morphology import binary_closing, binary_opening, binary_dilation, remove_small_objects, disk
+from skimage.morphology import closing, opening, dilation, remove_small_objects, disk
 from skimage.measure import label, regionprops
 
 
@@ -37,9 +37,9 @@ def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
     mask = (smoothed > thresh).astype(np.uint8)
 
     # Morphological cleanup
-    mask = binary_closing(mask, disk(5)).astype(np.uint8)
-    mask = binary_opening(mask, disk(3)).astype(np.uint8)
-    mask = remove_small_objects(mask.astype(bool), min_size=min_size).astype(np.uint8)
+    mask = closing(mask, disk(5)).astype(np.uint8)
+    mask = opening(mask, disk(3)).astype(np.uint8)
+    mask = remove_small_objects(mask.astype(bool), max_size=min_size - 1).astype(np.uint8)
 
     # Select component: second largest if largest is whole image
     labeled = label(mask)
@@ -60,6 +60,6 @@ def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
 
     # Dilate to include surrounding pixels
     if dilate_radius > 0:
-        mask = binary_dilation(mask, disk(dilate_radius)).astype(np.uint8)
+        mask = dilation(mask, disk(dilate_radius)).astype(np.uint8)
 
     return mask