Implemented batch processing for predictions to optimize memory usage and added progress tracking. Updated constants for batch inference size and adjusted contour drawing parameters for better visibility in UI elements.

S2FApp/app.py

@@ -18,6 +18,7 @@ if S2F_ROOT not in sys.path:
     sys.path.insert(0, S2F_ROOT)
 
 from config.constants import (
+    BATCH_INFERENCE_SIZE,
     BATCH_MAX_IMAGES,
     COLORMAPS,
     DEFAULT_SUBSTRATE,
@@ -464,11 +465,19 @@ if just_ran_batch:
         try:
             predictor = _load_predictor(model_type, checkpoint, ckp_folder)
             sub_val = substrate_val if model_type == "single_cell" and not use_manual else DEFAULT_SUBSTRATE
-            pred_results = predictor.predict_batch(
-                imgs_batch,
-                substrate=sub_val,
-                substrate_config=substrate_config if model_type == "single_cell" else None,
-            )
+            n_images = len(imgs_batch)
+            progress_bar = st.progress(0, text=f"Predicting 0 / {n_images} images")
+            pred_results = []
+            for start in range(0, n_images, BATCH_INFERENCE_SIZE):
+                chunk = imgs_batch[start : start + BATCH_INFERENCE_SIZE]
+                chunk_results = predictor.predict_batch(
+                    chunk,
+                    substrate=sub_val,
+                    substrate_config=substrate_config if model_type == "single_cell" else None,
+                )
+                pred_results.extend(chunk_results)
+                progress_bar.progress(min(start + len(chunk), n_images) / n_images,
+                                     text=f"Predicting {len(pred_results)} / {n_images} images")
             batch_results = [
                 {
                     "img": img_b.copy(),

S2FApp/config/constants.py

@@ -13,6 +13,8 @@ DEFAULT_SUBSTRATE = "Fibroblasts_Fibronectin_6KPa"
 CANVAS_SIZE = 320
 SAMPLE_THUMBNAIL_LIMIT = 8
 BATCH_MAX_IMAGES = 5
+# Max images per model forward pass (avoids OOM on Hugging Face free tier)
+BATCH_INFERENCE_SIZE = 2
 COLORMAP_N_SAMPLES = 64
 
 # Model type labels

S2FApp/predictor.py

@@ -13,7 +13,7 @@ S2F_ROOT = os.path.dirname(os.path.abspath(__file__))
 if S2F_ROOT not in sys.path:
     sys.path.insert(0, S2F_ROOT)
 
-from config.constants import DEFAULT_SUBSTRATE, MODEL_INPUT_SIZE
+from config.constants import BATCH_INFERENCE_SIZE, DEFAULT_SUBSTRATE, MODEL_INPUT_SIZE
 from models.s2f_model import create_s2f_model
 from utils.paths import get_ckp_base, model_subfolder
 from utils.substrate_settings import get_settings_of_category, compute_settings_normalization
@@ -178,18 +178,23 @@ class S2FPredictor:
 
         return heatmap, force, pixel_sum
 
-    def predict_batch(self, images, substrate=None, substrate_config=None):
+    def predict_batch(self, images, substrate=None, substrate_config=None, batch_size=None,
+                     on_progress=None):
         """
-        Run prediction on a batch of images (single forward pass).
+        Run prediction on a batch of images. Processes in chunks to avoid OOM on
+        memory-constrained environments (e.g. Hugging Face free tier).
 
         Args:
             images: List of (img_array, key) or list of img arrays. img_array: (H, W) or (H, W, C).
             substrate: Substrate name for single-cell mode (same for all images).
             substrate_config: Optional dict with 'pixelsize' and 'young' (same for all).
+            batch_size: Max images per forward pass (default: BATCH_INFERENCE_SIZE). Use 1 for minimal memory.
+            on_progress: Optional callback(processed: int, total: int) called after each forward pass.
 
         Returns:
             List of (heatmap, force, pixel_sum) tuples.
         """
+        batch_size = batch_size if batch_size is not None else BATCH_INFERENCE_SIZE
         imgs = []
         for item in images:
             img = item[0] if isinstance(item, tuple) else item
@@ -200,27 +205,34 @@ class S2FPredictor:
                 img = img / 255.0
             img = cv2.resize(img, (MODEL_INPUT_SIZE, MODEL_INPUT_SIZE))
             imgs.append(img)
-        x = torch.from_numpy(np.stack(imgs)).float().unsqueeze(1).to(self.device)  # [B, 1, H, W]
-
-        if self.model_type == "single_cell" and self.norm_params is not None:
-            sub = substrate if substrate is not None else DEFAULT_SUBSTRATE
-            settings_ch = create_settings_channels_single(
-                sub, self.device, x.shape[2], x.shape[3],
-                config_path=self.config_path, substrate_config=substrate_config
-            )
-            settings_batch = settings_ch.expand(x.shape[0], -1, -1, -1)
-            x = torch.cat([x, settings_batch], dim=1)  # [B, 3, H, W]
-
-        with torch.no_grad():
-            pred = self.generator(x)
-
-        if self._use_tanh_output:
-            pred = (pred + 1.0) / 2.0
 
         results = []
-        for i in range(pred.shape[0]):
-            heatmap = pred[i, 0].cpu().numpy()
-            force = sum_force_map(pred[i : i + 1]).item()
-            pixel_sum = float(np.sum(heatmap))
-            results.append((heatmap, force, pixel_sum))
+        for start in range(0, len(imgs), batch_size):
+            chunk = imgs[start : start + batch_size]
+            x = torch.from_numpy(np.stack(chunk)).float().unsqueeze(1).to(self.device)  # [B, 1, H, W]
+
+            if self.model_type == "single_cell" and self.norm_params is not None:
+                sub = substrate if substrate is not None else DEFAULT_SUBSTRATE
+                settings_ch = create_settings_channels_single(
+                    sub, self.device, x.shape[2], x.shape[3],
+                    config_path=self.config_path, substrate_config=substrate_config
+                )
+                settings_batch = settings_ch.expand(x.shape[0], -1, -1, -1)
+                x = torch.cat([x, settings_batch], dim=1)  # [B, 3, H, W]
+
+            with torch.no_grad():
+                pred = self.generator(x)
+
+            if self._use_tanh_output:
+                pred = (pred + 1.0) / 2.0
+
+            for i in range(pred.shape[0]):
+                heatmap = pred[i, 0].cpu().numpy()
+                force = sum_force_map(pred[i : i + 1]).item()
+                pixel_sum = float(np.sum(heatmap))
+                results.append((heatmap, force, pixel_sum))
+
+            if on_progress is not None:
+                on_progress(len(results), len(imgs))
+
         return results

S2FApp/static/s2f_styles.css

@@ -99,6 +99,9 @@ section[data-testid="stSidebar"] {
 /* === Sidebar === */
 section[data-testid="stSidebar"] {
     width: 360px !important;
+    max-height: 100vh !important;
+    overflow-x: hidden !important;
+    overflow-y: auto !important;
     background:
         linear-gradient(180deg, #f8fafc 0%, #f1f5f9 100%),
         /* Subtle vertical rhythm */
@@ -403,8 +406,26 @@ div[data-testid="stHorizontalBlock"]:has([data-testid="stDownloadButton"]):has([
 .stWarning { border-radius: 8px !important; }
 .stInfo { border-radius: 8px !important; }
 
-/* === Selectbox, toggle === */
-[data-testid="stSelectbox"] > div > div { border-radius: 8px !important; }
+/* === Selectbox, multiselect, toggle === */
+[data-testid="stSelectbox"] > div > div,
+[data-testid="stSelectbox"] input,
+[data-testid="stMultiSelect"] > div > div,
+[data-testid="stMultiSelect"] input {
+    color: #1e293b !important;
+    background-color: #ffffff !important;
+}
+[data-testid="stSelectbox"] > div > div,
+[data-testid="stMultiSelect"] > div > div {
+    border-radius: 8px !important;
+    border: 1px solid #cbd5e1 !important;
+}
+/* Dropdown options: ensure visible text on light background */
+[role="listbox"] [role="option"],
+[data-baseweb="menu"] li,
+ul[role="listbox"] li {
+    color: #1e293b !important;
+    background-color: #ffffff !important;
+}
 .selectbox-label {
     margin: 0;
     padding-top: 0.4rem;

S2FApp/ui/heatmaps.py

@@ -76,7 +76,7 @@ def make_annotated_heatmap_multi_regions(heatmap_rgb, masks, labels, cell_mask=N
     annotated = heatmap_rgb.copy()
     if cell_mask is not None and np.any(cell_mask > 0):
         contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cv2.drawContours(annotated, contours, -1, (255, 0, 0), 5)
+        cv2.drawContours(annotated, contours, -1, (255, 0, 0), 3)
     for i, mask in enumerate(masks):
         color = _REGION_COLORS[i % len(_REGION_COLORS)]
         _draw_region_overlay(annotated, mask, color, fill_alpha, stroke_width=2)
@@ -113,6 +113,6 @@ def add_cell_contour_to_fig(fig_pl, cell_mask, row=1, col=2):
             x.append(x[0])
             y.append(y[0])
         fig_pl.add_trace(
-            go.Scatter(x=x, y=y, mode="lines", line=dict(color="red", width=4), showlegend=False),
+            go.Scatter(x=x, y=y, mode="lines", line=dict(color="red", width=3), showlegend=False),
             row=row, col=col
         )

S2FApp/ui/measure_tool.py

@@ -156,7 +156,7 @@ def compute_region_metrics(raw_heatmap, mask, original_vals=None):
     }
 
 
-def _draw_contour_on_image(img_rgb, mask, stroke_color=(255, 0, 0), stroke_width=5):
+def _draw_contour_on_image(img_rgb, mask, stroke_color=(255, 0, 0), stroke_width=3):
     """Draw contour from mask on RGB image. Resizes mask to match img if needed."""
     h, w = img_rgb.shape[:2]
     if mask.shape[:2] != (h, w):
@@ -304,7 +304,7 @@ def render_region_canvas(display_heatmap, raw_heatmap=None, bf_img=None, origina
             st.caption("Bright-field")
             bf_display = bf_rgb.copy()
             if cell_mask is not None and np.any(cell_mask > 0):
-                bf_display = _draw_contour_on_image(bf_display, cell_mask, stroke_color=(255, 0, 0), stroke_width=5)
+                bf_display = _draw_contour_on_image(bf_display, cell_mask, stroke_color=(255, 0, 0), stroke_width=3)
             st.image(bf_display, width=CANVAS_SIZE)
     else:
         st.markdown("**Draw a region** on the heatmap.")

S2FApp/utils/report.py

@@ -68,7 +68,7 @@ def heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name="Jet", cell_mask=N
     if cell_mask is not None and np.any(cell_mask > 0):
         contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         if contours:
-            cv2.drawContours(heatmap_rgb, contours, -1, (255, 0, 0), 5)
+            cv2.drawContours(heatmap_rgb, contours, -1, (255, 0, 0), 3)
     return heatmap_rgb