"""기호 인식과 독립적으로 수식의 2D 배치 관계를 추론하는 연구용 레이어다.""" from __future__ import annotations from dataclasses import dataclass from typing import Any, Sequence import numpy as np from .math_tray import infer_math_trays RELATION_TYPES = ("RIGHT", "SUPERSCRIPT", "SUBSCRIPT", "ABOVE", "BELOW", "CONTAINS") STRUCTURAL_RELATION_TYPES = frozenset({"SUPERSCRIPT", "SUBSCRIPT", "ABOVE", "BELOW", "CONTAINS"}) @dataclass(frozen=True, slots=True) class RelationConfig: """필요 변수: 첨자 기하 임계값. 작동 원리: writer validation에서 고를 수 있도록 관계 판정값을 외부화한다.""" script_min_height_ratio: float = 0.18 script_max_height_ratio: float = 0.82 script_vertical_shift_ratio: float = 0.18 script_max_horizontal_gap_ratio: float = 0.75 def _label(segment: dict[str, Any]) -> str: """필요 변수: 인식된 segment. 작동 원리: 구조 힌트를 우선하고 첫 전문가의 top-1을 보조 라벨로 읽는다.""" hint = segment.get("structural_hint") if hint: return str(hint) for result in (segment.get("results") or {}).values(): candidates = result.get("candidates") or [] if candidates: return str(candidates[0]["label"]) return "" def _geometry(segment: dict[str, Any]) -> dict[str, float]: """필요 변수: segment bbox. 작동 원리: 관계 판정에 필요한 폭·높이·중심 좌표를 한 번 계산한다.""" box = segment["box"] left, top = float(box["left"]), float(box["top"]) right, bottom = float(box["right"]), float(box["bottom"]) return { "left": left, "top": top, "right": right, "bottom": bottom, "width": max(right - left, 1.0), "height": max(bottom - top, 1.0), "cx": (left + right) * 0.5, "cy": (top + bottom) * 0.5, } def infer_spatial_relations( segments: Sequence[dict[str, Any]], config: RelationConfig | None = None, *, trays: Sequence[dict[str, Any]] | None = None, ) -> list[dict[str, Any]]: """필요 변수: geometry group과 선택적 OCR 결과. 작동 원리: 각 기호를 한 번만 자식으로 소비해 2D 관계 그래프를 만든다.""" relation_config = config or RelationConfig() if len(segments) < 2: return [] boxes = [_geometry(segment) for segment in segments] labels = [_label(segment) for segment in segments] maximum_height = max(box["height"] for box in boxes) body_heights = [box["height"] for box in boxes if box["height"] >= maximum_height * 0.25] reference_height = float(np.median(body_heights)) relations: list[dict[str, Any]] = [] assigned: set[int] = set() structural_parents: set[int] = set() # 근호는 내부 기호를 symbol group에 합치지 않고 CONTAINS edge로만 소유한다. for parent, (box, label) in enumerate(zip(boxes, labels, strict=True)): root_like = label in {r"\sqrt", r"\sqrt{}"} if not root_like: continue structural_parents.add(parent) for child, inner in enumerate(boxes): if child == parent or child in assigned: continue inside_x = box["left"] + box["width"] * 0.20 <= inner["cx"] <= box["right"] + reference_height * 0.15 inside_y = inner["cy"] >= box["top"] and inner["top"] <= box["bottom"] + reference_height * 0.35 if inside_x and inside_y: relations.append({"parent": parent, "child": child, "type": "CONTAINS", "confidence": 0.95}) assigned.add(child) # 분수는 bar 단독이 아니라 분자·분모가 모두 점유된 원자적 Tray일 때만 관계를 만든다. tray_rows = list(trays) if trays is not None else infer_math_trays(segments) fraction_trays = [tray for tray in tray_rows if tray["type"] == "FRACTION"] for tray in fraction_trays: parent = int(tray["anchor"]) structural_parents.add(parent) for slot, relation in (("numerator", "ABOVE"), ("denominator", "BELOW")): for child in tray["slots"].get(slot, []): if child in assigned: continue relations.append({ "parent": parent, "child": child, "type": relation, "confidence": float(tray["constraint_score"]), }) assigned.add(child) # 적분·합·곱·극한의 상하한은 일반 최근접 첨자보다 수학 Tray의 anchor 소유권을 우선한다. operator_trays = [tray for tray in tray_rows if tray["type"] in {"INTEGRAL", "SUM", "PRODUCT", "LIM"}] for tray in operator_trays: parent = int(tray["anchor"]) structural_parents.add(parent) for slot, relation in (("lower", "SUBSCRIPT"), ("upper", "SUPERSCRIPT")): for child in tray["slots"].get(slot, []): if child in assigned: continue relations.append({ "parent": parent, "child": int(child), "type": relation, "confidence": float(tray["constraint_score"]), }) assigned.add(int(child)) # 남은 기호는 가장 가까운 왼쪽 기호를 기준으로 첨자 또는 RIGHT 관계를 갖는다. for child, box in enumerate(boxes): if child in assigned or child in structural_parents: continue candidates = [index for index, other in enumerate(boxes) if index != child and other["cx"] < box["cx"]] if not candidates: continue parent = min(candidates, key=lambda index: max(0.0, box["left"] - boxes[index]["right"])) base = boxes[parent] horizontal_gap = max(0.0, box["left"] - base["right"]) script_sized = ( reference_height * relation_config.script_min_height_ratio <= box["height"] <= base["height"] * relation_config.script_max_height_ratio ) script_near = horizontal_gap <= reference_height * relation_config.script_max_horizontal_gap_ratio vertical_shift = base["height"] * relation_config.script_vertical_shift_ratio if script_sized and script_near and box["cy"] < base["cy"] - vertical_shift: relation = "SUPERSCRIPT" confidence = 0.82 elif script_sized and script_near and box["cy"] > base["cy"] + vertical_shift: relation = "SUBSCRIPT" confidence = 0.82 else: relation = "RIGHT" confidence = 0.75 relations.append({"parent": parent, "child": child, "type": relation, "confidence": confidence}) assigned.add(child) return sorted(relations, key=lambda item: (item["child"], item["parent"], item["type"])) def serialize_relation_graph( labels: Sequence[str], segments: Sequence[dict[str, Any]], relations: Sequence[dict[str, Any]], ) -> str: """필요 변수: group별 라벨·bbox·관계 edge. 작동 원리: 구조 자식을 한 번만 소비해 2D graph를 LaTeX로 직렬화한다.""" if len(labels) != len(segments): raise ValueError("관계 serializer의 label과 segment 수가 다릅니다.") children: dict[int, dict[str, list[int]]] = {} structural_children: set[int] = set() for relation in relations: relation_type = str(relation["type"]) if relation_type not in STRUCTURAL_RELATION_TYPES: continue parent, child = int(relation["parent"]), int(relation["child"]) if parent < 0 or child < 0 or parent >= len(labels) or child >= len(labels) or parent == child: raise ValueError("관계 graph에 유효하지 않은 node index가 있습니다.") children.setdefault(parent, {}).setdefault(relation_type, []).append(child) structural_children.add(child) def left(index: int) -> float: """필요 변수: node index. 작동 원리: 같은 구조 슬롯 안의 기호를 원래 x 순서로 정렬한다.""" return float(segments[index]["box"]["left"]) def render_sequence(indices: Sequence[int], active: frozenset[int]) -> str: """필요 변수: node index 열·재귀 경로. 작동 원리: 같은 슬롯의 node를 좌→우로 한 번씩 출력한다.""" return "".join(render_node(index, active) for index in sorted(set(indices), key=left)) def render_node(index: int, active: frozenset[int]) -> str: """필요 변수: 현재 node·재귀 경로. 작동 원리: root/fraction/script 구조를 중첩 가능한 LaTeX node로 만든다.""" if index in active: raise ValueError("관계 graph에 cycle이 있습니다.") next_active = active | {index} slots = children.get(index, {}) above = slots.get("ABOVE", []) below = slots.get("BELOW", []) contained = slots.get("CONTAINS", []) if above and below: base = rf"\frac{{{render_sequence(above, next_active)}}}{{{render_sequence(below, next_active)}}}" elif contained and labels[index] in {r"\sqrt", r"\sqrt{}"}: base = rf"\sqrt{{{render_sequence(contained, next_active)}}}" else: base = labels[index] superscript = slots.get("SUPERSCRIPT", []) subscript = slots.get("SUBSCRIPT", []) if subscript: base += rf"_{{{render_sequence(subscript, next_active)}}}" if superscript: base += rf"^{{{render_sequence(superscript, next_active)}}}" return base roots = [index for index in range(len(labels)) if index not in structural_children] return render_sequence(roots, frozenset())