|
|
| import math |
| import numpy as np |
| from numba import njit, uint32, uint64 |
| from PIL import Image |
| from pbc3_trace import timed |
|
|
|
|
| RESAMPLE_FILTER = Image.Resampling.BICUBIC |
| RESAMPLE_REDUCING_GAP = None |
| PALETTE_GENERATED = 0 |
|
|
|
|
| _UINT64_MASK = (1 << 64) - 1 |
| _UINT32_SCALE = 1.0 / 4294967296.0 |
|
|
|
|
| @njit(inline="always", cache=True) |
| def pcg_step(state): |
| old_state = state |
| state = uint64(old_state * uint64(6364136223846793005) + uint64(1442695040888963407)) |
| xorshifted = uint32(((old_state >> uint64(18)) ^ old_state) >> uint64(27)) |
| rot = uint32(old_state >> uint64(59)) |
| out = uint32((xorshifted >> rot) | (xorshifted << ((-rot) & uint32(31)))) |
| return state, out |
|
|
|
|
| class PBC3Rng: |
| def __init__(self, seed: int): |
| self.state = (int(seed) + 0x853C49E6748FEA9B) & _UINT64_MASK |
| self.next_uint32() |
|
|
| def next_uint32(self) -> int: |
| old_state = self.state |
| self.state = (old_state * 6364136223846793005 + 1442695040888963407) & _UINT64_MASK |
| xorshifted = (((old_state >> 18) ^ old_state) >> 27) & 0xFFFFFFFF |
| rot = (old_state >> 59) & 31 |
| return ((xorshifted >> rot) | ((xorshifted << ((-rot) & 31)) & 0xFFFFFFFF)) & 0xFFFFFFFF |
|
|
| def uniform(self, low: float = 0.0, high: float = 1.0) -> float: |
| return float(low) + (float(high) - float(low)) * (self.next_uint32() * _UINT32_SCALE) |
|
|
| def integers(self, low: int, high: int = None) -> int: |
| if high is None: |
| low, high = 0, low |
| low, high = int(low), int(high) |
| if high <= low: |
| return low |
| return low + self.next_uint32() % (high - low) |
|
|
|
|
| def bicubic_resample_2d(values, out_h: int, out_w: int) -> np.ndarray: |
| """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() |
| 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.float32, |
| ) |
|
|
|
|
| def ceil_div(a: int, b: int) -> int: |
| """## Returns ceil(a / b) for integer cell/grid math""" |
| return (int(a) + int(b) - 1) // int(b) |
|
|
|
|
| def norm(values) -> np.ndarray: |
| """## Returns values normalized to 0..1, or all 1s when the range is flat""" |
| arr = np.asarray(values, dtype=np.float64) |
| rng = arr.max() - arr.min() |
| if rng <= 0: |
| return np.ones_like(arr) |
| return (arr - arr.min()) / rng |
|
|
|
|
| def interp(start: float, end: float, step: int, count: int) -> float: |
| """## Returns a linear interpolation between start and end for the current patch step""" |
| if count <= 1: |
| return float(end) |
| p = (step - 1) / max(1, count - 1) |
| return float(start) * (1 - p) + float(end) * p |
|
|
|
|
| def integral(a) -> np.ndarray: |
| """## Returns a padded integral image for fast rectangle sums""" |
| return np.pad(a.astype(np.int64).cumsum(0).cumsum(1), ((1, 0), (1, 0))) |
|
|
|
|
| @njit(cache=True) |
| def _box_cell_bound_kernel(integral_arr, x, y, bw, bh, cell_size): |
| nx = (bw + cell_size - 1) // cell_size |
| ny = (bh + cell_size - 1) // cell_size |
| total = 0.0 |
| for iy in range(ny): |
| y0 = y + iy * cell_size |
| y1 = y + bh if iy == ny - 1 else y + (iy + 1) * cell_size |
| for ix in range(nx): |
| x0 = x + ix * cell_size |
| x1 = x + bw if ix == nx - 1 else x + (ix + 1) * cell_size |
| s = float(integral_arr[y1, x1] - integral_arr[y0, x1] - integral_arr[y1, x0] + integral_arr[y0, x0]) |
| total += s * s / ((y1 - y0) * (x1 - x0)) |
| return total |
|
|
|
|
| @njit(cache=True) |
| def _base_cell_size_kernel(res, max_cell): |
| h, w = res.shape |
| s = 0.0 |
| for i in range(h): |
| for j in range(w): |
| v = res[i, j] |
| s += v if v >= 0 else -v |
| mean_abs = s / (h * w) |
| if mean_abs <= 0.0: |
| return max_cell |
|
|
| gx = 0.0 |
| if w > 1: |
| for i in range(h): |
| for j in range(w - 1): |
| d = res[i, j + 1] - res[i, j] |
| gx += d if d >= 0 else -d |
| gx /= h * (w - 1) |
|
|
| gy = 0.0 |
| if h > 1: |
| for i in range(h - 1): |
| for j in range(w): |
| d = res[i + 1, j] - res[i, j] |
| gy += d if d >= 0 else -d |
| gy /= (h - 1) * w |
|
|
| ratio = (gx + gy) / (mean_abs + 1.0) |
| if ratio < 0.25: |
| return 32 |
| if ratio < 0.5: |
| return 16 |
| if ratio < 1.0: |
| return 8 |
| return 4 |
|
|
|
|
| @njit(cache=True) |
| def _cell_mean_kernel(values, cell_size): |
| h, w = values.shape |
| ny = (h + cell_size - 1) // cell_size |
| nx = (w + cell_size - 1) // cell_size |
| out = np.empty((ny, nx), dtype=np.float32) |
| for iy in range(ny): |
| y0 = iy * cell_size |
| y1 = min(h, y0 + cell_size) |
| for ix in range(nx): |
| x0 = ix * cell_size |
| x1 = min(w, x0 + cell_size) |
| total = 0.0 |
| for y in range(y0, y1): |
| for x in range(x0, x1): |
| total += values[y, x] |
| out[iy, ix] = total / ((y1 - y0) * (x1 - x0)) |
| return out |
|
|
|
|
| @njit(cache=True) |
| def _cell_representability_kernel(values, cell_size): |
| h, w = values.shape |
| total_energy = 0.0 |
| for y in range(h): |
| for x in range(w): |
| total_energy += values[y, x] * values[y, x] |
| if total_energy <= 0.0: |
| return 1.0 |
|
|
| ny = (h + cell_size - 1) // cell_size |
| nx = (w + cell_size - 1) // cell_size |
| captured = 0.0 |
| for iy in range(ny): |
| y0 = iy * cell_size |
| y1 = min(h, y0 + cell_size) |
| for ix in range(nx): |
| x0 = ix * cell_size |
| x1 = min(w, x0 + cell_size) |
| total = 0.0 |
| for y in range(y0, y1): |
| for x in range(x0, x1): |
| total += values[y, x] |
| captured += total * total / ((y1 - y0) * (x1 - x0)) |
| return min(1.0, max(0.0, captured / total_energy)) |
|
|
|
|
| @njit(cache=True) |
| def _anchor_block_scores_kernel(err, block_size): |
| h, w = err.shape |
| ny = (h + block_size - 1) // block_size |
| nx = (w + block_size - 1) // block_size |
| n = ny * nx |
| scores = np.empty(n, dtype=np.float64) |
| ys = np.empty(n, dtype=np.int64) |
| xs = np.empty(n, dtype=np.int64) |
| k = 0 |
| for by in range(ny): |
| y0 = by * block_size |
| y1 = min(h, y0 + block_size) |
| for bx in range(nx): |
| x0 = bx * block_size |
| x1 = min(w, x0 + block_size) |
| s = 0.0 |
| for i in range(y0, y1): |
| for j in range(x0, x1): |
| s += err[i, j] |
| scores[k] = s / ((y1 - y0) * (x1 - x0)) |
| ys[k] = (y0 + y1 - 1) // 2 |
| xs[k] = (x0 + x1 - 1) // 2 |
| k += 1 |
| return scores, ys, xs |
|
|
|
|
| @njit(cache=True) |
| def _mse_kernel(a, b): |
| total = 0.0 |
| n = a.size |
| flat_a = a.ravel() |
| flat_b = b.ravel() |
| for i in range(n): |
| d = float(flat_a[i]) - float(flat_b[i]) |
| total += d * d |
| return total / n |
|
|
|
|
|
|
| @njit(cache=True) |
| def _palette_bounds_kernel(values): |
| flat = values.ravel() |
| mn = int(flat[0]) |
| mx = int(flat[0]) |
| for i in range(1, flat.size): |
| v = int(flat[i]) |
| if v < mn: |
| mn = v |
| if v > mx: |
| mx = v |
| neg = -mn if mn < 0 else 0 |
| pos = mx if mx > 0 else 0 |
| if neg > 255: |
| neg = 255 |
| if pos > 255: |
| pos = 255 |
| return neg, pos |
|
|
|
|
| @njit(cache=True) |
| def _mask_from_values_kernel(flat, mask_size, negative_max, positive_max, positive_bias): |
| mask = np.zeros(mask_size, dtype=np.uint8) |
| mask[0] = 1 |
| side_bits = mask_size - 1 |
| negative_max = max(0, int(negative_max)) |
| positive_max = max(0, int(positive_max)) |
| if side_bits == 0 or (negative_max == 0 and positive_max == 0): |
| return mask |
|
|
| if negative_max == 0: |
| pos_count = min(side_bits, positive_max) |
| neg_count = 0 |
| elif positive_max == 0: |
| pos_count = 0 |
| neg_count = min(side_bits, negative_max) |
| else: |
| raw_pos = side_bits * positive_max / (positive_max + negative_max) |
| pos_count = int(math.ceil(raw_pos)) if positive_bias else int(math.floor(raw_pos)) |
| pos_count = min(side_bits - 1, max(1, pos_count), positive_max) |
| neg_count = min(side_bits - pos_count, negative_max) |
| if neg_count == 0 and negative_max > 0 and side_bits > pos_count: |
| neg_count = 1 |
| pos_count = max(1, pos_count - 1) |
|
|
| if pos_count > 0 and positive_max > 0: |
| for i in range(flat.size): |
| v = int(flat[i]) |
| if v > 0: |
| if v > positive_max: |
| v = positive_max |
| b = 1 + min(((v - 1) * pos_count) // positive_max, pos_count - 1) |
| if b < mask_size: |
| mask[b] = 1 |
|
|
| if neg_count > 0 and negative_max > 0: |
| for i in range(flat.size): |
| v = int(flat[i]) |
| if v < 0: |
| mag = -v |
| if mag > negative_max: |
| mag = negative_max |
| b = 1 + pos_count + min(((mag - 1) * neg_count) // negative_max, neg_count - 1) |
| if b < mask_size: |
| mask[b] = 1 |
| return mask |
|
|
|
|
| @njit(cache=True) |
| def _quantize_signed_kernel(vals, pal): |
| out = np.empty(vals.size, dtype=np.uint16) |
| flat = vals.ravel() |
| for i in range(flat.size): |
| v = int(flat[i]) |
| best_i = 0 |
| best_d = abs(v - int(pal[0])) |
| for j in range(1, pal.size): |
| d = abs(v - int(pal[j])) |
| if d < best_d: |
| best_d = d |
| best_i = j |
| out[i] = best_i |
| return out |
|
|
|
|
| def palette_bounds(values) -> tuple[int, int]: |
| """## Returns negative max (min or 0) and positive max (true max) |
| The reason there is a "negative max" instead of just a "min" is that palette generation is symmetric, and the negative max is used to determine the number of negative bins in the mask.""" |
| return _palette_bounds_kernel(np.ascontiguousarray(values, dtype=np.int16)) |
|
|
|
|
|
|
| def range_counts(mask_size, negative_max=255, positive_max=255, positive_bias=True) -> tuple[int, int]: |
| """## Returns the number of positive and negative bins in the mask, given the mask size and the maximum positive and negative values. |
| - The number of bins is determined by the relative sizes of the positive and negative ranges, and positive_bias is a tiebreaker |
| - If for example mask size is 4, negative max is 10, positive max is 20, then the positive range is supposed to be twice as large as the negative range, so counting the 1 zero-bin, there should be 2 positive bins and 1 negative bin, representing ranges (-10,-1), (0), (1, 10), (11, 20) |
| - If in that example mask size was 3 though, to make sure there is still a negative and a positive range, the positive range would have to be 1 bin""" |
| side_bits = max(0, int(mask_size) - 1) |
| negative_max = max(0, int(negative_max)) |
| positive_max = max(0, int(positive_max)) |
| if side_bits == 0 or (negative_max == 0 and positive_max == 0): |
| return 0, 0 |
| if negative_max == 0: |
| return min(side_bits, positive_max), 0 |
| if positive_max == 0: |
| return 0, min(side_bits, negative_max) |
| raw_pos = side_bits * positive_max / (positive_max + negative_max) |
| pos_count = math.ceil(raw_pos) if positive_bias else math.floor(raw_pos) |
| pos_count = min(side_bits - 1, max(1, pos_count), positive_max) |
| neg_count = min(side_bits - pos_count, negative_max) |
| if neg_count == 0 and negative_max > 0 and side_bits > pos_count: |
| neg_count = 1 |
| pos_count = max(1, pos_count - 1) |
| return pos_count, neg_count |
|
|
|
|
| def range_for_mask_index(index, mask_size, negative_max=255, positive_max=255, positive_bias=True) -> tuple[int, int] | None: |
| """## Returns the range of values for a given index in a mask. |
| ### 0 is always the zero bin |
| PS: It might seem like negative max and positive max are useless here but it's needed for a bunch of functions so |
| it's best to just compute them once and pass them around instead of recomputing them every time inside the functions""" |
| pos_count, neg_count = range_counts(mask_size, negative_max, positive_max, positive_bias) |
| if index == 0: |
| return 0, 0 |
| if 1 <= index <= pos_count: |
| bin_i = index - 1 |
| start = 1 + (bin_i * positive_max) // pos_count |
| end = ((bin_i + 1) * positive_max) // pos_count |
| return (start, end) if start <= end else None |
| bin_i = index - 1 - pos_count |
| if 0 <= bin_i < neg_count: |
| low_mag = 1 + (bin_i * negative_max) // neg_count |
| high_mag = ((bin_i + 1) * negative_max) // neg_count |
| return (-high_mag, -low_mag) if high_mag >= low_mag else None |
| return None |
|
|
|
|
| def mask_from_values(values, mask_size, negative_max=255, positive_max=255, positive_bias=True) -> list[int]: |
| """## Returns a mask of the given size, with 1s for the bins that have values in them, and 0s for the bins that don't""" |
| flat = np.clip(np.rint(np.asarray(values)).astype(np.int32).ravel(), -negative_max, positive_max) |
| mask = _mask_from_values_kernel( |
| np.ascontiguousarray(flat, dtype=np.int32), |
| int(mask_size), |
| int(negative_max), |
| int(positive_max), |
| bool(positive_bias), |
| ) |
| return [int(x) for x in mask] |
|
|
|
|
|
|
| def active_value_count(mask, negative_max=255, positive_max=255, positive_bias=True) -> int: |
| """## Returns the number of values within the mask selected ranges""" |
| count = 0 |
| for i, bit in enumerate(mask): |
| if bit: |
| r = range_for_mask_index(i, len(mask), negative_max, positive_max, positive_bias) |
| if r is not None: |
| start, end = r |
| count += end - start + 1 |
| return max(1, count) |
|
|
|
|
| def resolve_palette_bitcount(mask, max_bitcount, negative_max=255, positive_max=255, positive_bias=True) -> int: |
| """## Returns the number of bits needed to represent the active values in the mask |
| This is so for example if the max bitcount is 8 but the mask only has 60 active values then the bitcount |
| for this palette can be reduced to 6 instead with no loss""" |
| value_count = active_value_count(mask, negative_max, positive_max, positive_bias) |
| needed = max(1, math.ceil(math.log2(value_count))) |
| return min(int(max_bitcount), needed) |
|
|
|
|
| def palette_generator(mask, max_bitcount, negative_max=255, positive_max=255, positive_bias=True) -> np.ndarray: |
| """## Returns a 1D array palette which is all the selected values for the patch""" |
| bitcount = resolve_palette_bitcount(mask, max_bitcount, negative_max, positive_max, positive_bias) |
| size = 1 << bitcount |
| active_ranges = [] |
| for i, bit in enumerate(mask): |
| if bit: |
| r = range_for_mask_index(i, len(mask), negative_max, positive_max, positive_bias) |
| if r is not None: |
| active_ranges.append(r) |
| palette = [] |
| if mask and mask[0]: |
| palette.append(0) |
| active_ranges = [r for r in active_ranges if r != (0, 0)] |
| value_count = active_value_count(mask, negative_max, positive_max, positive_bias) |
| if size >= value_count: |
| for start, end in active_ranges: |
| palette.extend(range(start, end + 1)) |
| if len(palette) < size: |
| palette.extend([palette[-1] if palette else 0] * (size - len(palette))) |
| return np.array(palette[:size], dtype=np.int16) |
| if not active_ranges: |
| return np.zeros(size, dtype=np.int16) |
| remaining = size - len(palette) |
| counts = [0] * len(active_ranges) |
| for i in range(remaining): |
| counts[i % len(active_ranges)] += 1 |
| for (start, end), count in zip(active_ranges, counts): |
| if count == 1: |
| palette.append(int(round((start + end) / 2))) |
| elif count > 1: |
| for j in range(count): |
| t = (j + 1) / (count + 1) |
| palette.append(int(round(start + (end - start) * t))) |
| if len(palette) < size: |
| palette.extend([palette[-1] if palette else 0] * (size - len(palette))) |
| return np.array(palette[:size], dtype=np.int16) |
|
|
|
|
| def quantize_signed(values, palette): |
| """## Returns the indices of the closest palette values for each value in the input array |
| `quantize_signed([8,13,20], [5,30]) == [0,0,1]`""" |
| vals = np.asarray(values, dtype=np.int16) |
| pal = np.asarray(palette, dtype=np.int16) |
| out = _quantize_signed_kernel(np.ascontiguousarray(vals), np.ascontiguousarray(pal)) |
| return out.reshape(vals.shape) |
|
|
|
|
|
|
| def signed_resample(values, out_h: int, out_w: int, trace=None, purpose: str = None) -> np.ndarray: |
| """## Resizes a signed grid to a patch-sized int16 delta image""" |
| if trace is not None and purpose: |
| trace.count(f"signed_resample.{purpose}") |
| return np.rint(bicubic_resample_2d(values, int(out_h), int(out_w))).astype(np.int16) |
|
|
|
|
| def signed_resample_cells(values, cell_size: int) -> np.ndarray: |
| """## Resamples a patch residual to its cell grid size""" |
| h, w = values.shape |
| return signed_resample(values, ceil_div(h, cell_size), ceil_div(w, cell_size)) |
|
|
|
|
| def project_residual_to_grid(values, cell_size: int, mode: str = "bicubic") -> np.ndarray: |
| """Projects a signed residual patch onto the stored coarse grid.""" |
| cell_size = int(cell_size) |
| if cell_size < 1: |
| raise ValueError("cell_size must be >= 1") |
| mode = str(mode).strip().lower() |
| if mode == "bicubic": |
| return signed_resample_cells(values, cell_size) |
| if mode == "cell_mean": |
| return _cell_mean_kernel(np.ascontiguousarray(values, dtype=np.float32), cell_size) |
| raise ValueError(f"unknown residual projection mode: {mode}") |
|
|
|
|
| def cell_representability(residual_patch, cell_size: int) -> float: |
| """Returns the fraction of residual energy representable by scalar cells.""" |
| cell_size = int(cell_size) |
| if cell_size < 1: |
| raise ValueError("cell_size must be >= 1") |
| return float(_cell_representability_kernel( |
| np.ascontiguousarray(residual_patch, dtype=np.float64), cell_size |
| )) |
|
|
|
|
| def apply_grid(canvas_layer, x: int, y: int, w: int, h: int, cell_size: int, values, trace=None) -> None: |
| """## Applies a signed patch grid into one canvas channel""" |
| delta = signed_resample(values, h, w, trace=trace, purpose="selected_patch_application") |
| canvas_layer[y:y + h, x:x + w] += delta |
|
|
|
|
| def apply_delta(canvas_layer, x: int, y: int, w: int, h: int, delta) -> None: |
| """## Applies an already reconstructed signed patch delta""" |
| delta = np.asarray(delta, dtype=np.int16) |
| if delta.shape != (int(h), int(w)): |
| raise ValueError(f"delta shape must be {(int(h), int(w))}, got {delta.shape}") |
| canvas_layer[y:y + h, x:x + w] += delta |
|
|
|
|
| def image_mse(reference, reconstructed) -> float: |
| """## Returns mean squared error between two same-shaped image arrays""" |
| a = np.ascontiguousarray(reference, dtype=np.float32) |
| b = np.ascontiguousarray(reconstructed, dtype=np.float32) |
| if a.shape != b.shape: |
| raise ValueError(f"MSE arrays must have the same shape, got {a.shape} and {b.shape}") |
| return float(_mse_kernel(a, b)) |
|
|
|
|
| def final_mse(reference_image, reconstructed_image) -> float: |
| """## Returns final image MSE from two PIL images or image-like arrays""" |
| return image_mse(np.asarray(reference_image, dtype=np.float32), np.asarray(reconstructed_image, dtype=np.float32)) |
|
|
|
|
| def box_cell_bound(integral_arr, x: int, y: int, bw: int, bh: int, cell_size: int) -> float: |
| """## Returns the best-case signed-error energy for a box/cell-size pair""" |
| return float(_box_cell_bound_kernel( |
| np.ascontiguousarray(integral_arr, dtype=np.int64), |
| int(x), |
| int(y), |
| int(bw), |
| int(bh), |
| int(cell_size), |
| )) |
|
|
|
|
| def top_anchors(visible_error_channel, top_k: int, block_size: int, channel: int) -> list[tuple[int, int, int]]: |
| """## Returns the strongest error anchors as (channel, y, x) tuples""" |
| h, w = visible_error_channel.shape |
| block_size = max(1, int(block_size)) |
| if block_size == 1: |
| flat = visible_error_channel.reshape(-1) |
| k = min(int(top_k), flat.size) |
| idx = np.argpartition(flat, -k)[-k:] |
| idx = idx[np.argsort(flat[idx])[::-1]] |
| return [(channel, int(i) // w, int(i) % w) for i in idx] |
|
|
| scores, ys, xs = _anchor_block_scores_kernel(np.ascontiguousarray(visible_error_channel), block_size) |
| if scores.size == 0: |
| return [] |
| k = min(int(top_k), scores.size) |
| idx = np.argpartition(scores, -k)[-k:] |
| order = idx[np.argsort(scores[idx])[::-1]] |
| return [(channel, int(ys[i]), int(xs[i])) for i in order] |
|
|
|
|
| def select_top_indices(scores, keep: int) -> np.ndarray: |
| """## Returns unordered indices of the top `keep` scores""" |
| scores = np.asarray(scores) |
| n = scores.size |
| keep = max(1, int(keep)) |
| if keep >= n: |
| return np.arange(n) |
| cut = n - keep |
| if cut < keep: |
| drop = np.argpartition(scores, cut)[:cut] |
| mask = np.ones(n, dtype=bool) |
| mask[drop] = False |
| return np.nonzero(mask)[0] |
| return np.argpartition(scores, -keep)[-keep:] |
|
|
|
|
| def sample_box(rng, anchor, image_w: int, image_h: int, config) -> tuple[int, int, int, int, int, int, int]: |
| """## Samples a candidate patch box around an anchor point""" |
| c, ay, ax = anchor |
| min_size = max(1, int(config.min_patch_size)) |
| max_w = max(min_size, min(int(config.max_patch_size), image_w)) |
| max_h = max(min_size, min(int(config.max_patch_size), image_h)) |
| w = int(round(2 ** rng.uniform(math.log2(min_size), math.log2(max_w)))) |
| h = int(round(2 ** rng.uniform(math.log2(min_size), math.log2(max_h)))) |
| w = min(max(1, w), image_w) |
| h = min(max(1, h), image_h) |
| x_min = max(0, ax - w + 1) |
| x_max = min(ax, image_w - w) |
| y_min = max(0, ay - h + 1) |
| y_max = min(ay, image_h - h) |
| x = int(rng.integers(x_min, x_max + 1)) if x_min <= x_max else max(0, min(ax, image_w - w)) |
| y = int(rng.integers(y_min, y_max + 1)) if y_min <= y_max else max(0, min(ay, image_h - h)) |
| return c, x, y, w, h, ax, ay |
|
|
|
|
| def search_boxes(target, canvas, config, rng, channel: int, depth: int, search_q: float, ordered=False, trace=None): |
| """Samples error-centered boxes and keeps the best prescores.""" |
| h_img, w_img, _ = target.shape |
| with timed(trace, "action.search.visible_error"): |
| visible = np.clip(canvas[:, :, channel], 0, 255).astype(np.int32) |
| abs_error = np.abs(target[:, :, channel] - visible) |
| with timed(trace, "action.search.integral_error"): |
| integral_abs = integral(abs_error) |
| with timed(trace, "action.search.anchors"): |
| anchors = top_anchors(abs_error, config.top_k, config.anchor_block_size, channel) |
| if not anchors: |
| return [] |
|
|
| specs, sums, areas = [], [], [] |
| with timed(trace, "action.search.sample_boxes", depth=int(depth)): |
| for i in range(max(1, int(depth))): |
| c, x, y, bw, bh, _, _ = sample_box(rng, anchors[i % len(anchors)], w_img, h_img, config) |
| total = integral_abs[y + bh, x + bw] - integral_abs[y, x + bw] - integral_abs[y + bh, x] + integral_abs[y, x] |
| if total <= 0: |
| continue |
| sums.append(float(total)) |
| areas.append(float(bw * bh)) |
| specs.append((c, x, y, bw, bh)) |
| if not specs: |
| return [] |
|
|
| with timed(trace, "action.search.prescore"): |
| scores = search_q * norm(sums) - (1.0 - search_q) * norm(areas) |
| with timed(trace, "action.search.select_top_boxes"): |
| keep = select_top_indices(scores, config.proposal_depth) |
| if ordered: |
| keep = sorted(keep, key=lambda i: scores[i], reverse=True) |
| return [specs[i] for i in keep] |
|
|
|
|
| def base_cell_size(residual_patch, config) -> int: |
| """## Returns the rough cell size suggested by local residual frequency""" |
| if residual_patch.size == 0: |
| return int(config.max_cell_size) |
| return int(_base_cell_size_kernel(np.ascontiguousarray(residual_patch, dtype=np.float64), int(config.max_cell_size))) |
|
|
|
|
| def patch_bits_for(patch, channel_bits: int) -> int: |
| """## Returns the serialized bit cost of a complete patch""" |
| grid_bits = ceil_div(patch["w"], patch["cell_size"]) * ceil_div(patch["h"], patch["cell_size"]) * patch["bitcount"] |
| return channel_bits + 64 + 16 + 1 + 10 + len(patch["mask"]) + 8 + 8 + 4 + grid_bits |
|
|
|
|
| def patch_header_bits(channel_bits: int, mask_size: int) -> int: |
| """## Returns the fixed patch header bit cost before grid indices""" |
| return channel_bits + 64 + 10 + mask_size + 8 + 8 + 4 + 16 + 1 |
|
|
|
|
| def make_patch(channel: int, x: int, y: int, w: int, h: int, cell_size: int, residual, config, max_bitcount: int, trace=None): |
| """## Builds a generated-palette patch and the signed values it will apply""" |
| mode = str(getattr(config, "residual_projection_mode", "bicubic")).strip().lower() |
| if trace is not None: |
| trace.count("candidates_built") |
| trace.count("grid_coefficients_processed", ceil_div(h, cell_size) * ceil_div(w, cell_size)) |
| trace.count(f"projection.{mode}") |
| if mode == "bicubic": |
| trace.count("signed_resample.projection") |
| small = np.rint(project_residual_to_grid(residual, cell_size, mode)).astype(np.int16) |
| negative_max, positive_max = palette_bounds(small) |
| mask = mask_from_values(small, config.mask_size, negative_max, positive_max, config.positive_bias) |
| pal = palette_generator(mask, max_bitcount, negative_max, positive_max, config.positive_bias) |
| indices = quantize_signed(np.clip(small, -negative_max, positive_max), pal) |
| values = pal[indices] |
| bitcount = resolve_palette_bitcount(mask, max_bitcount, negative_max, positive_max, config.positive_bias) |
| return { |
| "channel": channel, |
| "x": x, |
| "y": y, |
| "w": w, |
| "h": h, |
| "cell_size": cell_size, |
| "indices": indices, |
| "palette_mode": PALETTE_GENERATED, |
| "palette": None, |
| "bitcount": bitcount, |
| "mask": mask, |
| "neg": negative_max, |
| "pos": positive_max, |
| "max_bitcount": max_bitcount, |
| }, values |
|
|
|
|
|
|
| def debug_line(kind: str, **items) -> str: |
| """## Returns one plain-text debug line with key=value fields""" |
| return kind + " " + " ".join(f"{k}={v}" for k, v in items.items()) |
|
|