cleaned code

app.py

@@ -20,6 +20,7 @@ from config.constants import (
     MODEL_TYPE_LABELS,
     SAMPLE_EXTENSIONS,
 )
+from utils.paths import get_ckp_base, model_subfolder
 from utils.segmentation import estimate_cell_mask
 from utils.substrate_settings import list_substrates
 from utils.display import apply_display_scale
@@ -32,11 +33,6 @@ from ui.components import (
     HAS_DRAWABLE_CANVAS,
 )
 
-try:
-    from streamlit_drawable_canvas import st_canvas
-except (ImportError, AttributeError):
-    pass  # HAS_DRAWABLE_CANVAS from ui.components
-
 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."
@@ -73,6 +69,21 @@ def _get_measure_dialog_fn():
     return measure_region_dialog if (HAS_DRAWABLE_CANVAS and ST_DIALOG) else None
 
 
+def _populate_measure_session_state(heatmap, img, pixel_sum, force, key_img, colormap_name,
+                                    display_mode, auto_cell_boundary):
+    """Populate session state for the measure tool."""
+    cell_mask = estimate_cell_mask(heatmap)
+    st.session_state["measure_raw_heatmap"] = heatmap.copy()
+    st.session_state["measure_display_mode"] = display_mode
+    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)
+    st.session_state["measure_colormap"] = colormap_name
+    st.session_state["measure_auto_cell_on"] = auto_cell_boundary
+    st.session_state["measure_cell_vals"] = build_cell_vals(heatmap, cell_mask, pixel_sum, force) if auto_cell_boundary else None
+    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")
 
 # Theme CSS (inject based on sidebar selection)
@@ -123,11 +134,7 @@ st.title("🦠 Shape2Force (S2F)")
 st.caption("Predict traction force maps from bright-field microscopy images of cells or spheroids")
 
 # Folders
-ckp_base = os.path.join(S2F_ROOT, "ckp")
-if not os.path.isdir(ckp_base):
-    project_root = os.path.dirname(S2F_ROOT)
-    if os.path.isdir(os.path.join(project_root, "ckp")):
-        ckp_base = os.path.join(project_root, "ckp")
+ckp_base = get_ckp_base(S2F_ROOT)
 ckp_single_cell = os.path.join(ckp_base, "single_cell")
 ckp_spheroid = os.path.join(ckp_base, "spheroid")
 sample_base = os.path.join(S2F_ROOT, "samples")
@@ -179,7 +186,7 @@ with st.sidebar:
 
     ckp_files = get_ckp_files_for_model(model_type)
     ckp_folder = ckp_single_cell if model_type == "single_cell" else ckp_spheroid
-    ckp_subfolder_name = "single_cell" if model_type == "single_cell" else "spheroid"
+    ckp_subfolder_name = model_subfolder(model_type)
 
     if ckp_files:
         checkpoint = st.selectbox(
@@ -265,7 +272,7 @@ if img_source == "Upload":
 else:
     sample_files = get_sample_files_for_model(model_type)
     sample_folder = sample_single_cell if model_type == "single_cell" else sample_spheroid
-    sample_subfolder_name = "single_cell" if model_type == "single_cell" else "spheroid"
+    sample_subfolder_name = model_subfolder(model_type)
     if sample_files:
         selected_sample = st.selectbox(
             f"Select example image (from `samples/{sample_subfolder_name}/`)",
@@ -345,16 +352,10 @@ if just_ran:
                 "pixel_sum": pixel_sum,
                 "cache_key": cache_key,
             }
-            st.session_state["measure_raw_heatmap"] = heatmap.copy()
-            st.session_state["measure_display_mode"] = display_mode
-            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)
-            st.session_state["measure_colormap"] = colormap_name
-            cell_mask = estimate_cell_mask(heatmap)
-            st.session_state["measure_auto_cell_on"] = auto_cell_boundary
-            st.session_state["measure_cell_vals"] = build_cell_vals(heatmap, cell_mask, pixel_sum, force) if auto_cell_boundary else None
-            st.session_state["measure_cell_mask"] = cell_mask if auto_cell_boundary else None
+            _populate_measure_session_state(
+                heatmap, img, pixel_sum, force, key_img, colormap_name,
+                display_mode, auto_cell_boundary,
+            )
 
             render_result_display(
                 img, heatmap, display_heatmap, pixel_sum, force, key_img,
@@ -373,16 +374,10 @@ elif has_cached:
     img, heatmap, force, pixel_sum = r["img"], r["heatmap"], r["force"], r["pixel_sum"]
     display_heatmap = apply_display_scale(heatmap, display_mode)
 
-    st.session_state["measure_raw_heatmap"] = heatmap.copy()
-    st.session_state["measure_display_mode"] = display_mode
-    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)
-    st.session_state["measure_colormap"] = colormap_name
-    cell_mask = estimate_cell_mask(heatmap)
-    st.session_state["measure_auto_cell_on"] = auto_cell_boundary
-    st.session_state["measure_cell_vals"] = build_cell_vals(heatmap, cell_mask, pixel_sum, force) if auto_cell_boundary else None
-    st.session_state["measure_cell_mask"] = cell_mask if auto_cell_boundary else None
+    _populate_measure_session_state(
+        heatmap, img, pixel_sum, force, key_img, colormap_name,
+        display_mode, auto_cell_boundary,
+    )
 
     if st.session_state.pop("open_measure_dialog", False):
         measure_region_dialog()

predictor.py

@@ -14,6 +14,7 @@ if S2F_ROOT not in sys.path:
     sys.path.insert(0, S2F_ROOT)
 
 from models.s2f_model import create_s2f_model
+from utils.paths import get_ckp_base, model_subfolder
 from utils.substrate_settings import get_settings_of_category, compute_settings_normalization
 from utils import config
 
@@ -89,12 +90,8 @@ class S2FPredictor:
         """
         self.model_type = model_type
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        ckp_base = os.path.join(S2F_ROOT, "ckp")
-        if not os.path.isdir(ckp_base):
-            project_root = os.path.dirname(S2F_ROOT)
-            if os.path.isdir(os.path.join(project_root, "ckp")):
-                ckp_base = os.path.join(project_root, "ckp")
-        subfolder = "single_cell" if model_type == "single_cell" else "spheroid"
+        ckp_base = get_ckp_base(S2F_ROOT)
+        subfolder = model_subfolder(model_type)
         ckp_dir = ckp_folder if ckp_folder else os.path.join(ckp_base, subfolder)
         if not os.path.isdir(ckp_dir):
             ckp_dir = ckp_base  # fallback if subfolders not used

ui/components.py

@@ -17,7 +17,7 @@ from config.constants import (
     TOOL_LABELS,
 )
 from utils.display import apply_display_scale, cv_colormap_to_plotly_colorscale
-from utils.report import heatmap_to_rgb, heatmap_to_png_bytes, create_pdf_report
+from utils.report import heatmap_to_rgb, heatmap_to_rgb_with_contour, heatmap_to_png_bytes, create_pdf_report
 from utils.segmentation import estimate_cell_mask
 
 try:
@@ -102,17 +102,6 @@ def _obj_to_pts(obj, scale_x, scale_y, heatmap_w, heatmap_h):
     return pts
 
 
-def parse_canvas_shapes_to_mask(json_data, canvas_h, canvas_w, heatmap_h, heatmap_w):
-    """Parse drawn shapes from streamlit-drawable-canvas json_data and create binary mask (combined)."""
-    masks, _ = parse_canvas_shapes_to_masks(json_data, canvas_h, canvas_w, heatmap_h, heatmap_w)
-    if not masks:
-        return None, 0
-    combined = np.zeros((heatmap_h, heatmap_w), dtype=np.uint8)
-    for m in masks:
-        combined = np.maximum(combined, m)
-    return combined, len(masks)
-
-
 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"]:
@@ -237,9 +226,7 @@ def render_region_canvas(display_heatmap, raw_heatmap=None, bf_img=None, origina
     """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(display_heatmap, colormap_name)
-    if cell_mask is not None and np.any(cell_mask > 0):
-        heatmap_rgb = _draw_contour_on_image(heatmap_rgb.copy(), cell_mask, stroke_color=(255, 0, 0), stroke_width=2)
+    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("""

utils/paths.py

@@ -0,0 +1,17 @@
+"""Path resolution utilities for S2F App."""
+import os
+
+
+def get_ckp_base(root):
+    """Resolve checkpoint base directory (S2FApp/ckp or project/ckp)."""
+    ckp_base = os.path.join(root, "ckp")
+    if not os.path.isdir(ckp_base):
+        project_root = os.path.dirname(root)
+        if os.path.isdir(os.path.join(project_root, "ckp")):
+            ckp_base = os.path.join(project_root, "ckp")
+    return ckp_base
+
+
+def model_subfolder(model_type):
+    """Return subfolder name for model type: 'single_cell' or 'spheroid'."""
+    return "single_cell" if model_type == "single_cell" else "spheroid"

utils/report.py

@@ -21,13 +21,19 @@ def heatmap_to_rgb(scaled_heatmap, colormap_name="Jet"):
     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."""
+def heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name="Jet", cell_mask=None):
+    """Convert heatmap to RGB, optionally drawing red cell contour. Mask must match heatmap shape."""
     heatmap_rgb = heatmap_to_rgb(scaled_heatmap, colormap_name)
     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)
+    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)
     buf = io.BytesIO()
     Image.fromarray(heatmap_rgb).save(buf, format="PNG")
     buf.seek(0)
@@ -62,11 +68,7 @@ def create_pdf_report(img, display_heatmap, raw_heatmap, pixel_sum, force, base_
     c.setFont("Helvetica", 9)
     c.drawString(72, img_top - img_h - 12, "Input: Bright-field")
 
-    heatmap_rgb = heatmap_to_rgb(display_heatmap, colormap_name)
-    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)
+    heatmap_rgb = heatmap_to_rgb_with_contour(display_heatmap, colormap_name, cell_mask)
     hm_buf = io.BytesIO()
     Image.fromarray(heatmap_rgb).save(hm_buf, format="PNG")
     hm_buf.seek(0)