Spaces:
Running on Zero
Running on Zero
| """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} | |