Spaces:
Running on Zero
Running on Zero
File size: 5,932 Bytes
85300f6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Nine-way conditional trees and connected-component multi-point readout."""
import numpy as np
from .hierarchy import POINT_RC
def softmax9(logits):
values = np.asarray(logits, dtype=np.float64)
if values.shape != (9,) or not np.isfinite(values).all():
raise ValueError("Expected nine finite logits")
probabilities = np.exp(values - values.max())
return probabilities / probabilities.sum()
def grid_location(path):
row = col = 0
for choice in path:
dr, dc = POINT_RC[choice]
row, col = row * 3 + dr, col * 3 + dc
return row, col
def branch_spec(builder, parent, choice):
# Copy mutable histories, but share immutable PIL images and input metadata.
child = dict(parent)
for key in ("images", "image_keys", "counts", "mapping", "path"):
child[key] = list(parent[key])
if 'roi_boxes' in parent:
child['roi_boxes'] = list(parent['roi_boxes'])
child["path"].append(choice)
builder.advance(child, choice)
return child
def probability_map(builder, request, evaluate_nodes, mode="full", root_relative_cutoff=0.01):
"""Three levels, 729 leaves; labels/masks/counts are not accepted as arguments.
root_pruned uses only the first decision to omit low-mass root subtrees.
Unvisited leaves remain explicitly unknown, with their missing mass reported.
No probability renormalization is applied after pruning.
"""
if mode not in ("full", "root_pruned"):
raise ValueError("Unknown tree policy")
if not 0 <= root_relative_cutoff <= 1:
raise ValueError("Invalid root cutoff")
if request["input"]["answer_space"]["kind"] != "point":
raise ValueError("Point input required")
grid = np.zeros((27, 27), dtype=np.float64)
known = np.zeros((27, 27), dtype=bool)
nodes, pruned = [], []
frontier = [(builder.initial_spec(request["input"], request["id"]), 1.0)]
for depth in range(3):
logits = evaluate_nodes([state for state, _ in frontier])
if len(logits) != len(frontier):
raise ValueError("Node evaluator returned a different batch size")
following = []
for (state, mass), scores in zip(frontier, logits):
probabilities = softmax9(scores)
path = list(state["path"])
nodes.append({"path": path, "mass": mass, "logits": list(map(float, scores)),
"probabilities": probabilities.tolist()})
cutoff = probabilities.max() * root_relative_cutoff
for choice, probability in enumerate(probabilities):
child_path = [*path, choice]
child_mass = mass * float(probability)
if depth == 2:
row, col = grid_location(child_path)
if known[row, col]:
raise ValueError("Duplicate leaf location")
grid[row, col], known[row, col] = child_mass, True
elif depth == 0 and mode == "root_pruned" and probability < cutoff:
pruned.append({"path": child_path, "mass": child_mass})
else:
following.append((branch_spec(builder, state, choice), child_mass))
frontier = following
retained, omitted = float(grid.sum()), sum(r["mass"] for r in pruned)
if not np.isclose(retained + omitted, 1, atol=1e-12, rtol=0):
raise ValueError("Probability mass was lost or duplicated")
return {"probabilities": grid, "known": known, "nodes": nodes,
"evaluated_nodes": len(nodes), "evaluated_leaves": int(known.sum()),
"retained_mass": retained, "pruned_mass": omitted, "pruned_branches": pruned}
def component_points(probabilities, known=None, relative_threshold=0.1, connectivity=8):
"""Pick each superlevel component's maximum cell, with deterministic tie breaks."""
grid = np.asarray(probabilities, dtype=np.float64)
if grid.ndim != 2 or not np.isfinite(grid).all() or (grid < 0).any():
raise ValueError("Invalid spatial probability grid")
if not 0 < relative_threshold <= 1 or connectivity not in (4, 8):
raise ValueError("Invalid connected-component settings")
known = np.ones_like(grid, dtype=bool) if known is None else np.asarray(known, dtype=bool)
if known.shape != grid.shape:
raise ValueError("Known-cell mask shape mismatch")
peak = float(grid[known].max()) if known.any() else 0.0
threshold = peak * relative_threshold
mask = known & (grid >= threshold) & (grid > 0)
visited = np.zeros_like(mask)
neighbors = [(-1, 0), (1, 0), (0, -1), (0, 1)]
if connectivity == 8:
neighbors += [(-1, -1), (-1, 1), (1, -1), (1, 1)]
height, width = grid.shape
components, points = [], []
for row, col in zip(*np.where(mask)):
if visited[row, col]:
continue
visited[row, col] = True
stack, cells = [(int(row), int(col))], []
while stack:
y, x = stack.pop()
cells.append((y, x))
for dy, dx in neighbors:
yy, xx = y + dy, x + dx
if 0 <= yy < height and 0 <= xx < width and mask[yy, xx] and not visited[yy, xx]:
visited[yy, xx] = True
stack.append((yy, xx))
# Prefer the highest probability, then topmost/leftmost on exact ties.
y, x = min(cells, key=lambda cell: (-grid[cell], cell[0], cell[1]))
point = [(x + 0.5) / width, (y + 0.5) / height]
points.append(point)
components.append({"cells": [list(cell) for cell in cells], "peak_cell": [y, x],
"peak_probability": float(grid[y, x]),
"mass": sum(float(grid[cell]) for cell in cells)})
return {"points": points, "components": components, "threshold": threshold,
"relative_threshold": relative_threshold, "connectivity": connectivity}
|