"""OCR top-label과 Tray 슬롯 계약을 lattice candidate score로 변환한다.""" from __future__ import annotations from dataclasses import dataclass import math from typing import Any, Sequence import numpy as np from .cross_visual import CrossVisualModel, cross_pair_feature_rows from .expression_contracts import horizontal_sides, segments_complete from .equality_visual import EqualityVisualModel, equality_pair_feature_rows from .math_tray import infer_math_trays from .segmentation_lattice import LATTICE_FEATURE_NAMES from .structure_relations import STRUCTURAL_RELATION_TYPES, infer_spatial_relations ANCHOR_TYPES = { r"\sqrt": "ROOT", r"\sqrt{}": "ROOT", r"\int": "INTEGRAL", r"\oint": "INTEGRAL", r"\iint": "INTEGRAL", r"\iiint": "INTEGRAL", r"\sum": "LARGE_OPERATOR", r"\prod": "LARGE_OPERATOR", r"\lim": "LARGE_OPERATOR", } RELIABLE_MULTISTROKE = {"i", r"\pi", r"\rightarrow", r"\pm", r"\leq", r"\geq"} MULTISTROKE_FAMILY_LABELS = frozenset({ r"\sum", r"\Sigma", r"\pi", r"\Pi", r"\rightarrow", r"\shortrightarrow", r"\longrightarrow", r"\Rightarrow", r"\neq", r"\not\equiv", r"\pm", }) MULTISTROKE_FAMILY_NAMES = frozenset({ r"\pi", r"\rightarrow", r"\shortrightarrow", r"\longrightarrow", r"\Rightarrow", r"\neq", r"\not\equiv", r"\pm", }) def _multistroke_family_geometry_guard( label: str, family: str, feature_row: np.ndarray, ) -> bool: """필요 변수: 합친 OCR label/family·lattice geometry. 작동 원리: 완성 단일기호 열을 먹는 넓은 alias 후보를 차단한다.""" width_ref = float(feature_row[2]) height_ref = float(feature_row[3]) pair_gap_max = float(feature_row[12]) if label == r"\Sigma" and width_ref < 2.0: return False if (label == r"\pm" or family == r"\pm") and height_ref > 2.0: return False arrow_family = { r"\rightarrow", r"\shortrightarrow", r"\longrightarrow", r"\Rightarrow", } if (label in arrow_family or family in arrow_family) and pair_gap_max > 0.50: return False return True @dataclass(frozen=True, slots=True) class TrayJointWeights: """필요 변수: validation 선택 가중치. 작동 원리: 구조 보너스와 두 must-not-link 감점을 한 계약으로 고정한다.""" tray: float = 4.0 symbol: float = 4.0 fraction: float = 8.0 infix: float = 8.0 competition: float = 0.0 local_baseline: float = 0.0 group_bias: float = -2.0 def component_competition_penalties( candidates: Sequence[dict[str, Any]], features: np.ndarray, *, mode: str = "joint", ) -> np.ndarray: """필요 변수: lattice 후보·OCR merge gain. 작동 원리: 합친 OCR이 구성획보다 약한 다획 후보에 must-not-link 값을 준다.""" if len(candidates) != len(features): raise ValueError("component competition candidate/feature 수가 다릅니다.") if mode not in {"top1", "joint"}: raise ValueError(f"지원하지 않는 component competition mode입니다: {mode}") merge_index = len(LATTICE_FEATURE_NAMES) + 6 entropy_index = len(LATTICE_FEATURE_NAMES) + 7 penalties = np.zeros(len(candidates), dtype=np.float32) for index, candidate in enumerate(candidates): if len(candidate["source_indices"]) <= 1: continue merge_gain = float(features[index][merge_index]) deficit = -merge_gain if mode == "joint": deficit -= 0.5 * float(features[index][entropy_index]) penalties[index] = min(1.0, max(0.0, deficit)) return penalties def _group_box(group: frozenset[int], strokes: Sequence[dict[str, Any]]) -> dict[str, float]: """필요 변수: stroke group·좌표. 작동 원리: candidate와 context symbol의 bbox를 계산한다.""" points = [point for index in group for point in strokes[index]["points"]] xs, ys = [float(point[0]) for point in points], [float(point[1]) for point in points] return {"left": min(xs), "top": min(ys), "right": max(xs), "bottom": max(ys)} def _recognized_segment(box: dict[str, float], label: str, score: float, family: str = "") -> dict[str, Any]: """필요 변수: 후보 bbox·top label/family·확률. 작동 원리: Tray와 부분식 parser용 최소 OCR segment를 만든다.""" return {"box": box, "results": {"expanded": {"candidates": [{ "label": label, "score": score, "visual_family": family, }]}}} def local_baseline_boundary_penalties( strokes: Sequence[dict[str, Any]], candidates: Sequence[dict[str, Any]], ocr_labels: Sequence[str], features: np.ndarray, base_partition: Sequence[frozenset[int]], *, ocr_families: Sequence[str] | None = None, ) -> np.ndarray: """필요 변수: 초기 partition·OCR segment·공간 관계. 작동 원리: 첨자·분수 등 구조 경계 양쪽을 동시에 먹는 병합만 감점한다.""" count = len(candidates) if len(ocr_labels) != count or len(features) != count: raise ValueError("local baseline candidate/label/feature 수가 다릅니다.") families = list(ocr_families or [""] * count) if len(families) != count: raise ValueError("local baseline candidate/family 수가 다릅니다.") index_by_group = { frozenset(int(value) for value in candidate["source_indices"]): index for index, candidate in enumerate(candidates) } partition_groups = [frozenset(group) for group in base_partition if frozenset(group) in index_by_group] top1_index = len(LATTICE_FEATURE_NAMES) segments = [] for group in partition_groups: index = index_by_group[group] segments.append(_recognized_segment( _group_box(group, strokes), ocr_labels[index], float(features[index][top1_index]), families[index], )) trays = infer_math_trays(segments) relations = infer_spatial_relations(segments, trays=trays) structural_edges = [ ( partition_groups[int(relation["parent"])], partition_groups[int(relation["child"])], float(relation["confidence"]), ) for relation in relations if str(relation["type"]) in STRUCTURAL_RELATION_TYPES ] penalties = np.zeros(count, dtype=np.float32) for index, candidate in enumerate(candidates): group = frozenset(int(value) for value in candidate["source_indices"]) if len(group) <= 1: continue penalties[index] = max( ( confidence for parent, child, confidence in structural_edges if group & parent and group & child ), default=0.0, ) return penalties def _candidate_box( candidate: dict[str, Any], group: frozenset[int], strokes: Sequence[dict[str, Any]], ) -> dict[str, float]: """필요 변수: lattice 후보·원 stroke. 작동 원리: cache bbox를 우선 재사용하고 최소 입력은 계산으로 호환한다.""" box = candidate.get("box") return dict(box) if isinstance(box, dict) else _group_box(group, strokes) def _point_xy(point: Any) -> tuple[float, float]: """필요 변수: 배열 또는 객체형 point. 작동 원리: 획 형태 계산용 좌표를 공통 형식으로 읽는다.""" if isinstance(point, dict): return float(point.get("x", 0.0)), float(point.get("y", 0.0)) return float(point[0]), float(point[1]) def _stroke_direction(stroke: dict[str, Any]) -> tuple[float, float]: """필요 변수: 두 점 이상인 stroke. 작동 원리: 첫·끝점 벡터를 방향 정규화해 평행·교차 형태를 판정한다.""" points = stroke.get("points") or [] if len(points) < 2: return 0.0, 0.0 start, end = _point_xy(points[0]), _point_xy(points[-1]) dx, dy = end[0] - start[0], end[1] - start[1] length = math.hypot(dx, dy) return (dx / length, dy / length) if length > 1e-6 else (0.0, 0.0) def _completed_operand_context_score( group: frozenset[int], group_box: dict[str, float], context: Sequence[tuple[frozenset[int], dict[str, Any]]], ) -> float: """필요 변수: 중위 후보 bbox·초기 partition. 작동 원리: 좌변·우변이 모두 완료된 표현일 때만 구조 점수를 연다.""" usable = [(other_group, segment) for other_group, segment in context if not other_group & group] left, right = horizontal_sides(group_box, usable) return 1.0 if segments_complete(left) and segments_complete(right) else 0.0 def _occupied_operand_context_score( group: frozenset[int], group_box: dict[str, float], context: Sequence[tuple[frozenset[int], dict[str, Any]]], ) -> float: """필요 변수: 교차선 후보·현재 partition. 작동 원리: 곱셈은 OCR 문법을 강제하지 않고 같은 행의 양쪽 점유만 확인한다.""" usable = [(other_group, segment) for other_group, segment in context if not other_group & group] left, right = horizontal_sides(group_box, usable) return 1.0 if left and right else 0.0 def _infix_shape_score( group: frozenset[int], strokes: Sequence[dict[str, Any]], stroke_boxes: Sequence[dict[str, float]], ) -> tuple[str, float]: """필요 변수: 두 획 후보·사전 계산 bbox. 작동 원리: 평행선 equality와 교차선 multiply family를 분리해 반환한다.""" if len(group) != 2: return "", 0.0 first_index, second_index = sorted(group) first, second = strokes[first_index], strokes[second_index] first_direction, second_direction = _stroke_direction(first), _stroke_direction(second) horizontal = min(abs(first_direction[0]), abs(second_direction[0])) vertical_leak = max(abs(first_direction[1]), abs(second_direction[1])) first_box, second_box = stroke_boxes[first_index], stroke_boxes[second_index] first_width = max(first_box["right"] - first_box["left"], 1e-6) second_width = max(second_box["right"] - second_box["left"], 1e-6) x_overlap = max(0.0, min(first_box["right"], second_box["right"]) - max(first_box["left"], second_box["left"])) overlap_ratio = x_overlap / min(first_width, second_width) length_ratio = min(first_width, second_width) / max(first_width, second_width) first_y = (first_box["top"] + first_box["bottom"]) * 0.5 second_y = (second_box["top"] + second_box["bottom"]) * 0.5 separation_ratio = abs(first_y - second_y) / max((first_width + second_width) * 0.5, 1e-6) parallel = abs(first_direction[0] * second_direction[0] + first_direction[1] * second_direction[1]) if ( horizontal >= 0.92 and vertical_leak <= 0.38 and parallel >= 0.94 and overlap_ratio >= 0.60 and length_ratio >= 0.60 and 0.04 <= separation_ratio <= 0.65 ): return "equality", min(1.0, 0.35 + 0.25 * overlap_ratio + 0.20 * length_ratio + 0.20 * parallel) diagonal = min(abs(first_direction[0] * first_direction[1]), abs(second_direction[0] * second_direction[1])) * 2.0 opposite_slopes = first_direction[0] * first_direction[1] * second_direction[0] * second_direction[1] < 0.0 x_intersects = max(first_box["left"], second_box["left"]) <= min(first_box["right"], second_box["right"]) y_intersects = max(first_box["top"], second_box["top"]) <= min(first_box["bottom"], second_box["bottom"]) if opposite_slopes and x_intersects and y_intersects and diagonal >= 0.55: return "multiply", min(1.0, diagonal) return "", 0.0 def candidate_signals( strokes: Sequence[dict[str, Any]], candidates: Sequence[dict[str, Any]], ocr_labels: Sequence[str], features: np.ndarray, base_partition: Sequence[frozenset[int]], *, ocr_families: Sequence[str] | None = None, strict_equality: bool = False, equality_model: EqualityVisualModel | None = None, cross_model: CrossVisualModel | None = None, cross_gap_ratio: float = 0.40, multistroke_family_boost: float = 6.0, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """필요 변수: 원 stroke·후보/label/feature·cross/family 설정. 작동 원리: Tray·신뢰 label·중위식 evidence를 계산한다.""" count = len(candidates) if len(ocr_labels) != count or len(features) != count: raise ValueError("Tray joint candidate/label/feature 수가 다릅니다.") if cross_gap_ratio < 0.0: raise ValueError("cross gap ratio는 0 이상이어야 합니다.") if multistroke_family_boost < 0.0: raise ValueError("multistroke family boost는 0 이상이어야 합니다.") families = list(ocr_families or [""] * count) if len(families) != count: raise ValueError("Tray joint candidate/family 수가 다릅니다.") tray_signal = np.zeros(count, dtype=np.float32) symbol_signal = np.zeros(count, dtype=np.float32) infix_signal = np.zeros(count, dtype=np.float32) top1 = features[:, len(LATTICE_FEATURE_NAMES)] stroke_boxes = [_group_box(frozenset({index}), strokes) for index in range(len(strokes))] index_by_group = { frozenset(int(value) for value in candidate["source_indices"]): index for index, candidate in enumerate(candidates) } pair_groups = [ frozenset(int(value) for value in candidate["source_indices"]) for candidate in candidates if len(candidate["source_indices"]) == 2 ] equality_probability: dict[frozenset[int], float] = {} cross_probability: dict[frozenset[int], float] = {} if equality_model is not None and pair_groups: ordered_pairs = [tuple(sorted(group)) for group in pair_groups] pair_rows = equality_pair_feature_rows(ordered_pairs, strokes) equality_probability = { group: equality_model.probability(row) for group, row in zip(pair_groups, pair_rows, strict=True) } if cross_model is not None and pair_groups: nonzero_heights = [box["bottom"] - box["top"] for box in stroke_boxes if box["bottom"] - box["top"] > 1e-5] reference = float(np.median(nonzero_heights)) if nonzero_heights else 1.0 plausible_groups = [] for group in pair_groups: first_index, second_index = sorted(group) first_box, second_box = stroke_boxes[first_index], stroke_boxes[second_index] x_gap = max(first_box["left"] - second_box["right"], second_box["left"] - first_box["right"], 0.0) y_gap = max(first_box["top"] - second_box["bottom"], second_box["top"] - first_box["bottom"], 0.0) if x_gap <= reference * cross_gap_ratio and y_gap <= reference * cross_gap_ratio: plausible_groups.append(group) ordered_pairs = [tuple(sorted(group)) for group in plausible_groups] pair_rows = cross_pair_feature_rows(ordered_pairs, strokes) cross_probability = { group: cross_model.probability(row) for group, row in zip(plausible_groups, pair_rows, strict=True) } context = [ ( group, _recognized_segment( _candidate_box(candidates[index_by_group[group]], group, strokes), ocr_labels[index_by_group[group]], float(top1[index_by_group[group]]), families[index_by_group[group]], ), ) for group in base_partition ] for index, candidate in enumerate(candidates): group = frozenset(int(value) for value in candidate["source_indices"]) label = ocr_labels[index] group_box = _candidate_box(candidate, group, strokes) shape_family, shape_score = _infix_shape_score(group, strokes, stroke_boxes) learned_equality = equality_probability.get(group, 0.0) learned_cross = cross_probability.get(group, 0.0) if equality_model is not None and learned_equality >= equality_model.threshold: shape_family, shape_score = "equality", learned_equality if cross_model is not None and learned_cross >= cross_model.threshold and learned_cross > learned_equality: # threshold 직상단의 불확실한 pair가 큰 병합 보너스를 받지 않도록 초과 margin만 사용한다. margin = (learned_cross - cross_model.threshold) / max(1.0 - cross_model.threshold, 1e-6) shape_family, shape_score = "cross_visual", min(1.0, max(0.0, margin)) if shape_family == "equality": # 등호 인식은 수식 완성 여부와 독립적이다. parser는 shadow 분석에서만 gating한다. context_score = _completed_operand_context_score(group, group_box, context) if strict_equality else 1.0 elif shape_family == "cross_visual": context_score = 1.0 else: context_score = _occupied_operand_context_score(group, group_box, context) infix_signal[index] = shape_score * context_score if len(group) > 1: if label in RELIABLE_MULTISTROKE: symbol_signal[index] = float(top1[index]) if ( label in MULTISTROKE_FAMILY_LABELS or families[index] in MULTISTROKE_FAMILY_NAMES ) and _multistroke_family_geometry_guard( label, families[index], features[index], ): # 합친 OCR family가 이미 살아 있을 때만 추가한다. 인접 획이라는 이유만으로 병합하지 않는다. symbol_signal[index] += float(top1[index]) * multistroke_family_boost tray_type = ANCHOR_TYPES.get(label) if tray_type is None: continue anchor = _recognized_segment(_group_box(group, strokes), label, float(top1[index]), families[index]) segments = [anchor, *(segment for other_group, segment in context if not other_group & group)] matches = ( tray for tray in infer_math_trays(segments) if tray["anchor"] == 0 and tray["type"] == tray_type and tray["required_satisfied"] ) tray_signal[index] = max((float(tray["constraint_score"]) for tray in matches), default=0.0) return tray_signal, symbol_signal, infix_signal def adjusted_logits( base_logits: np.ndarray, tray_signal: np.ndarray, symbol_signal: np.ndarray, fraction_penalty: np.ndarray, weights: TrayJointWeights, infix_signal: np.ndarray | None = None, competition_penalty: np.ndarray | None = None, local_baseline_penalty: np.ndarray | None = None, ) -> np.ndarray: """필요 변수: base logit·구조 signal·세 penalty. 작동 원리: joint selector 목적함수의 candidate logit을 만든다.""" effective_infix = np.zeros_like(base_logits) if infix_signal is None else infix_signal effective_competition = ( np.zeros_like(base_logits) if competition_penalty is None else competition_penalty ) effective_local_baseline = ( np.zeros_like(base_logits) if local_baseline_penalty is None else local_baseline_penalty ) arrays = ( base_logits, tray_signal, symbol_signal, fraction_penalty, effective_infix, effective_competition, effective_local_baseline, ) if len({len(values) for values in arrays}) != 1: raise ValueError("Tray joint logit과 signal 길이가 다릅니다.") return ( base_logits + weights.tray * tray_signal + weights.symbol * symbol_signal + weights.infix * effective_infix - weights.fraction * fraction_penalty - weights.competition * effective_competition - weights.local_baseline * effective_local_baseline )