added

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+*.tif filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+Archive/
+.DS_Store

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "S2FApp"]
+    path = S2FApp
+    url = git@hf.co:spaces/kaveh/Shape2force

README.md

@@ -0,0 +1,75 @@
+# Shape2Force (S2F)
+
+Predict force maps from bright-field microscopy images of single-cell or spheroid using deep learning. 
+
+**Web App:** The app is published to [Hugging Face Spaces](https://huggingface.co/spaces/kaveh/Shape2force). To work on it locally: `git clone git@hf.co:spaces/kaveh/Shape2force S2FApp`
+
+---
+
+## Quick Start
+
+**Web app (local):**
+```bash
+cd S2FApp
+pip install -r requirements.txt
+streamlit run app.py
+```
+
+Or use the [online app](https://huggingface.co/spaces/kaveh/Shape2force) on Hugging Face. Place checkpoints (`.pth`) in `S2FApp/ckp/` for __local use__; the Space downloads them automatically.
+
+---
+
+## Ways to Use S2F
+
+### 1. Web App
+
+Run the Streamlit GUI from `S2FApp/`:
+
+```bash
+cd S2FApp && streamlit run app.py
+```
+
+1. Choose **Model type**: Single cell or Spheroid
+2. Select a **Checkpoint** from `ckp/`
+3. For single-cell: pick **Substrate** (e.g. fibroblasts_PDMS)
+4. Upload an image or pick from `sample/`
+5. Click **Run prediction**
+
+Output: heatmap, cell force (sum), and basic stats.
+
+----
+
+### 2. Jupyter Notebook
+
+For interactive usage and custom analysis, you may use the notebook:
+
+- **`notebooks/evaluate_model.ipynb`** – Load data, run evaluation, plot predictions, and save per-sample metrics.
+
+Once cloned the repo. open the notebook in Jupyter and adjust the configuration cell (paths, model type, substrate).
+
+---
+
+### 3. Training & Fine-Tuning
+
+**Dataset layout:** A folder with `train/` and `test/` subfolders. Each subfolder has:
+- `BF_001.tif` (bright-field image)
+- `*_gray.jpg` (force map / heatmap)
+- Optional `.txt` (cell_area, sum_force)
+
+**Single-cell:**
+```bash
+python -m training.train --data path/to/dataset --model single_cell --epochs 100 --substrate fibroblasts_PDMS
+```
+
+**Spheroid:**
+```bash
+python -m training.train --data path/to/dataset --model spheroid --epochs 100
+```
+
+**Resume / fine-tune from checkpoint:**
+```bash
+python -m training.train --data path/to/dataset --model single_cell --resume ckp/last_checkpoint.pth --epochs 150
+```
+
+---
+

S2FApp/.dockerignore

@@ -0,0 +1,8 @@
+.git
+.gitignore
+__pycache__
+*.py[cod]
+.venv
+venv
+.DS_Store
+ckp/*.pth

S2FApp/.gitattributes

@@ -0,0 +1,2 @@
+*.tif filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text

S2FApp/.gitignore

@@ -0,0 +1,8 @@
+__pycache__
+*.py[cod]
+.venv
+venv
+.DS_Store
+ckp/*.pth
+sample/*.tif
+sample/*.tiff

S2FApp/Dockerfile

@@ -0,0 +1,62 @@
+# Shape2Force (S2F) - Hugging Face Spaces
+FROM python:3.10-slim
+
+# Create user for HF Spaces (runs as UID 1000)
+RUN useradd -m -u 1000 user
+
+WORKDIR /app
+
+# Install system deps for OpenCV
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Install Python dependencies (exclude heavy training deps for smaller image)
+RUN pip install --no-cache-dir \
+    torch torchvision \
+    numpy opencv-python streamlit matplotlib Pillow plotly \
+    huggingface_hub
+
+# Copy app code (chown for HF Spaces permissions)
+COPY --chown=user:user app.py predictor.py ./
+COPY --chown=user:user models/ models/
+COPY --chown=user:user utils/ utils/
+COPY --chown=user:user config/ config/
+COPY --chown=user:user sample/ sample/
+RUN mkdir -p ckp && chown user:user ckp
+
+# Download checkpoints from Hugging Face if ckp is empty (for Space deployment)
+# Set HF_MODEL_REPO env to your model repo, e.g. kaveh/Shape2Force
+ARG HF_MODEL_REPO=kaveh/Shape2Force
+ENV HF_MODEL_REPO=${HF_MODEL_REPO}
+
+RUN python -c "
+import os
+from pathlib import Path
+ckp = Path('/app/ckp')
+if not list(ckp.glob('*.pth')):
+    try:
+        from huggingface_hub import hf_hub_download, list_repo_files
+        repo = os.environ.get('HF_MODEL_REPO', 'kaveh/Shape2Force')
+        files = list_repo_files(repo)
+        pth_files = [f for f in files if f.startswith('ckp/') and f.endswith('.pth')]
+        for f in pth_files:
+            hf_hub_download(repo_id=repo, filename=f, local_dir='/app')
+        print('Downloaded checkpoints from', repo)
+    except Exception as e:
+        print('Could not download checkpoints:', e)
+else:
+    print('Checkpoints already present')
+"
+
+# Ensure ckp contents are readable by user
+RUN chown -R user:user ckp
+
+USER user
+
+EXPOSE 8501
+CMD ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]

S2FApp/README.md

@@ -0,0 +1,37 @@
+---
+tags:
+- cell-mechanobiology
+- microscopy
+- image-to-image
+- pytorch
+license: cc-by-4.0
+sdk: docker
+app_port: 8501
+---
+
+# Shape2Force (S2F) App
+
+Predict force maps from bright-field microscopy images using deep learning. 
+
+## Quick Start
+
+```bash
+pip install -r requirements.txt
+streamlit run app.py
+```
+
+Checkpoints are downloaded automatically from the [Shape2Force model repo](https://huggingface.co/kaveh/Shape2Force) when running in Docker. For local use, place `.pth` files in `ckp/`.
+
+## Usage
+
+1. Choose **Model type**: Single cell or Spheroid
+2. Select a **Checkpoint** from `ckp/`
+3. For single-cell: pick **Substrate** (e.g. fibroblasts_PDMS)
+4. Upload an image or pick from `sample/`
+5. Click **Run prediction**
+
+Output: heatmap, cell force (sum), and basic stats.
+
+## Full Project
+
+For training, evaluation, and notebooks, see the main [Shape2Force repository](https://github.com/Angione-Lab/Shape2Force).

S2FApp/app.py

@@ -0,0 +1,205 @@
+"""
+Shape2Force (S2F) - GUI for force map prediction from bright field microscopy images.
+"""
+import os
+import sys
+import io
+import cv2
+cv2.utils.logging.setLogLevel(cv2.utils.logging.LOG_LEVEL_ERROR)
+
+import numpy as np
+import streamlit as st
+from PIL import Image
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+# Ensure S2F is in path
+S2F_ROOT = os.path.dirname(os.path.abspath(__file__))
+if S2F_ROOT not in sys.path:
+    sys.path.insert(0, S2F_ROOT)
+
+from predictor import S2FPredictor
+from utils.substrate_settings import list_substrates
+
+st.set_page_config(page_title="Shape2Force (S2F)", page_icon="🔬", layout="centered")
+st.markdown("""
+    <style>
+    section[data-testid="stSidebar"] { width: 380px !important; }
+    </style>
+    """, unsafe_allow_html=True)
+st.title("🔬 Shape2Force (S2F)")
+st.caption("Predict force maps from bright field microscopy images")
+
+# Folders
+ckp_folder = os.path.join(S2F_ROOT, "ckp")
+sample_folder = os.path.join(S2F_ROOT, "sample")
+ckp_files = []
+if os.path.isdir(ckp_folder):
+    ckp_files = sorted([f for f in os.listdir(ckp_folder) if f.endswith(".pth")])
+SAMPLE_EXTENSIONS = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
+sample_files = []
+if os.path.isdir(sample_folder):
+    sample_files = sorted([f for f in os.listdir(sample_folder)
+                          if f.lower().endswith(SAMPLE_EXTENSIONS)])
+
+# Sidebar: model configuration
+with st.sidebar:
+    st.header("Model configuration")
+    model_type = st.radio(
+        "Model type",
+        ["single_cell", "spheroid"],
+        format_func=lambda x: "Single cell" if x == "single_cell" else "Spheroid",
+        horizontal=False,
+    )
+
+    if ckp_files:
+        checkpoint = st.selectbox(
+            "Checkpoint",
+            ckp_files,
+            help="Select a .pth file from the ckp folder",
+        )
+    else:
+        st.warning("No .pth files in ckp/ folder. Add checkpoints to load.")
+        checkpoint = None
+
+    substrate_config = None
+    substrate_val = "fibroblasts_PDMS"
+    use_manual = False
+    if model_type == "single_cell":
+        try:
+            substrates = list_substrates()
+            substrate_val = st.selectbox(
+                "Substrate (from config)",
+                substrates,
+                help="Select a preset from config/substrate_settings.json",
+            )
+            use_manual = st.checkbox("Enter substrate values manually", value=False)
+            if use_manual:
+                st.caption("Enter pixelsize (µm/px) and Young's modulus (Pa)")
+                manual_pixelsize = st.number_input("Pixelsize (µm/px)", min_value=0.1, max_value=50.0,
+                                                   value=3.0769, step=0.1, format="%.4f")
+                manual_young = st.number_input("Young's modulus (Pa)", min_value=100.0, max_value=100000.0,
+                                               value=6000.0, step=100.0, format="%.0f")
+                substrate_config = {"pixelsize": manual_pixelsize, "young": manual_young}
+            else:
+                substrate_config = None
+        except FileNotFoundError:
+            st.error("config/substrate_settings.json not found")
+
+    st.divider()
+    st.subheader("Display")
+    display_size = st.slider("Image size (px)", min_value=200, max_value=800, value=350, step=50,
+                             help="Adjust display size. Drag to pan, scroll to zoom.")
+
+    st.divider()
+
+# Main area: image input
+img_source = st.radio("Image source", ["Upload", "Sample"], horizontal=True, label_visibility="collapsed")
+img = None
+uploaded = None
+selected_sample = None
+
+if img_source == "Upload":
+    uploaded = st.file_uploader(
+        "Upload bright field image",
+        type=["tif", "tiff", "png", "jpg", "jpeg"],
+        help="Bright field microscopy image (grayscale or RGB)",
+    )
+    if uploaded:
+        bytes_data = uploaded.read()
+        nparr = np.frombuffer(bytes_data, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
+        uploaded.seek(0)  # reset for potential re-read
+else:
+    if sample_files:
+        selected_sample = st.selectbox(
+            "Select sample image",
+            sample_files,
+            format_func=lambda x: x,
+        )
+        if selected_sample:
+            sample_path = os.path.join(sample_folder, selected_sample)
+            img = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)
+        # Show sample thumbnails
+        st.caption("Sample images (add more to the `sample/` folder)")
+        n_cols = min(4, len(sample_files))
+        cols = st.columns(n_cols)
+        for i, fname in enumerate(sample_files[:8]):  # show up to 8
+            with cols[i % n_cols]:
+                path = os.path.join(sample_folder, fname)
+                sample_img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+                if sample_img is not None:
+                    st.image(sample_img, caption=fname, width='content')
+    else:
+        st.info("No sample images found. Add images to the `sample/` folder, or use Upload.")
+
+run = st.button("Run prediction", type="primary")
+has_image = img is not None
+
+if run and checkpoint and has_image:
+    with st.spinner("Loading model and predicting..."):
+        try:
+            predictor = S2FPredictor(
+                model_type=model_type,
+                checkpoint_path=checkpoint,
+                ckp_folder=ckp_folder,
+            )
+            if img is not None:
+                sub_val = substrate_val if model_type == "single_cell" and not use_manual else "fibroblasts_PDMS"
+                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!")
+
+                # Metrics
+                col1, col2, col3, col4 = st.columns(4)
+                with col1:
+                    st.metric("Sum of all pixels", f"{pixel_sum:.2f}")
+                with col2:
+                    st.metric("Cell force (scaled)", f"{force:.2f}")
+                with col3:
+                    st.metric("Heatmap max", f"{np.max(heatmap):.4f}")
+                with col4:
+                    st.metric("Heatmap mean", f"{np.mean(heatmap):.4f}")
+
+                # Visualization - Plotly with zoom/pan
+                fig_pl = make_subplots(rows=1, cols=2, subplot_titles=["", ""])
+                fig_pl.add_trace(go.Heatmap(z=img, colorscale="gray", showscale=False), row=1, col=1)
+                fig_pl.add_trace(go.Heatmap(z=heatmap, colorscale="Jet", zmin=0, zmax=1, showscale=True), row=1, col=2)
+                fig_pl.update_layout(
+                    height=display_size,
+                    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)
+
+                # Download
+                heatmap_uint8 = (np.clip(heatmap, 0, 1) * 255).astype(np.uint8)
+                heatmap_rgb = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
+                heatmap_rgb = cv2.cvtColor(heatmap_rgb, cv2.COLOR_BGR2RGB)
+                pil_heatmap = Image.fromarray(heatmap_rgb)
+                buf_hm = io.BytesIO()
+                pil_heatmap.save(buf_hm, format="PNG")
+                buf_hm.seek(0)
+                st.download_button("Download Heatmap", data=buf_hm.getvalue(),
+                                   file_name="s2f_heatmap.png", mime="image/png")
+
+        except Exception as e:
+            st.error(f"Prediction failed: {e}")
+            import traceback
+            st.code(traceback.format_exc())
+
+elif run and not checkpoint:
+    st.warning("Please add checkpoint files to the ckp/ folder and select one.")
+elif run and not has_image:
+    st.warning("Please upload an image or select a sample.")
+
+# Footer
+st.sidebar.divider()
+st.sidebar.caption("Place .pth checkpoints in the ckp/ folder")

S2FApp/config/substrate_settings.json

@@ -0,0 +1 @@
+{"substrates":{"fibroblasts_PDMS":{"name":"Fibroblasts on PDMS (6 kPa)","pixelsize":3.0769,"young":6000},"U2OS_PDMS":{"name":"U2OS cells on PDMS (6 kPa)","pixelsize":6.1538,"young":6000},"PDMS_1kPa":{"name":"PDMS soft hydrogel (1 kPa, 10 µm/px)","pixelsize":9.8138,"young":1000},"PDMS_10kPa":{"name":"PDMS stiff hydrogel (10 kPa, 10 µm/px)","pixelsize":9.8138,"young":10000},"PDMS_1kPa_3um":{"name":"PDMS soft hydrogel (1 kPa, 3 µm/px)","pixelsize":3.0769,"young":1000},"PDMS_10kPa_3um":{"name":"PDMS stiff hydrogel (10 kPa, 3 µm/px)","pixelsize":3.0769,"young":10000}},"default_substrate":"fibroblasts_PDMS"}

S2FApp/models/__init__.py

@@ -0,0 +1 @@
+from .s2f_model import create_s2f_model, S2FGenerator

S2FApp/models/blocks.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+import torch
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = nn.Conv2d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else None
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        if self.downsample:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        return out

S2FApp/models/cbam.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes, ratio=16):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc = nn.Sequential(
+            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
+            nn.ReLU(),
+            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(self.max_pool(x))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+        padding = kernel_size // 2
+        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv(x)
+        return self.sigmoid(x)
+
+
+class CBAM(nn.Module):
+    def __init__(self, in_planes, ratio=16, kernel_size=7):
+        super(CBAM, self).__init__()
+        self.channel_attention = ChannelAttention(in_planes, ratio)
+        self.spatial_attention = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        x = x * self.channel_attention(x)
+        x = x * self.spatial_attention(x)
+        return x

S2FApp/models/s2f_model.py

@@ -0,0 +1,435 @@
+"""
+S2F (Shape2Force) model for force map prediction (inference only).
+Supports single-cell and spheroid modes.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .blocks import ResidualBlock
+from .cbam import CBAM
+
+from utils import config
+from utils.substrate_settings import (
+    get_settings_of_category,
+    compute_settings_normalization,
+    load_substrate_config,
+)
+
+
+def normalize_settings(substrate_name, normalization_params, config=None, config_path=None):
+    """
+    Normalize settings for a given substrate.
+
+    Args:
+        substrate_name (str): Name of the substrate
+        normalization_params (dict): Normalization parameters
+
+    Returns:
+        tuple: (normalized_pixelsize, normalized_young)
+    """
+    settings = get_settings_of_category(substrate_name, config=config, config_path=config_path)
+
+    # Min-max normalization to [0, 1]
+    pixelsize_norm = (settings['pixelsize'] - normalization_params['pixelsize']['min']) / \
+                     (normalization_params['pixelsize']['max'] - normalization_params['pixelsize']['min'])
+
+    young_norm = (settings['young'] - normalization_params['young']['min']) / \
+                 (normalization_params['young']['max'] - normalization_params['young']['min'])
+
+    return pixelsize_norm, young_norm
+
+
+def create_settings_channels(metadata, normalization_params, device, image_shape, config_path=None):
+    """
+    Create settings channels for a batch of images.
+
+    Args:
+        metadata (dict): Batch metadata containing substrate information
+        normalization_params (dict): Normalization parameters
+        device: Device to create tensors on
+        image_shape (tuple): Shape of input images (B, C, H, W)
+
+    Returns:
+        torch.Tensor: Settings channels [B, 2, H, W] where channels are [pixelsize, young]
+    """
+    batch_size, _, height, width = image_shape
+
+    # Create settings channels
+    pixelsize_channel = torch.zeros(batch_size, 1, height, width, device=device)
+    young_channel = torch.zeros(batch_size, 1, height, width, device=device)
+
+    for i in range(batch_size):
+        substrate = metadata['substrate'][i]
+        pixelsize_norm, young_norm = normalize_settings(
+            substrate, normalization_params, config_path=config_path
+        )
+
+        # Fill entire channel with normalized value
+        pixelsize_channel[i, 0] = pixelsize_norm
+        young_channel[i, 0] = young_norm
+
+    # Concatenate channels
+    settings_channels = torch.cat([pixelsize_channel, young_channel], dim=1)  # [B, 2, H, W]
+
+    return settings_channels
+
+
+class GlobalContextModule(nn.Module):
+    """Global context module for capturing cell shape information"""
+    def __init__(self, in_channels):
+        super().__init__()
+        self.global_pool = nn.AdaptiveAvgPool2d(1)
+        self.global_conv = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels//4, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels//4, in_channels, 1),
+            nn.Sigmoid()
+        )
+        self.large_kernel = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels, 3, padding=1, groups=in_channels),
+            nn.Conv2d(in_channels, in_channels, 1),
+            nn.BatchNorm2d(in_channels),
+            nn.ReLU(inplace=True)
+        )
+        self.multi_scale = nn.ModuleList([
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=1, dilation=1),
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=2, dilation=2),
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=4, dilation=4),
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=8, dilation=8)
+        ])
+        self.fusion = nn.Conv2d(in_channels, in_channels, 1)
+
+    def forward(self, x):
+        global_ctx = self.global_pool(x)
+        global_weight = self.global_conv(global_ctx)
+        large_features = self.large_kernel(x)
+        multi_scale_features = []
+        for conv in self.multi_scale:
+            multi_scale_features.append(conv(x))
+        multi_scale_out = torch.cat(multi_scale_features, dim=1)
+        multi_scale_out = self.fusion(multi_scale_out)
+        return x + (large_features * global_weight) + multi_scale_out
+
+
+class HierarchicalAttention(nn.Module):
+    """Hierarchical attention combining spatial and channel attention"""
+    def __init__(self, channels):
+        super().__init__()
+        self.spatial_att = nn.Sequential(
+            nn.Conv2d(channels, channels//8, 1),
+            nn.Conv2d(channels//8, 1, 3, padding=1),
+            nn.Sigmoid()
+        )
+        self.channel_att = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(channels, channels//16, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels//16, channels, 1),
+            nn.Sigmoid()
+        )
+        self.cross_att = nn.Sequential(
+            nn.Conv2d(channels, channels//4, 1),
+            nn.BatchNorm2d(channels//4),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels//4, channels, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        spatial_weight = self.spatial_att(x)
+        channel_weight = self.channel_att(x)
+        attended = x * spatial_weight * channel_weight
+        cross_weight = self.cross_att(attended)
+        return x + (attended * cross_weight)
+
+
+class EnhancedAttentionGate(nn.Module):
+    """Enhanced attention gate with global context"""
+    def __init__(self, F_g, F_l, F_int):
+        super().__init__()
+        self.W_g = nn.Sequential(
+            nn.Conv2d(F_g, F_int, kernel_size=1),
+            nn.BatchNorm2d(F_int)
+        )
+        self.W_x = nn.Sequential(
+            nn.Conv2d(F_l, F_int, kernel_size=1),
+            nn.BatchNorm2d(F_int)
+        )
+        self.psi = nn.Sequential(
+            nn.ReLU(inplace=True),
+            nn.Conv2d(F_int, F_int//2, kernel_size=3, padding=1),
+            nn.BatchNorm2d(F_int//2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(F_int//2, 1, kernel_size=1),
+            nn.Sigmoid()
+        )
+        self.global_context = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(F_l, F_int//4, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(F_int//4, 1, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, g, x):
+        g1 = self.W_g(g)
+        x1 = self.W_x(x)
+        if g1.shape[2:] != x1.shape[2:]:
+            g1 = F.interpolate(g1, size=x1.shape[2:], mode='bilinear', align_corners=False)
+        psi = self.psi(g1 + x1)
+        global_weight = self.global_context(x)
+        psi = psi * global_weight
+        if psi.shape[2:] != x.shape[2:]:
+            psi = F.interpolate(psi, size=x.shape[2:], mode='bilinear', align_corners=False)
+        return x * psi
+
+
+class S2FGenerator(nn.Module):
+    """
+    S2F (Shape2Force) model: U-Net generator for force map prediction.
+    Supports substrate-specific settings as additional input channels.
+    """
+    def __init__(self,
+                 in_channels=1,
+                 out_channels=1,
+                 img_size=1024,
+                 bridge_type='cbam',
+                 use_multi_scale_input=True):
+        super().__init__()
+
+        self.img_size = img_size
+        self.bridge_type = bridge_type
+        self.use_multi_scale_input = use_multi_scale_input
+
+        if self.use_multi_scale_input:
+            self.scale_pyramid = nn.ModuleList([
+                nn.Conv2d(in_channels, 32, 3, padding=1),
+                nn.Sequential(
+                    nn.AvgPool2d(2, stride=2),
+                    nn.Conv2d(in_channels, 32, 3, padding=1)
+                ),
+                nn.Sequential(
+                    nn.AvgPool2d(4, stride=4),
+                    nn.Conv2d(in_channels, 32, 3, padding=1)
+                )
+            ])
+            self.initial_conv = nn.Conv2d(96, 64, 1)
+        else:
+            self.initial_conv = nn.Conv2d(in_channels, 64, 3, padding=1)
+
+        def enhanced_conv_block(in_c, out_c, use_attention=True):
+            layers = [
+                nn.Conv2d(in_c, out_c, 3, padding=1),
+                nn.BatchNorm2d(out_c),
+                nn.ReLU(inplace=True),
+                ResidualBlock(out_c, out_c)
+            ]
+            if use_attention:
+                layers.append(HierarchicalAttention(out_c))
+            return nn.Sequential(*layers)
+
+        def dilated_conv_block(in_c, out_c, use_global_context=False):
+            layers = [
+                nn.Conv2d(in_c, out_c, 3, padding=2, dilation=2),
+                nn.BatchNorm2d(out_c),
+                nn.ReLU(inplace=True),
+                ResidualBlock(out_c, out_c)
+            ]
+            if use_global_context:
+                layers.append(GlobalContextModule(out_c))
+            return nn.Sequential(*layers)
+
+        self.encoder1 = enhanced_conv_block(64, 64, use_attention=False)
+        self.pool1 = nn.MaxPool2d(2)
+        self.encoder2 = enhanced_conv_block(64, 128, use_attention=True)
+        self.pool2 = nn.MaxPool2d(2)
+        self.encoder3 = dilated_conv_block(128, 256, use_global_context=True)
+        self.pool3 = nn.MaxPool2d(2)
+        self.encoder4 = dilated_conv_block(256, 512, use_global_context=True)
+        self.pool4 = nn.MaxPool2d(2)
+
+        if bridge_type == 'cbam':
+            self.bridge = nn.Sequential(
+                dilated_conv_block(512, 1024, use_global_context=True),
+                CBAM(1024),
+                GlobalContextModule(1024),
+                HierarchicalAttention(1024)
+            )
+        else:
+            self.bridge = nn.Sequential(
+                dilated_conv_block(512, 1024, use_global_context=True),
+                GlobalContextModule(1024),
+                HierarchicalAttention(1024)
+            )
+
+        self.att4 = EnhancedAttentionGate(512, 512, 256)
+        self.att3 = EnhancedAttentionGate(256, 256, 128)
+        self.att2 = EnhancedAttentionGate(128, 128, 64)
+        self.att1 = EnhancedAttentionGate(64, 64, 32)
+
+        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
+        self.dec4 = enhanced_conv_block(1024, 512, use_attention=True)
+        self.refine4 = HierarchicalAttention(512)
+        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.dec3 = enhanced_conv_block(512, 256, use_attention=True)
+        self.refine3 = HierarchicalAttention(256)
+        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.dec2 = enhanced_conv_block(256, 128, use_attention=True)
+        self.refine2 = HierarchicalAttention(128)
+        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.dec1 = enhanced_conv_block(128, 64, use_attention=True)
+        self.refine1 = HierarchicalAttention(64)
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(64, 32, 3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, out_channels, 1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        if self.use_multi_scale_input:
+            scale_features = []
+            for i, scale_conv in enumerate(self.scale_pyramid):
+                if i == 0:
+                    scale_features.append(scale_conv(x))
+                else:
+                    scale_out = scale_conv(x)
+                    scale_out = F.interpolate(scale_out, size=x.shape[2:], mode='bilinear', align_corners=False)
+                    scale_features.append(scale_out)
+            fused = torch.cat(scale_features, dim=1)
+            initial_features = self.initial_conv(fused)
+        else:
+            initial_features = self.initial_conv(x)
+
+        e1 = self.encoder1(initial_features)
+        e2 = self.encoder2(self.pool1(e1))
+        e3 = self.encoder3(self.pool2(e2))
+        e4 = self.encoder4(self.pool3(e3))
+        b = self.bridge(self.pool4(e4))
+
+        g4 = self.up4(b)
+        x4 = self.att4(g4, e4)
+        d4 = self.dec4(torch.cat([g4, x4], dim=1))
+        d4 = self.refine4(d4)
+        g3 = self.up3(d4)
+        x3 = self.att3(g3, e3)
+        d3 = self.dec3(torch.cat([g3, x3], dim=1))
+        d3 = self.refine3(d3)
+        g2 = self.up2(d3)
+        x2 = self.att2(g2, e2)
+        d2 = self.dec2(torch.cat([g2, x2], dim=1))
+        d2 = self.refine2(d2)
+        g1 = self.up1(d2)
+        x1 = self.att1(g1, e1)
+        d1 = self.dec1(torch.cat([g1, x1], dim=1))
+        d1 = self.refine1(d1)
+        out = self.final_conv(d1)
+        return out
+
+    def load_checkpoint_with_expansion(self, checkpoint_path, strict=False):
+        """Load checkpoint and expand from 1-channel to 3-channel if needed."""
+        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+        generator_state = checkpoint['generator_state_dict']
+        needs_expansion = False
+
+        if 'scale_pyramid.0.weight' in generator_state:
+            old_shape = generator_state['scale_pyramid.0.weight'].shape
+            current_shape = self.scale_pyramid[0].weight.shape
+            if old_shape[1] != current_shape[1]:
+                needs_expansion = True
+        elif 'initial_conv.weight' in generator_state:
+            old_shape = generator_state['initial_conv.weight'].shape
+            current_shape = self.initial_conv.weight.shape
+            if old_shape[1] != current_shape[1]:
+                needs_expansion = True
+
+        if needs_expansion:
+            generator_state = self._expand_generator_state(generator_state)
+
+        self.load_state_dict(generator_state, strict=strict)
+        return checkpoint
+
+    def _expand_generator_state(self, generator_state):
+        """Expand generator state dict from 1-channel to 3-channel input."""
+        expanded_state = generator_state.copy()
+        if 'scale_pyramid.0.weight' in generator_state:
+            for i in range(3):
+                key = f'scale_pyramid.{i}.weight' if i == 0 else f'scale_pyramid.{i}.1.weight'
+                if key in generator_state:
+                    old_weight = generator_state[key]
+                    new_weight = torch.zeros(32, 3, 3, 3)
+                    new_weight[:, 0:1, :, :] = old_weight
+                    expanded_state[key] = new_weight
+        elif 'initial_conv.weight' in generator_state:
+            old_weight = generator_state['initial_conv.weight']
+            new_weight = torch.zeros(64, 3, 3, 3)
+            new_weight[:, 0:1, :, :] = old_weight
+            expanded_state['initial_conv.weight'] = new_weight
+        return expanded_state
+
+
+class PatchGANDiscriminator(nn.Module):
+    """PatchGAN Discriminator (included for create_s2f_model compatibility)."""
+    def __init__(self, in_channels=2, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        super().__init__()
+        use_bias = norm_layer == nn.InstanceNorm2d
+        self.initial_conv = nn.Sequential(
+            nn.Conv2d(in_channels, ndf, kernel_size=4, stride=2, padding=1, bias=use_bias),
+            nn.LeakyReLU(0.2, inplace=True)
+        )
+        self.layers = nn.ModuleList()
+        nf_mult, nf_mult_prev = 1, 1
+        for n in range(1, n_layers):
+            nf_mult_prev, nf_mult = nf_mult, min(2 ** n, 8)
+            self.layers.append(nn.Sequential(
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=4, stride=2, padding=1, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, inplace=True)
+            ))
+        nf_mult_prev, nf_mult = nf_mult, min(2 ** n_layers, 8)
+        self.layers.append(nn.Sequential(
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=4, stride=1, padding=1, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, inplace=True)
+        ))
+        self.output_conv = nn.Conv2d(ndf * nf_mult, 1, kernel_size=4, stride=1, padding=1)
+        self.attention = nn.Sequential(
+            nn.Conv2d(ndf * nf_mult, ndf * nf_mult // 4, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(ndf * nf_mult // 4, ndf * nf_mult, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, input):
+        x = self.initial_conv(input)
+        for layer in self.layers:
+            x = layer(x)
+        x = x * self.attention(x)
+        return self.output_conv(x)
+
+
+def create_s2f_model(
+    in_channels=1,
+    out_channels=1,
+    img_size=1024,
+    bridge_type='cbam',
+    use_multi_scale_input=True,
+    ndf=64,
+    n_layers=3,
+):
+    """Create S2F model with generator and discriminator."""
+    generator = S2FGenerator(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        img_size=img_size,
+        bridge_type=bridge_type,
+        use_multi_scale_input=use_multi_scale_input,
+    )
+    discriminator = PatchGANDiscriminator(
+        in_channels=in_channels + out_channels,
+        ndf=ndf,
+        n_layers=n_layers
+    )
+    return generator, discriminator

S2FApp/predictor.py

@@ -0,0 +1,164 @@
+"""
+Core inference logic for S2F (Shape2Force).
+Predicts force maps from bright field microscopy images.
+"""
+import os
+import sys
+import cv2
+import torch
+import numpy as np
+
+# Ensure S2F is in path when running from project root or S2F
+S2F_ROOT = os.path.dirname(os.path.abspath(__file__))
+if S2F_ROOT not in sys.path:
+    sys.path.insert(0, S2F_ROOT)
+
+from models.s2f_model import create_s2f_model
+from utils.substrate_settings import get_settings_of_category, compute_settings_normalization
+from utils import config
+
+
+def load_image(filepath, target_size=1024):
+    """Load and preprocess a bright field image."""
+    img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
+    if img is None:
+        raise ValueError(f"Could not load image: {filepath}")
+    if isinstance(target_size, int):
+        target_size = (target_size, target_size)
+    img = cv2.resize(img, target_size)
+    img = img.astype(np.float32) / 255.0
+    return img
+
+
+def sum_force_map(force_map):
+    """Compute cell force as sum of pixel values scaled by SCALE_FACTOR_FORCE."""
+    if isinstance(force_map, np.ndarray):
+        force_map = torch.from_numpy(force_map.astype(np.float32))
+    if force_map.dim() == 2:
+        force_map = force_map.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
+    elif force_map.dim() == 3:
+        force_map = force_map.unsqueeze(0)  # [1, 1, H, W]
+    # force_map: [B, 1, H, W], sum over spatial dims (2, 3)
+    return torch.sum(force_map, dim=(2, 3)) * config.SCALE_FACTOR_FORCE
+
+
+def create_settings_channels_single(substrate_name, device, height, width, config_path=None,
+                                    substrate_config=None):
+    """
+    Create settings channels for a single image (single-cell mode).
+
+    Args:
+        substrate_name: Substrate name (used if substrate_config is None)
+        device: torch device
+        height, width: spatial dimensions
+        config_path: Path to substrate config JSON
+        substrate_config: Optional dict with 'pixelsize' and 'young'. If provided, overrides substrate_name.
+    """
+    norm_params = compute_settings_normalization(config_path=config_path)
+    if substrate_config is not None and 'pixelsize' in substrate_config and 'young' in substrate_config:
+        settings = substrate_config
+    else:
+        settings = get_settings_of_category(substrate_name, config_path=config_path)
+    pmin, pmax = norm_params['pixelsize']['min'], norm_params['pixelsize']['max']
+    ymin, ymax = norm_params['young']['min'], norm_params['young']['max']
+    pixelsize_norm = (settings['pixelsize'] - pmin) / (pmax - pmin) if pmax > pmin else 0.5
+    young_norm = (settings['young'] - ymin) / (ymax - ymin) if ymax > ymin else 0.5
+    pixelsize_norm = max(0.0, min(1.0, pixelsize_norm))
+    young_norm = max(0.0, min(1.0, young_norm))
+    pixelsize_ch = torch.full(
+        (1, 1, height, width), pixelsize_norm, device=device, dtype=torch.float32
+    )
+    young_ch = torch.full(
+        (1, 1, height, width), young_norm, device=device, dtype=torch.float32
+    )
+    return torch.cat([pixelsize_ch, young_ch], dim=1)
+
+
+class S2FPredictor:
+    """
+    Shape2Force predictor for single-cell or spheroid force map prediction.
+    """
+
+    def __init__(self, model_type="single_cell", checkpoint_path=None, ckp_folder=None, device=None):
+        """
+        Args:
+            model_type: "single_cell" or "spheroid"
+            checkpoint_path: Path to .pth checkpoint (relative to ckp_folder or absolute)
+            ckp_folder: Folder containing checkpoints (default: S2F/ckp)
+            device: "cuda" or "cpu" (auto-detected if None)
+        """
+        self.model_type = model_type
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        ckp_folder = ckp_folder or os.path.join(S2F_ROOT, "ckp")
+
+        in_channels = 3 if model_type == "single_cell" else 1
+        generator, _ = create_s2f_model(in_channels=in_channels)
+        self.generator = generator
+
+        if checkpoint_path:
+            full_path = checkpoint_path
+            if not os.path.isabs(checkpoint_path):
+                full_path = os.path.join(ckp_folder, checkpoint_path)
+            if not os.path.exists(full_path):
+                raise FileNotFoundError(f"Checkpoint not found: {full_path}")
+
+            # Single-cell: use load_checkpoint_with_expansion (handles 1ch->3ch if needed)
+            if model_type == "single_cell":
+                self.generator.load_checkpoint_with_expansion(full_path, strict=True)
+            else:
+                checkpoint = torch.load(full_path, map_location="cpu", weights_only=False)
+                state = checkpoint.get("generator_state_dict", checkpoint)
+                self.generator.load_state_dict(state, strict=True)
+
+        self.generator = self.generator.to(self.device)
+        self.generator.eval()
+
+        self.norm_params = compute_settings_normalization() if model_type == "single_cell" else None
+        self.config_path = os.path.join(S2F_ROOT, "config", "substrate_settings.json")
+
+    def predict(self, image_path=None, image_array=None, substrate="fibroblasts_PDMS",
+                substrate_config=None):
+        """
+        Run prediction on an image.
+
+        Args:
+            image_path: Path to bright field image (tif, png, jpg)
+            image_array: numpy array (H, W) or (H, W, C) in [0, 255] or [0, 1]
+            substrate: Substrate name for single-cell mode (used if substrate_config is None)
+            substrate_config: Optional dict with 'pixelsize' and 'young'. Overrides substrate lookup.
+
+        Returns:
+            heatmap: numpy array (1024, 1024) in [0, 1]
+            force: scalar cell force (sum of heatmap * SCALE_FACTOR_FORCE)
+            pixel_sum: raw sum of all pixel values in heatmap
+        """
+        if image_path is not None:
+            img = load_image(image_path)
+        elif image_array is not None:
+            img = np.asarray(image_array, 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, (1024, 1024))
+        else:
+            raise ValueError("Provide image_path or image_array")
+
+        x = torch.from_numpy(img).float().unsqueeze(0).unsqueeze(0).to(self.device)  # [1,1,H,W]
+
+        if self.model_type == "single_cell" and self.norm_params is not None:
+            settings_ch = create_settings_channels_single(
+                substrate, self.device, x.shape[2], x.shape[3],
+                config_path=self.config_path, substrate_config=substrate_config
+            )
+            x = torch.cat([x, settings_ch], dim=1)  # [1,3,H,W]
+
+        with torch.no_grad():
+            pred = self.generator(x)
+
+        pred = (pred + 1.0) / 2.0  # Tanh to [0, 1]
+        heatmap = pred[0, 0].cpu().numpy()
+        force = sum_force_map(pred).item()
+        pixel_sum = float(np.sum(heatmap))
+
+        return heatmap, force, pixel_sum

S2FApp/requirements.txt

@@ -0,0 +1,10 @@
+# Shape2Force App - inference only
+torch>=2.0.0
+torchvision>=0.15.0
+numpy>=1.20.0
+opencv-python>=4.5.0
+streamlit>=1.28.0
+matplotlib>=3.5.0
+Pillow>=9.0.0
+plotly>=5.14.0
+huggingface_hub>=0.20.0

S2FApp/utils/__init__.py

@@ -0,0 +1 @@
+from . import config

S2FApp/utils/config.py

@@ -0,0 +1,3 @@
+# Constants for force and area scaling (used in force map prediction)
+SCALE_FACTOR_FORCE = 1e-3
+SCALE_FACTOR_AREA = 1e-4

S2FApp/utils/metrics.py

@@ -0,0 +1,447 @@
+"""Metrics for S2F training and evaluation.
+
+Includes: MSE, MS-SSIM, Pixel Correlation (Pearson), Relative Magnitude Error (WFM),
+and evaluation helpers for notebooks and scripts.
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from skimage.metrics import structural_similarity as ssim
+from scipy.stats import pearsonr
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+try:
+    from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure
+    from torchmetrics import MeanSquaredError
+    HAS_TORCHMETRICS = True
+except ImportError:
+    HAS_TORCHMETRICS = False
+
+
+def calculate_mse(y_true, y_pred):
+    if isinstance(y_true, torch.Tensor):
+        return F.mse_loss(y_pred, y_true).item()
+    return float(np.mean((np.asarray(y_true) - np.asarray(y_pred)) ** 2))
+
+
+def calculate_psnr(y_true, y_pred, max_pixel_value=1.0):
+    mse = calculate_mse(y_true, y_pred)
+    if mse == 0:
+        return float('inf')
+    return 20 * np.log10(max_pixel_value / np.sqrt(mse))
+
+
+def calculate_ssim_tensor(y_true, y_pred, data_range=1.0):
+    if isinstance(y_true, torch.Tensor):
+        y_true = y_true.detach().cpu().numpy()
+    if isinstance(y_pred, torch.Tensor):
+        y_pred = y_pred.detach().cpu().numpy()
+    ssim_values = []
+    batch_size = y_true.shape[0]
+    for i in range(batch_size):
+        if len(y_true.shape) == 4:
+            true_img = y_true[i, 0] if y_true.shape[1] == 1 else y_true[i, 0]
+            pred_img = y_pred[i, 0] if y_pred.shape[1] == 1 else y_pred[i, 0]
+        else:
+            true_img, pred_img = y_true[i], y_pred[i]
+        ssim_values.append(ssim(true_img, pred_img, data_range=data_range))
+    return np.mean(ssim_values)
+
+
+def calculate_pearson_correlation(y_true, y_pred):
+    if isinstance(y_true, torch.Tensor):
+        y_true = y_true.cpu().numpy()
+    if isinstance(y_pred, torch.Tensor):
+        y_pred = y_pred.cpu().numpy()
+    correlation, _ = pearsonr(y_true.flatten(), y_pred.flatten())
+    return correlation
+
+
+def calculate_individual_pixel_correlation(y_true, y_pred):
+    """Pixel-wise Pearson correlation per sample in batch."""
+    if isinstance(y_true, torch.Tensor):
+        y_true = y_true.cpu().numpy()
+    if isinstance(y_pred, torch.Tensor):
+        y_pred = y_pred.cpu().numpy()
+    correlations = []
+    batch_size = y_true.shape[0]
+    for i in range(batch_size):
+        true_flat = y_true[i].flatten()
+        pred_flat = y_pred[i].flatten()
+        r, _ = pearsonr(true_flat, pred_flat)
+        correlations.append(r)
+    return correlations
+
+
+# --- WFM (Wrinkle Force Microscopy) metrics for heatmap as magnitude ---
+
+def _to_numpy_wfm(x):
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy()
+    return np.asarray(x)
+
+
+def _ensure_shape_wfm(f):
+    """Ensure (N, 2, H, W). Heatmap -> fx=magnitude, fy=0."""
+    if f.ndim == 3:
+        if f.shape[-1] == 2:
+            f = np.transpose(f, (2, 0, 1))[None, ...]
+        elif f.shape[0] == 2:
+            f = f[None, ...]
+        else:
+            raise ValueError(f"Unsupported 3D shape {f.shape}")
+    elif f.ndim == 4:
+        if f.shape[-1] == 2:
+            f = np.transpose(f, (0, 3, 1, 2))
+    else:
+        raise ValueError(f"Unsupported ndim={f.ndim}")
+    return f
+
+
+def _force_mag_wfm(f):
+    fx, fy = f[:, 0], f[:, 1]
+    return np.sqrt(fx**2 + fy**2)
+
+
+def wfm_correlation(y_true, y_pred, mode="magnitude"):
+    """Pearson correlation between prediction and ground truth (magnitude mode for heatmaps)."""
+    t = _ensure_shape_wfm(_to_numpy_wfm(y_true))
+    p = _ensure_shape_wfm(_to_numpy_wfm(y_pred))
+    if t.shape != p.shape:
+        raise ValueError(f"Shape mismatch: true {t.shape} vs pred {p.shape}")
+    if mode == "magnitude":
+        tv = _force_mag_wfm(t).ravel()
+        pv = _force_mag_wfm(p).ravel()
+    else:
+        raise ValueError(f"Unknown mode '{mode}'")
+    tv, pv = tv.astype(np.float64), pv.astype(np.float64)
+    if np.allclose(tv.std(), 0) or np.allclose(pv.std(), 0):
+        return 0.0
+    return float(np.corrcoef(tv, pv)[0, 1])
+
+
+def wfm_relative_magnitude_error(y_true, y_pred, eps=1e-8):
+    """Relative magnitude error for heatmap-as-magnitude."""
+    t = _ensure_shape_wfm(_to_numpy_wfm(y_true))
+    p = _ensure_shape_wfm(_to_numpy_wfm(y_pred))
+    if t.shape != p.shape:
+        raise ValueError(f"Shape mismatch: true {t.shape} vs pred {p.shape}")
+    mag_t = _force_mag_wfm(t)
+    mag_p = _force_mag_wfm(p)
+    fbar = np.mean(mag_t)
+    if np.isclose(fbar, 0):
+        return 0.0
+    rel = np.abs(mag_p - mag_t) / (mag_t + eps)
+    w = mag_t / fbar
+    return float(np.mean(rel * w))
+
+
+def apply_threshold_mask(tensor, threshold=0.0):
+    return tensor * (tensor >= threshold).float()
+
+
+def detect_tanh_output_model(model):
+    """Detect if model outputs [-1, 1] (Tanh)."""
+    if hasattr(model, 'use_sigmoid') and not model.use_sigmoid:
+        return True
+    if hasattr(model, 'use_tanh_output') and model.use_tanh_output:
+        return True
+    if hasattr(model, 'final_conv'):
+        fc = model.final_conv
+        if isinstance(fc, nn.Sequential):
+            if isinstance(fc[-1], nn.Tanh):
+                return True
+        elif isinstance(fc, nn.Tanh):
+            return True
+    return False
+
+
+def convert_tanh_to_sigmoid_range(tensor):
+    return (tensor + 1.0) / 2.0
+
+
+# --- TorchMetrics wrapper for MS-SSIM ---
+
+class TorchMetricsWrapper:
+    def __init__(self, device='cpu'):
+        self.device = device
+        self.reset_metrics()
+
+    def reset_metrics(self):
+        if HAS_TORCHMETRICS:
+            self.ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
+            self.mse = MeanSquaredError().to(self.device)
+        else:
+            self.ms_ssim = None
+            self.mse = None
+
+    def compute_ms_ssim(self, y_true, y_pred):
+        if not HAS_TORCHMETRICS:
+            return float(calculate_ssim_tensor(y_true, y_pred))  # fallback to SSIM
+        y_true = y_true.to(self.device)
+        y_pred = y_pred.to(self.device)
+        if y_true.shape[1] == 1:
+            pass
+        else:
+            y_true, y_pred = y_true[:, 0:1], y_pred[:, 0:1]
+        return self.ms_ssim(y_pred, y_true).item()
+
+    def compute_mse(self, y_true, y_pred):
+        if not HAS_TORCHMETRICS:
+            return calculate_mse(y_true, y_pred)
+        y_true = y_true.to(self.device)
+        y_pred = y_pred.to(self.device)
+        return self.mse(y_pred, y_true).item()
+
+
+# --- Full evaluation on dataset ---
+
+def evaluate_metrics_on_dataset(generator, data_loader, device=None, description="Evaluating",
+                               save_predictions=False, threshold=0.0, use_settings=False,
+                               normalization_params=None, config_path=None, substrate_override=None):
+    """
+    Evaluate S2F generator on a dataset. Returns MSE, MS-SSIM, Pixel Correlation,
+    Relative Magnitude Error, and force sum/mean correlations.
+    """
+    if device is None:
+        device = torch.device('mps' if torch.backends.mps.is_available() else
+                              'cuda' if torch.cuda.is_available() else 'cpu')
+
+    generator = generator.to(device)
+    generator.eval()
+    metrics_wrapper = TorchMetricsWrapper(device=device)
+
+    heatmap_mse = []
+    heatmap_ms_ssim = []
+    heatmap_pixel_corr = []
+    wfm_corr_mag = []
+    wfm_rel_mag_err = []
+    force_sum_gt, force_sum_pred = [], []
+    force_mean_gt, force_mean_pred = [], []
+    individual_predictions = [] if save_predictions else None
+
+    with torch.no_grad():
+        for batch_idx, batch_data in enumerate(tqdm(data_loader, desc=description)):
+            if len(batch_data) == 5:
+                images, heatmaps, _, _, metadata = batch_data
+                has_metadata = True
+            else:
+                images, heatmaps, _, _ = batch_data
+                has_metadata = False
+
+            images = images.to(device, dtype=torch.float32)
+            heatmaps = heatmaps.to(device, dtype=torch.float32)
+
+            if use_settings and normalization_params is not None:
+                from models.s2f_model import create_settings_channels
+                meta = metadata if has_metadata else {'substrate': [substrate_override or 'fibroblasts_PDMS'] * images.size(0)}
+                settings_ch = create_settings_channels(meta, normalization_params, device, images.shape, config_path=config_path)
+                images = torch.cat([images, settings_ch], dim=1)
+
+            pred = generator(images)
+            if detect_tanh_output_model(generator):
+                pred = convert_tanh_to_sigmoid_range(pred)
+
+            gt_thresh = apply_threshold_mask(heatmaps, threshold)
+            pred_thresh = pred  # no threshold on pred for metrics
+
+            heatmap_mse.append(metrics_wrapper.compute_mse(gt_thresh, pred_thresh))
+            heatmap_ms_ssim.append(metrics_wrapper.compute_ms_ssim(gt_thresh, pred_thresh))
+            heatmap_pixel_corr.extend(calculate_individual_pixel_correlation(gt_thresh, pred_thresh))
+
+            # WFM: heatmap as magnitude (fx=magnitude, fy=0)
+            B, _, H, W = gt_thresh.shape
+            gt_ff = torch.zeros(B, 2, H, W, device=device)
+            pred_ff = torch.zeros(B, 2, H, W, device=device)
+            gt_ff[:, 0], pred_ff[:, 0] = gt_thresh[:, 0], pred_thresh[:, 0]
+            try:
+                wfm_corr_mag.append(wfm_correlation(gt_ff, pred_ff, mode="magnitude"))
+                wfm_rel_mag_err.append(wfm_relative_magnitude_error(gt_ff, pred_ff))
+            except Exception:
+                wfm_corr_mag.append(float('nan'))
+                wfm_rel_mag_err.append(float('nan'))
+
+            force_sum_gt.extend(torch.sum(gt_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+            force_sum_pred.extend(torch.sum(pred_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+            force_mean_gt.extend(torch.mean(gt_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+            force_mean_pred.extend(torch.mean(pred_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+
+            if save_predictions:
+                for i in range(images.size(0)):
+                    p, t = pred_thresh[i:i+1], gt_thresh[i:i+1]
+                    gt_ff_i = torch.zeros(1, 2, H, W, device=device)
+                    pred_ff_i = torch.zeros(1, 2, H, W, device=device)
+                    gt_ff_i[0, 0], pred_ff_i[0, 0] = t[0, 0], p[0, 0]
+                    try:
+                        rme = wfm_relative_magnitude_error(gt_ff_i, pred_ff_i)
+                    except Exception:
+                        rme = float('nan')
+                    individual_predictions.append({
+                        'batch_idx': batch_idx,
+                        'sample_idx': i,
+                        'original_image': images[i].cpu().numpy(),
+                        'ground_truth': heatmaps[i].cpu().numpy(),
+                        'ground_truth_thresholded': gt_thresh[i].cpu().numpy(),
+                        'prediction': pred[i].cpu().numpy(),
+                        'prediction_thresholded': pred_thresh[i].cpu().numpy(),
+                        'mse': metrics_wrapper.compute_mse(t, p),
+                        'ms_ssim': metrics_wrapper.compute_ms_ssim(t, p),
+                        'pixel_correlation': calculate_pearson_correlation(t, p),
+                        'wfm_relative_magnitude_error': rme,
+                        'force_sum_gt': torch.sum(gt_thresh[i]).item(),
+                        'force_sum_pred': torch.sum(pred_thresh[i]).item(),
+                        'force_mean_gt': torch.mean(gt_thresh[i]).item(),
+                        'force_mean_pred': torch.mean(pred_thresh[i]).item(),
+                    })
+
+    valid_wfm_corr = [x for x in wfm_corr_mag if not np.isnan(x)]
+    valid_wfm_rme = [x for x in wfm_rel_mag_err if not np.isnan(x)]
+    try:
+        force_sum_corr, _ = pearsonr(force_sum_gt, force_sum_pred)
+        force_mean_corr, _ = pearsonr(force_mean_gt, force_mean_pred)
+    except Exception:
+        force_sum_corr = force_mean_corr = 0.0
+    if force_sum_corr is None or (isinstance(force_sum_corr, float) and np.isnan(force_sum_corr)):
+        force_sum_corr = 0.0
+    if force_mean_corr is None or (isinstance(force_mean_corr, float) and np.isnan(force_mean_corr)):
+        force_mean_corr = 0.0
+
+    results = {
+        'heatmap': {
+            'mse': np.mean(heatmap_mse),
+            'mse_std': np.std(heatmap_mse),
+            'ms_ssim': np.mean(heatmap_ms_ssim),
+            'ms_ssim_std': np.std(heatmap_ms_ssim),
+            'pixel_correlation': np.mean(heatmap_pixel_corr),
+            'pixel_correlation_std': np.std(heatmap_pixel_corr),
+        },
+        'wfm': {
+            'correlation_magnitude': np.mean(valid_wfm_corr) if valid_wfm_corr else float('nan'),
+            'correlation_magnitude_std': np.std(valid_wfm_corr) if valid_wfm_corr else float('nan'),
+            'relative_magnitude_error': np.mean(valid_wfm_rme) if valid_wfm_rme else float('nan'),
+            'relative_magnitude_error_std': np.std(valid_wfm_rme) if valid_wfm_rme else float('nan'),
+        },
+        'force_sum': {
+            'correlation': float(force_sum_corr),
+            'gt_mean': np.mean(force_sum_gt),
+            'pred_mean': np.mean(force_sum_pred),
+            'gt_std': np.std(force_sum_gt),
+            'pred_std': np.std(force_sum_pred),
+        },
+        'force_mean': {
+            'correlation': float(force_mean_corr),
+            'gt_mean': np.mean(force_mean_gt),
+            'pred_mean': np.mean(force_mean_pred),
+        },
+    }
+
+    if save_predictions:
+        results['individual_predictions'] = individual_predictions
+    return results
+
+
+def print_metrics_report(report, threshold=0.0, uses_tanh=False):
+    """Print formatted metrics report."""
+    for name, metrics in report.items():
+        print(f"\n🔸 {name.upper()} SET METRICS" + (f" (threshold={threshold})" if threshold > 0 else ""))
+        print("-" * 60)
+        print("HEATMAP METRICS:")
+        print(f" MSE:             {metrics['heatmap']['mse']:.6f} ± {metrics['heatmap']['mse_std']:.6f}")
+        print(f" MS-SSIM:         {metrics['heatmap']['ms_ssim']:.4f} ± {metrics['heatmap']['ms_ssim_std']:.4f}")
+        print(f" Pixel Corr:      {metrics['heatmap']['pixel_correlation']:.4f} ± {metrics['heatmap']['pixel_correlation_std']:.4f}")
+        print("WFM METRICS (heatmap as magnitude):")
+        print(f" Correlation (Magnitude): {metrics['wfm']['correlation_magnitude']:.4f} ± {metrics['wfm']['correlation_magnitude_std']:.4f}")
+        print(f" Relative Magnitude Error: {metrics['wfm']['relative_magnitude_error']:.4f} ± {metrics['wfm']['relative_magnitude_error_std']:.4f}")
+        print("FORCE SUM CORRELATION:")
+        print(f" Correlation: {metrics['force_sum']['correlation']:.4f}")
+        print(f" GT Mean: {metrics['force_sum']['gt_mean']:.2f} ± {metrics['force_sum']['gt_std']:.2f}")
+        print(f" Pred Mean: {metrics['force_sum']['pred_mean']:.2f} ± {metrics['force_sum']['pred_std']:.2f}")
+        if uses_tanh:
+            print(" Note: Model outputs [-1,1], converted to [0,1] for evaluation")
+    print("=" * 60)
+
+
+def gen_prediction_plots(individual_predictions, save_dir, sort_by='ms_ssim', sort_order='desc', threshold=0.0):
+    """Generate prediction plots (BF | GT | Pred) sorted by metric."""
+    os.makedirs(save_dir, exist_ok=True)
+    reverse = (sort_order.lower() == 'desc') if sort_by.lower() not in ['mse', 'wfm_relative_magnitude_error'] else (sort_order.lower() == 'desc')
+    valid = [p for p in individual_predictions if not np.isnan(p.get(sort_by.lower(), 0))]
+    sorted_preds = sorted(valid, key=lambda x: x[sort_by.lower()], reverse=reverse)
+    print(f"Sorting {len(sorted_preds)} predictions by {sort_by} ({sort_order})")
+    for rank, p in enumerate(tqdm(sorted_preds, desc="Saving plots"), 1):
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+        img = p['original_image']
+        axes[0].imshow(img[0] if img.ndim == 3 else img, cmap='gray')
+        axes[0].set_title('Bright Field')
+        axes[0].axis('off')
+        gt = p['ground_truth']
+        axes[1].imshow(gt[0] if gt.ndim == 3 else gt, cmap='jet', vmin=0, vmax=1)
+        axes[1].set_title('Ground Truth')
+        axes[1].axis('off')
+        pr = p['prediction']
+        axes[2].imshow(pr[0] if pr.ndim == 3 else pr, cmap='jet', vmin=0, vmax=1)
+        axes[2].set_title('Prediction')
+        axes[2].axis('off')
+        m = (f"MSE: {p['mse']:.4f} | MS-SSIM: {p['ms_ssim']:.4f} | "
+             f"Pixel Corr: {p['pixel_correlation']:.4f} | Rel Mag Err: {p.get('wfm_relative_magnitude_error', 'N/A')}")
+        fig.suptitle(f"Rank {rank} (by {sort_by})\n{m}", fontsize=10, y=0.02)
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_dir, f"rank{rank:03d}_batch{p['batch_idx']:03d}_sample{p['sample_idx']:02d}.png"), dpi=150, bbox_inches='tight')
+        plt.close()
+
+
+def plot_predictions(loader, generator, n_samples, device, threshold=0.0,
+                     use_settings=False, normalization_params=None, config_path=None, substrate_override=None):
+    """Plot BF | GT | Pred for first n_samples from loader."""
+    generator = generator.to(device)
+    generator.eval()
+    bf_list, gt_list, meta_list = [], [], []
+    it = iter(loader)
+    while len(bf_list) < n_samples:
+        try:
+            batch = next(it)
+        except StopIteration:
+            break
+        if len(batch) == 5:
+            images, heatmaps, _, _, meta = batch
+        else:
+            images, heatmaps = batch[0], batch[1]
+            meta = None
+        for i in range(images.shape[0]):
+            if len(bf_list) >= n_samples:
+                break
+            bf_list.append(images[i])
+            gt_list.append(heatmaps[i])
+            meta_list.append(meta)
+    n = min(n_samples, len(bf_list))
+    bf_batch = torch.stack(bf_list[:n]).to(device, dtype=torch.float32)
+    if use_settings and normalization_params:
+        from models.s2f_model import create_settings_channels
+        sub = substrate_override or 'fibroblasts_PDMS'
+        meta_dict = {'substrate': [sub] * n}
+        settings_ch = create_settings_channels(meta_dict, normalization_params, device, bf_batch.shape, config_path=config_path)
+        bf_batch = torch.cat([bf_batch, settings_ch], dim=1)
+    with torch.no_grad():
+        pred = generator(bf_batch)
+        if detect_tanh_output_model(generator):
+            pred = convert_tanh_to_sigmoid_range(pred)
+    if threshold > 0:
+        pred = pred * (pred >= threshold).float()
+    fig, axes = plt.subplots(n, 3, figsize=(12, 4 * n))
+    if n == 1:
+        axes = axes.reshape(1, -1)
+    for i in range(n):
+        axes[i, 0].imshow(bf_list[i].squeeze().cpu().numpy(), cmap='gray')
+        axes[i, 0].set_title('Bright Field')
+        axes[i, 0].axis('off')
+        axes[i, 1].imshow(gt_list[i].squeeze().cpu().numpy(), cmap='jet', vmin=0, vmax=1)
+        axes[i, 1].set_title('Ground Truth')
+        axes[i, 1].axis('off')
+        axes[i, 2].imshow(pred[i].squeeze().cpu().numpy(), cmap='jet', vmin=0, vmax=1)
+        axes[i, 2].set_title('Prediction')
+        axes[i, 2].axis('off')
+    plt.tight_layout()
+    plt.show()

S2FApp/utils/substrate_settings.py

@@ -0,0 +1,129 @@
+"""
+Substrate settings for force map prediction.
+Loads from config/substrate_settings.json - users can edit this file to add/modify substrates.
+"""
+import os
+import json
+
+
+def _default_config_path():
+    """Default path to substrate settings config (S2F/config/substrate_settings.json)."""
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    project_root = os.path.dirname(this_dir)  # S2F root
+    return os.path.join(project_root, 'config', 'substrate_settings.json')
+
+
+def load_substrate_config(config_path=None):
+    """
+    Load substrate settings from config file.
+
+    Args:
+        config_path: Path to JSON config. If None, uses config/substrate_settings.json in S2F root.
+
+    Returns:
+        dict: Config with 'substrates', 'default_substrate'
+    """
+    path = config_path or _default_config_path()
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"Substrate config not found at {path}. "
+            "Create config/substrate_settings.json or pass config_path."
+        )
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def resolve_substrate(name, config=None, config_path=None):
+    """
+    Resolve substrate name to a canonical substrate key.
+
+    Args:
+        name: Substrate key (e.g. 'fibroblasts_PDMS', 'PDMS_10kPa')
+        config: Pre-loaded config dict. If None, loads from config_path.
+        config_path: Path to config file (used if config is None).
+
+    Returns:
+        str: Canonical substrate key
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+
+    s = (name or '').strip()
+    if not s:
+        return config.get('default_substrate', 'fibroblasts_PDMS')
+
+    substrates = config.get('substrates', {})
+    s_lower = s.lower()
+    for key in substrates:
+        if key.lower() == s_lower:
+            return key
+    for key in substrates:
+        if s_lower.startswith(key.lower()) or key.lower().startswith(s_lower):
+            return key
+
+    return config.get('default_substrate', 'fibroblasts_PDMS')
+
+
+def get_settings_of_category(substrate_name, config=None, config_path=None):
+    """
+    Get pixelsize and young's modulus for a substrate.
+
+    Args:
+        substrate_name: Substrate or folder name (case-insensitive)
+        config: Pre-loaded config dict. If None, loads from config_path.
+        config_path: Path to config file (used if config is None).
+
+    Returns:
+        dict: {'name': str, 'pixelsize': float, 'young': float}
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+
+    substrate_key = resolve_substrate(substrate_name, config=config)
+    substrates = config.get('substrates', {})
+    default = config.get('default_substrate', 'fibroblasts_PDMS')
+
+    if substrate_key in substrates:
+        return substrates[substrate_key].copy()
+
+    default_settings = substrates.get(default, {'name': 'Fibroblasts on PDMS', 'pixelsize': 3.0769, 'young': 6000})
+    return default_settings.copy()
+
+
+def list_substrates(config=None, config_path=None):
+    """
+    Return list of available substrate keys for user selection.
+
+    Returns:
+        list: Substrate keys
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+    return list(config.get('substrates', {}).keys())
+
+
+def compute_settings_normalization(config=None, config_path=None):
+    """
+    Compute min-max normalization parameters from all substrates in config.
+
+    Returns:
+        dict: {'pixelsize': {'min', 'max'}, 'young': {'min', 'max'}}
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+
+    substrates = config.get('substrates', {})
+    all_pixelsizes = [s['pixelsize'] for s in substrates.values()]
+    all_youngs = [s['young'] for s in substrates.values()]
+
+    if not all_pixelsizes or not all_youngs:
+        pixelsize_min, pixelsize_max = 3.0769, 9.8138
+        young_min, young_max = 1000.0, 10000.0
+    else:
+        pixelsize_min, pixelsize_max = min(all_pixelsizes), max(all_pixelsizes)
+        young_min, young_max = min(all_youngs), max(all_youngs)
+
+    return {
+        'pixelsize': {'min': pixelsize_min, 'max': pixelsize_max},
+        'young': {'min': young_min, 'max': young_max}
+    }

ckp/ckp_singlecell.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e4daea628c86ab8dc04e1adff47584a48c3a0104abbc6a79a4296f55c571698
+size 959538864

ckp/ckp_spheroid.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5764cd3e29209dcd49e186b1ae3432ed46822de14bf3e47443fc25873b5c096d
+size 959520496

config/substrate_settings.json

@@ -0,0 +1 @@
+{"substrates":{"fibroblasts_PDMS":{"name":"Fibroblasts on PDMS (6 kPa)","pixelsize":3.0769,"young":6000},"U2OS_PDMS":{"name":"U2OS cells on PDMS (6 kPa)","pixelsize":6.1538,"young":6000},"PDMS_1kPa":{"name":"PDMS soft hydrogel (1 kPa, 10 µm/px)","pixelsize":9.8138,"young":1000},"PDMS_10kPa":{"name":"PDMS stiff hydrogel (10 kPa, 10 µm/px)","pixelsize":9.8138,"young":10000},"PDMS_1kPa_3um":{"name":"PDMS soft hydrogel (1 kPa, 3 µm/px)","pixelsize":3.0769,"young":1000},"PDMS_10kPa_3um":{"name":"PDMS stiff hydrogel (10 kPa, 3 µm/px)","pixelsize":3.0769,"young":10000}},"default_substrate":"fibroblasts_PDMS"}

data/__init__.py

@@ -0,0 +1 @@
+from .cell_dataset import prepare_data, load_folder_data, ImageDataset

data/augmentations.py

@@ -0,0 +1,101 @@
+import numpy as np
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+import torch
+from scipy.ndimage import gaussian_filter, map_coordinates
+
+
+class AdvancedAugmentations:
+    def __init__(self, target_size=(1024, 1024)):
+        self.target_size = target_size
+
+    def __call__(self, image, heatmap):
+        image = TF.to_pil_image(image)
+        heatmap = TF.to_pil_image(heatmap)
+
+        if np.random.rand() > 0.5:
+            image = TF.hflip(image)
+            heatmap = TF.hflip(heatmap)
+        if np.random.rand() > 0.5:
+            image = TF.vflip(image)
+            heatmap = TF.vflip(heatmap)
+
+        if np.random.rand() > 0.5:
+            angle = np.random.uniform(-45, 45)
+            image = TF.rotate(image, angle, interpolation=TF.InterpolationMode.BILINEAR)
+            heatmap = TF.rotate(heatmap, angle, interpolation=TF.InterpolationMode.BILINEAR)
+
+        if np.random.rand() > 0.5:
+            width, height = image.size
+            crop_size = int(min(width, height) * np.random.uniform(0.8, 1.0))
+            i, j, h, w = T.RandomCrop.get_params(image, (crop_size, crop_size))
+            image = TF.crop(image, i, j, h, w)
+            heatmap = TF.crop(heatmap, i, j, h, w)
+            image = TF.resize(image, self.target_size, interpolation=TF.InterpolationMode.BILINEAR)
+            heatmap = TF.resize(heatmap, self.target_size, interpolation=TF.InterpolationMode.BILINEAR)
+
+        if np.random.rand() > 0.5:
+            image, heatmap = self.random_affine(image, heatmap)
+
+        if not isinstance(image, torch.Tensor):
+            image = TF.to_tensor(image)
+        if not isinstance(heatmap, torch.Tensor):
+            heatmap = TF.to_tensor(heatmap)
+
+        if np.random.rand() > 0.5:
+            brightness_factor = np.random.uniform(0.8, 1.2)
+            image = TF.adjust_brightness(image, brightness_factor)
+        if np.random.rand() > 0.5:
+            contrast_factor = np.random.uniform(0.8, 1.2)
+            image = TF.adjust_contrast(image, contrast_factor)
+
+        if np.random.rand() > 0.5:
+            noise_level = np.random.uniform(0.01, 0.05)
+            noise = torch.randn_like(image) * noise_level
+            image = torch.clamp(image + noise, 0, 1)
+
+        if np.random.rand() > 0.5:
+            image, heatmap = self.elastic_transform(image, heatmap)
+
+        return image, heatmap
+
+    def random_affine(self, image, heatmap):
+        degrees = [-10.0, 10.0]
+        translate = [0.05, 0.05]
+        scale = [0.95, 1.05]
+        shear = [-5.0, 5.0]
+        params = T.RandomAffine.get_params(degrees, translate, scale, shear, image.size)
+        angle, translate, scale, shear = params
+        translate = list(translate)
+        shear = list(shear)
+        image = TF.affine(image, angle, translate, scale, shear, interpolation=TF.InterpolationMode.BILINEAR)
+        heatmap = TF.affine(heatmap, angle, translate, scale, shear, interpolation=TF.InterpolationMode.BILINEAR)
+        return image, heatmap
+
+    def elastic_transform(self, image, heatmap, alpha=50, sigma=4):
+        if isinstance(image, torch.Tensor):
+            image_np = image.permute(1, 2, 0).numpy()
+            heatmap_np = heatmap.permute(1, 2, 0).numpy()
+        else:
+            image_np = np.asarray(image)
+            heatmap_np = np.asarray(heatmap)
+            if image_np.ndim == 2:
+                image_np = image_np[:, :, np.newaxis]
+            if heatmap_np.ndim == 2:
+                heatmap_np = heatmap_np[:, :, np.newaxis]
+
+        shape = image_np.shape[:2]
+        dx = gaussian_filter((np.random.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
+        dy = gaussian_filter((np.random.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
+        x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1))
+
+        image_transformed = np.zeros_like(image_np)
+        heatmap_transformed = np.zeros_like(heatmap_np)
+        for i in range(image_np.shape[2]):
+            image_transformed[..., i] = map_coordinates(image_np[..., i], indices, order=1).reshape(shape)
+        for i in range(heatmap_np.shape[2]):
+            heatmap_transformed[..., i] = map_coordinates(heatmap_np[..., i], indices, order=1).reshape(shape)
+
+        return torch.from_numpy(image_transformed).float().permute(2, 0, 1), \
+               torch.from_numpy(heatmap_transformed).float().permute(2, 0, 1)

data/cell_dataset.py

@@ -0,0 +1,188 @@
+"""
+Dataset and data loading for S2F training.
+Expects folder structure: each subfolder has BF_001.tif (bright field), *_gray.jpg (heatmap), and optionally .txt (cell_area, sum_force).
+"""
+import os
+import cv2
+import torch
+from torch.utils.data import Dataset, DataLoader
+from sklearn.model_selection import train_test_split
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+
+from utils import config
+
+
+def blur_force_map(force_map, ksize=25, sigma=10):
+    if ksize % 2 == 0:
+        ksize += 1
+    if force_map.dim() == 3:
+        force_map = force_map.unsqueeze(0)
+    device = force_map.device
+    force_map = force_map.cpu()
+    blurred_maps = []
+    for i in range(force_map.size(0)):
+        force_np = force_map[i, 0].numpy().astype(np.float32)
+        blurred = cv2.GaussianBlur(force_np, (ksize, ksize), sigmaX=sigma)
+        blurred_maps.append(blurred)
+    return torch.from_numpy(np.stack(blurred_maps)).to(device)
+
+
+class ImageDataset(Dataset):
+    def __init__(self, image_pairs, transform=None, channel_first=True,
+                 blur_heatmap=False, threshold=0.0, return_metadata=False):
+        self.image_pairs = image_pairs
+        self.transform = transform
+        self.channel_first = channel_first
+        self.blur_heatmap = blur_heatmap
+        self.threshold = threshold
+        self.return_metadata = return_metadata
+
+    def __len__(self):
+        return len(self.image_pairs)
+
+    def __getitem__(self, idx):
+        if self.return_metadata:
+            bf_image, hm_image, numbers, metadata = self.image_pairs[idx]
+        else:
+            bf_image, hm_image, numbers = self.image_pairs[idx]
+        if isinstance(numbers, tuple):
+            cell_area, sum_force = numbers
+        else:
+            cell_area = 0
+            sum_force = numbers
+
+        image = torch.from_numpy(bf_image).float().unsqueeze(0)
+        heatmap = torch.from_numpy(hm_image).float().unsqueeze(0)
+        if self.transform:
+            image, heatmap = self.transform(image, heatmap)
+        cell_area = torch.tensor(cell_area, dtype=torch.float32)
+        sum_force = torch.tensor(sum_force, dtype=torch.float32)
+        heatmap[heatmap <= self.threshold] = 0
+        if self.blur_heatmap:
+            heatmap = blur_force_map(heatmap)
+        if not self.channel_first:
+            image = image.permute(2, 1, 0)
+            heatmap = heatmap.permute(2, 1, 0)
+        if self.return_metadata:
+            return image, heatmap, cell_area, sum_force, metadata
+        return image, heatmap, cell_area, sum_force
+
+
+def load_image(filepath, target_size):
+    img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
+    if isinstance(target_size, int):
+        target_size = (target_size, target_size)
+    img = cv2.resize(img, target_size)
+    img = img / 255.0
+    return img.astype(np.float32)
+
+
+def load_text_data(filepath):
+    with open(filepath, 'r') as f:
+        lines = [line.strip() for line in f if line.strip()]
+        cell_area_diff = float(lines[0].split(":")[1].strip()) * config.SCALE_FACTOR_AREA
+        sum_force_diff = float(lines[1].split(":")[1].strip()) * config.SCALE_FACTOR_FORCE
+        return (cell_area_diff, sum_force_diff)
+
+
+def load_images_from_subfolders(root_folder, target_size, load_numerical_data=True,
+                                load_force_sum=False, return_metadata=False, substrate=None):
+    paired_images = []
+    numerical_data = []
+    metadata = []
+    for subfolder in os.listdir(root_folder):
+        subfolder_path = os.path.join(root_folder, subfolder)
+        if not os.path.isdir(subfolder_path):
+            continue
+        bf_image_path = hm_image_path = txt_file_path = None
+        for filename in os.listdir(subfolder_path):
+            if filename.endswith("BF_001.tif"):
+                bf_image_path = os.path.join(subfolder_path, filename)
+            elif filename.endswith("_gray.jpg"):
+                hm_image_path = os.path.join(subfolder_path, filename)
+            elif filename.endswith(".txt"):
+                txt_file_path = os.path.join(subfolder_path, filename)
+
+        if return_metadata:
+            if substrate is None:
+                from utils.substrate_settings import list_substrates
+                raise ValueError("substrate must be passed when return_metadata=True. Options: " +
+                                 ", ".join(list_substrates()))
+            metadata.append({'folder_name': subfolder, 'substrate': substrate, 'root_folder': root_folder})
+
+        if load_numerical_data:
+            if bf_image_path and hm_image_path and txt_file_path:
+                paired_images.append((bf_image_path, hm_image_path))
+                numerical_data.append(load_text_data(txt_file_path))
+        elif load_force_sum:
+            if bf_image_path and hm_image_path:
+                paired_images.append((bf_image_path, hm_image_path))
+                hm = load_image(hm_image_path, target_size)
+                numerical_data.append((0, float(np.sum(hm)) * config.SCALE_FACTOR_FORCE))
+        else:
+            if bf_image_path and hm_image_path:
+                paired_images.append((bf_image_path, hm_image_path))
+
+    with ThreadPoolExecutor() as executor:
+        bf_loaded = list(executor.map(lambda p: load_image(p[0], target_size), paired_images))
+        hm_loaded = list(executor.map(lambda p: load_image(p[1], target_size), paired_images))
+    if not numerical_data:
+        numerical_data = [(0, 0)] * len(bf_loaded)
+    if return_metadata:
+        return list(zip(bf_loaded, hm_loaded, numerical_data, metadata))
+    return list(zip(bf_loaded, hm_loaded, numerical_data))
+
+
+def prepare_data(input_folder, batch_size=8, target_size=(1024, 1024), split_size=0.2,
+                 use_augmentations=True, train_test_sep_folder=True, channel_first=True,
+                 load_numerical_data=False, load_force_sum=False, blur_heatmap=False,
+                 threshold=0.0, return_metadata=False, substrate=None):
+    if load_numerical_data and load_force_sum:
+        raise ValueError("load_numerical_data and load_force_sum cannot be True at the same time")
+
+    if train_test_sep_folder:
+        train_folder = os.path.join(input_folder, 'train')
+        test_folder = os.path.join(input_folder, 'test')
+        if not (os.path.exists(train_folder) and os.path.exists(test_folder)):
+            raise ValueError(f"train/test folders not found in {input_folder}")
+        train_pairs = load_images_from_subfolders(train_folder, target_size=target_size,
+                                                  load_numerical_data=load_numerical_data,
+                                                  load_force_sum=load_force_sum,
+                                                  return_metadata=return_metadata, substrate=substrate)
+        val_pairs = load_images_from_subfolders(test_folder, target_size=target_size,
+                                                load_numerical_data=load_numerical_data,
+                                                load_force_sum=load_force_sum,
+                                                return_metadata=return_metadata, substrate=substrate)
+    else:
+        image_pairs = load_images_from_subfolders(input_folder, target_size=target_size,
+                                                  load_numerical_data=load_numerical_data,
+                                                  load_force_sum=load_force_sum,
+                                                  return_metadata=return_metadata, substrate=substrate)
+        train_pairs, val_pairs = train_test_split(image_pairs, test_size=split_size, random_state=42)
+
+    train_transform = None
+    if use_augmentations:
+        from .augmentations import AdvancedAugmentations
+        train_transform = AdvancedAugmentations(target_size)
+
+    train_dataset = ImageDataset(train_pairs, transform=train_transform, channel_first=channel_first,
+                                 blur_heatmap=blur_heatmap, threshold=threshold, return_metadata=return_metadata)
+    train_dataset.name = os.path.basename(input_folder)
+    val_dataset = ImageDataset(val_pairs, channel_first=channel_first,
+                               blur_heatmap=blur_heatmap, threshold=threshold, return_metadata=return_metadata)
+    val_dataset.name = os.path.basename(input_folder)
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+    return train_loader, val_loader
+
+
+def load_folder_data(folder_path, substrate=None, img_size=1024, blur_heatmap=False,
+                     batch_size=2, threshold=0.0, return_metadata=False):
+    val_pairs = load_images_from_subfolders(folder_path, target_size=img_size,
+                                            load_numerical_data=False, load_force_sum=False,
+                                            return_metadata=return_metadata, substrate=substrate)
+    val_dataset = ImageDataset(val_pairs, channel_first=True, blur_heatmap=blur_heatmap,
+                               threshold=threshold, return_metadata=return_metadata)
+    val_dataset.name = os.path.basename(folder_path)
+    return DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

models/__init__.py

@@ -0,0 +1 @@
+from .s2f_model import create_s2f_model, S2FGenerator

models/blocks.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+import torch
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = nn.Conv2d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else None
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        if self.downsample:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        return out

models/cbam.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes, ratio=16):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc = nn.Sequential(
+            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
+            nn.ReLU(),
+            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(self.max_pool(x))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+        padding = kernel_size // 2
+        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv(x)
+        return self.sigmoid(x)
+
+
+class CBAM(nn.Module):
+    def __init__(self, in_planes, ratio=16, kernel_size=7):
+        super(CBAM, self).__init__()
+        self.channel_attention = ChannelAttention(in_planes, ratio)
+        self.spatial_attention = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        x = x * self.channel_attention(x)
+        x = x * self.spatial_attention(x)
+        return x

models/s2f_model.py

@@ -0,0 +1,435 @@
+"""
+S2F (Shape2Force) model for force map prediction (inference only).
+Supports single-cell and spheroid modes.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .blocks import ResidualBlock
+from .cbam import CBAM
+
+from utils import config
+from utils.substrate_settings import (
+    get_settings_of_category,
+    compute_settings_normalization,
+    load_substrate_config,
+)
+
+
+def normalize_settings(substrate_name, normalization_params, config=None, config_path=None):
+    """
+    Normalize settings for a given substrate.
+
+    Args:
+        substrate_name (str): Name of the substrate
+        normalization_params (dict): Normalization parameters
+
+    Returns:
+        tuple: (normalized_pixelsize, normalized_young)
+    """
+    settings = get_settings_of_category(substrate_name, config=config, config_path=config_path)
+
+    # Min-max normalization to [0, 1]
+    pixelsize_norm = (settings['pixelsize'] - normalization_params['pixelsize']['min']) / \
+                     (normalization_params['pixelsize']['max'] - normalization_params['pixelsize']['min'])
+
+    young_norm = (settings['young'] - normalization_params['young']['min']) / \
+                 (normalization_params['young']['max'] - normalization_params['young']['min'])
+
+    return pixelsize_norm, young_norm
+
+
+def create_settings_channels(metadata, normalization_params, device, image_shape, config_path=None):
+    """
+    Create settings channels for a batch of images.
+
+    Args:
+        metadata (dict): Batch metadata containing substrate information
+        normalization_params (dict): Normalization parameters
+        device: Device to create tensors on
+        image_shape (tuple): Shape of input images (B, C, H, W)
+
+    Returns:
+        torch.Tensor: Settings channels [B, 2, H, W] where channels are [pixelsize, young]
+    """
+    batch_size, _, height, width = image_shape
+
+    # Create settings channels
+    pixelsize_channel = torch.zeros(batch_size, 1, height, width, device=device)
+    young_channel = torch.zeros(batch_size, 1, height, width, device=device)
+
+    for i in range(batch_size):
+        substrate = metadata['substrate'][i]
+        pixelsize_norm, young_norm = normalize_settings(
+            substrate, normalization_params, config_path=config_path
+        )
+
+        # Fill entire channel with normalized value
+        pixelsize_channel[i, 0] = pixelsize_norm
+        young_channel[i, 0] = young_norm
+
+    # Concatenate channels
+    settings_channels = torch.cat([pixelsize_channel, young_channel], dim=1)  # [B, 2, H, W]
+
+    return settings_channels
+
+
+class GlobalContextModule(nn.Module):
+    """Global context module for capturing cell shape information"""
+    def __init__(self, in_channels):
+        super().__init__()
+        self.global_pool = nn.AdaptiveAvgPool2d(1)
+        self.global_conv = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels//4, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels//4, in_channels, 1),
+            nn.Sigmoid()
+        )
+        self.large_kernel = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels, 3, padding=1, groups=in_channels),
+            nn.Conv2d(in_channels, in_channels, 1),
+            nn.BatchNorm2d(in_channels),
+            nn.ReLU(inplace=True)
+        )
+        self.multi_scale = nn.ModuleList([
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=1, dilation=1),
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=2, dilation=2),
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=4, dilation=4),
+            nn.Conv2d(in_channels, in_channels//4, 3, padding=8, dilation=8)
+        ])
+        self.fusion = nn.Conv2d(in_channels, in_channels, 1)
+
+    def forward(self, x):
+        global_ctx = self.global_pool(x)
+        global_weight = self.global_conv(global_ctx)
+        large_features = self.large_kernel(x)
+        multi_scale_features = []
+        for conv in self.multi_scale:
+            multi_scale_features.append(conv(x))
+        multi_scale_out = torch.cat(multi_scale_features, dim=1)
+        multi_scale_out = self.fusion(multi_scale_out)
+        return x + (large_features * global_weight) + multi_scale_out
+
+
+class HierarchicalAttention(nn.Module):
+    """Hierarchical attention combining spatial and channel attention"""
+    def __init__(self, channels):
+        super().__init__()
+        self.spatial_att = nn.Sequential(
+            nn.Conv2d(channels, channels//8, 1),
+            nn.Conv2d(channels//8, 1, 3, padding=1),
+            nn.Sigmoid()
+        )
+        self.channel_att = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(channels, channels//16, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels//16, channels, 1),
+            nn.Sigmoid()
+        )
+        self.cross_att = nn.Sequential(
+            nn.Conv2d(channels, channels//4, 1),
+            nn.BatchNorm2d(channels//4),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels//4, channels, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        spatial_weight = self.spatial_att(x)
+        channel_weight = self.channel_att(x)
+        attended = x * spatial_weight * channel_weight
+        cross_weight = self.cross_att(attended)
+        return x + (attended * cross_weight)
+
+
+class EnhancedAttentionGate(nn.Module):
+    """Enhanced attention gate with global context"""
+    def __init__(self, F_g, F_l, F_int):
+        super().__init__()
+        self.W_g = nn.Sequential(
+            nn.Conv2d(F_g, F_int, kernel_size=1),
+            nn.BatchNorm2d(F_int)
+        )
+        self.W_x = nn.Sequential(
+            nn.Conv2d(F_l, F_int, kernel_size=1),
+            nn.BatchNorm2d(F_int)
+        )
+        self.psi = nn.Sequential(
+            nn.ReLU(inplace=True),
+            nn.Conv2d(F_int, F_int//2, kernel_size=3, padding=1),
+            nn.BatchNorm2d(F_int//2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(F_int//2, 1, kernel_size=1),
+            nn.Sigmoid()
+        )
+        self.global_context = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(F_l, F_int//4, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(F_int//4, 1, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, g, x):
+        g1 = self.W_g(g)
+        x1 = self.W_x(x)
+        if g1.shape[2:] != x1.shape[2:]:
+            g1 = F.interpolate(g1, size=x1.shape[2:], mode='bilinear', align_corners=False)
+        psi = self.psi(g1 + x1)
+        global_weight = self.global_context(x)
+        psi = psi * global_weight
+        if psi.shape[2:] != x.shape[2:]:
+            psi = F.interpolate(psi, size=x.shape[2:], mode='bilinear', align_corners=False)
+        return x * psi
+
+
+class S2FGenerator(nn.Module):
+    """
+    S2F (Shape2Force) model: U-Net generator for force map prediction.
+    Supports substrate-specific settings as additional input channels.
+    """
+    def __init__(self,
+                 in_channels=1,
+                 out_channels=1,
+                 img_size=1024,
+                 bridge_type='cbam',
+                 use_multi_scale_input=True):
+        super().__init__()
+
+        self.img_size = img_size
+        self.bridge_type = bridge_type
+        self.use_multi_scale_input = use_multi_scale_input
+
+        if self.use_multi_scale_input:
+            self.scale_pyramid = nn.ModuleList([
+                nn.Conv2d(in_channels, 32, 3, padding=1),
+                nn.Sequential(
+                    nn.AvgPool2d(2, stride=2),
+                    nn.Conv2d(in_channels, 32, 3, padding=1)
+                ),
+                nn.Sequential(
+                    nn.AvgPool2d(4, stride=4),
+                    nn.Conv2d(in_channels, 32, 3, padding=1)
+                )
+            ])
+            self.initial_conv = nn.Conv2d(96, 64, 1)
+        else:
+            self.initial_conv = nn.Conv2d(in_channels, 64, 3, padding=1)
+
+        def enhanced_conv_block(in_c, out_c, use_attention=True):
+            layers = [
+                nn.Conv2d(in_c, out_c, 3, padding=1),
+                nn.BatchNorm2d(out_c),
+                nn.ReLU(inplace=True),
+                ResidualBlock(out_c, out_c)
+            ]
+            if use_attention:
+                layers.append(HierarchicalAttention(out_c))
+            return nn.Sequential(*layers)
+
+        def dilated_conv_block(in_c, out_c, use_global_context=False):
+            layers = [
+                nn.Conv2d(in_c, out_c, 3, padding=2, dilation=2),
+                nn.BatchNorm2d(out_c),
+                nn.ReLU(inplace=True),
+                ResidualBlock(out_c, out_c)
+            ]
+            if use_global_context:
+                layers.append(GlobalContextModule(out_c))
+            return nn.Sequential(*layers)
+
+        self.encoder1 = enhanced_conv_block(64, 64, use_attention=False)
+        self.pool1 = nn.MaxPool2d(2)
+        self.encoder2 = enhanced_conv_block(64, 128, use_attention=True)
+        self.pool2 = nn.MaxPool2d(2)
+        self.encoder3 = dilated_conv_block(128, 256, use_global_context=True)
+        self.pool3 = nn.MaxPool2d(2)
+        self.encoder4 = dilated_conv_block(256, 512, use_global_context=True)
+        self.pool4 = nn.MaxPool2d(2)
+
+        if bridge_type == 'cbam':
+            self.bridge = nn.Sequential(
+                dilated_conv_block(512, 1024, use_global_context=True),
+                CBAM(1024),
+                GlobalContextModule(1024),
+                HierarchicalAttention(1024)
+            )
+        else:
+            self.bridge = nn.Sequential(
+                dilated_conv_block(512, 1024, use_global_context=True),
+                GlobalContextModule(1024),
+                HierarchicalAttention(1024)
+            )
+
+        self.att4 = EnhancedAttentionGate(512, 512, 256)
+        self.att3 = EnhancedAttentionGate(256, 256, 128)
+        self.att2 = EnhancedAttentionGate(128, 128, 64)
+        self.att1 = EnhancedAttentionGate(64, 64, 32)
+
+        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
+        self.dec4 = enhanced_conv_block(1024, 512, use_attention=True)
+        self.refine4 = HierarchicalAttention(512)
+        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.dec3 = enhanced_conv_block(512, 256, use_attention=True)
+        self.refine3 = HierarchicalAttention(256)
+        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.dec2 = enhanced_conv_block(256, 128, use_attention=True)
+        self.refine2 = HierarchicalAttention(128)
+        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.dec1 = enhanced_conv_block(128, 64, use_attention=True)
+        self.refine1 = HierarchicalAttention(64)
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(64, 32, 3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, out_channels, 1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        if self.use_multi_scale_input:
+            scale_features = []
+            for i, scale_conv in enumerate(self.scale_pyramid):
+                if i == 0:
+                    scale_features.append(scale_conv(x))
+                else:
+                    scale_out = scale_conv(x)
+                    scale_out = F.interpolate(scale_out, size=x.shape[2:], mode='bilinear', align_corners=False)
+                    scale_features.append(scale_out)
+            fused = torch.cat(scale_features, dim=1)
+            initial_features = self.initial_conv(fused)
+        else:
+            initial_features = self.initial_conv(x)
+
+        e1 = self.encoder1(initial_features)
+        e2 = self.encoder2(self.pool1(e1))
+        e3 = self.encoder3(self.pool2(e2))
+        e4 = self.encoder4(self.pool3(e3))
+        b = self.bridge(self.pool4(e4))
+
+        g4 = self.up4(b)
+        x4 = self.att4(g4, e4)
+        d4 = self.dec4(torch.cat([g4, x4], dim=1))
+        d4 = self.refine4(d4)
+        g3 = self.up3(d4)
+        x3 = self.att3(g3, e3)
+        d3 = self.dec3(torch.cat([g3, x3], dim=1))
+        d3 = self.refine3(d3)
+        g2 = self.up2(d3)
+        x2 = self.att2(g2, e2)
+        d2 = self.dec2(torch.cat([g2, x2], dim=1))
+        d2 = self.refine2(d2)
+        g1 = self.up1(d2)
+        x1 = self.att1(g1, e1)
+        d1 = self.dec1(torch.cat([g1, x1], dim=1))
+        d1 = self.refine1(d1)
+        out = self.final_conv(d1)
+        return out
+
+    def load_checkpoint_with_expansion(self, checkpoint_path, strict=False):
+        """Load checkpoint and expand from 1-channel to 3-channel if needed."""
+        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+        generator_state = checkpoint['generator_state_dict']
+        needs_expansion = False
+
+        if 'scale_pyramid.0.weight' in generator_state:
+            old_shape = generator_state['scale_pyramid.0.weight'].shape
+            current_shape = self.scale_pyramid[0].weight.shape
+            if old_shape[1] != current_shape[1]:
+                needs_expansion = True
+        elif 'initial_conv.weight' in generator_state:
+            old_shape = generator_state['initial_conv.weight'].shape
+            current_shape = self.initial_conv.weight.shape
+            if old_shape[1] != current_shape[1]:
+                needs_expansion = True
+
+        if needs_expansion:
+            generator_state = self._expand_generator_state(generator_state)
+
+        self.load_state_dict(generator_state, strict=strict)
+        return checkpoint
+
+    def _expand_generator_state(self, generator_state):
+        """Expand generator state dict from 1-channel to 3-channel input."""
+        expanded_state = generator_state.copy()
+        if 'scale_pyramid.0.weight' in generator_state:
+            for i in range(3):
+                key = f'scale_pyramid.{i}.weight' if i == 0 else f'scale_pyramid.{i}.1.weight'
+                if key in generator_state:
+                    old_weight = generator_state[key]
+                    new_weight = torch.zeros(32, 3, 3, 3)
+                    new_weight[:, 0:1, :, :] = old_weight
+                    expanded_state[key] = new_weight
+        elif 'initial_conv.weight' in generator_state:
+            old_weight = generator_state['initial_conv.weight']
+            new_weight = torch.zeros(64, 3, 3, 3)
+            new_weight[:, 0:1, :, :] = old_weight
+            expanded_state['initial_conv.weight'] = new_weight
+        return expanded_state
+
+
+class PatchGANDiscriminator(nn.Module):
+    """PatchGAN Discriminator (included for create_s2f_model compatibility)."""
+    def __init__(self, in_channels=2, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        super().__init__()
+        use_bias = norm_layer == nn.InstanceNorm2d
+        self.initial_conv = nn.Sequential(
+            nn.Conv2d(in_channels, ndf, kernel_size=4, stride=2, padding=1, bias=use_bias),
+            nn.LeakyReLU(0.2, inplace=True)
+        )
+        self.layers = nn.ModuleList()
+        nf_mult, nf_mult_prev = 1, 1
+        for n in range(1, n_layers):
+            nf_mult_prev, nf_mult = nf_mult, min(2 ** n, 8)
+            self.layers.append(nn.Sequential(
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=4, stride=2, padding=1, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, inplace=True)
+            ))
+        nf_mult_prev, nf_mult = nf_mult, min(2 ** n_layers, 8)
+        self.layers.append(nn.Sequential(
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=4, stride=1, padding=1, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, inplace=True)
+        ))
+        self.output_conv = nn.Conv2d(ndf * nf_mult, 1, kernel_size=4, stride=1, padding=1)
+        self.attention = nn.Sequential(
+            nn.Conv2d(ndf * nf_mult, ndf * nf_mult // 4, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(ndf * nf_mult // 4, ndf * nf_mult, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, input):
+        x = self.initial_conv(input)
+        for layer in self.layers:
+            x = layer(x)
+        x = x * self.attention(x)
+        return self.output_conv(x)
+
+
+def create_s2f_model(
+    in_channels=1,
+    out_channels=1,
+    img_size=1024,
+    bridge_type='cbam',
+    use_multi_scale_input=True,
+    ndf=64,
+    n_layers=3,
+):
+    """Create S2F model with generator and discriminator."""
+    generator = S2FGenerator(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        img_size=img_size,
+        bridge_type=bridge_type,
+        use_multi_scale_input=use_multi_scale_input,
+    )
+    discriminator = PatchGANDiscriminator(
+        in_channels=in_channels + out_channels,
+        ndf=ndf,
+        n_layers=n_layers
+    )
+    return generator, discriminator

notebooks/evaluate_model.ipynb

@@ -0,0 +1,248 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# S2F Model Evaluation\n",
+    "\n",
+    "This notebook shows how to evaluate a trained Shape2Force (S2F) model on your dataset.\n",
+    "\n",
+    "**Metrics computed:**\n",
+    "- **MSE** – Mean Squared Error\n",
+    "- **MS-SSIM** – Multi-Scale Structural Similarity\n",
+    "- **Pixel Correlation** – Pearson correlation between predicted and ground-truth heatmaps\n",
+    "- **Relative Magnitude Error** – WFM-style weighted relative error\n",
+    "- **Force Sum/Mean Correlation** – Correlation of total force per sample\n",
+    "\n",
+    "Run the cells below after adjusting paths and settings."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import warnings\n",
+    "warnings.filterwarnings('ignore')\n",
+    "import sys\n",
+    "import os\n",
+    "import cv2\n",
+    "cv2.utils.logging.setLogLevel(cv2.utils.logging.LOG_LEVEL_ERROR)\n",
+    "\n",
+    "cwd = os.getcwd()\n",
+    "S2F_ROOT = cwd if os.path.exists(os.path.join(cwd, 'models')) else os.path.dirname(cwd)\n",
+    "sys.path.insert(0, S2F_ROOT)\n",
+    "\n",
+    "import torch\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from data.cell_dataset import prepare_data, load_folder_data\n",
+    "from models.s2f_model import create_s2f_model, compute_settings_normalization\n",
+    "from utils.metrics import (\n",
+    "    evaluate_metrics_on_dataset,\n",
+    "    print_metrics_report,\n",
+    "    gen_prediction_plots,\n",
+    "    plot_predictions,\n",
+    ")\n",
+    "\n",
+    "print(f\"S2F root: {S2F_ROOT}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Adjust these paths ---\n",
+    "USE_SINGLE_CELL = True  # True = single-cell (with substrate), False = spheroid\n",
+    "DATA_FOLDER = os.path.join(S2F_ROOT, 'sample')  # or path to your dataset\n",
+    "CHECKPOINT_PATH = os.path.join(S2F_ROOT, 'ckp', 'best_checkpoint.pth')\n",
+    "\n",
+    "IMAGE_SIZE = 1024\n",
+    "BATCH_SIZE = 2\n",
+    "THRESHOLD = 0.0  # Threshold for heatmap metrics (0 = no threshold)\n",
+    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "SUBSTRATE = 'fibroblasts_PDMS'  # For single-cell mode"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Option A: Load from folder with train/test subfolders\n",
+    "train_folder = os.path.join(DATA_FOLDER, 'train')\n",
+    "test_folder = os.path.join(DATA_FOLDER, 'test')\n",
+    "\n",
+    "if os.path.exists(train_folder) and os.path.exists(test_folder):\n",
+    "    train_loader, val_loader = prepare_data(\n",
+    "        DATA_FOLDER,\n",
+    "        batch_size=BATCH_SIZE,\n",
+    "        target_size=(IMAGE_SIZE, IMAGE_SIZE),\n",
+    "        use_augmentations=False,\n",
+    "        train_test_sep_folder=True,\n",
+    "        return_metadata=USE_SINGLE_CELL,\n",
+    "        substrate=SUBSTRATE if USE_SINGLE_CELL else None,\n",
+    "    )\n",
+    "    print(f\"Loaded train: {len(train_loader.dataset)} samples, val: {len(val_loader.dataset)} samples\")\n",
+    "else:\n",
+    "    # Option B: Load from a single folder (e.g. test only)\n",
+    "    val_loader = load_folder_data(\n",
+    "        DATA_FOLDER,\n",
+    "        substrate=SUBSTRATE if USE_SINGLE_CELL else None,\n",
+    "        img_size=IMAGE_SIZE,\n",
+    "        batch_size=BATCH_SIZE,\n",
+    "        return_metadata=USE_SINGLE_CELL,\n",
+    "    )\n",
+    "    train_loader = None\n",
+    "    print(f\"Loaded {len(val_loader.dataset)} samples from {DATA_FOLDER}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "in_channels = 3 if USE_SINGLE_CELL else 1\n",
+    "generator, _ = create_s2f_model(in_channels=in_channels)\n",
+    "checkpoint = torch.load(CHECKPOINT_PATH, map_location='cpu', weights_only=False)\n",
+    "generator.load_state_dict(checkpoint.get('generator_state_dict', checkpoint), strict=True)\n",
+    "generator = generator.to(DEVICE)\n",
+    "generator.eval()\n",
+    "\n",
+    "print(f\"Loaded checkpoint from {CHECKPOINT_PATH}\")\n",
+    "if 'epoch' in checkpoint:\n",
+    "    print(f\"  Epoch: {checkpoint['epoch']}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run Evaluation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config_path = os.path.join(S2F_ROOT, 'config', 'substrate_settings.json')\n",
+    "normalization_params = compute_settings_normalization(config_path=config_path) if USE_SINGLE_CELL else None\n",
+    "\n",
+    "val_results = evaluate_metrics_on_dataset(\n",
+    "    generator,\n",
+    "    val_loader,\n",
+    "    device=DEVICE,\n",
+    "    description=\"Evaluating\",\n",
+    "    save_predictions=True,\n",
+    "    threshold=THRESHOLD,\n",
+    "    use_settings=USE_SINGLE_CELL,\n",
+    "    normalization_params=normalization_params,\n",
+    "    config_path=config_path,\n",
+    "    substrate_override=SUBSTRATE,\n",
+    ")\n",
+    "\n",
+    "report = {'validation': val_results}\n",
+    "if train_loader is not None:\n",
+    "    train_results = evaluate_metrics_on_dataset(\n",
+    "        generator, train_loader, device=DEVICE, description=\"Training\",\n",
+    "        threshold=THRESHOLD, use_settings=USE_SINGLE_CELL,\n",
+    "        normalization_params=normalization_params, config_path=config_path,\n",
+    "        substrate_override=SUBSTRATE,\n",
+    "    )\n",
+    "    report = {'train': train_results, 'validation': val_results}\n",
+    "\n",
+    "print_metrics_report(report, threshold=THRESHOLD)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Visualize Predictions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Quick preview: plot first few samples\n",
+    "plot_predictions(\n",
+    "    val_loader,\n",
+    "    generator,\n",
+    "    n_samples=3,\n",
+    "    device=DEVICE,\n",
+    "    use_settings=USE_SINGLE_CELL,\n",
+    "    normalization_params=normalization_params,\n",
+    "    config_path=config_path,\n",
+    "    substrate_override=SUBSTRATE,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Save individual prediction plots (sorted by MSE, best first)\n",
+    "plot_output_dir = os.path.join(S2F_ROOT, 'outputs', 'evaluation_plots')\n",
+    "if 'individual_predictions' in val_results:\n",
+    "    gen_prediction_plots(\n",
+    "        val_results['individual_predictions'],\n",
+    "        save_dir=plot_output_dir,\n",
+    "        sort_by='mse',\n",
+    "        sort_order='asc',\n",
+    "        threshold=THRESHOLD,\n",
+    "    )\n",
+    "    print(f\"Saved plots to {plot_output_dir}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

requirements.txt

@@ -0,0 +1,18 @@
+# Shape2Force (S2F) - GUI and training
+torch>=2.0.0
+torchvision>=0.15.0
+numpy>=1.20.0
+opencv-python>=4.5.0
+streamlit>=1.28.0
+matplotlib>=3.5.0
+Pillow>=9.0.0
+plotly>=5.14.0
+
+# Training & evaluation
+torchmetrics>=1.0.0
+diffusers>=0.20.0
+tqdm>=4.65.0
+scikit-image>=0.19.0
+scipy>=1.9.0
+scikit-learn>=1.0.0
+pandas>=1.3.0

training/__init__.py

@@ -0,0 +1 @@
+from .s2f_trainer import train_s2f, S2FLoss, calculate_soft_dice_loss

training/evaluate.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""
+S2F evaluation script.
+Metrics: MSE, MS-SSIM, Pixel Correlation, Relative Magnitude Error, Force Sum/Mean correlation.
+
+Usage:
+  python -m training.evaluate --model single_cell --checkpoint ckp/best_checkpoint.pth --data path/to/test
+  python -m training.evaluate --model spheroid --checkpoint ckp/best_checkpoint.pth --data path/to/test
+"""
+import os
+import sys
+import argparse
+
+S2F_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if S2F_ROOT not in sys.path:
+    sys.path.insert(0, S2F_ROOT)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Evaluate S2F model')
+    parser.add_argument('--data', required=True, help='Path to test folder (subfolders with BF_001.tif, *_gray.jpg)')
+    parser.add_argument('--model', choices=['single_cell', 'spheroid'], default='single_cell')
+    parser.add_argument('--checkpoint', required=True, help='Path to .pth checkpoint')
+    parser.add_argument('--substrate', default='fibroblasts_PDMS', help='Substrate for single_cell')
+    parser.add_argument('--batch_size', type=int, default=2)
+    parser.add_argument('--img_size', type=int, default=1024)
+    parser.add_argument('--threshold', type=float, default=0.0, help='Threshold for heatmap metrics')
+    parser.add_argument('--output', default=None, help='Optional CSV path for per-sample metrics')
+    parser.add_argument('--save_plots', default=None, help='Directory to save prediction plots')
+    parser.add_argument('--device', default='cuda')
+    args = parser.parse_args()
+
+    from data.cell_dataset import load_folder_data
+    from models.s2f_model import create_s2f_model, compute_settings_normalization
+    from utils.metrics import (
+        evaluate_metrics_on_dataset,
+        print_metrics_report,
+        gen_prediction_plots,
+        detect_tanh_output_model,
+    )
+    import torch
+    import pandas as pd
+
+    use_settings = args.model == 'single_cell'
+    config_path = os.path.join(S2F_ROOT, 'config', 'substrate_settings.json')
+
+    print(f"Loading data from {args.data}")
+    val_loader = load_folder_data(
+        args.data,
+        substrate=args.substrate if use_settings else None,
+        img_size=args.img_size,
+        batch_size=args.batch_size,
+        return_metadata=use_settings,
+    )
+
+    in_channels = 3 if use_settings else 1
+    generator, _ = create_s2f_model(in_channels=in_channels)
+    ckpt = torch.load(args.checkpoint, map_location='cpu', weights_only=False)
+    generator.load_state_dict(ckpt.get('generator_state_dict', ckpt), strict=True)
+
+    norm_params = compute_settings_normalization(config_path=config_path) if use_settings else None
+    uses_tanh = detect_tanh_output_model(generator)
+
+    results = evaluate_metrics_on_dataset(
+        generator,
+        val_loader,
+        device=args.device,
+        description="Evaluating",
+        save_predictions=(args.save_plots is not None or args.output is not None),
+        threshold=args.threshold,
+        use_settings=use_settings,
+        normalization_params=norm_params,
+        config_path=config_path,
+        substrate_override=args.substrate,
+    )
+
+    report = {'validation': results}
+    print_metrics_report(report, threshold=args.threshold, uses_tanh=uses_tanh)
+    print(f"Samples: {len(val_loader.dataset)}")
+
+    if args.save_plots and 'individual_predictions' in results:
+        gen_prediction_plots(
+            results['individual_predictions'],
+            args.save_plots,
+            sort_by='mse',
+            sort_order='asc',
+            threshold=args.threshold,
+        )
+        print(f"Saved prediction plots to {args.save_plots}")
+
+    if args.output:
+        preds = results.get('individual_predictions', [])
+        if preds:
+            df = pd.DataFrame([{
+                'mse': p['mse'],
+                'ms_ssim': p['ms_ssim'],
+                'pixel_correlation': p['pixel_correlation'],
+                'relative_magnitude_error': p.get('wfm_relative_magnitude_error'),
+                'force_sum_gt': p['force_sum_gt'],
+                'force_sum_pred': p['force_sum_pred'],
+            } for p in preds])
+            df.to_csv(args.output, index=False)
+            print(f"Saved per-sample metrics to {args.output}")
+        else:
+            # Fallback: write summary only
+            with open(args.output.replace('.csv', '_summary.txt'), 'w') as f:
+                f.write(f"MSE: {results['heatmap']['mse']:.6f}\n")
+                f.write(f"MS-SSIM: {results['heatmap']['ms_ssim']:.4f}\n")
+                f.write(f"Pixel Corr: {results['heatmap']['pixel_correlation']:.4f}\n")
+                f.write(f"Rel Mag Error: {results['wfm']['relative_magnitude_error']:.4f}\n")
+            print(f"Saved summary to {args.output.replace('.csv', '_summary.txt')}")
+
+
+if __name__ == '__main__':
+    main()

training/s2f_trainer.py

@@ -0,0 +1,278 @@
+"""
+S2F training logic: loss, metrics, and training loop.
+"""
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+from tqdm.auto import tqdm
+
+# Add S2F root to path
+S2F_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if S2F_ROOT not in sys.path:
+    sys.path.insert(0, S2F_ROOT)
+
+from models.s2f_model import create_settings_channels, compute_settings_normalization
+from utils.metrics import calculate_psnr, calculate_ssim_tensor, calculate_pearson_correlation
+from scipy.stats import pearsonr
+
+
+class S2FLoss(nn.Module):
+    """S2F loss: reconstruction (L1) + GAN + optional force consistency."""
+    def __init__(self, lambda_L1=100.0, lambda_gan=1.0, lambda_force=1.0,
+                 gan_mode='vanilla', custom_loss=None, use_force_consistency=False,
+                 force_consistency_target='mean'):
+        super().__init__()
+        self.lambda_L1 = lambda_L1
+        self.lambda_gan = lambda_gan
+        self.lambda_force = lambda_force
+        self.gan_mode = gan_mode
+        self.use_force_consistency = use_force_consistency
+        self.force_consistency_target = force_consistency_target
+        self.reconstruction_loss = custom_loss if custom_loss is not None else nn.L1Loss()
+        self.force_consistency_loss = nn.MSELoss() if use_force_consistency else None
+        self.gan_loss = nn.BCEWithLogitsLoss() if gan_mode == 'vanilla' else nn.MSELoss()
+
+    def forward(self, pred, target, disc_pred=None, disc_target=None):
+        recon_loss = self.reconstruction_loss(pred, target)
+        gan_loss = 0.0
+        if disc_pred is not None and disc_target is not None:
+            gan_loss = self.gan_loss(disc_pred, disc_target)
+        force_loss = 0.0
+        if self.use_force_consistency and self.force_consistency_loss is not None:
+            if self.force_consistency_target == 'mean':
+                pred_global = torch.mean(pred.view(pred.size(0), -1), dim=1, keepdim=True)
+                target_global = torch.mean(target.view(target.size(0), -1), dim=1, keepdim=True)
+            else:
+                pred_global = torch.sum(pred.view(pred.size(0), -1), dim=1, keepdim=True)
+                target_global = torch.sum(target.view(target.size(0), -1), dim=1, keepdim=True)
+            force_loss = self.force_consistency_loss(pred_global, target_global)
+        total = self.lambda_L1 * recon_loss + self.lambda_gan * gan_loss + self.lambda_force * force_loss
+        return total, recon_loss, gan_loss, force_loss
+
+
+def calculate_soft_dice_loss(pred, target, smooth=1e-6):
+    """Dice score (higher is better)."""
+    pred_flat = pred.view(pred.size(0), -1)
+    target_flat = target.view(target.size(0), -1)
+    intersection = (pred_flat * target_flat).sum(dim=1)
+    dice_scores = (2.0 * intersection + smooth) / (pred_flat.sum(dim=1) + target_flat.sum(dim=1) + smooth)
+    return dice_scores.mean().item()
+
+
+def train_s2f(generator, discriminator, train_loader, val_loader, device='cuda',
+              num_epochs=100, g_lr=2e-4, d_lr=2e-4, beta1=0.5, beta2=0.999,
+              save_dir='ckp', lambda_L1=100.0, lambda_gan=1.0, lambda_force=1.0,
+              gan_mode='vanilla', save_predictions_every=5, custom_loss=None,
+              loaded_metadata=False, use_settings=False, use_force_consistency=False,
+              force_consistency_target='mean', config_path=None):
+    """
+    Train S2F model.
+    """
+    from diffusers.optimization import get_cosine_schedule_with_warmup
+
+    config_path = config_path or os.path.join(S2F_ROOT, 'config', 'substrate_settings.json')
+    normalization_params = None
+    if use_settings:
+        if not loaded_metadata:
+            raise ValueError("loaded_metadata must be True when use_settings=True")
+        normalization_params = compute_settings_normalization(config_path=config_path)
+
+    history = {'g_loss': [], 'd_loss': [], 'g_recon_loss': [], 'g_gan_loss': [], 'g_force_loss': [],
+               'train_loss': [], 'val_loss': [], 'train_ssim': [], 'val_ssim': [],
+               'train_psnr': [], 'val_psnr': [], 'train_mse': [], 'val_mse': [],
+               'train_dice_score': [], 'val_dice_score': []}
+
+    if not torch.cuda.is_available() and device == 'cuda':
+        device = 'cpu'
+    generator = generator.to(device)
+    discriminator = discriminator.to(device)
+    criterion = S2FLoss(lambda_L1=lambda_L1, lambda_gan=lambda_gan, lambda_force=lambda_force,
+                        gan_mode=gan_mode, custom_loss=custom_loss,
+                        use_force_consistency=use_force_consistency,
+                        force_consistency_target=force_consistency_target)
+    g_optimizer = torch.optim.Adam(generator.parameters(), lr=g_lr, betas=(beta1, beta2))
+    d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=d_lr, betas=(beta1, beta2))
+    num_steps = len(train_loader) * num_epochs
+    g_scheduler = get_cosine_schedule_with_warmup(g_optimizer, int(num_steps * 0.1), num_steps)
+    d_scheduler = get_cosine_schedule_with_warmup(d_optimizer, int(num_steps * 0.1), num_steps)
+
+    os.makedirs(save_dir, exist_ok=True)
+    vis_dir = os.path.join(save_dir, 'visualizations')
+    os.makedirs(vis_dir, exist_ok=True)
+    best_val_loss = float('inf')
+    disc_output_shape = None
+
+    for epoch in range(num_epochs):
+        generator.train()
+        discriminator.train()
+        g_loss_total = d_loss_total = g_recon_total = g_gan_total = g_force_total = 0.0
+        train_ssim = train_psnr = train_mse = train_dice = 0.0
+        pbar = tqdm(train_loader, desc=f'Epoch {epoch}')
+
+        for batch_data in pbar:
+            if loaded_metadata:
+                input_images, target_images, _, _, metadata = batch_data
+            else:
+                input_images, target_images, _, _ = batch_data
+            input_images = input_images.to(device, dtype=torch.float32)
+            target_images = target_images.to(device, dtype=torch.float32)
+            batch_size = input_images.size(0)
+
+            if use_settings and normalization_params is not None:
+                settings_channels = create_settings_channels(
+                    metadata, normalization_params, device, input_images.shape,
+                    config_path=config_path
+                )
+                input_images = torch.cat([input_images, settings_channels], dim=1)
+
+            target_scaled = target_images * 2.0 - 1.0
+            if disc_output_shape is None:
+                with torch.no_grad():
+                    dummy = torch.cat([input_images[:1], target_scaled[:1]], dim=1)
+                    disc_output_shape = discriminator(dummy).shape[2:]
+            real_labels = torch.ones(batch_size, 1, *disc_output_shape).to(device)
+            fake_labels = torch.zeros(batch_size, 1, *disc_output_shape).to(device)
+
+            g_optimizer.zero_grad()
+            fake_images = generator(input_images)
+            fake_for_loss = (fake_images + 1.0) / 2.0
+            fake_input = torch.cat([input_images, fake_images], dim=1)
+            fake_pred = discriminator(fake_input)
+            g_loss, g_recon, g_gan, g_force = criterion(fake_for_loss, target_images, fake_pred, real_labels)
+            g_loss.backward()
+            g_optimizer.step()
+
+            d_optimizer.zero_grad()
+            real_input = torch.cat([input_images, target_scaled], dim=1)
+            real_pred = discriminator(real_input)
+            d_real = criterion.gan_loss(real_pred, real_labels)
+            fake_input_d = torch.cat([input_images, fake_images.detach()], dim=1)
+            fake_pred_d = discriminator(fake_input_d)
+            d_fake = criterion.gan_loss(fake_pred_d, fake_labels)
+            d_loss = (d_real + d_fake) * 0.5
+            d_loss.backward()
+            d_optimizer.step()
+            g_scheduler.step()
+            d_scheduler.step()
+
+            g_loss_total += g_loss.item()
+            d_loss_total += d_loss.item()
+            g_recon_total += g_recon.item()
+            g_gan_total += g_gan.item()
+            g_force_total += g_force.item() if isinstance(g_force, torch.Tensor) else g_force
+            train_ssim += calculate_ssim_tensor(fake_for_loss, target_images)
+            train_psnr += calculate_psnr(fake_for_loss, target_images)
+            train_mse += F.mse_loss(fake_for_loss, target_images).item()
+            train_dice += calculate_soft_dice_loss(fake_for_loss, target_images)
+            pbar.set_postfix({'G': g_loss.item(),
+                'D': d_loss.item(), 'Dice': train_dice / (pbar.n + 1)})
+
+        n_train = len(train_loader)
+        g_loss_total /= n_train
+        d_loss_total /= n_train
+        train_ssim /= n_train
+        train_psnr /= n_train
+        train_mse /= n_train
+        train_dice /= n_train
+
+        generator.eval()
+        val_loss = val_ssim = val_psnr = val_mse = val_dice = 0.0
+        with torch.no_grad():
+            for batch_data in val_loader:
+                if loaded_metadata:
+                    input_images, target_images, _, _, metadata = batch_data
+                else:
+                    input_images, target_images, _, _ = batch_data
+                input_images = input_images.to(device, dtype=torch.float32)
+                target_images = target_images.to(device, dtype=torch.float32)
+                if use_settings and normalization_params is not None:
+                    settings_channels = create_settings_channels(
+                        metadata, normalization_params, device, input_images.shape,
+                        config_path=config_path
+                    )
+                    input_images = torch.cat([input_images, settings_channels], dim=1)
+                fake_images = generator(input_images)
+                fake_for_loss = (fake_images + 1.0) / 2.0
+                _, recon_loss, _, force_loss = criterion(fake_for_loss, target_images)
+                val_loss += recon_loss.item()
+                val_ssim += calculate_ssim_tensor(fake_for_loss, target_images)
+                val_psnr += calculate_psnr(fake_for_loss, target_images)
+                val_mse += F.mse_loss(fake_for_loss, target_images).item()
+                val_dice += calculate_soft_dice_loss(fake_for_loss, target_images)
+        n_val = len(val_loader)
+        val_loss /= n_val
+        val_ssim /= n_val
+        val_psnr /= n_val
+        val_mse /= n_val
+        val_dice /= n_val
+
+        history['g_loss'].append(g_loss_total)
+        history['d_loss'].append(d_loss_total)
+        history['train_loss'].append(g_loss_total)
+        history['val_loss'].append(val_loss)
+        history['train_ssim'].append(train_ssim)
+        history['val_ssim'].append(val_ssim)
+        history['train_psnr'].append(train_psnr)
+        history['val_psnr'].append(val_psnr)
+        history['train_mse'].append(train_mse)
+        history['val_mse'].append(val_mse)
+        history['train_dice_score'].append(train_dice)
+        history['val_dice_score'].append(val_dice)
+
+        best_mark = "✓" if val_loss < best_val_loss else ""
+        print(f"Train: G_Loss:{g_loss_total:.4f} D_Loss:{d_loss_total:.4f} "
+              f"MSE:{train_mse:.4f} SSIM:{train_ssim:.4f} Dice:{train_dice:.4f}")
+        print(f"Valid: Loss:{val_loss:.4f} MSE:{val_mse:.4f} SSIM:{val_ssim:.4f} Dice:{val_dice:.4f} {best_mark}")
+
+        checkpoint = {
+            'epoch': epoch,
+            'generator_state_dict': generator.state_dict(),
+            'discriminator_state_dict': discriminator.state_dict(),
+            'g_optimizer_state_dict': g_optimizer.state_dict(),
+            'd_optimizer_state_dict': d_optimizer.state_dict(),
+            'val_loss': val_loss,
+            'history': history
+        }
+        torch.save(checkpoint, os.path.join(save_dir, 'last_checkpoint.pth'))
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            torch.save(checkpoint, os.path.join(save_dir, 'best_checkpoint.pth'))
+
+        if epoch % save_predictions_every == 0:
+            generator.eval()
+            with torch.no_grad():
+                batch_data = next(iter(val_loader))
+                if loaded_metadata:
+                    input_images, target_images, _, _, metadata = batch_data
+                else:
+                    input_images, target_images, _, _ = batch_data
+                input_images = input_images.to(device, dtype=torch.float32)
+                target_images = target_images.to(device, dtype=torch.float32)
+                if use_settings and normalization_params is not None:
+                    settings_channels = create_settings_channels(
+                        metadata, normalization_params, device, input_images.shape,
+                        config_path=config_path
+                    )
+                    input_images = torch.cat([input_images, settings_channels], dim=1)
+                fake_images = generator(input_images)
+                fake_vis = (fake_images + 1.0) / 2.0
+                n_vis = min(4, input_images.size(0))
+                fig, axes = plt.subplots(3, n_vis, figsize=(4 * n_vis, 12))
+                if n_vis == 1:
+                    axes = axes.reshape(3, 1)
+                for i in range(n_vis):
+                    axes[0, i].imshow(input_images[i, 0].cpu().numpy(), cmap='gray')
+                    axes[0, i].axis('off')
+                    axes[1, i].imshow(fake_vis[i, 0].cpu().numpy(), cmap='jet', vmin=0, vmax=1)
+                    axes[1, i].axis('off')
+                    axes[2, i].imshow(target_images[i, 0].cpu().numpy(), cmap='jet', vmin=0, vmax=1)
+                    axes[2, i].axis('off')
+                plt.tight_layout()
+                plt.savefig(os.path.join(vis_dir, f'predictions_epoch_{epoch:02d}.png'), dpi=150, bbox_inches='tight')
+                plt.close()
+                print(f"Saved visualization for epoch {epoch}")
+
+    return history

training/train.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""
+S2F training script.
+Usage:
+  python -m training.train --data path/to/dataset --model single_cell --epochs 100
+  python -m training.train --data path/to/dataset --model spheroid --epochs 50
+"""
+import os
+import sys
+import argparse
+
+S2F_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if S2F_ROOT not in sys.path:
+    sys.path.insert(0, S2F_ROOT)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Train S2F model')
+    parser.add_argument('--data', required=True, help='Path to dataset (must have train/ and test/ subfolders)')
+    parser.add_argument('--model', choices=['single_cell', 'spheroid'], default='single_cell',
+                        help='Model type: single_cell (with substrate) or spheroid')
+    parser.add_argument('--substrate', default=None,
+                        help='Substrate name for single_cell when metadata not in dataset (e.g. fibroblasts_PDMS)')
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=4)
+    parser.add_argument('--img_size', type=int, default=1024)
+    parser.add_argument('--save_dir', default='ckp', help='Checkpoint save directory')
+    parser.add_argument('--g_lr', type=float, default=2e-4)
+    parser.add_argument('--d_lr', type=float, default=2e-4)
+    parser.add_argument('--resume', type=str, default=None, help='Path to checkpoint to resume from')
+    parser.add_argument('--device', default='cuda')
+    parser.add_argument('--no_augment', action='store_true', help='Disable augmentations')
+    parser.add_argument('--use_force_consistency', action='store_true')
+    parser.add_argument('--force_target', choices=['mean', 'sum'], default='mean')
+    args = parser.parse_args()
+
+    from data.cell_dataset import prepare_data
+    from models.s2f_model import create_s2f_model
+    from training.s2f_trainer import train_s2f
+
+    use_settings = args.model == 'single_cell'
+    substrate = args.substrate or 'fibroblasts_PDMS'
+    return_metadata = use_settings
+
+    print(f"Loading data from {args.data} (model={args.model})")
+    train_loader, val_loader = prepare_data(
+        args.data,
+        batch_size=args.batch_size,
+        target_size=(args.img_size, args.img_size),
+        use_augmentations=not args.no_augment,
+        train_test_sep_folder=True,
+        return_metadata=return_metadata,
+        substrate=substrate if use_settings else None,
+    )
+
+    in_channels = 3 if use_settings else 1
+    generator, discriminator = create_s2f_model(in_channels=in_channels)
+
+    if args.resume:
+        ckpt = __import__('torch').load(args.resume, map_location='cpu', weights_only=False)
+        generator.load_state_dict(ckpt.get('generator_state_dict', ckpt), strict=True)
+        print(f"Resumed from {args.resume}")
+
+    history = train_s2f(
+        generator, discriminator,
+        train_loader, val_loader,
+        device=args.device,
+        num_epochs=args.epochs,
+        g_lr=args.g_lr, d_lr=args.d_lr,
+        save_dir=args.save_dir,
+        loaded_metadata=return_metadata,
+        use_settings=use_settings,
+        use_force_consistency=args.use_force_consistency,
+        force_consistency_target=args.force_target,
+    )
+    print(f"Training complete. Checkpoints saved to {args.save_dir}")
+
+
+if __name__ == '__main__':
+    main()

utils/__init__.py

@@ -0,0 +1 @@
+from . import config

utils/config.py

@@ -0,0 +1,3 @@
+# Constants for force and area scaling (used in force map prediction)
+SCALE_FACTOR_FORCE = 1e-3
+SCALE_FACTOR_AREA = 1e-4

utils/metrics.py

@@ -0,0 +1,447 @@
+"""Metrics for S2F training and evaluation.
+
+Includes: MSE, MS-SSIM, Pixel Correlation (Pearson), Relative Magnitude Error (WFM),
+and evaluation helpers for notebooks and scripts.
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from skimage.metrics import structural_similarity as ssim
+from scipy.stats import pearsonr
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+try:
+    from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure
+    from torchmetrics import MeanSquaredError
+    HAS_TORCHMETRICS = True
+except ImportError:
+    HAS_TORCHMETRICS = False
+
+
+def calculate_mse(y_true, y_pred):
+    if isinstance(y_true, torch.Tensor):
+        return F.mse_loss(y_pred, y_true).item()
+    return float(np.mean((np.asarray(y_true) - np.asarray(y_pred)) ** 2))
+
+
+def calculate_psnr(y_true, y_pred, max_pixel_value=1.0):
+    mse = calculate_mse(y_true, y_pred)
+    if mse == 0:
+        return float('inf')
+    return 20 * np.log10(max_pixel_value / np.sqrt(mse))
+
+
+def calculate_ssim_tensor(y_true, y_pred, data_range=1.0):
+    if isinstance(y_true, torch.Tensor):
+        y_true = y_true.detach().cpu().numpy()
+    if isinstance(y_pred, torch.Tensor):
+        y_pred = y_pred.detach().cpu().numpy()
+    ssim_values = []
+    batch_size = y_true.shape[0]
+    for i in range(batch_size):
+        if len(y_true.shape) == 4:
+            true_img = y_true[i, 0] if y_true.shape[1] == 1 else y_true[i, 0]
+            pred_img = y_pred[i, 0] if y_pred.shape[1] == 1 else y_pred[i, 0]
+        else:
+            true_img, pred_img = y_true[i], y_pred[i]
+        ssim_values.append(ssim(true_img, pred_img, data_range=data_range))
+    return np.mean(ssim_values)
+
+
+def calculate_pearson_correlation(y_true, y_pred):
+    if isinstance(y_true, torch.Tensor):
+        y_true = y_true.cpu().numpy()
+    if isinstance(y_pred, torch.Tensor):
+        y_pred = y_pred.cpu().numpy()
+    correlation, _ = pearsonr(y_true.flatten(), y_pred.flatten())
+    return correlation
+
+
+def calculate_individual_pixel_correlation(y_true, y_pred):
+    """Pixel-wise Pearson correlation per sample in batch."""
+    if isinstance(y_true, torch.Tensor):
+        y_true = y_true.cpu().numpy()
+    if isinstance(y_pred, torch.Tensor):
+        y_pred = y_pred.cpu().numpy()
+    correlations = []
+    batch_size = y_true.shape[0]
+    for i in range(batch_size):
+        true_flat = y_true[i].flatten()
+        pred_flat = y_pred[i].flatten()
+        r, _ = pearsonr(true_flat, pred_flat)
+        correlations.append(r)
+    return correlations
+
+
+# --- WFM (Wrinkle Force Microscopy) metrics for heatmap as magnitude ---
+
+def _to_numpy_wfm(x):
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy()
+    return np.asarray(x)
+
+
+def _ensure_shape_wfm(f):
+    """Ensure (N, 2, H, W). Heatmap -> fx=magnitude, fy=0."""
+    if f.ndim == 3:
+        if f.shape[-1] == 2:
+            f = np.transpose(f, (2, 0, 1))[None, ...]
+        elif f.shape[0] == 2:
+            f = f[None, ...]
+        else:
+            raise ValueError(f"Unsupported 3D shape {f.shape}")
+    elif f.ndim == 4:
+        if f.shape[-1] == 2:
+            f = np.transpose(f, (0, 3, 1, 2))
+    else:
+        raise ValueError(f"Unsupported ndim={f.ndim}")
+    return f
+
+
+def _force_mag_wfm(f):
+    fx, fy = f[:, 0], f[:, 1]
+    return np.sqrt(fx**2 + fy**2)
+
+
+def wfm_correlation(y_true, y_pred, mode="magnitude"):
+    """Pearson correlation between prediction and ground truth (magnitude mode for heatmaps)."""
+    t = _ensure_shape_wfm(_to_numpy_wfm(y_true))
+    p = _ensure_shape_wfm(_to_numpy_wfm(y_pred))
+    if t.shape != p.shape:
+        raise ValueError(f"Shape mismatch: true {t.shape} vs pred {p.shape}")
+    if mode == "magnitude":
+        tv = _force_mag_wfm(t).ravel()
+        pv = _force_mag_wfm(p).ravel()
+    else:
+        raise ValueError(f"Unknown mode '{mode}'")
+    tv, pv = tv.astype(np.float64), pv.astype(np.float64)
+    if np.allclose(tv.std(), 0) or np.allclose(pv.std(), 0):
+        return 0.0
+    return float(np.corrcoef(tv, pv)[0, 1])
+
+
+def wfm_relative_magnitude_error(y_true, y_pred, eps=1e-8):
+    """Relative magnitude error for heatmap-as-magnitude."""
+    t = _ensure_shape_wfm(_to_numpy_wfm(y_true))
+    p = _ensure_shape_wfm(_to_numpy_wfm(y_pred))
+    if t.shape != p.shape:
+        raise ValueError(f"Shape mismatch: true {t.shape} vs pred {p.shape}")
+    mag_t = _force_mag_wfm(t)
+    mag_p = _force_mag_wfm(p)
+    fbar = np.mean(mag_t)
+    if np.isclose(fbar, 0):
+        return 0.0
+    rel = np.abs(mag_p - mag_t) / (mag_t + eps)
+    w = mag_t / fbar
+    return float(np.mean(rel * w))
+
+
+def apply_threshold_mask(tensor, threshold=0.0):
+    return tensor * (tensor >= threshold).float()
+
+
+def detect_tanh_output_model(model):
+    """Detect if model outputs [-1, 1] (Tanh)."""
+    if hasattr(model, 'use_sigmoid') and not model.use_sigmoid:
+        return True
+    if hasattr(model, 'use_tanh_output') and model.use_tanh_output:
+        return True
+    if hasattr(model, 'final_conv'):
+        fc = model.final_conv
+        if isinstance(fc, nn.Sequential):
+            if isinstance(fc[-1], nn.Tanh):
+                return True
+        elif isinstance(fc, nn.Tanh):
+            return True
+    return False
+
+
+def convert_tanh_to_sigmoid_range(tensor):
+    return (tensor + 1.0) / 2.0
+
+
+# --- TorchMetrics wrapper for MS-SSIM ---
+
+class TorchMetricsWrapper:
+    def __init__(self, device='cpu'):
+        self.device = device
+        self.reset_metrics()
+
+    def reset_metrics(self):
+        if HAS_TORCHMETRICS:
+            self.ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
+            self.mse = MeanSquaredError().to(self.device)
+        else:
+            self.ms_ssim = None
+            self.mse = None
+
+    def compute_ms_ssim(self, y_true, y_pred):
+        if not HAS_TORCHMETRICS:
+            return float(calculate_ssim_tensor(y_true, y_pred))  # fallback to SSIM
+        y_true = y_true.to(self.device)
+        y_pred = y_pred.to(self.device)
+        if y_true.shape[1] == 1:
+            pass
+        else:
+            y_true, y_pred = y_true[:, 0:1], y_pred[:, 0:1]
+        return self.ms_ssim(y_pred, y_true).item()
+
+    def compute_mse(self, y_true, y_pred):
+        if not HAS_TORCHMETRICS:
+            return calculate_mse(y_true, y_pred)
+        y_true = y_true.to(self.device)
+        y_pred = y_pred.to(self.device)
+        return self.mse(y_pred, y_true).item()
+
+
+# --- Full evaluation on dataset ---
+
+def evaluate_metrics_on_dataset(generator, data_loader, device=None, description="Evaluating",
+                               save_predictions=False, threshold=0.0, use_settings=False,
+                               normalization_params=None, config_path=None, substrate_override=None):
+    """
+    Evaluate S2F generator on a dataset. Returns MSE, MS-SSIM, Pixel Correlation,
+    Relative Magnitude Error, and force sum/mean correlations.
+    """
+    if device is None:
+        device = torch.device('mps' if torch.backends.mps.is_available() else
+                              'cuda' if torch.cuda.is_available() else 'cpu')
+
+    generator = generator.to(device)
+    generator.eval()
+    metrics_wrapper = TorchMetricsWrapper(device=device)
+
+    heatmap_mse = []
+    heatmap_ms_ssim = []
+    heatmap_pixel_corr = []
+    wfm_corr_mag = []
+    wfm_rel_mag_err = []
+    force_sum_gt, force_sum_pred = [], []
+    force_mean_gt, force_mean_pred = [], []
+    individual_predictions = [] if save_predictions else None
+
+    with torch.no_grad():
+        for batch_idx, batch_data in enumerate(tqdm(data_loader, desc=description)):
+            if len(batch_data) == 5:
+                images, heatmaps, _, _, metadata = batch_data
+                has_metadata = True
+            else:
+                images, heatmaps, _, _ = batch_data
+                has_metadata = False
+
+            images = images.to(device, dtype=torch.float32)
+            heatmaps = heatmaps.to(device, dtype=torch.float32)
+
+            if use_settings and normalization_params is not None:
+                from models.s2f_model import create_settings_channels
+                meta = metadata if has_metadata else {'substrate': [substrate_override or 'fibroblasts_PDMS'] * images.size(0)}
+                settings_ch = create_settings_channels(meta, normalization_params, device, images.shape, config_path=config_path)
+                images = torch.cat([images, settings_ch], dim=1)
+
+            pred = generator(images)
+            if detect_tanh_output_model(generator):
+                pred = convert_tanh_to_sigmoid_range(pred)
+
+            gt_thresh = apply_threshold_mask(heatmaps, threshold)
+            pred_thresh = pred  # no threshold on pred for metrics
+
+            heatmap_mse.append(metrics_wrapper.compute_mse(gt_thresh, pred_thresh))
+            heatmap_ms_ssim.append(metrics_wrapper.compute_ms_ssim(gt_thresh, pred_thresh))
+            heatmap_pixel_corr.extend(calculate_individual_pixel_correlation(gt_thresh, pred_thresh))
+
+            # WFM: heatmap as magnitude (fx=magnitude, fy=0)
+            B, _, H, W = gt_thresh.shape
+            gt_ff = torch.zeros(B, 2, H, W, device=device)
+            pred_ff = torch.zeros(B, 2, H, W, device=device)
+            gt_ff[:, 0], pred_ff[:, 0] = gt_thresh[:, 0], pred_thresh[:, 0]
+            try:
+                wfm_corr_mag.append(wfm_correlation(gt_ff, pred_ff, mode="magnitude"))
+                wfm_rel_mag_err.append(wfm_relative_magnitude_error(gt_ff, pred_ff))
+            except Exception:
+                wfm_corr_mag.append(float('nan'))
+                wfm_rel_mag_err.append(float('nan'))
+
+            force_sum_gt.extend(torch.sum(gt_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+            force_sum_pred.extend(torch.sum(pred_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+            force_mean_gt.extend(torch.mean(gt_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+            force_mean_pred.extend(torch.mean(pred_thresh, dim=[1, 2, 3]).cpu().numpy().tolist())
+
+            if save_predictions:
+                for i in range(images.size(0)):
+                    p, t = pred_thresh[i:i+1], gt_thresh[i:i+1]
+                    gt_ff_i = torch.zeros(1, 2, H, W, device=device)
+                    pred_ff_i = torch.zeros(1, 2, H, W, device=device)
+                    gt_ff_i[0, 0], pred_ff_i[0, 0] = t[0, 0], p[0, 0]
+                    try:
+                        rme = wfm_relative_magnitude_error(gt_ff_i, pred_ff_i)
+                    except Exception:
+                        rme = float('nan')
+                    individual_predictions.append({
+                        'batch_idx': batch_idx,
+                        'sample_idx': i,
+                        'original_image': images[i].cpu().numpy(),
+                        'ground_truth': heatmaps[i].cpu().numpy(),
+                        'ground_truth_thresholded': gt_thresh[i].cpu().numpy(),
+                        'prediction': pred[i].cpu().numpy(),
+                        'prediction_thresholded': pred_thresh[i].cpu().numpy(),
+                        'mse': metrics_wrapper.compute_mse(t, p),
+                        'ms_ssim': metrics_wrapper.compute_ms_ssim(t, p),
+                        'pixel_correlation': calculate_pearson_correlation(t, p),
+                        'wfm_relative_magnitude_error': rme,
+                        'force_sum_gt': torch.sum(gt_thresh[i]).item(),
+                        'force_sum_pred': torch.sum(pred_thresh[i]).item(),
+                        'force_mean_gt': torch.mean(gt_thresh[i]).item(),
+                        'force_mean_pred': torch.mean(pred_thresh[i]).item(),
+                    })
+
+    valid_wfm_corr = [x for x in wfm_corr_mag if not np.isnan(x)]
+    valid_wfm_rme = [x for x in wfm_rel_mag_err if not np.isnan(x)]
+    try:
+        force_sum_corr, _ = pearsonr(force_sum_gt, force_sum_pred)
+        force_mean_corr, _ = pearsonr(force_mean_gt, force_mean_pred)
+    except Exception:
+        force_sum_corr = force_mean_corr = 0.0
+    if force_sum_corr is None or (isinstance(force_sum_corr, float) and np.isnan(force_sum_corr)):
+        force_sum_corr = 0.0
+    if force_mean_corr is None or (isinstance(force_mean_corr, float) and np.isnan(force_mean_corr)):
+        force_mean_corr = 0.0
+
+    results = {
+        'heatmap': {
+            'mse': np.mean(heatmap_mse),
+            'mse_std': np.std(heatmap_mse),
+            'ms_ssim': np.mean(heatmap_ms_ssim),
+            'ms_ssim_std': np.std(heatmap_ms_ssim),
+            'pixel_correlation': np.mean(heatmap_pixel_corr),
+            'pixel_correlation_std': np.std(heatmap_pixel_corr),
+        },
+        'wfm': {
+            'correlation_magnitude': np.mean(valid_wfm_corr) if valid_wfm_corr else float('nan'),
+            'correlation_magnitude_std': np.std(valid_wfm_corr) if valid_wfm_corr else float('nan'),
+            'relative_magnitude_error': np.mean(valid_wfm_rme) if valid_wfm_rme else float('nan'),
+            'relative_magnitude_error_std': np.std(valid_wfm_rme) if valid_wfm_rme else float('nan'),
+        },
+        'force_sum': {
+            'correlation': float(force_sum_corr),
+            'gt_mean': np.mean(force_sum_gt),
+            'pred_mean': np.mean(force_sum_pred),
+            'gt_std': np.std(force_sum_gt),
+            'pred_std': np.std(force_sum_pred),
+        },
+        'force_mean': {
+            'correlation': float(force_mean_corr),
+            'gt_mean': np.mean(force_mean_gt),
+            'pred_mean': np.mean(force_mean_pred),
+        },
+    }
+
+    if save_predictions:
+        results['individual_predictions'] = individual_predictions
+    return results
+
+
+def print_metrics_report(report, threshold=0.0, uses_tanh=False):
+    """Print formatted metrics report."""
+    for name, metrics in report.items():
+        print(f"\n🔸 {name.upper()} SET METRICS" + (f" (threshold={threshold})" if threshold > 0 else ""))
+        print("-" * 60)
+        print("HEATMAP METRICS:")
+        print(f" MSE:             {metrics['heatmap']['mse']:.6f} ± {metrics['heatmap']['mse_std']:.6f}")
+        print(f" MS-SSIM:         {metrics['heatmap']['ms_ssim']:.4f} ± {metrics['heatmap']['ms_ssim_std']:.4f}")
+        print(f" Pixel Corr:      {metrics['heatmap']['pixel_correlation']:.4f} ± {metrics['heatmap']['pixel_correlation_std']:.4f}")
+        print("WFM METRICS (heatmap as magnitude):")
+        print(f" Correlation (Magnitude): {metrics['wfm']['correlation_magnitude']:.4f} ± {metrics['wfm']['correlation_magnitude_std']:.4f}")
+        print(f" Relative Magnitude Error: {metrics['wfm']['relative_magnitude_error']:.4f} ± {metrics['wfm']['relative_magnitude_error_std']:.4f}")
+        print("FORCE SUM CORRELATION:")
+        print(f" Correlation: {metrics['force_sum']['correlation']:.4f}")
+        print(f" GT Mean: {metrics['force_sum']['gt_mean']:.2f} ± {metrics['force_sum']['gt_std']:.2f}")
+        print(f" Pred Mean: {metrics['force_sum']['pred_mean']:.2f} ± {metrics['force_sum']['pred_std']:.2f}")
+        if uses_tanh:
+            print(" Note: Model outputs [-1,1], converted to [0,1] for evaluation")
+    print("=" * 60)
+
+
+def gen_prediction_plots(individual_predictions, save_dir, sort_by='ms_ssim', sort_order='desc', threshold=0.0):
+    """Generate prediction plots (BF | GT | Pred) sorted by metric."""
+    os.makedirs(save_dir, exist_ok=True)
+    reverse = (sort_order.lower() == 'desc') if sort_by.lower() not in ['mse', 'wfm_relative_magnitude_error'] else (sort_order.lower() == 'desc')
+    valid = [p for p in individual_predictions if not np.isnan(p.get(sort_by.lower(), 0))]
+    sorted_preds = sorted(valid, key=lambda x: x[sort_by.lower()], reverse=reverse)
+    print(f"Sorting {len(sorted_preds)} predictions by {sort_by} ({sort_order})")
+    for rank, p in enumerate(tqdm(sorted_preds, desc="Saving plots"), 1):
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+        img = p['original_image']
+        axes[0].imshow(img[0] if img.ndim == 3 else img, cmap='gray')
+        axes[0].set_title('Bright Field')
+        axes[0].axis('off')
+        gt = p['ground_truth']
+        axes[1].imshow(gt[0] if gt.ndim == 3 else gt, cmap='jet', vmin=0, vmax=1)
+        axes[1].set_title('Ground Truth')
+        axes[1].axis('off')
+        pr = p['prediction']
+        axes[2].imshow(pr[0] if pr.ndim == 3 else pr, cmap='jet', vmin=0, vmax=1)
+        axes[2].set_title('Prediction')
+        axes[2].axis('off')
+        m = (f"MSE: {p['mse']:.4f} | MS-SSIM: {p['ms_ssim']:.4f} | "
+             f"Pixel Corr: {p['pixel_correlation']:.4f} | Rel Mag Err: {p.get('wfm_relative_magnitude_error', 'N/A')}")
+        fig.suptitle(f"Rank {rank} (by {sort_by})\n{m}", fontsize=10, y=0.02)
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_dir, f"rank{rank:03d}_batch{p['batch_idx']:03d}_sample{p['sample_idx']:02d}.png"), dpi=150, bbox_inches='tight')
+        plt.close()
+
+
+def plot_predictions(loader, generator, n_samples, device, threshold=0.0,
+                     use_settings=False, normalization_params=None, config_path=None, substrate_override=None):
+    """Plot BF | GT | Pred for first n_samples from loader."""
+    generator = generator.to(device)
+    generator.eval()
+    bf_list, gt_list, meta_list = [], [], []
+    it = iter(loader)
+    while len(bf_list) < n_samples:
+        try:
+            batch = next(it)
+        except StopIteration:
+            break
+        if len(batch) == 5:
+            images, heatmaps, _, _, meta = batch
+        else:
+            images, heatmaps = batch[0], batch[1]
+            meta = None
+        for i in range(images.shape[0]):
+            if len(bf_list) >= n_samples:
+                break
+            bf_list.append(images[i])
+            gt_list.append(heatmaps[i])
+            meta_list.append(meta)
+    n = min(n_samples, len(bf_list))
+    bf_batch = torch.stack(bf_list[:n]).to(device, dtype=torch.float32)
+    if use_settings and normalization_params:
+        from models.s2f_model import create_settings_channels
+        sub = substrate_override or 'fibroblasts_PDMS'
+        meta_dict = {'substrate': [sub] * n}
+        settings_ch = create_settings_channels(meta_dict, normalization_params, device, bf_batch.shape, config_path=config_path)
+        bf_batch = torch.cat([bf_batch, settings_ch], dim=1)
+    with torch.no_grad():
+        pred = generator(bf_batch)
+        if detect_tanh_output_model(generator):
+            pred = convert_tanh_to_sigmoid_range(pred)
+    if threshold > 0:
+        pred = pred * (pred >= threshold).float()
+    fig, axes = plt.subplots(n, 3, figsize=(12, 4 * n))
+    if n == 1:
+        axes = axes.reshape(1, -1)
+    for i in range(n):
+        axes[i, 0].imshow(bf_list[i].squeeze().cpu().numpy(), cmap='gray')
+        axes[i, 0].set_title('Bright Field')
+        axes[i, 0].axis('off')
+        axes[i, 1].imshow(gt_list[i].squeeze().cpu().numpy(), cmap='jet', vmin=0, vmax=1)
+        axes[i, 1].set_title('Ground Truth')
+        axes[i, 1].axis('off')
+        axes[i, 2].imshow(pred[i].squeeze().cpu().numpy(), cmap='jet', vmin=0, vmax=1)
+        axes[i, 2].set_title('Prediction')
+        axes[i, 2].axis('off')
+    plt.tight_layout()
+    plt.show()

utils/substrate_settings.py

@@ -0,0 +1,129 @@
+"""
+Substrate settings for force map prediction.
+Loads from config/substrate_settings.json - users can edit this file to add/modify substrates.
+"""
+import os
+import json
+
+
+def _default_config_path():
+    """Default path to substrate settings config (S2F/config/substrate_settings.json)."""
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    project_root = os.path.dirname(this_dir)  # S2F root
+    return os.path.join(project_root, 'config', 'substrate_settings.json')
+
+
+def load_substrate_config(config_path=None):
+    """
+    Load substrate settings from config file.
+
+    Args:
+        config_path: Path to JSON config. If None, uses config/substrate_settings.json in S2F root.
+
+    Returns:
+        dict: Config with 'substrates', 'default_substrate'
+    """
+    path = config_path or _default_config_path()
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"Substrate config not found at {path}. "
+            "Create config/substrate_settings.json or pass config_path."
+        )
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def resolve_substrate(name, config=None, config_path=None):
+    """
+    Resolve substrate name to a canonical substrate key.
+
+    Args:
+        name: Substrate key (e.g. 'fibroblasts_PDMS', 'PDMS_10kPa')
+        config: Pre-loaded config dict. If None, loads from config_path.
+        config_path: Path to config file (used if config is None).
+
+    Returns:
+        str: Canonical substrate key
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+
+    s = (name or '').strip()
+    if not s:
+        return config.get('default_substrate', 'fibroblasts_PDMS')
+
+    substrates = config.get('substrates', {})
+    s_lower = s.lower()
+    for key in substrates:
+        if key.lower() == s_lower:
+            return key
+    for key in substrates:
+        if s_lower.startswith(key.lower()) or key.lower().startswith(s_lower):
+            return key
+
+    return config.get('default_substrate', 'fibroblasts_PDMS')
+
+
+def get_settings_of_category(substrate_name, config=None, config_path=None):
+    """
+    Get pixelsize and young's modulus for a substrate.
+
+    Args:
+        substrate_name: Substrate or folder name (case-insensitive)
+        config: Pre-loaded config dict. If None, loads from config_path.
+        config_path: Path to config file (used if config is None).
+
+    Returns:
+        dict: {'name': str, 'pixelsize': float, 'young': float}
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+
+    substrate_key = resolve_substrate(substrate_name, config=config)
+    substrates = config.get('substrates', {})
+    default = config.get('default_substrate', 'fibroblasts_PDMS')
+
+    if substrate_key in substrates:
+        return substrates[substrate_key].copy()
+
+    default_settings = substrates.get(default, {'name': 'Fibroblasts on PDMS', 'pixelsize': 3.0769, 'young': 6000})
+    return default_settings.copy()
+
+
+def list_substrates(config=None, config_path=None):
+    """
+    Return list of available substrate keys for user selection.
+
+    Returns:
+        list: Substrate keys
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+    return list(config.get('substrates', {}).keys())
+
+
+def compute_settings_normalization(config=None, config_path=None):
+    """
+    Compute min-max normalization parameters from all substrates in config.
+
+    Returns:
+        dict: {'pixelsize': {'min', 'max'}, 'young': {'min', 'max'}}
+    """
+    if config is None:
+        config = load_substrate_config(config_path)
+
+    substrates = config.get('substrates', {})
+    all_pixelsizes = [s['pixelsize'] for s in substrates.values()]
+    all_youngs = [s['young'] for s in substrates.values()]
+
+    if not all_pixelsizes or not all_youngs:
+        pixelsize_min, pixelsize_max = 3.0769, 9.8138
+        young_min, young_max = 1000.0, 10000.0
+    else:
+        pixelsize_min, pixelsize_max = min(all_pixelsizes), max(all_pixelsizes)
+        young_min, young_max = min(all_youngs), max(all_youngs)
+
+    return {
+        'pixelsize': {'min': pixelsize_min, 'max': pixelsize_max},
+        'young': {'min': young_min, 'max': young_max}
+    }