Update spm.py

spm.py

@@ -32,9 +32,7 @@ def _to_divisible_by(img, N):
 def _edgelogic(i, j, ph, pw, N, overlap):
     """
     Base (no-overlap) patch is [i*ph:(i+1)*ph, j*pw:(j+1)*pw].
-    Extend with overlap, biasing inward.
-    Uses 2*overlap for edges to keep patch areas roughly comparable.
-    Returns (start_h, end_h, start_w, end_w) BEFORE clamping to image bounds.
+    Extend with overlap, biasing inward. Uses 2*overlap at borders.
     """
     start_h = i * ph
     start_w = j * pw
@@ -70,21 +68,20 @@ def spm_augment(
     mix_prob=0.5,
     beta_a=2.0,
     beta_b=2.0,
-    overlap_px=0,
+    overlap_pct=0.0,   # percentage of patch size (0..49 typically)
     seed=None
 ):
     """
     SPM-style augmentation with optional overlap + feathered blending.
 
-    When overlap_px <= 0:
+    When overlap_pct <= 0:
       - Standard global shuffle over N×N patches;
       - Per-patch mixing with a single alpha ~ Beta(a,b) for the image.
 
-    When overlap_px > 0:
-      - Each base cell (N×N grid) expands by +/-overlap_px (2*overlap at borders),
-        clipped to the image. Patches are mixed per location and alpha sampled per-patch
-        for a bit more stochasticity (can be changed to per-image alpha by editing below).
-      - Patches are blended into the canvas with a feather mask of size `overlap_px`.
+    When overlap_pct > 0:
+      - Each base cell (N×N grid) expands by ±overlap_px (derived from percentage),
+        clipped to the image. Patches are mixed per location with per-patch alpha.
+      - Patches are blended into the canvas with a feather mask of size overlap_px.
     """
     # Normalize to PIL and ensure divisibility
     if isinstance(image, np.ndarray):
@@ -100,10 +97,10 @@ def spm_augment(
     ph = H // N
     pw = W // N
 
-    # Clamp overlap to < half patch size
-    if overlap_px is None:
-        overlap_px = 0
-    overlap_px = int(overlap_px)
+    # Convert percentage to pixel overlap; clamp to < half patch size
+    pct = float(overlap_pct)
+    pct = max(0.0, min(pct, 49.0))  # keep below 50% for stability
+    overlap_px = int(round((pct / 100.0) * min(ph, pw)))
     max_ov = max(0, min(ph, pw) // 2 - 1)
     ov = int(np.clip(overlap_px, 0, max_ov))
 
@@ -167,7 +164,6 @@ def spm_augment(
     total = len(patches)
     perm = rng.permutation(total)
 
-    # We'll sample alpha per-patch to echo your overlap snippet
     def sample_alpha():
         if beta_a > 0 and beta_b > 0:
             return float(rng.beta(beta_a, beta_b))
@@ -178,9 +174,7 @@ def spm_augment(
 
     for k, (sh, eh, sw, ew) in enumerate(coords):
         if rng.random() >= float(mix_prob):
-            # keep original content in that region
-            src = patches[k]
-            patch = src
+            patch = patches[k]
         else:
             lam = sample_alpha()
             src = patches[k].astype(np.float32)
@@ -188,18 +182,12 @@ def spm_augment(
             patch = lam * shf + (1.0 - lam) * src
 
         ph_k, pw_k, _ = patch.shape
-        # Slice feather mask down if needed (near borders)
         mask2d = feather_full[:ph_k, :pw_k]
-        if arr.shape[2] == 1:
-            mask3d = mask2d[..., None]
-        else:
-            mask3d = np.repeat(mask2d[..., None], arr.shape[2], axis=2)
+        mask3d = mask2d[..., None] if arr.shape[2] == 1 else np.repeat(mask2d[..., None], arr.shape[2], axis=2)
 
-        # Accumulate
         canvas[sh:eh, sw:ew] += patch * mask3d
         weight[sh:eh, sw:ew] += mask2d
 
-    # Normalize
     weight = np.clip(weight, 1e-8, None)
     out = (canvas / weight[..., None])
     out = np.clip(out, 0, 255).astype(np.uint8)