""" postprocess.py ------------------------------------------------------------------- Purpose : Merge fragmented same-class detections that represent a single continuous defect — line-shaped defects such as a stitched "seam" or a "Warp_Weft" thread run — which standard NMS does not catch. The set of classes this applies to is configurable via PostprocessConfig.merge_classes (see api/main.py for the active config). Background ---------- Standard YOLO NMS (governed by the `iou` parameter at inference time) only suppresses boxes that overlap each other ABOVE the iou threshold. For a long, curved, or diagonal defect like a stitched seam or a warp/weft thread line, the model sometimes proposes two (or more) spatially adjacent but barely-overlapping boxes covering different segments of the same continuous line. Their IoU can be well below 0.5, so lowering `iou` at inference time does not merge them — and lowering it too far risks merging genuinely separate nearby defects (false suppression), which is a worse failure mode for a QC system than an occasional duplicate label. This module implements a class-aware, proximity-AND-orientation based merge step that runs AFTER NMS, as a separate, explicit, auditable stage. It only merges boxes of the SAME class that are (a) touching or within `gap_threshold` pixels of an already-grouped ORIGINAL box (not a grown union — see merge_fragmented_detections docstring for why), and (b) consistently oriented/elongated in the same general direction, when both boxes are elongated enough to have a confident direction. Boxes from different classes (e.g. a "hole" box overlapping a "seam" box, as seen in the diagnostic images) are never merged, since they represent genuinely different defect types occupying nearby regions. Why orientation matters (not just gap distance): - Two same-class boxes can be geometrically close (small gap) while being fragments of two DIFFERENT, unrelated defects — e.g. two distinct stitch lines in different parts of a garment that happen to be within gap_threshold_px of each other via a chain of intermediate boxes. Gap distance alone can't tell these apart from genuine fragments of one continuous line. - A true fragmented line-defect (seam or warp/weft thread) tends to keep a consistent elongation direction across its pieces (they're slices of the same line). Unrelated nearby defects usually don't share that direction. - When a box isn't elongated enough to have a confident direction (close to square/blob), the orientation check is skipped for that pair and gap distance alone decides — this keeps the check from blocking merges it has no real signal about. Why a separate module instead of tuning `iou` further: - iou is a single global knob shared by ALL 5 classes. Lowering it to fix seam fragmentation risks under-suppressing legitimately separate Stain/hole/Warp_Weft detections. - This merge step is class-scoped and distance-bounded, so it only acts on the specific failure pattern observed (seam fragmentation), leaving other classes' NMS behavior untouched. ------------------------------------------------------------------- """ from dataclasses import dataclass, field @dataclass class PostprocessConfig: """Tunable parameters for the merge step. Keep all magic numbers here instead of inline in business logic.""" # Classes eligible for proximity merging. Only classes prone to # being split into multiple boxes along a continuous line # should be listed here (currently just "seam"). merge_classes: set[str] = field(default_factory=lambda: {"seam"}) # Max gap (in pixels, in the ORIGINAL image's coordinate space) # between two same-class boxes for them to be considered part # of the same continuous defect. Boxes that already overlap # have a gap of 0 and are always merged. gap_threshold_px: float = 20.0 # Minimum elongation ratio (long side / short side) for a box # to be treated as "line-like" for orientation purposes. Boxes # below this ratio are closer to square/blob-shaped and skip # the orientation check (handled as a safe default — see # _orientation_compatible docstring). elongation_ratio_min: float = 1.5 # Max allowed angular difference (in degrees) between two # line-like boxes' long-axis directions for them to be # considered "the same line" and eligible for merging. max_angle_diff_deg: float = 35.0 def _box_gap(box_a: list[float], box_b: list[float]) -> float: """ Returns the Chebyshev-style gap between two axis-aligned boxes: 0 if they overlap or touch, otherwise the max of the x-gap and y-gap between their nearest edges. """ ax1, ay1, ax2, ay2 = box_a bx1, by1, bx2, by2 = box_b gap_x = max(0.0, max(bx1, ax1) - min(bx2, ax2)) gap_y = max(0.0, max(by1, ay1) - min(by2, ay2)) return max(gap_x, gap_y) def _box_orientation(box: list[float]) -> tuple[float, float] | None: """ Estimate the long-axis "direction" of a box as (elongation_ratio, angle_deg). angle_deg is measured from the horizontal axis, in [0, 180): 0 -> long axis is horizontal 90 -> long axis is vertical Returns None if the box is not elongated enough (close to square) to have a meaningful direction — see `elongation_ratio_min` in PostprocessConfig. A box without a clear direction is treated as direction-agnostic (handled by the caller, not here). """ x1, y1, x2, y2 = box w = max(x2 - x1, 1e-6) h = max(y2 - y1, 1e-6) long_side = max(w, h) short_side = min(w, h) ratio = long_side / short_side # Axis-aligned bbox: angle is simply 0 (wide) or 90 (tall). # This is a coarse proxy for true stitch-line orientation # (we don't have the actual line/contour, only its bbox), but # it is enough to distinguish "this fragment runs roughly # left-right" from "this one runs roughly top-bottom or is in # a totally different part of the image at a different angle". angle_deg = 0.0 if w >= h else 90.0 return ratio, angle_deg def _orientation_compatible( box_a: list[float], box_b: list[float], config: "PostprocessConfig", ) -> bool: """ Decide whether two boxes are plausibly fragments of the SAME continuous line-shaped defect, based on shape + alignment — not just proximity. Rationale: two same-class boxes can be geometrically close (small gap) yet belong to clearly different defects — e.g. two separate, distant stitch lines on a garment that happen to be within `gap_threshold_px` of each other via a chain of intermediate boxes. Proximity alone can't tell them apart. Requiring similar elongation + similar long-axis direction catches the common real-world pattern: a single seam fragmented by the detector keeps a consistent, strongly elongated (line-like) shape across its pieces, while spurious or imprecise fragments of unrelated nearby defects tend to be closer to square. Default-to-BLOCK for non-elongated boxes (validated against real failure cases, not assumed): a box that is NOT clearly elongated (ratio below `elongation_ratio_min`) is treated as untrustworthy shape evidence for "this is a slice of a long line" and merging is blocked for that pair, even if the gap is small. This is intentionally the conservative direction for a QC system — failing to merge two genuine fragments produces an extra box (a minor annoyance, caught by a human reviewer over- splitting a defect into two regions); wrongly merging produces one giant box that can swallow most of the frame and obscure or misrepresent the actual defect location/extent, which is the worse failure mode we're specifically trying to eliminate. """ ratio_a, angle_a = _box_orientation(box_a) ratio_b, angle_b = _box_orientation(box_b) if ratio_a < config.elongation_ratio_min or ratio_b < config.elongation_ratio_min: return False # not clearly line-like — don't merge on shape grounds # Compare long-axis direction. Since our coarse angle is only # ever 0 or 90 (axis-aligned proxy), "compatible" simply means # both fragments run the same general way (both wide, or both # tall). A genuinely diagonal seam tends to produce fragments # that are each elongated in a consistent way relative to one # another even under this coarse proxy; a wide fragment paired # with a tall, unrelated fragment elsewhere in the frame is the # pattern we want to block. angle_diff = abs(angle_a - angle_b) angle_diff = min(angle_diff, 180.0 - angle_diff) return angle_diff <= config.max_angle_diff_deg def merge_fragmented_detections( defects: list[dict], config: PostprocessConfig | None = None, ) -> list[dict]: """ Merge same-class detections that are likely fragments of a single continuous defect (e.g. a seam split across 2+ boxes). Args: defects: list of dicts as produced by run_inference(), each with keys "type", "confidence", "bbox" ([x1,y1,x2,y2]). config: PostprocessConfig. Uses defaults if not provided. Returns: A new list of defect dicts, with eligible same-class fragments merged into a single bounding box (the union of the merged boxes) and the confidence set to the max of the merged group. Defects of non-merge-eligible classes, or with no nearby same-class neighbor, are returned unchanged. """ if config is None: config = PostprocessConfig() n = len(defects) used = [False] * n merged_defects: list[dict] = [] for i in range(n): if used[i]: continue current = defects[i] used[i] = True if current["type"] not in config.merge_classes: merged_defects.append(current) continue # group = indices of all defects merged into this cluster so far. # Gap and orientation are always measured against the ORIGINAL # boxes in this group (never against the growing union box), # so one large/loose fragment can't "tow in" a distant box just # because the union happens to have grown close to it. This # prevents the transitive chain-growth failure mode where # A merges with B, the union then reaches C, and a single # detection ends up swallowing the whole image. group = [i] merged_box = list(current["bbox"]) merged_conf = current["confidence"] changed = True while changed: changed = False for j in range(n): if used[j] or defects[j]["type"] != current["type"]: continue candidate_box = defects[j]["bbox"] # Must be close to (and oriented consistently with) # at least one ORIGINAL box already in the group — # not just the accumulated union. is_candidate = False for k in group: member_box = defects[k]["bbox"] if _box_gap(member_box, candidate_box) > config.gap_threshold_px: continue if not _orientation_compatible(member_box, candidate_box, config): continue is_candidate = True break if not is_candidate: continue bx1, by1, bx2, by2 = candidate_box merged_box = [ min(merged_box[0], bx1), min(merged_box[1], by1), max(merged_box[2], bx2), max(merged_box[3], by2), ] merged_conf = max(merged_conf, defects[j]["confidence"]) used[j] = True group.append(j) changed = True merged_defects.append({ "type": current["type"], "confidence": merged_conf, "bbox": [round(v, 1) for v in merged_box], }) return merged_defects