"""实体分布热力图:把实体 AABB 聚合为「占用网格 + 簇」,供自适应瓦片规划使用。 纯函数模块(仅依赖 `math`);可视化 `render_overlay` 内部按需 import Pillow, 使本模块在无 Pillow 的环境下仍可被单测与主流程引用。 坐标约定(与 CAD 一致): - 世界坐标 x 向右、y **向上**;bbox = (min_x, min_y, max_x, max_y)。 - 网格 row 0 在**最上方**、col 0 在**最左**(与截图 / 图片一致),便于直接对照画面。 """ import math # 退化 bbox(线/点)的膨胀比例(相对格尺寸),保证零宽/零高的实体也能被计入 _EPS_FRAC = 1e-3 def norm_box(b): """把任意 bbox 归一为 (minx,miny,maxx,maxy) 的 float 元组;非法返回 None。""" if not isinstance(b, (list, tuple)) or len(b) != 4: return None try: x1, y1, x2, y2 = (float(v) for v in b) except (TypeError, ValueError): return None if not all(math.isfinite(v) for v in (x1, y1, x2, y2)): return None if x1 > x2: x1, x2 = x2, x1 if y1 > y2: y1, y2 = y2, y1 return (x1, y1, x2, y2) def entity_boxes(envelopes): """从若干查询信封收集合法实体 bbox。 envelopes 可以是单个信封 dict,也可以是信封的可迭代集合;每个信封形如 `{"items": [{"bbox": [..]}, ...], "total": N}`(其它字段忽略)。 返回 [(minx,miny,maxx,maxy), ...];缺字段 / 非法的条目直接跳过。 """ if isinstance(envelopes, dict): envelopes = [envelopes] out = [] for env in envelopes or []: if not isinstance(env, dict): continue items = env.get("items") if not isinstance(items, list): continue for it in items: if not isinstance(it, dict): continue box = norm_box(it.get("bbox")) if box is not None: out.append(box) return out def box_area(b): """bbox 面积。""" return max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1]) def union_box(boxes): """若干 bbox 的并集包围盒;空集合返回 None。""" boxes = [b for b in boxes if b] if not boxes: return None return ( min(b[0] for b in boxes), min(b[1] for b in boxes), max(b[2] for b in boxes), max(b[3] for b in boxes), ) def drop_oversized(boxes, window, max_frac=0.5): """丢弃 bbox 面积超过图框面积 `max_frac` 的实体。 这类实体多是**图框边框 / 图签块**等"结构"(实测每个图框恰有 1 个 bbox 等于整幅图框), 其 bbox 会铺满整张热力图、掩盖真正的细节分布;它们无细节可读,整体图已能覆盖。 """ area = (float(window["max_x"]) - float(window["min_x"])) * ( float(window["max_y"]) - float(window["min_y"]) ) if area <= 0: return list(boxes) limit = area * float(max_frac) return [b for b in boxes if box_area(b) <= limit] def div_to_cell(window, div): """按「长边切 div 份」求格子世界尺寸。""" w = float(window["max_x"]) - float(window["min_x"]) h = float(window["max_y"]) - float(window["min_y"]) div = max(1, int(div)) return max(w, h) / div def build_grid(window, boxes, cell): """把 window 按 cell 尺寸切成网格,统计每格命中的实体数与覆盖比例。 返回 dict: {"window", "cols", "rows", "cell_w", "cell_h", "counts": [[int]], "coverage": [[float]]} - counts[r][c] :与格 (r,c) 相交的实体数(同一实体跨多格则各计一次); - coverage[r][c]:该格被实体 bbox 覆盖的面积比例(0~1,多实体叠加后夹到 1)。 """ min_x, min_y = float(window["min_x"]), float(window["min_y"]) max_x, max_y = float(window["max_x"]), float(window["max_y"]) cell = float(cell) if cell <= 0 or max_x <= min_x or max_y <= min_y: return { "window": window, "cols": 0, "rows": 0, "cell_w": 0.0, "cell_h": 0.0, "counts": [], "coverage": [], } cols = max(1, int(math.ceil((max_x - min_x) / cell))) rows = max(1, int(math.ceil((max_y - min_y) / cell))) cw = (max_x - min_x) / cols ch = (max_y - min_y) / rows counts = [[0] * cols for _ in range(rows)] coverage = [[0.0] * cols for _ in range(rows)] for b in boxes: bx1, by1, bx2, by2 = b # 退化 bbox(水平/垂直线、点)宽或高为 0,按格尺寸的千分之一膨胀,保证能被计入 if bx2 - bx1 <= 0: bx1 -= cw * _EPS_FRAC * 0.5 bx2 += cw * _EPS_FRAC * 0.5 if by2 - by1 <= 0: by1 -= ch * _EPS_FRAC * 0.5 by2 += ch * _EPS_FRAC * 0.5 # 与 window 求交,完全在外的直接跳过 bx1, by1 = max(bx1, min_x), max(by1, min_y) bx2, by2 = min(bx2, max_x), min(by2, max_y) if bx1 > bx2 or by1 > by2: continue c0 = min(max(int((bx1 - min_x) / cw), 0), cols - 1) c1 = min(max(int((bx2 - min_x) / cw), 0), cols - 1) # y 向上 → row 0 在最上方 r0 = min(max(int((max_y - by2) / ch), 0), rows - 1) r1 = min(max(int((max_y - by1) / ch), 0), rows - 1) for r in range(r0, r1 + 1): cy_top = max_y - r * ch cy_bot = cy_top - ch for c in range(c0, c1 + 1): cx1 = min_x + c * cw cx2 = cx1 + cw ix = max(0.0, min(bx2, cx2) - max(bx1, cx1)) iy = max(0.0, min(by2, cy_top) - max(by1, cy_bot)) # 用真实交面积判定命中,避免恰好落在格边界上的实体被计入相邻空格 if ix <= 0 or iy <= 0: continue counts[r][c] += 1 coverage[r][c] = min(1.0, coverage[r][c] + (ix * iy) / (cw * ch)) return { "window": window, "cols": cols, "rows": rows, "cell_w": cw, "cell_h": ch, "counts": counts, "coverage": coverage, } def occupied_cells(grid, min_count=1): """占用格坐标列表 [(row, col), ...](counts >= min_count),行优先。""" counts = grid.get("counts") or [] out = [] for r, row in enumerate(counts): for c, n in enumerate(row): if n >= min_count: out.append((r, c)) return out def _dilate(mask, radius): """把布尔占用掩码向外膨胀 radius 格(半径 0 原样返回)。""" if radius <= 0: return mask rows, cols = len(mask), len(mask[0]) if mask else 0 out = [[False] * cols for _ in range(rows)] for r in range(rows): for c in range(cols): if not mask[r][c]: continue for dr in range(-radius, radius + 1): for dc in range(-radius, radius + 1): rr, cc = r + dr, c + dc if 0 <= rr < rows and 0 <= cc < cols: out[rr][cc] = True return out def cluster_cells(grid, min_count=1, dilate=0): """把占用格按 4-连通聚成若干簇(可选先膨胀 dilate 格,合并相近内容)。 返回按「实体计数降序」排列的簇列表,每项: {"bbox": (minx,miny,maxx,maxy), "cells": [(r,c),...], "count": int} bbox 为该簇**实际占用格**的并集(不含膨胀出来的空格),故不会虚增范围。 """ rows, cols = grid.get("rows", 0), grid.get("cols", 0) counts = grid.get("counts") or [] if not rows or not cols or not counts: return [] raw = [[counts[r][c] >= min_count for c in range(cols)] for r in range(rows)] mask = _dilate(raw, int(dilate)) seen = [[False] * cols for _ in range(rows)] clusters = [] for r in range(rows): for c in range(cols): if not mask[r][c] or seen[r][c]: continue stack = [(r, c)] seen[r][c] = True cells = [] while stack: cr, cc = stack.pop() if raw[cr][cc]: cells.append((cr, cc)) for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)): nr, nc = cr + dr, cc + dc if 0 <= nr < rows and 0 <= nc < cols and mask[nr][nc] and not seen[nr][nc]: seen[nr][nc] = True stack.append((nr, nc)) if not cells: continue # 膨胀出来的孤立空格,无实体,丢弃 boxes = [_cell_box(grid, cr, cc) for cr, cc in cells] clusters.append({ "bbox": union_box(boxes), "cells": sorted(cells), "count": sum(counts[cr][cc] for cr, cc in cells), }) clusters.sort(key=lambda cl: (-cl["count"], -cl["bbox"][3], cl["bbox"][0])) return clusters def _cell_box(grid, r, c): """格 (r,c) 的世界 bbox(row 0 在上)。""" win = grid["window"] min_x, max_y = float(win["min_x"]), float(win["max_y"]) cw, ch = grid["cell_w"], grid["cell_h"] x1 = min_x + c * cw y_top = max_y - r * ch return (x1, y_top - ch, x1 + cw, y_top) def grid_to_windows(clusters, margin=0.0): """把簇 bbox 转成 zoom 窗口列表(保持输入顺序,即实体计数降序)。""" out = [] for cl in clusters: x1, y1, x2, y2 = cl["bbox"] w, h = x2 - x1, y2 - y1 mx, my = w * margin, h * margin out.append({ "min_x": x1 - mx, "min_y": y1 - my, "max_x": x2 + mx, "max_y": y2 + my, }) return out def grid_stats(grid, clusters=None): """统计摘要(供日志 / 目视对照):占用率、实体数分布、簇数与最大簇。""" rows, cols = grid.get("rows", 0), grid.get("cols", 0) total = rows * cols occ = occupied_cells(grid) counts = grid.get("counts") or [] hits = [n for row in counts for n in row if n > 0] stats = { "rows": rows, "cols": cols, "occupied": len(occ), "occupied_ratio": (len(occ) / total) if total else 0.0, "entity_cell_hits": len(hits), "max_count": max(hits) if hits else 0, } if clusters is not None: stats["clusters"] = len(clusters) stats["max_cluster_cells"] = max((len(cl["cells"]) for cl in clusters), default=0) stats["max_cluster_count"] = max((cl["count"] for cl in clusters), default=0) return stats def render_overlay(path, window, boxes, grid, clusters=None, *, size=1200, bg=(18, 18, 18)): """把热力图渲染成 PNG 供人工目视:窗口边框 + 实体 bbox 轮廓 + 热力着色 + 簇框。 - 长边缩放到 size 像素,保持世界坐标长宽比; - 热力色:占用越密越亮(绿→黄→红); - 簇框用青色粗线标出。 需要 Pillow;缺失时抛 RuntimeError(调用方自行决定是否忽略)。 """ try: from PIL import Image, ImageDraw except ImportError as exc: # pragma: no cover - 依赖缺失 raise RuntimeError("渲染热力图需要 Pillow") from exc min_x, min_y = float(window["min_x"]), float(window["min_y"]) max_x, max_y = float(window["max_x"]), float(window["max_y"]) w, h = max_x - min_x, max_y - min_y if w <= 0 or h <= 0: raise ValueError("window 尺寸非法") scale = size / max(w, h) iw, ih = max(1, int(round(w * scale))), max(1, int(round(h * scale))) img = Image.new("RGB", (iw, ih), bg) draw = ImageDraw.Draw(img) def to_px(x, y): return (int(round((x - min_x) * scale)), int(round((max_y - y) * scale))) # 热力格子 counts = grid.get("counts") or [] max_count = max((n for row in counts for n in row), default=0) or 1 for r, row in enumerate(counts): for c, n in enumerate(row): if n <= 0: continue x1, y1, x2, y2 = _cell_box(grid, r, c) p1, p2 = to_px(x1, y2), to_px(x2, y1) t = min(1.0, n / max_count) color = (int(60 + 195 * t), int(200 - 140 * t), int(60 - 40 * t)) draw.rectangle([p1, p2], fill=color) # 实体 bbox 轮廓(细线) for b in boxes: draw.rectangle([to_px(b[0], b[3]), to_px(b[2], b[1])], outline=(90, 90, 90)) # 簇框(青色粗线) for cl in clusters or []: b = cl["bbox"] draw.rectangle([to_px(b[0], b[3]), to_px(b[2], b[1])], outline=(0, 220, 220), width=2) # 窗口边框(白色) draw.rectangle([to_px(min_x, max_y), to_px(max_x, min_y)], outline=(255, 255, 255), width=1) img.save(str(path)) return path