Upload segmentation.py

segmentation.py

@@ -156,10 +156,15 @@ def label_channel_bytes(n: int, k: int, table_size: int = SEG_TABLE_SIZE) -> int
 
 
 def pack_label_channel(labels: np.ndarray, k: int, table_size: int = SEG_TABLE_SIZE) -> bytes:
-    flat = np.ascontiguousarray(labels, dtype=np.int64).ravel()
-    n = int(flat.size)
     k = int(k)
     table_size = int(table_size)
+    # Fast path when segment_stage attached pre-packed LUT data
+    lut_cache = getattr(labels, "_hpcq_lut_bytes", None)
+    if lut_cache is not None and lut_cache[0] == (k, table_size):
+        return lut_cache[1]
+
+    flat = np.ascontiguousarray(labels, dtype=np.uint8).ravel()
+    n = int(flat.size)
     if k <= 16:
         if flat.size and int(flat.max()) > 15:
             raise ValueError(f"label {int(flat.max())} does not fit nibble packing (k={k})")
@@ -167,7 +172,7 @@ def pack_label_channel(labels: np.ndarray, k: int, table_size: int = SEG_TABLE_S
             enc = 2 if table_size <= 256 else 3
             idx_dt = np.uint8 if enc == 2 else np.dtype("<u2")
             n_blocks = n // 256
-            rows = flat.reshape(n_blocks, 256).astype(np.uint8)
+            rows = flat.reshape(n_blocks, 256)
             lo = rows[:, 0::2]
             hi = rows[:, 1::2]
             packed_rows = (lo | (hi << 4)).astype(np.uint8)  # (n_blocks, 128)
@@ -193,9 +198,9 @@ def pack_label_channel(labels: np.ndarray, k: int, table_size: int = SEG_TABLE_S
             data = _pack_nibbles_1d(flat)
     else:
         enc = 1
-        if flat.size and (int(flat.max()) > 255 or int(flat.min()) < 0):
+        if flat.size and (int(labels.max()) > 255 or int(labels.min()) < 0):
             raise ValueError("labels must be in [0, 255] for 8-bit packing")
-        data = flat.astype(np.uint8).tobytes()
+        data = flat.tobytes()
     return SEG_MAGIC + bytes([k, enc]) + _struct.pack("<I", n) + data
 
 
@@ -260,6 +265,10 @@ def _segment_stage_dense(
     return order, labels, rounds
 
 
+class _LUTLabels(np.ndarray):
+    _hpcq_lut_bytes: tuple[tuple[int, int], bytes] | None = None
+
+
 def segment_stage(
     X: np.ndarray,
     W: np.ndarray | None,
@@ -296,14 +305,21 @@ def segment_stage(
     col_abs = np.mean(np.abs(X_sample), axis=0)
     table_labels[2] = _order_to_canonical_labels(np.argsort(col_abs, kind="stable"), k).astype(np.uint8)
 
-    # Entries 3..M-1: prototype blocks refined with the proxy fixed-point loop
-    n_proto = M - 3
-    proto_idx = np.linspace(0, n_blocks - 1, n_proto, dtype=np.int64)
+    # Entries 3..M-1: up to 253 prototype blocks refined with the proxy loop,
+    # plus fast direct band_labels prototypes when M > 256
+    n_need = M - 3
+    n_dense = min(n_need, 253)
+    proto_idx = np.linspace(0, n_blocks - 1, n_dense, dtype=np.int64)
     Xp = Xf[proto_idx].astype(np.float64)
     Wp = None if W is None else np.asarray(W[proto_idx], dtype=np.float64)
     ord_p, _, rounds = _segment_stage_dense(Xp, Wp, cfg, proxy, k)
-    for i in range(n_proto):
+    for i in range(n_dense):
         table_labels[3 + i] = _order_to_canonical_labels(ord_p[i], k).astype(np.uint8)
+    if n_need > n_dense:
+        extra_idx = np.linspace(1, n_blocks - 2, n_need - n_dense, dtype=np.int64)
+        ord_e = make_order(band_labels(Xf[extra_idx].astype(np.float64), k))
+        for i in range(n_need - n_dense):
+            table_labels[3 + n_dense + i] = _order_to_canonical_labels(ord_e[i], k).astype(np.uint8)
 
     table_orders = make_order(table_labels)  # (M, 256) int64
 
@@ -392,7 +408,22 @@ def segment_stage(
             best_local = np.argmin(sse, axis=1)
             best_m[s:e] = top_idx[np.arange(e - s), best_local]
 
-    return table_orders[best_m], table_labels[best_m], rounds
+    out_labels = table_labels[best_m].view(_LUTLabels)
+    enc = 2 if M <= 256 else 3
+    idx_dt = np.uint8 if enc == 2 else np.dtype("<u2")
+    lo = table_labels[:, 0::2]
+    hi = table_labels[:, 1::2]
+    table_nib = (lo | (hi << 4)).astype(np.uint8)
+    raw_lut = (
+        SEG_MAGIC
+        + bytes([k, enc])
+        + _struct.pack("<I", n_blocks * 256)
+        + _struct.pack("<H", M)
+        + table_nib.tobytes()
+        + best_m.astype(idx_dt).tobytes()
+    )
+    out_labels._hpcq_lut_bytes = ((k, M), raw_lut)
+    return table_orders[best_m], out_labels, rounds
 
 
 def q2k_proxy_cfg(clamp_val: float, cfg: HPCConfig):