20086551557

Committed 10 Dec 2025 03:43AM UTC coverage: 90.303% (+0.3%) from 90.041%

Build # 20086551557

Build Type

push

github

Committed by

bramp

Commit Message

Refactor arrow shaft detection: unified method, stroked line support, multi-head grouping

- Merge _find_simple_shaft, _find_stroked_line_shaft, and _find_cornered_shaft
  into a single unified _find_shaft method that handles all shaft types by
  extracting points and finding closest/furthest from the arrowhead tip
- Add support for stroked line shafts (stroke_color instead of fill_color)
- Add tail_grouping_tolerance config for grouping arrowheads with nearby tails
- Group arrowheads that share the same shaft_block (L-shaped arrows with
  multiple heads at different ends)
- Use union-find algorithm to group arrowheads by shared shaft or tail proximity
- Extract colors_match to shared utils module
- Add comprehensive tests for stroked line shafts, tail correctness, and
  multi-head arrow grouping
- Update golden files for pages 011, 013, 015, 017 with corrected arrow detection

Coverage Stats

204 of 206 new or added lines in 5 files covered. (99.03%)

252 existing lines in 14 files now uncovered.

11855 of 13128 relevant lines covered (90.3%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

83.59

/src/build_a_long/pdf_extract/classifier/steps/subassembly_classifier.py

"""
SubAssembly classifier.

Purpose
-------
Identify sub-assembly callout boxes on LEGO instruction pages. SubAssemblies
typically:
- Are white/light-colored rectangular boxes with black borders
- Are larger than individual parts (to contain a small build diagram)
- May contain a count label (e.g., "2x") indicating how many to build
- May contain step numbers for multi-step subassemblies
- May have an arrow pointing from them to the main diagram

Scoring is based on intrinsic properties of the box:
- Fill color (white/light)
- Size (larger than minimum threshold)

Child element discovery (step_count, step_numbers, diagrams, arrows) is
deferred to build time per DESIGN.md principles.

Debugging
---------
Set environment variables to aid investigation without code changes:

- LOG_LEVEL=DEBUG
    Enables DEBUG-level logging (if not already configured by caller).
"""

from __future__ import annotations

import logging
from typing import ClassVar

from build_a_long.pdf_extract.classifier.candidate import Candidate
from build_a_long.pdf_extract.classifier.classification_result import (
    CandidateFailedError,
    ClassificationResult,
)
from build_a_long.pdf_extract.classifier.config import SubAssemblyConfig
from build_a_long.pdf_extract.classifier.label_classifier import LabelClassifier
from build_a_long.pdf_extract.classifier.score import (
    Score,
    Weight,
    find_best_scoring,
)
from build_a_long.pdf_extract.classifier.text import extract_step_number_value
from build_a_long.pdf_extract.classifier.utils import score_white_fill
from build_a_long.pdf_extract.extractor.bbox import (
    BBox,
    filter_contained,
    filter_overlapping,
    group_by_similar_bbox,
)
from build_a_long.pdf_extract.extractor.lego_page_elements import (
    Diagram,
    StepCount,
    StepNumber,
    SubAssembly,
    SubAssemblyStep,
)
from build_a_long.pdf_extract.extractor.page_blocks import Blocks, Drawing, Image, Text

log = logging.getLogger(__name__)


class _SubAssemblyScore(Score):
    """Internal score representation for subassembly classification.

    Scores based on intrinsic box properties only. Child element discovery
    (step_count, step_numbers, diagrams, arrows) is deferred to build time.
    """

    box_score: float
    """Score based on box having white/light fill (0.0-1.0)."""

    has_step_count: bool
    """Whether a step_count candidate exists inside (for scoring bonus)."""

    has_diagram_or_images: bool
    """Whether diagram candidates or images exist inside (for scoring bonus)."""

    has_step_numbers: bool
    """Whether step_number candidates exist inside (for multi-step subassemblies)."""

    config: SubAssemblyConfig
    """Configuration containing weights for score calculation."""

    def score(self) -> Weight:
        """Calculate final weighted score from components."""
        count_score = 1.0 if self.has_step_count else 0.0
        diagram_score = 1.0 if self.has_diagram_or_images else 0.0

        return (
            self.box_score * self.config.box_shape_weight
            + count_score * self.config.count_weight
            + diagram_score * self.config.diagram_weight
        )


class SubAssemblyClassifier(LabelClassifier):
    """Classifier for subassembly callout boxes."""

    output: ClassVar[str] = "subassembly"
    requires: ClassVar[frozenset[str]] = frozenset(
        {"step_count", "step_number", "diagram"}
    )

    def _score(self, result: ClassificationResult) -> None:
        """Score Drawing blocks as potential subassembly boxes."""
        page_data = result.page_data
        subassembly_config = self.config.subassembly

        # Get step_count, step_number, and diagram candidates
        step_count_candidates = result.get_scored_candidates(
            "step_count", valid_only=False, exclude_failed=True
        )
        step_number_candidates = result.get_scored_candidates(
            "step_number", valid_only=False, exclude_failed=True
        )
        diagram_candidates = result.get_scored_candidates(
            "diagram", valid_only=False, exclude_failed=True
        )

        # Find rectangular drawing blocks that could be subassembly boxes
        # Filter by size constraints first
        max_width = page_data.bbox.width * subassembly_config.max_page_width_ratio
        max_height = page_data.bbox.height * subassembly_config.max_page_height_ratio

        valid_drawings: list[Drawing] = []
        for block in page_data.blocks:
            if not isinstance(block, Drawing):
                continue

            bbox = block.bbox

            # Skip boxes smaller than minimum subassembly size
            if (
                bbox.width < subassembly_config.min_subassembly_width
                or bbox.height < subassembly_config.min_subassembly_height
            ):
                continue

            # Skip boxes larger than maximum subassembly size
            if bbox.width > max_width or bbox.height > max_height:
                log.debug(
                    "[subassembly] Skipping oversized box at %s "
                    "(%.1f x %.1f > max %.1f x %.1f)",
                    bbox,
                    bbox.width,
                    bbox.height,
                    max_width,
                    max_height,
                )
                continue

            valid_drawings.append(block)

        # Group drawings with similar bboxes (e.g., white-filled box and
        # black-bordered box for the same subassembly)
        groups = group_by_similar_bbox(valid_drawings, tolerance=2.0)

        # Process each group - create one candidate per unique bbox region
        for group in groups:
            # Use union of all grouped drawings' bboxes
            bbox = BBox.union_all([d.bbox for d in group])

            # Score each drawing's colors and pick the best
            best_box_score = max(score_white_fill(d) for d in group)
            if best_box_score < 0.3:
                continue

            # Check for child elements inside the box (for scoring only)
            # Actual candidate discovery happens at build time
            has_step_count = bool(
                self._find_candidate_inside(bbox, step_count_candidates)
            )
            has_step_numbers = bool(
                self._find_all_candidates_inside(bbox, step_number_candidates)
            )
            diagrams_inside = self._find_all_diagrams_inside(bbox, diagram_candidates)
            images_inside = self._find_images_inside(bbox, page_data.blocks)
            has_diagram_or_images = bool(diagrams_inside or images_inside)

            # We need at least a box - count and diagram are optional
            score_details = _SubAssemblyScore(
                box_score=best_box_score,
                has_step_count=has_step_count,
                has_diagram_or_images=has_diagram_or_images,
                has_step_numbers=has_step_numbers,
                config=subassembly_config,
            )

            if score_details.score() < subassembly_config.min_score:
                log.debug(
                    "[subassembly] Rejected box at %s: score=%.2f < min_score=%.2f",
                    bbox,
                    score_details.score(),
                    subassembly_config.min_score,
                )
                continue

            result.add_candidate(
                Candidate(
                    bbox=bbox,
                    label="subassembly",
                    score=score_details.score(),
                    score_details=score_details,
                    source_blocks=list(group),
                )
            )
            log.debug(
                "[subassembly] Candidate at %s: has_count=%s, "
                "has_steps=%s, has_diagrams_or_images=%s, score=%.2f",
                bbox,
                has_step_count,
                has_step_numbers,
                has_diagram_or_images,
                score_details.score(),
            )

    def _find_candidate_inside(
        self, bbox: BBox, candidates: list[Candidate]
    ) -> Candidate | None:
        """Find the best candidate that is fully inside the given box.

        Args:
            bbox: The bounding box of the subassembly container
            candidates: Candidates to search

        Returns:
            The best candidate inside the box, or None
        """
        return find_best_scoring(filter_contained(candidates, bbox))

    def _find_all_candidates_inside(
        self, bbox: BBox, candidates: list[Candidate]
    ) -> list[Candidate]:
        """Find all candidates that are fully inside the given box.

        Args:
            bbox: The bounding box of the subassembly container
            candidates: Candidates to search

        Returns:
            List of candidates inside the box, sorted by score (highest first)
        """
        inside = filter_contained(candidates, bbox)

        # Sort by score (highest first)
        inside.sort(key=lambda c: c.score, reverse=True)
        return inside

    def _find_diagram_inside(
        self, bbox: BBox, diagram_candidates: list[Candidate]
    ) -> Candidate | None:
        """Find the best diagram candidate that overlaps with the box.

        Args:
            bbox: The bounding box of the subassembly container
            diagram_candidates: Diagram candidates to search

        Returns:
            The best diagram candidate overlapping the box, or None
        """
        best_candidate = None
        best_overlap = 0.0

        # Use filter_overlapping to narrow down candidates
        overlapping_candidates = filter_overlapping(diagram_candidates, bbox)

        for candidate in overlapping_candidates:
            # Calculate overlap area
            # TODO Should this use bbox.intersection_area(candidate.bbox)?
            overlap = bbox.intersect(candidate.bbox)
            overlap_area = overlap.width * overlap.height
            if overlap_area > best_overlap:
                best_candidate = candidate
                best_overlap = overlap_area

        return best_candidate

    def _find_all_diagrams_inside(
        self, bbox: BBox, diagram_candidates: list[Candidate]
    ) -> list[Candidate]:
        """Find all diagram candidates that are fully inside the box.

        Args:
            bbox: The bounding box of the subassembly container
            diagram_candidates: Diagram candidates to search

        Returns:
            List of diagram candidates inside the box, sorted by area (largest first)
        """
        diagrams = filter_contained(diagram_candidates, bbox)
        # Sort by area (largest first)
        diagrams.sort(key=lambda c: c.bbox.area, reverse=True)
        return diagrams

    def _find_images_inside(self, bbox: BBox, blocks: list[Blocks]) -> list[Image]:
        """Find Image blocks that are fully inside the given box.

        This directly looks at Image blocks, bypassing the diagram clustering.
        Images inside subassembly boxes often get clustered with larger diagrams
        outside the box, so we need to find them directly.

        Args:
            bbox: The bounding box of the subassembly container
            blocks: All blocks on the page

        Returns:
            List of Image blocks fully inside the box, sorted by area (largest first)
        """
        min_area = 100.0  # Skip very small images (decorative elements)

        potential_images = [
            b for b in blocks if isinstance(b, Image) and b.bbox.area >= min_area
        ]
        images = filter_contained(potential_images, bbox)

        # Sort by area (largest first) - larger images are more likely to be diagrams
        images.sort(key=lambda img: img.bbox.area, reverse=True)
        return images

    def build(self, candidate: Candidate, result: ClassificationResult) -> SubAssembly:
        """Construct a SubAssembly element from a candidate.

        Child element discovery happens here at build time:
        - Find step_count inside the box
        - Find step_numbers inside the box
        - Find diagrams and images inside the box
        - Match step_numbers with diagrams/images
        """
        bbox = candidate.bbox
        page_data = result.page_data

        # Get candidates for child element discovery
        step_count_candidates = result.get_scored_candidates(
            "step_count", valid_only=False, exclude_failed=True
        )
        step_number_candidates = result.get_scored_candidates(
            "step_number", valid_only=False, exclude_failed=True
        )
        diagram_candidates = result.get_scored_candidates(
            "diagram", valid_only=False, exclude_failed=True
        )

        # Find step_count inside the box and build it
        count = None
        step_count_candidate = self._find_candidate_inside(bbox, step_count_candidates)
        if step_count_candidate:
            count_elem = result.build(step_count_candidate)
            assert isinstance(count_elem, StepCount)
            count = count_elem

        # Find step_numbers inside the box
        step_nums_inside = self._find_all_candidates_inside(
            bbox, step_number_candidates
        )

        # Find diagrams and images inside the box
        diagrams_inside = self._find_all_diagrams_inside(bbox, diagram_candidates)
        images_inside = self._find_images_inside(bbox, page_data.blocks)

        # Build steps if we have step numbers inside
        steps: list[SubAssemblyStep] = []
        if step_nums_inside:
            # Build step numbers and match them with diagrams or images
            # Pass a slightly inset bbox to constrain diagram clustering.
            # The inset avoids capturing the white border of the subassembly
            # box.
            # TODO Turn the -3 into a config
            inset_bbox = bbox.expand(-3.0)  # Shrink by 3 points on all sides
            steps = self._build_subassembly_steps(
                step_nums_inside,
                diagrams_inside,
                images_inside,
                result,
                constraint_bbox=inset_bbox,
            )

        # Build a single diagram if present and no steps were built
        # Pass a slightly inset bbox as a constraint to prevent diagram
        # clustering from expanding beyond the subassembly bounds.
        # The inset avoids capturing the white border of the subassembly box.
        # TODO Turn the -3 into a config
        diagram = None
        if not steps:
            # TODO Check all diagrams are built inside the box
            if diagrams_inside:
                inset_bbox = bbox.expand(-3.0)  # Shrink by 3 points on all sides
                diagram_elem = result.build(
                    diagrams_inside[0], constraint_bbox=inset_bbox
                )
                assert isinstance(diagram_elem, Diagram)
                diagram = diagram_elem
            elif images_inside:
                # Fall back to using an Image directly as the diagram
                diagram = Diagram(bbox=images_inside[0].bbox)

        # Subassemblies must contain at least one diagram
        # (either in steps or standalone)
        has_diagram = diagram is not None or any(s.diagram is not None for s in steps)
        if not has_diagram:
            raise ValueError(
                f"SubAssembly at {bbox} has no diagram - "
                "subassemblies must contain at least one diagram"
            )

        return SubAssembly(
            bbox=bbox,
            steps=steps,
            diagram=diagram,
            count=count,
        )

    def _build_subassembly_steps(
        self,
        step_number_candidates: list[Candidate],
        diagram_candidates: list[Candidate],
        images_inside: list[Image],
        result: ClassificationResult,
        constraint_bbox: BBox,
    ) -> list[SubAssemblyStep]:
        """Build SubAssemblyStep elements by matching step numbers with diagrams.

        Uses a simple heuristic: diagrams are typically to the right of and/or
        below the step number. For each step number, find the best matching
        diagram based on position. If no diagram candidates are available,
        uses Image blocks found directly inside the subassembly box.

        Args:
            step_number_candidates: Step number candidates inside the subassembly
            diagram_candidates: Diagram candidates inside the subassembly
            images_inside: Image blocks found directly inside the subassembly
            result: Classification result for building elements
            constraint_bbox: Bounding box to constrain diagram clustering

        Returns:
            List of SubAssemblyStep elements, sorted by step number value
        """
        steps: list[SubAssemblyStep] = []
        used_diagram_ids: set[int] = set()
        used_image_ids: set[int] = set()

        # Sort step numbers by their value for consistent ordering
        sorted_step_nums = sorted(
            step_number_candidates,
            key=lambda c: self._extract_step_value(c),
        )

        for step_num_candidate in sorted_step_nums:
            # Build the step number element
            step_num_elem = result.build(step_num_candidate)
            assert isinstance(step_num_elem, StepNumber)

            # Find the best matching diagram for this step
            best_diagram: Diagram | None = None
            best_diagram_id: int | None = None
            best_score = -float("inf")

            # First try diagram candidates
            for diagram_candidate in diagram_candidates:
                diagram_id = id(diagram_candidate)
                if diagram_id in used_diagram_ids:
                    continue

                # Skip candidates that are already failed (e.g., from a previous
                # diagram build that clustered and claimed shared images)
                if diagram_candidate.failure_reason:
                    continue

                # Score this diagram for this step
                score = self._score_step_diagram_match(
                    step_num_candidate.bbox, diagram_candidate.bbox
                )
                if score > best_score:
                    # Build the diagram with constraint to prevent clustering
                    # beyond the subassembly bounds
                    # TODO maybe pass images_inside to constrain diagram clustering?
                    try:
                        diagram_elem = result.build(
                            diagram_candidate, constraint_bbox=constraint_bbox
                        )

                        assert isinstance(diagram_elem, Diagram)
                        best_score = score
                        best_diagram_id = diagram_id
                        best_diagram = diagram_elem
                    except CandidateFailedError:
                        # This candidate was claimed by another diagram during
                        # clustering - skip it and try the next one
                        continue

            # If no diagram candidate found, try Image blocks directly
            if best_diagram is None:
                best_image: Image | None = None
                best_image_id: int | None = None
                best_score = -float("inf")

                for image in images_inside:
                    image_id = id(image)
                    if image_id in used_image_ids:
                        continue

                    # Score this image for this step
                    score = self._score_step_diagram_match(
                        step_num_candidate.bbox, image.bbox
                    )
                    if score > best_score:
                        best_score = score
                        best_image_id = image_id
                        best_image = image

                if best_image is not None and best_image_id is not None:
                    used_image_ids.add(best_image_id)
                    # Create a Diagram from the Image
                    best_diagram = Diagram(bbox=best_image.bbox)

            if best_diagram_id is not None:
                used_diagram_ids.add(best_diagram_id)

            # Compute bbox for the step
            step_bbox = step_num_elem.bbox
            if best_diagram:
                step_bbox = step_bbox.union(best_diagram.bbox)

            steps.append(
                SubAssemblyStep(
                    bbox=step_bbox,
                    step_number=step_num_elem,
                    diagram=best_diagram,
                )
            )

        return steps

    def _extract_step_value(self, candidate: Candidate) -> int:
        """Extract the step number value from a candidate.

        Args:
            candidate: A step_number candidate

        Returns:
            The step number value, or 0 if not extractable
        """
        if candidate.source_blocks and isinstance(candidate.source_blocks[0], Text):
            text_block = candidate.source_blocks[0]
            value = extract_step_number_value(text_block.text)
            return value if value is not None else 0
        return 0

    def _score_step_diagram_match(self, step_bbox: BBox, diagram_bbox: BBox) -> float:
        """Score how well a diagram matches a step number in a subassembly.

        In subassemblies, diagrams are typically positioned to the right of
        and/or below the step number.

        Args:
            step_bbox: The step number bounding box
            diagram_bbox: The diagram bounding box

        Returns:
            Score (higher is better match)
        """
        # Prefer diagrams that are:
        # 1. To the right of the step number (positive x_offset)
        # 2. Below or at same level (positive or small negative y_offset)
        # 3. Close by (small distance)

        x_offset = diagram_bbox.x0 - step_bbox.x1
        y_offset = diagram_bbox.y0 - step_bbox.y0

        # X score: prefer diagrams to the right
        if x_offset >= 0:
            x_score = 1.0 - min(x_offset / 200.0, 0.5)
        else:
            x_score = 0.5 + x_offset / 100.0  # Penalize left position

        # Y score: prefer diagrams at same level or below
        if abs(y_offset) < 50:
            y_score = 1.0
        elif y_offset >= 0:
            y_score = 0.8 - min(y_offset / 200.0, 0.3)
        else:
            y_score = 0.5 + y_offset / 100.0  # Penalize above position

        return x_score + y_score

1	"""
2	SubAssembly classifier.
3
4	Purpose
5	-------
6	Identify sub-assembly callout boxes on LEGO instruction pages. SubAssemblies
7	typically:
8	- Are white/light-colored rectangular boxes with black borders
9	- Are larger than individual parts (to contain a small build diagram)
10	- May contain a count label (e.g., "2x") indicating how many to build
11	- May contain step numbers for multi-step subassemblies
12	- May have an arrow pointing from them to the main diagram
13
14	Scoring is based on intrinsic properties of the box:
15	- Fill color (white/light)
16	- Size (larger than minimum threshold)
17
18	Child element discovery (step_count, step_numbers, diagrams, arrows) is
19	deferred to build time per DESIGN.md principles.
20
21	Debugging
22	---------
23	Set environment variables to aid investigation without code changes:
24
25	- LOG_LEVEL=DEBUG
26	Enables DEBUG-level logging (if not already configured by caller).
27	"""
28
29	from __future__ import annotations	1✔
30
31	import logging	1✔
32	from typing import ClassVar	1✔
33
34	from build_a_long.pdf_extract.classifier.candidate import Candidate	1✔
35	from build_a_long.pdf_extract.classifier.classification_result import (	1✔
36	CandidateFailedError,
37	ClassificationResult,
38	)
39	from build_a_long.pdf_extract.classifier.config import SubAssemblyConfig	1✔
40	from build_a_long.pdf_extract.classifier.label_classifier import LabelClassifier	1✔
41	from build_a_long.pdf_extract.classifier.score import (	1✔
42	Score,
43	Weight,
44	find_best_scoring,
45	)
46	from build_a_long.pdf_extract.classifier.text import extract_step_number_value	1✔
47	from build_a_long.pdf_extract.classifier.utils import score_white_fill	1✔
48	from build_a_long.pdf_extract.extractor.bbox import (	1✔
49	BBox,
50	filter_contained,
51	filter_overlapping,
52	group_by_similar_bbox,
53	)
54	from build_a_long.pdf_extract.extractor.lego_page_elements import (	1✔
55	Diagram,
56	StepCount,
57	StepNumber,
58	SubAssembly,
59	SubAssemblyStep,
60	)
61	from build_a_long.pdf_extract.extractor.page_blocks import Blocks, Drawing, Image, Text	1✔
62
63	log = logging.getLogger(__name__)	1✔
64
65
66	class _SubAssemblyScore(Score):	1✔
67	"""Internal score representation for subassembly classification.
68
69	Scores based on intrinsic box properties only. Child element discovery
70	(step_count, step_numbers, diagrams, arrows) is deferred to build time.
71	"""
72
73	box_score: float	1✔
74	"""Score based on box having white/light fill (0.0-1.0)."""	1✔
75
76	has_step_count: bool	1✔
77	"""Whether a step_count candidate exists inside (for scoring bonus)."""	1✔
78
79	has_diagram_or_images: bool	1✔
80	"""Whether diagram candidates or images exist inside (for scoring bonus)."""	1✔
81
82	has_step_numbers: bool	1✔
83	"""Whether step_number candidates exist inside (for multi-step subassemblies)."""	1✔
84
85	config: SubAssemblyConfig	1✔
86	"""Configuration containing weights for score calculation."""	1✔
87
88	def score(self) -> Weight:	1✔
89	"""Calculate final weighted score from components."""
90	count_score = 1.0 if self.has_step_count else 0.0	1✔
91	diagram_score = 1.0 if self.has_diagram_or_images else 0.0	1✔
92
93	return (	1✔
94	self.box_score * self.config.box_shape_weight
95	+ count_score * self.config.count_weight
96	+ diagram_score * self.config.diagram_weight
97	)
98
99
100	class SubAssemblyClassifier(LabelClassifier):	1✔
101	"""Classifier for subassembly callout boxes."""
102
103	output: ClassVar[str] = "subassembly"	1✔
104	requires: ClassVar[frozenset[str]] = frozenset(	1✔
105	{"step_count", "step_number", "diagram"}
106	)
107
108	def _score(self, result: ClassificationResult) -> None:	1✔
109	"""Score Drawing blocks as potential subassembly boxes."""
110	page_data = result.page_data	1✔
111	subassembly_config = self.config.subassembly	1✔
112
113	# Get step_count, step_number, and diagram candidates
114	step_count_candidates = result.get_scored_candidates(	1✔
115	"step_count", valid_only=False, exclude_failed=True
116	)
117	step_number_candidates = result.get_scored_candidates(	1✔
118	"step_number", valid_only=False, exclude_failed=True
119	)
120	diagram_candidates = result.get_scored_candidates(	1✔
121	"diagram", valid_only=False, exclude_failed=True
122	)
123
124	# Find rectangular drawing blocks that could be subassembly boxes
125	# Filter by size constraints first
126	max_width = page_data.bbox.width * subassembly_config.max_page_width_ratio	1✔
127	max_height = page_data.bbox.height * subassembly_config.max_page_height_ratio	1✔
128
129	valid_drawings: list[Drawing] = []	1✔
130	for block in page_data.blocks:	1✔
131	if not isinstance(block, Drawing):	1✔
132	continue	1✔
133
134	bbox = block.bbox	1✔
135
136	# Skip boxes smaller than minimum subassembly size
137	if (	1✔
138	bbox.width < subassembly_config.min_subassembly_width
139	or bbox.height < subassembly_config.min_subassembly_height
140	):
141	continue	1✔
142
143	# Skip boxes larger than maximum subassembly size
144	if bbox.width > max_width or bbox.height > max_height:	1✔
145	log.debug(	1✔
146	"[subassembly] Skipping oversized box at %s "
147	"(%.1f x %.1f > max %.1f x %.1f)",
148	bbox,
149	bbox.width,
150	bbox.height,
151	max_width,
152	max_height,
153	)
154	continue	1✔
155
156	valid_drawings.append(block)	1✔
157
158	# Group drawings with similar bboxes (e.g., white-filled box and
159	# black-bordered box for the same subassembly)
160	groups = group_by_similar_bbox(valid_drawings, tolerance=2.0)	1✔
161
162	# Process each group - create one candidate per unique bbox region
163	for group in groups:	1✔
164	# Use union of all grouped drawings' bboxes
165	bbox = BBox.union_all([d.bbox for d in group])	1✔
166
167	# Score each drawing's colors and pick the best
168	best_box_score = max(score_white_fill(d) for d in group)	1✔
169	if best_box_score < 0.3:	1✔
170	continue	1✔
171
172	# Check for child elements inside the box (for scoring only)
173	# Actual candidate discovery happens at build time
174	has_step_count = bool(	1✔
175	self._find_candidate_inside(bbox, step_count_candidates)
176	)
177	has_step_numbers = bool(	1✔
178	self._find_all_candidates_inside(bbox, step_number_candidates)
179	)
180	diagrams_inside = self._find_all_diagrams_inside(bbox, diagram_candidates)	1✔
181	images_inside = self._find_images_inside(bbox, page_data.blocks)	1✔
182	has_diagram_or_images = bool(diagrams_inside or images_inside)	1✔
183
184	# We need at least a box - count and diagram are optional
185	score_details = _SubAssemblyScore(	1✔
186	box_score=best_box_score,
187	has_step_count=has_step_count,
188	has_diagram_or_images=has_diagram_or_images,
189	has_step_numbers=has_step_numbers,
190	config=subassembly_config,
191	)
192
193	if score_details.score() < subassembly_config.min_score:	1✔
194	log.debug(	1✔
195	"[subassembly] Rejected box at %s: score=%.2f < min_score=%.2f",
196	bbox,
197	score_details.score(),
198	subassembly_config.min_score,
199	)
200	continue	1✔
201
202	result.add_candidate(	1✔
203	Candidate(
204	bbox=bbox,
205	label="subassembly",
206	score=score_details.score(),
207	score_details=score_details,
208	source_blocks=list(group),
209	)
210	)
211	log.debug(	1✔
212	"[subassembly] Candidate at %s: has_count=%s, "
213	"has_steps=%s, has_diagrams_or_images=%s, score=%.2f",
214	bbox,
215	has_step_count,
216	has_step_numbers,
217	has_diagram_or_images,
218	score_details.score(),
219	)
220
221	def _find_candidate_inside(	1✔
222	self, bbox: BBox, candidates: list[Candidate]
223	) -> Candidate \| None:
224	"""Find the best candidate that is fully inside the given box.
225
226	Args:
227	bbox: The bounding box of the subassembly container
228	candidates: Candidates to search
229
230	Returns:
231	The best candidate inside the box, or None
232	"""
233	return find_best_scoring(filter_contained(candidates, bbox))	1✔
234
235	def _find_all_candidates_inside(	1✔
236	self, bbox: BBox, candidates: list[Candidate]
237	) -> list[Candidate]:
238	"""Find all candidates that are fully inside the given box.
239
240	Args:
241	bbox: The bounding box of the subassembly container
242	candidates: Candidates to search
243
244	Returns:
245	List of candidates inside the box, sorted by score (highest first)
246	"""
247	inside = filter_contained(candidates, bbox)	1✔
248
249	# Sort by score (highest first)
250	inside.sort(key=lambda c: c.score, reverse=True)	1✔
251	return inside	1✔
252
253	def _find_diagram_inside(	1✔
254	self, bbox: BBox, diagram_candidates: list[Candidate]
255	) -> Candidate \| None:
256	"""Find the best diagram candidate that overlaps with the box.
257
258	Args:
259	bbox: The bounding box of the subassembly container
260	diagram_candidates: Diagram candidates to search
261
262	Returns:
263	The best diagram candidate overlapping the box, or None
264	"""
UNCOV 265	best_candidate = None	×
UNCOV 266	best_overlap = 0.0	×
267
268	# Use filter_overlapping to narrow down candidates
UNCOV 269	overlapping_candidates = filter_overlapping(diagram_candidates, bbox)	×
270
UNCOV 271	for candidate in overlapping_candidates:	×
272	# Calculate overlap area
273	# TODO Should this use bbox.intersection_area(candidate.bbox)?
UNCOV 274	overlap = bbox.intersect(candidate.bbox)	×
UNCOV 275	overlap_area = overlap.width * overlap.height	×
UNCOV 276	if overlap_area > best_overlap:	×
UNCOV 277	best_candidate = candidate	×
UNCOV 278	best_overlap = overlap_area	×
279
UNCOV 280	return best_candidate	×
281
282	def _find_all_diagrams_inside(	1✔
283	self, bbox: BBox, diagram_candidates: list[Candidate]
284	) -> list[Candidate]:
285	"""Find all diagram candidates that are fully inside the box.
286
287	Args:
288	bbox: The bounding box of the subassembly container
289	diagram_candidates: Diagram candidates to search
290
291	Returns:
292	List of diagram candidates inside the box, sorted by area (largest first)
293	"""
294	diagrams = filter_contained(diagram_candidates, bbox)	1✔
295	# Sort by area (largest first)
296	diagrams.sort(key=lambda c: c.bbox.area, reverse=True)	1✔
297	return diagrams	1✔
298
299	def _find_images_inside(self, bbox: BBox, blocks: list[Blocks]) -> list[Image]:	1✔
300	"""Find Image blocks that are fully inside the given box.
301
302	This directly looks at Image blocks, bypassing the diagram clustering.
303	Images inside subassembly boxes often get clustered with larger diagrams
304	outside the box, so we need to find them directly.
305
306	Args:
307	bbox: The bounding box of the subassembly container
308	blocks: All blocks on the page
309
310	Returns:
311	List of Image blocks fully inside the box, sorted by area (largest first)
312	"""
313	min_area = 100.0 # Skip very small images (decorative elements)	1✔
314
315	potential_images = [	1✔
316	b for b in blocks if isinstance(b, Image) and b.bbox.area >= min_area
317	]
318	images = filter_contained(potential_images, bbox)	1✔
319
320	# Sort by area (largest first) - larger images are more likely to be diagrams
321	images.sort(key=lambda img: img.bbox.area, reverse=True)	1✔
322	return images	1✔
323
324	def build(self, candidate: Candidate, result: ClassificationResult) -> SubAssembly:	1✔
325	"""Construct a SubAssembly element from a candidate.
326
327	Child element discovery happens here at build time:
328	- Find step_count inside the box
329	- Find step_numbers inside the box
330	- Find diagrams and images inside the box
331	- Match step_numbers with diagrams/images
332	"""
333	bbox = candidate.bbox	1✔
334	page_data = result.page_data	1✔
335
336	# Get candidates for child element discovery
337	step_count_candidates = result.get_scored_candidates(	1✔
338	"step_count", valid_only=False, exclude_failed=True
339	)
340	step_number_candidates = result.get_scored_candidates(	1✔
341	"step_number", valid_only=False, exclude_failed=True
342	)
343	diagram_candidates = result.get_scored_candidates(	1✔
344	"diagram", valid_only=False, exclude_failed=True
345	)
346
347	# Find step_count inside the box and build it
348	count = None	1✔
349	step_count_candidate = self._find_candidate_inside(bbox, step_count_candidates)	1✔
350	if step_count_candidate:	1✔
351	count_elem = result.build(step_count_candidate)	1✔
352	assert isinstance(count_elem, StepCount)	1✔
353	count = count_elem	1✔
354
355	# Find step_numbers inside the box
356	step_nums_inside = self._find_all_candidates_inside(	1✔
357	bbox, step_number_candidates
358	)
359
360	# Find diagrams and images inside the box
361	diagrams_inside = self._find_all_diagrams_inside(bbox, diagram_candidates)	1✔
362	images_inside = self._find_images_inside(bbox, page_data.blocks)	1✔
363
364	# Build steps if we have step numbers inside
365	steps: list[SubAssemblyStep] = []	1✔
366	if step_nums_inside:	1✔
367	# Build step numbers and match them with diagrams or images
368	# Pass a slightly inset bbox to constrain diagram clustering.
369	# The inset avoids capturing the white border of the subassembly
370	# box.
371	# TODO Turn the -3 into a config
372	inset_bbox = bbox.expand(-3.0) # Shrink by 3 points on all sides	1✔
373	steps = self._build_subassembly_steps(	1✔
374	step_nums_inside,
375	diagrams_inside,
376	images_inside,
377	result,
378	constraint_bbox=inset_bbox,
379	)
380
381	# Build a single diagram if present and no steps were built
382	# Pass a slightly inset bbox as a constraint to prevent diagram
383	# clustering from expanding beyond the subassembly bounds.
384	# The inset avoids capturing the white border of the subassembly box.
385	# TODO Turn the -3 into a config
386	diagram = None	1✔
387	if not steps:	1✔
388	# TODO Check all diagrams are built inside the box
389	if diagrams_inside:	1✔
390	inset_bbox = bbox.expand(-3.0) # Shrink by 3 points on all sides	1✔
391	diagram_elem = result.build(	1✔
392	diagrams_inside[0], constraint_bbox=inset_bbox
393	)
394	assert isinstance(diagram_elem, Diagram)	1✔
395	diagram = diagram_elem	1✔
396	elif images_inside:	1✔
397	# Fall back to using an Image directly as the diagram
398	diagram = Diagram(bbox=images_inside[0].bbox)	1✔
399
400	# Subassemblies must contain at least one diagram
401	# (either in steps or standalone)
402	has_diagram = diagram is not None or any(s.diagram is not None for s in steps)	1✔
403	if not has_diagram:	1✔
UNCOV 404	raise ValueError(	×
405	f"SubAssembly at {bbox} has no diagram - "
406	"subassemblies must contain at least one diagram"
407	)
408
409	return SubAssembly(	1✔
410	bbox=bbox,
411	steps=steps,
412	diagram=diagram,
413	count=count,
414	)
415
416	def _build_subassembly_steps(	1✔
417	self,
418	step_number_candidates: list[Candidate],
419	diagram_candidates: list[Candidate],
420	images_inside: list[Image],
421	result: ClassificationResult,
422	constraint_bbox: BBox,
423	) -> list[SubAssemblyStep]:
424	"""Build SubAssemblyStep elements by matching step numbers with diagrams.
425
426	Uses a simple heuristic: diagrams are typically to the right of and/or
427	below the step number. For each step number, find the best matching
428	diagram based on position. If no diagram candidates are available,
429	uses Image blocks found directly inside the subassembly box.
430
431	Args:
432	step_number_candidates: Step number candidates inside the subassembly
433	diagram_candidates: Diagram candidates inside the subassembly
434	images_inside: Image blocks found directly inside the subassembly
435	result: Classification result for building elements
436	constraint_bbox: Bounding box to constrain diagram clustering
437
438	Returns:
439	List of SubAssemblyStep elements, sorted by step number value
440	"""
441	steps: list[SubAssemblyStep] = []	1✔
442	used_diagram_ids: set[int] = set()	1✔
443	used_image_ids: set[int] = set()	1✔
444
445	# Sort step numbers by their value for consistent ordering
446	sorted_step_nums = sorted(	1✔
447	step_number_candidates,
448	key=lambda c: self._extract_step_value(c),
449	)
450
451	for step_num_candidate in sorted_step_nums:	1✔
452	# Build the step number element
453	step_num_elem = result.build(step_num_candidate)	1✔
454	assert isinstance(step_num_elem, StepNumber)	1✔
455
456	# Find the best matching diagram for this step
457	best_diagram: Diagram \| None = None	1✔
458	best_diagram_id: int \| None = None	1✔
459	best_score = -float("inf")	1✔
460
461	# First try diagram candidates
462	for diagram_candidate in diagram_candidates:	1✔
463	diagram_id = id(diagram_candidate)	1✔
464	if diagram_id in used_diagram_ids:	1✔
465	continue	1✔
466
467	# Skip candidates that are already failed (e.g., from a previous
468	# diagram build that clustered and claimed shared images)
469	if diagram_candidate.failure_reason:	1✔
470	continue	1✔
471
472	# Score this diagram for this step
473	score = self._score_step_diagram_match(	1✔
474	step_num_candidate.bbox, diagram_candidate.bbox
475	)
476	if score > best_score:	1✔
477	# Build the diagram with constraint to prevent clustering
478	# beyond the subassembly bounds
479	# TODO maybe pass images_inside to constrain diagram clustering?
480	try:	1✔
481	diagram_elem = result.build(	1✔
482	diagram_candidate, constraint_bbox=constraint_bbox
483	)
484
485	assert isinstance(diagram_elem, Diagram)	1✔
486	best_score = score	1✔
487	best_diagram_id = diagram_id	1✔
488	best_diagram = diagram_elem	1✔
UNCOV 489	except CandidateFailedError:	×
490	# This candidate was claimed by another diagram during
491	# clustering - skip it and try the next one
UNCOV 492	continue	×
493
494	# If no diagram candidate found, try Image blocks directly
495	if best_diagram is None:	1✔
UNCOV 496	best_image: Image \| None = None	×
UNCOV 497	best_image_id: int \| None = None	×
UNCOV 498	best_score = -float("inf")	×
499
UNCOV 500	for image in images_inside:	×
UNCOV 501	image_id = id(image)	×
UNCOV 502	if image_id in used_image_ids:	×
UNCOV 503	continue	×
504
505	# Score this image for this step
UNCOV 506	score = self._score_step_diagram_match(	×
507	step_num_candidate.bbox, image.bbox
508	)
UNCOV 509	if score > best_score:	×
UNCOV 510	best_score = score	×
UNCOV 511	best_image_id = image_id	×
512	best_image = image	×
513
UNCOV 514	if best_image is not None and best_image_id is not None:	×
515	used_image_ids.add(best_image_id)	×
516	# Create a Diagram from the Image
UNCOV 517	best_diagram = Diagram(bbox=best_image.bbox)	×
518
519	if best_diagram_id is not None:	1✔
520	used_diagram_ids.add(best_diagram_id)	1✔
521
522	# Compute bbox for the step
523	step_bbox = step_num_elem.bbox	1✔
524	if best_diagram:	1✔
525	step_bbox = step_bbox.union(best_diagram.bbox)	1✔
526
527	steps.append(	1✔
528	SubAssemblyStep(
529	bbox=step_bbox,
530	step_number=step_num_elem,
531	diagram=best_diagram,
532	)
533	)
534
535	return steps	1✔
536
537	def _extract_step_value(self, candidate: Candidate) -> int:	1✔
538	"""Extract the step number value from a candidate.
539
540	Args:
541	candidate: A step_number candidate
542
543	Returns:
544	The step number value, or 0 if not extractable
545	"""
546	if candidate.source_blocks and isinstance(candidate.source_blocks[0], Text):	1✔
547	text_block = candidate.source_blocks[0]	1✔
548	value = extract_step_number_value(text_block.text)	1✔
549	return value if value is not None else 0	1✔
UNCOV 550	return 0	×
551
552	def _score_step_diagram_match(self, step_bbox: BBox, diagram_bbox: BBox) -> float:	1✔
553	"""Score how well a diagram matches a step number in a subassembly.
554
555	In subassemblies, diagrams are typically positioned to the right of
556	and/or below the step number.
557
558	Args:
559	step_bbox: The step number bounding box
560	diagram_bbox: The diagram bounding box
561
562	Returns:
563	Score (higher is better match)
564	"""
565	# Prefer diagrams that are:
566	# 1. To the right of the step number (positive x_offset)
567	# 2. Below or at same level (positive or small negative y_offset)
568	# 3. Close by (small distance)
569
570	x_offset = diagram_bbox.x0 - step_bbox.x1	1✔
571	y_offset = diagram_bbox.y0 - step_bbox.y0	1✔
572
573	# X score: prefer diagrams to the right
574	if x_offset >= 0:	1✔
575	x_score = 1.0 - min(x_offset / 200.0, 0.5)	1✔
576	else:
577	x_score = 0.5 + x_offset / 100.0 # Penalize left position	1✔
578
579	# Y score: prefer diagrams at same level or below
580	if abs(y_offset) < 50:	1✔
581	y_score = 1.0	1✔
UNCOV 582	elif y_offset >= 0:	×
UNCOV 583	y_score = 0.8 - min(y_offset / 200.0, 0.3)	×
584	else:
UNCOV 585	y_score = 0.5 + y_offset / 100.0 # Penalize above position	×
586
587	return x_score + y_score	1✔

bramp / build-along / 20086551557

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous