added

.dockerignore

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

.gitattributes

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

.gitignore

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

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"]

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).

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")

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"}

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

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

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

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}
+    }