Fix cumulative patches chart sampling baselines

app.py

@@ -8,11 +8,12 @@ high local complexity = roughly what the encoder would spend bits on).
 
 Pipeline (mirrors codec_tools/pipeline/process_video_bitcost_readiness.py):
     1. Uniformly sample N frames from the input video.
-    2. smart_resize each frame so dims are multiples of `patch` and the
-       total pixel count <= max_pixels.
+    2. Resize each sampled frame to a fixed square patch grid driven by
+       `patch_size`.
     3. Slice every frame into a patch grid; score each patch by its
        Sobel gradient magnitude mean.
-    4. Pick the top-K highest-scoring patches per frame.
+    4. Pick the top-K highest-scoring patches under the selected GOP
+       grouping.
     5. Render a "selection visualization" video: kept patches stay in
        full color, dropped patches are faded to a gray-white wash so the
        viewer can see exactly which patches the codec stage chose.
@@ -21,7 +22,6 @@ Pipeline (mirrors codec_tools/pipeline/process_video_bitcost_readiness.py):
 """
 
 import json
-import math
 import os
 import shutil
 import subprocess
@@ -48,7 +48,7 @@ DEMO_PRESET = (
     DEMO_VIDEO_PATH,  # video_in
     16,               # sample_frames
     14,               # patch_size
-    1024,             # total_patches
+    3136,             # total_patches  (= 16 * 14^2)
     150000,           # max_pixels
     "sbs",            # viz_mode
     0.55,             # heatmap_alpha
@@ -63,16 +63,22 @@ DEMO_PRESET = (
 
 
 def smart_resize(frame: np.ndarray, max_pixels: int, factor: int) -> np.ndarray:
-    """Resize so h,w are multiples of `factor` and h*w <= max_pixels."""
-    h, w = frame.shape[:2]
-    pixels = h * w
-    if pixels > max_pixels:
-        scale = math.sqrt(max_pixels / pixels)
-        h = max(factor, int(h * scale))
-        w = max(factor, int(w * scale))
-    h = max(factor, (h // factor) * factor)
-    w = max(factor, (w // factor) * factor)
-    return cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA)
+    """Resize each frame to a square patch grid.
+
+    The demo uses `factor` as both:
+      - patch size in pixels
+      - patches per side in the resized frame
+
+    So patch_size=14 means:
+      - each patch is 14 x 14 pixels
+      - each frame is resized to 14 x 14 patches
+      - each frame therefore contributes 14^2 = 196 patch slots
+
+    `max_pixels` is kept for API compatibility with earlier revisions, but
+    the frame token count is now controlled by `factor` directly.
+    """
+    side_px = int(factor) * int(factor)
+    return cv2.resize(frame, (side_px, side_px), interpolation=cv2.INTER_AREA)
 
 
 def sample_frame_ids(total: int, n: int) -> List[int]:
@@ -83,6 +89,27 @@ def sample_frame_ids(total: int, n: int) -> List[int]:
     return [int(round(i)) for i in np.linspace(0, total - 1, n)]
 
 
+def split_budget_evenly(total_k: int, n_parts: int) -> List[int]:
+    total = max(0, int(total_k))
+    n = max(0, int(n_parts))
+    if n == 0:
+        return []
+    base, rem = divmod(total, n)
+    return [base + (1 if i < rem else 0) for i in range(n)]
+
+
+def sample_window_frame_ids(start: int, end: int, n: int) -> List[int]:
+    start_i = int(start)
+    end_i = int(end)
+    count = max(0, int(n))
+    if end_i < start_i or count <= 0:
+        return []
+    total = end_i - start_i + 1
+    if count >= total:
+        return list(range(start_i, end_i + 1))
+    return [start_i + x for x in sample_frame_ids(total, count)]
+
+
 def decode_frames(video_path: str, frame_ids: List[int]) -> List[np.ndarray]:
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
@@ -200,8 +227,10 @@ def topk_mask(score: np.ndarray, k: int) -> np.ndarray:
         return np.ones_like(score, dtype=np.uint8)
     if k <= 0:
         return np.zeros_like(score, dtype=np.uint8)
-    thresh = np.partition(flat, -k)[-k]
-    return (score >= thresh).astype(np.uint8)
+    out = np.zeros(flat.size, dtype=np.uint8)
+    keep_idx = np.argpartition(flat, -k)[-k:]
+    out[keep_idx] = 1
+    return out.reshape(score.shape)
 
 
 def global_topk_masks(
@@ -218,15 +247,17 @@ def global_topk_masks(
     arr = np.stack(grids, axis=0).astype(np.float32)  # [N, hb, wb]
     N, hb, wb = arr.shape
     flat = arr.reshape(-1)
-    if total_k >= flat.size:
+    k = int(total_k)
+    if k >= flat.size:
         masks = [np.ones((hb, wb), dtype=np.uint8) for _ in range(N)]
         return masks, int(flat.size)
-    if total_k <= 0:
+    if k <= 0:
         return [np.zeros((hb, wb), dtype=np.uint8) for _ in range(N)], 0
-    thresh = np.partition(flat, -total_k)[-total_k]
-    bool_mask = (arr >= thresh)
-    actual = int(bool_mask.sum())
-    return [bool_mask[i].astype(np.uint8) for i in range(N)], actual
+    mask_flat = np.zeros(flat.size, dtype=np.uint8)
+    keep_idx = np.argpartition(flat, -k)[-k:]
+    mask_flat[keep_idx] = 1
+    bool_mask = mask_flat.reshape(N, hb, wb)
+    return [bool_mask[i].astype(np.uint8) for i in range(N)], k
 
 
 def build_dynamic_groups(
@@ -323,14 +354,35 @@ def grouped_topk_masks(
             cursor = end + 1
 
     num_groups = max(1, len(groups))
-    per_group_budget = max(1, int(total_k) // num_groups)
+    target_k = max(0, int(total_k))
+
+    capacities = [
+        sum(int(g.size) for g in grids[s:e + 1])
+        for (s, e) in groups
+    ]
+    alloc = split_budget_evenly(target_k, num_groups)
+
+    leftover = 0
+    for i, cap in enumerate(capacities):
+        if alloc[i] > cap:
+            leftover += alloc[i] - cap
+            alloc[i] = cap
+    while leftover > 0:
+        progressed = False
+        for i, cap in enumerate(capacities):
+            if alloc[i] < cap and leftover > 0:
+                alloc[i] += 1
+                leftover -= 1
+                progressed = True
+        if not progressed:
+            break
 
     # Initialize empty masks, then fill per-group selections.
     out_masks = [np.zeros(g.shape, dtype=np.uint8) for g in grids]
     actual_total = 0
-    for (s, e) in groups:
+    for (s, e), group_k in zip(groups, alloc):
         sub = grids[s:e + 1]
-        sub_masks, sub_actual = global_topk_masks(sub, per_group_budget)
+        sub_masks, sub_actual = global_topk_masks(sub, group_k)
         for i, sm in enumerate(sub_masks):
             out_masks[s + i] = sm
         actual_total += sub_actual
@@ -573,7 +625,9 @@ def pack_canvases_per_group(
 def make_charts(
     grids: List[np.ndarray],
     masks: List[np.ndarray],
-    frame_ids: List[int],
+    codec_frame_ids: List[int],
+    uniform_frame_ids: List[int],
+    uniform_requested_frames: int,
     fps: float,
     total_duration_sec: float,
     total_patches_budget: int,
@@ -582,14 +636,13 @@ def make_charts(
     gop_label: str = "global",
 ):
     """One overlaid step chart: cumulative patches selected vs time, for
-    the codec saliency curve and a uniform-sampling baseline at the same
-    total budget.
+    the codec saliency curve and a uniform full-frame sampling baseline.
 
     X = time (s)
     Y = cumulative count of selected patches
-    Both curves end near the budget (codec: == total selected; uniform:
-    n_uniform_frames × grid_size, ≤ budget). The codec curve rises in
-    bursts where saliency is high; uniform rises in equal steps."""
+    The codec curve rises in bursts where saliency is high; the uniform
+    baseline rises in equal steps because every sampled full frame
+    contributes one complete patch grid."""
     fig, ax = plt.subplots(figsize=(9.2, 3.6), constrained_layout=True)
 
     fps_safe = float(fps) if fps and fps > 0 else 25.0
@@ -598,8 +651,9 @@ def make_charts(
     else:
         hb = wb = 1
     grid_size = hb * wb
+    all_frame_ids = list(codec_frame_ids) + list(uniform_frame_ids)
     duration = float(total_duration_sec) if total_duration_sec and total_duration_sec > 0 else (
-        (max(frame_ids) / fps_safe) if frame_ids else 1.0
+        (max(all_frame_ids) / fps_safe) if all_frame_ids else 1.0
     )
 
     # ─── Build step curves ──────────────────────────────────────────────
@@ -615,27 +669,22 @@ def make_charts(
         xx.append(duration); yy.append(prev)
         return xx, yy
 
-    times = [fid / fps_safe for fid in frame_ids]
+    times = [fid / fps_safe for fid in codec_frame_ids]
     counts = [int(m.sum()) for m in masks]
     codec_cum = list(np.cumsum(counts)) if counts else []
     codec_total = int(codec_cum[-1]) if codec_cum else 0
     xx_c, yy_c = _step(times, codec_cum)
 
-    # Uniform baseline: same N frames as codec (at the same timestamps),
-    # but the patch budget is split equally across them. Both curves now
-    # reach the same budget — what differs is *which* patches each method
-    # picks within each frame (saliency vs equal-allocation).
-    n_uniform = len(times) if times else 1
+    # Uniform baseline: evenly sample COMPLETE frames from the same time
+    # window, no codec saliency involved. Each sampled frame contributes a
+    # whole patch grid.
     budget_int = int(total_patches_budget)
-    if n_uniform > 0 and budget_int > 0:
-        base = budget_int // n_uniform
-        rem = budget_int - base * n_uniform
-        uni_per_step = [base + (1 if i < rem else 0) for i in range(n_uniform)]
-    else:
-        uni_per_step = []
+    requested_uniform = max(0, int(uniform_requested_frames))
+    n_uniform = len(uniform_frame_ids)
+    uni_per_step = [grid_size for _ in uniform_frame_ids]
     uni_cum = list(np.cumsum(uni_per_step)) if uni_per_step else []
     uni_total = int(uni_cum[-1]) if uni_cum else 0
-    uni_times = times if times else [duration * 0.5]
+    uni_times = [fid / fps_safe for fid in uniform_frame_ids]
     xx_u, yy_u = _step(uni_times, uni_cum)
 
     # ─── Plot ───────────────────────────────────────────────────────────
@@ -650,17 +699,22 @@ def make_charts(
     else:
         codec_lbl = f"Codec · {saliency_signal}  ({codec_total:,} patches)"
     if uni_per_step:
-        u_per = uni_per_step[0]
-        u_extra = sum(1 for x in uni_per_step if x != u_per)
-        if u_extra == 0:
-            uni_lbl = f"Uniform baseline  ({uni_total:,} total · {u_per}/frame)"
-        else:
-            uni_lbl = (
-                f"Uniform baseline  ({uni_total:,} total · "
-                f"~{budget_int // max(1, n_uniform)}/frame, ±1)"
-            )
+        unused = max(0, budget_int - uni_total)
+        frame_part = (
+            f"{n_uniform}/{requested_uniform} frames fit budget"
+            if requested_uniform != n_uniform else f"{n_uniform} frames"
+        )
+        uni_lbl = (
+            f"Uniform full frames  ({frame_part} · {grid_size}/frame · "
+            f"{uni_total:,} total"
+            + (f" · {unused:,} budget unused" if unused else "")
+            + ")"
+        )
     else:
-        uni_lbl = f"Uniform baseline  ({uni_total:,} patches)"
+        uni_lbl = (
+            f"Uniform full frames  (0/{requested_uniform} frames fit budget "
+            f"{budget_int:,}; need {grid_size} patches/frame)"
+        )
 
     ax.fill_between(xx_c, yy_c, step=None, alpha=0.12, color="#4f46e5")
     ax.plot(xx_c, yy_c, color="#4f46e5", linewidth=2.2, label=codec_lbl)
@@ -838,12 +892,16 @@ def process(
 
     hb, wb = grids[0].shape
     grid_size = int(grids[0].shape[0] * grids[0].shape[1]) if grids else 0
-    # Uniform baseline samples the SAME number of frames as codec, evenly
-    # spaced in time; the budget is split equally across them.
-    n_uniform = max(1, len(fids))
-    uniform_per_frame = (
-        int(int(total_patches)) // n_uniform if n_uniform > 0 else 0
+    # Uniform full-frame sampling baseline: evenly sample complete frames
+    # from the same time window, independent of codec saliency.
+    requested_budget = int(total_patches)
+    uniform_requested_frames = len(fids)
+    uniform_frame_count = min(
+        uniform_requested_frames,
+        requested_budget // max(1, grid_size),
     )
+    uniform_frame_ids = sample_window_frame_ids(f_start, f_end, uniform_frame_count)
+    uniform_total = int(len(uniform_frame_ids) * grid_size)
     info = {
         "input": meta,
         "params": {
@@ -872,22 +930,25 @@ def process(
         "frame_window": {
             "first_decoded": int(f_start),
             "last_decoded": int(f_end),
-            "actual_frame_ids": [int(x) for x in fids],
+            "codec_frame_ids": [int(x) for x in fids],
+            "uniform_full_frame_ids": [int(x) for x in uniform_frame_ids],
         },
         "codec_per_frame_patches": [int(m.sum()) for m in masks],
         "uniform_baseline": {
-            "frames": int(n_uniform),
-            "patches_per_frame": int(uniform_per_frame),
-            "total_patches": int(uniform_per_frame * n_uniform),
+            "mode": "uniform_full_frame_sampling",
+            "requested_frames": int(uniform_requested_frames),
+            "frames": int(len(uniform_frame_ids)),
+            "patches_per_frame": int(grid_size),
+            "frame_ids": [int(x) for x in uniform_frame_ids],
+            "requested_budget": requested_budget,
+            "unused_budget": int(max(0, requested_budget - uniform_total)),
+            "total_patches": uniform_total,
             "explanation": (
-                "Same N frames as codec, evenly spaced in time. The patch "
-                "budget is split equally per frame ({budget} ÷ {n} = "
-                "{per}); the codec, by contrast, concentrates the same "
-                "budget on high-saliency patches across those frames."
-            ).format(
-                budget=int(total_patches),
-                n=int(n_uniform),
-                per=int(uniform_per_frame),
+                "Uniformly sample complete frames from the same time window. "
+                f"The baseline targets the same sampled-frame count as codec "
+                f"({uniform_requested_frames}), but each full frame costs "
+                f"{grid_size} patches, so only {len(uniform_frame_ids)} full "
+                "frames may fit inside the requested budget."
             ),
         },
         "resized_frame_size": f"{tw}x{th}",
@@ -916,7 +977,8 @@ def process(
     progress(0.95, desc="Building charts")
     duration_sec = (total / fps) if fps > 0 else 0.0
     chart_fig = make_charts(
-        grids, masks, fids, fps, duration_sec,
+        grids, masks, fids, uniform_frame_ids, uniform_requested_frames,
+        fps, duration_sec,
         int(total_patches), saliency_signal,
         groups=groups, gop_label=gop_resolved,
     )
@@ -1337,12 +1399,14 @@ with gr.Blocks(**_BLOCK_KW) as demo:
                     4, 64, value=16, step=1, label="Sampled frames",
                 )
                 top_k = gr.Slider(
-                    64, 8192, value=1024, step=32,
+                    16, 16384, value=3136, step=16,
                     label="Total patches budget (whole video)",
-                    info="The single budget shared across the whole video. "
-                         "The codec saliency picks these patches GLOBALLY — "
-                         "high-energy frames may contribute many, low-energy "
-                         "frames may contribute zero.",
+                    info="Default = sample_frames x patch_size^2 "
+                    "(16 x 14^2 = 3136). The uniform baseline spends "
+                    "this budget on evenly sampled complete frames; the "
+                    "codec path spends it on saliency-selected patches. If "
+                    "budget < sample_frames x patch_size^2, the full-frame "
+                    "baseline will use fewer frames than codec.",
                 )
                 patch_size = gr.Radio(
                     PATCH_CHOICES, value=14, label="Patch size (px)",
@@ -1422,13 +1486,13 @@ with gr.Blocks(**_BLOCK_KW) as demo:
             with gr.Group(elem_classes="ovc-card ovc-card-primary"):
                 gr.Markdown("### Cumulative patches over time")
                 gr.Markdown(
-                    "<small>Same number of sampled frames and the same total "
-                    "patch budget for both methods. <b>Indigo</b>: codec "
-                    "saliency — rises in bursts where the frames carry more "
-                    "information. <b>Cyan (dashed)</b>: uniform baseline — "
-                    "the same budget split equally per frame, so each step "
-                    "has the same height. Both curves end exactly at the "
-                    "dotted <b>budget</b> reference line.</small>"
+                    "<small><b>Indigo</b>: codec method — selects patches "
+                    "within frames according to saliency, so the curve rises "
+                    "in bursts. <b>Cyan (dashed)</b>: uniform full-frame "
+                    "sampling — evenly samples complete frames from the same "
+                    "time window, targeting the same sampled-frame count as "
+                    "codec when the budget allows. Each step is one full "
+                    "patch grid. The dotted line marks the requested budget.</small>"
                 )
                 chart_out = gr.Plot(label="", show_label=False)