pbc3: restore PIL resampling and fix benchmarks

.gitignore

@@ -13,4 +13,6 @@ hpt_data_val/
 *.pt
 training_space/results/
 training_space/runs/*
-!training_space/runs/RL_checkpoint.pt
+!training_space/runs/RL_checkpoint.pt
+pbc3_resampling_benchmark.*
+pbc3_resampling_top_variants_20.*

PBC3.py

@@ -162,7 +162,10 @@ class PBC3:
             return canvas
         out = np.empty((new_h, new_w, ch), dtype=np.int32)
         for c in range(ch):
-            out[:, :, c] = np.rint(ops.bicubic_resample_2d(canvas[:, :, c], new_h, new_w)).astype(np.int32)
+            layer_arr = np.ascontiguousarray(canvas[:, :, c], dtype=np.float32)
+            layer = Image.frombuffer("F", (w, h), layer_arr, "raw", "F", 0, 1)
+            layer = layer.resize((new_w, new_h), cls.RESAMPLE_FILTER)
+            out[:, :, c] = np.rint(np.asarray(layer, dtype=np.float32)).astype(np.int32)
         return out
 
     @classmethod
@@ -477,12 +480,22 @@ class PBC3:
         return image
 
 
-def preload_numba() -> None:
-    """## Runs a tiny encode so numba compiles the hot kernels once"""
-    arr = ((np.arange(16 * 16 * 3, dtype=np.uint16).reshape(16, 16, 3) * 37) % 256).astype(np.uint8)
-    img = Image.fromarray(arr)
-    PBC3.compress(img, PBC3Config(patch_count=10, auto_downsample_init=True, learned_filler_enabled=False))
-    print("[preload] numba kernels compiled")
+def preload_numba(model_path: str = "patch_policy.npz") -> None:
+    """## Warms the production learned RGB/RGBA encode paths."""
+    h, w = 512, 768
+    base = np.arange(h * w * 4, dtype=np.uint32).reshape(h, w, 4)
+    config = PBC3Config.quality(
+        patch_count=50,
+        learned_filler_enabled=True,
+        learned_filler_model_path=model_path,
+        auto_downsample_max_pixels=250_000,
+        use_lzma=True,
+        compute_final_mse=True,
+    )
+    for channels in (3, 4):
+        arr = ((base[:, :, :channels] * 37 + channels * 19) % 256).astype(np.uint8)
+        PBC3.compress(Image.fromarray(arr), config=config)
+    print("[preload] production PBC3 paths warmed")
 
 
 if __name__ == "__main__":

PBC3_animation.py

@@ -2,6 +2,7 @@ import os
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 from PBC3 import PBC3, BitReader
+import pbc3_ops as ops
 
 
 def _font(size):
@@ -345,7 +346,7 @@ def animate_pbc3(
 
         tc = target.shape[2] if target is not None else 0
         prev_mse = float(np.mean((canvas[:, :, :tc] - target) ** 2)) if target is not None else None
-        PBC3.apply_grid(canvas[:, :, channel], x, y, pw, ph, cell_size, values)
+        ops.apply_grid(canvas[:, :, channel], x, y, pw, ph, cell_size, values)
         now_mse = float(np.mean((canvas[:, :, :tc] - target) ** 2)) if target is not None else None
         parts = [f"Patch {i}/{patch_count}", f"Stream Size: {current_kb:.2f} KB", f"(+{delta_kb:.2f} KB)"]
         if now_mse is not None:
@@ -364,4 +365,4 @@ def animate_pbc3(
             output_size,
         ))
 
-    return _write_frames(frames, output_path, fps, fallback_to_gif=fallback_to_gif)
+    return _write_frames(frames, output_path, fps, fallback_to_gif=fallback_to_gif)

pbc3_benchmark.py

@@ -9,7 +9,7 @@ import time
 import numpy as np
 from PIL import Image, ImageOps
 
-from PBC3 import PBC3
+from PBC3 import PBC3, preload_numba
 from pbc3_types import PBC3Config
 
 try:
@@ -216,6 +216,7 @@ def _run(codecs, n_trials):
         images = dataset()
         if not images:
             raise RuntimeError("No images found in hpt_data/")
+        preload_numba(os.path.join(PROJECT_DIR, "patch_policy.npz"))
         jobs = _jobs(codecs)
         total = n_trials * len(jobs) * len(images)
         run_id = time.strftime("%Y%m%d_%H%M%S")

pbc3_heads.py

@@ -153,7 +153,9 @@ class FillerHead:
             hidden_residual = target[y:y + bh, x:x + bw, c] - canvas[y:y + bh, x:x + bw, c]
             before = target[y:y + bh, x:x + bw, c] - np.clip(canvas[y:y + bh, x:x + bw, c], 0, 255)
             before_sse = float(np.sum(before.astype(np.int64) ** 2))
-            patch, values = ops.make_patch(c, x, y, bw, bh, cell_size, hidden_residual, config, bitcount)
+            patch, values = ops.make_patch(
+                c, x, y, bw, bh, cell_size, hidden_residual, config, bitcount,
+            )
             delta = ops.signed_resample(values, bh, bw).astype(np.int32)
             after = target[y:y + bh, x:x + bw, c] - np.clip(canvas[y:y + bh, x:x + bw, c] + delta, 0, 255)
             reduction = before_sse - float(np.sum(after.astype(np.int64) ** 2))

pbc3_ops.py

@@ -12,7 +12,6 @@ PALETTE_GENERATED = 0
 
 _UINT64_MASK = (1 << 64) - 1
 _UINT32_SCALE = 1.0 / 4294967296.0
-_RESAMPLE_COEFF_CACHE = {}
 
 
 @njit(inline="always", cache=True)
@@ -49,58 +48,17 @@ class PBC3Rng:
         return low + self.next_uint32() % (high - low)
 
 
-def _bicubic_kernel(x):
-    x = np.abs(np.asarray(x, dtype=np.float64))
-    out = np.zeros_like(x)
-    m1 = x < 1.0
-    m2 = (x >= 1.0) & (x < 2.0)
-    x1 = x[m1]
-    x2 = x[m2]
-    a = -0.5
-    out[m1] = ((a + 2.0) * x1 - (a + 3.0)) * x1 * x1 + 1.0
-    out[m2] = (((x2 - 5.0) * x2 + 8.0) * x2 - 4.0) * a
-    return out
-
-
-def _resample_coeffs(in_size: int, out_size: int) -> np.ndarray:
-    in_size, out_size = int(in_size), int(out_size)
-    key = (in_size, out_size)
-    coeffs = _RESAMPLE_COEFF_CACHE.get(key)
-    if coeffs is not None:
-        return coeffs
-
-    scale = in_size / out_size
-    filterscale = max(scale, 1.0)
-    support = 2.0 * filterscale
-    inv_filterscale = 1.0 / filterscale
-    coeffs = np.zeros((out_size, in_size), dtype=np.float64)
-
-    for xx in range(out_size):
-        center = (xx + 0.5) * scale
-        xmin = max(0, int(center - support + 0.5))
-        xmax = min(in_size, int(center + support + 0.5))
-        xs = np.arange(xmin, xmax, dtype=np.float64)
-        weights = _bicubic_kernel((xs - center + 0.5) * inv_filterscale)
-        total = weights.sum()
-        if total != 0.0:
-            weights /= total
-        coeffs[xx, xmin:xmax] = weights
-
-    _RESAMPLE_COEFF_CACHE[key] = coeffs
-    return coeffs
-
-
 def bicubic_resample_2d(values, out_h: int, out_w: int) -> np.ndarray:
-    arr = np.asarray(values, dtype=np.float64)
+    """Resamples a signed grid with the production PIL bicubic filter."""
+    arr = np.ascontiguousarray(values, dtype=np.float32)
     out_h, out_w = int(out_h), int(out_w)
     if arr.shape == (out_h, out_w):
-        return arr.copy()
-    out = arr
-    if arr.shape[1] != out_w:
-        out = out @ _resample_coeffs(arr.shape[1], out_w).T
-    if arr.shape[0] != out_h:
-        out = _resample_coeffs(arr.shape[0], out_h) @ out
-    return out
+        return arr.astype(np.float64, copy=True)
+    return np.asarray(
+        Image.frombuffer("F", (arr.shape[1], arr.shape[0]), arr, "raw", "F", 0, 1)
+        .resize((out_w, out_h), Image.Resampling.BICUBIC),
+        dtype=np.float64,
+    )
 
 
 def ceil_div(a: int, b: int) -> int:

pbc3_quick_rd.py

@@ -7,10 +7,10 @@ import time
 
 import numpy as np
 import optuna
-from PIL import Image
+from PIL import Image, ImageOps
 
 import pbc3_sweep
-from PBC3 import PBC3
+from PBC3 import PBC3, preload_numba
 from pbc3_types import PBC3Config
 
 RESULTS_PATH = os.path.join(pbc3_sweep.PROJECT_DIR, "pbc3_quick_rd.json")
@@ -110,19 +110,21 @@ def _eval_old_preset_subprocess(name, images, trial):
     code = r'''
 import json, os, sys, time
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageOps
 sys.path.insert(0, sys.argv[1])
 from PBC3 import PBC3
 try:
     from pbc3_types import PBC3Config
 except Exception:
     from pbc_types import PBC3Config
+from PBC3 import preload_numba
 preset = sys.argv[2]
 trial = int(sys.argv[3])
 images = json.loads(sys.stdin.read())
+preload_numba()
 out = []
 for im in images:
-    img = Image.open(im["path"]).convert("RGB")
+    img = ImageOps.exif_transpose(Image.open(im["path"])).convert("RGB")
     arr = np.asarray(img)
     pixels = arr.shape[0] * arr.shape[1]
     cfg = getattr(PBC3Config, preset)()
@@ -154,6 +156,7 @@ def _run(n_trials, run_old):
         images = pbc3_sweep.dataset()
         if not images:
             raise RuntimeError("No images found in hpt_data/")
+        preload_numba(os.path.join(pbc3_sweep.PROJECT_DIR, "patch_policy.npz"))
         _log(f"loaded {len(images)} images · {n_trials} trial(s)" + (" · live old enabled" if run_old else ""))
         out = {"created": time.time(), "n_trials": n_trials, "presets": []}
         for trial in range(1, n_trials + 1):
@@ -260,4 +263,4 @@ def results(mp_min=None, mp_max=None):
         "new_exists": os.path.exists(RESULTS_PATH),
         "old_pbc_exists": os.path.exists(os.path.join(OLD_PBC_DIR, "PBC3.py")),
         "rows": old["rows"] + live,
-    }
+    }

server.py

@@ -929,8 +929,8 @@ BOOT_ID = f"{time.time():.3f}"
 
 @app.on_event("startup")
 def _startup():
-    print("[startup] preloading Numba JITs...", flush=True)
-    preload_numba()
+    print("[startup] warming production PBC3 paths...", flush=True)
+    preload_numba(os.path.join(_HERE, "patch_policy.npz"))
     print("[startup] done.", flush=True)
     print(f"[startup] BOOT_ID={BOOT_ID}", flush=True)