Update wall_pipeline.py

wall_pipeline.py

@@ -397,99 +397,241 @@ class WallPipeline:
         return result
 
     # ══════════════════════════════════════════════════════════════════════════
-    # Stage 4 — Extract walls  (original)
+    # Stage 4 — Extract walls  (exact GeometryAgent.extract_walls_adaptive)
     # ══════════════════════════════════════════════════════════════════════════
     def _extract_walls(self, img: np.ndarray) -> np.ndarray:
+        """
+        Exact port of GeometryAgent.extract_walls_adaptive().
+        Uses analyze_image_characteristics() for the threshold, then:
+          H/V morph-open → body dilate → collision resolve → distance gate
+          → _remove_thin_lines → small-CC noise filter → _filter_double_lines_and_thick
+        """
         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
         h, w = gray.shape
 
-        otsu_val, _ = cv2.threshold(gray, 0, 255,
-                                     cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        # ── adaptive threshold (identical to analyze_image_characteristics) ──
         brightness = float(np.mean(gray))
+        contrast   = float(np.std(gray))
+        otsu_thr, _ = cv2.threshold(gray, 0, 255,
+                                     cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        wall_pct    = np.sum(_ > 0) / _.size * 100
         if brightness > 220:
-            thr = max(200, int(otsu_val * 1.1))
+            wall_threshold = max(200, int(otsu_thr * 1.1))
         elif brightness < 180:
-            thr = max(150, int(otsu_val * 0.9))
+            wall_threshold = max(150, int(otsu_thr * 0.9))
         else:
-            thr = int(otsu_val)
+            wall_threshold = int(otsu_thr)
 
-        _, binary = cv2.threshold(gray, thr, 255, cv2.THRESH_BINARY_INV)
+        _, binary = cv2.threshold(gray, wall_threshold, 255, cv2.THRESH_BINARY_INV)
 
-        min_line       = max(8, int(0.012 * w))
-        body_thickness = self._estimate_wall_thickness(binary)
+        min_line_len   = max(8, int(0.012 * w))
+        body_thickness = self._estimate_wall_body_thickness(binary, fallback=12)
         body_thickness = int(np.clip(body_thickness, 9, 30))
-        self._wall_thickness = body_thickness
 
-        k_h = cv2.getStructuringElement(cv2.MORPH_RECT, (min_line, 1))
-        k_v = cv2.getStructuringElement(cv2.MORPH_RECT, (1, min_line))
+        print(f"    min_line={min_line_len}px  body={body_thickness}px  (w={w}px)")
 
-        if _GPU:
-            # GPU morphology via cupy — simulate with erosion+dilation
-            long_h = _to_cpu(cp.asarray(
-                cv2.morphologyEx(binary, cv2.MORPH_OPEN, k_h)))
-            long_v = _to_cpu(cp.asarray(
-                cv2.morphologyEx(binary, cv2.MORPH_OPEN, k_v)))
-        else:
-            long_h = cv2.morphologyEx(binary, cv2.MORPH_OPEN, k_h)
-            long_v = cv2.morphologyEx(binary, cv2.MORPH_OPEN, k_v)
+        k_h    = cv2.getStructuringElement(cv2.MORPH_RECT, (min_line_len, 1))
+        k_v    = cv2.getStructuringElement(cv2.MORPH_RECT, (1, min_line_len))
+        long_h = cv2.morphologyEx(binary, cv2.MORPH_OPEN, k_h)
+        long_v = cv2.morphologyEx(binary, cv2.MORPH_OPEN, k_v)
 
         orig_walls = cv2.bitwise_or(long_h, long_v)
-        k_bh = cv2.getStructuringElement(cv2.MORPH_RECT, (1, body_thickness))
-        k_bv = cv2.getStructuringElement(cv2.MORPH_RECT, (body_thickness, 1))
-        dilated_h  = cv2.dilate(long_h, k_bh)
-        dilated_v  = cv2.dilate(long_v, k_bv)
-        walls      = cv2.bitwise_or(dilated_h, dilated_v)
-        collision  = cv2.bitwise_and(dilated_h, dilated_v)
-        safe_zone  = cv2.bitwise_and(collision, orig_walls)
-        walls      = cv2.bitwise_or(
-            cv2.bitwise_and(walls, cv2.bitwise_not(collision)), safe_zone)
-        dist       = cv2.distanceTransform(
-            cv2.bitwise_not(orig_walls), cv2.DIST_L2, 5)
-        keep_mask  = (dist <= (body_thickness / 2)).astype(np.uint8) * 255
-        walls      = cv2.bitwise_and(walls, keep_mask)
-        walls      = self._thin_line_filter(walls, body_thickness)
-        n, labels, stats, _ = cv2.connectedComponentsWithStats(walls, connectivity=8)
-        if n > 1:
+
+        k_bh      = cv2.getStructuringElement(cv2.MORPH_RECT, (1, body_thickness))
+        k_bv      = cv2.getStructuringElement(cv2.MORPH_RECT, (body_thickness, 1))
+        dilated_h = cv2.dilate(long_h, k_bh)
+        dilated_v = cv2.dilate(long_v, k_bv)
+        walls     = cv2.bitwise_or(dilated_h, dilated_v)
+
+        collision = cv2.bitwise_and(dilated_h, dilated_v)
+        safe_zone = cv2.bitwise_and(collision, orig_walls)
+        walls     = cv2.bitwise_or(
+            cv2.bitwise_and(walls, cv2.bitwise_not(collision)),
+            safe_zone
+        )
+
+        dist      = cv2.distanceTransform(cv2.bitwise_not(orig_walls), cv2.DIST_L2, 5)
+        keep_mask = (dist <= (body_thickness / 2)).astype(np.uint8) * 255
+        walls     = cv2.bitwise_and(walls, keep_mask)
+        walls     = self._remove_thin_lines(walls, min_thickness=body_thickness)
+
+        n_lbl, labels, stats, _ = cv2.connectedComponentsWithStats(walls, connectivity=8)
+        if n_lbl > 1:
             areas     = stats[1:, cv2.CC_STAT_AREA]
             min_noise = max(20, int(np.median(areas) * 0.0001))
-            lut       = np.zeros(n, np.uint8)
-            lut[1:]   = (areas >= min_noise).astype(np.uint8)
-            walls     = (lut[labels] * 255).astype(np.uint8)
+            keep_lut  = np.zeros(n_lbl, dtype=np.uint8)
+            keep_lut[1:] = (areas >= min_noise).astype(np.uint8)
+            walls = (keep_lut[labels] * 255).astype(np.uint8)
+
+        walls = self._filter_double_lines_and_thick(walls)
+
+        self._wall_thickness = body_thickness
+        print(f"    Walls: {np.count_nonzero(walls)} px "
+              f"({100*np.count_nonzero(walls)/walls.size:.1f}%)")
         return walls
 
-    def _estimate_wall_thickness(self, binary: np.ndarray, fallback: int = 12) -> int:
-        h, w    = binary.shape
-        n_cols  = min(200, w)
-        col_idx = np.linspace(0, w-1, n_cols, dtype=int)
-        runs    = []
-        max_run = max(2, int(h * 0.05))
-        for ci in col_idx:
-            col = (binary[:, ci] > 0).astype(np.int8)
-            pad = np.concatenate([[0], col, [0]])
-            d   = np.diff(pad.astype(np.int16))
-            s   = np.where(d ==  1)[0]
-            e   = np.where(d == -1)[0]
-            n_  = min(len(s), len(e))
-            r   = (e[:n_] - s[:n_]).astype(int)
-            runs.extend(r[(r >= 2) & (r <= max_run)].tolist())
-        if runs:
-            return int(np.median(runs))
+    def _estimate_wall_body_thickness(self, binary: np.ndarray,
+                                       fallback: int = 12) -> int:
+        """Exact GeometryAgent._estimate_wall_body_thickness — vectorised column scan."""
+        try:
+            h, w        = binary.shape
+            n_cols      = min(200, w)
+            col_indices = np.linspace(0, w - 1, n_cols, dtype=int)
+
+            cols   = (binary[:, col_indices] > 0).astype(np.int8)
+            padded = np.concatenate(
+                [np.zeros((1, n_cols), dtype=np.int8), cols,
+                 np.zeros((1, n_cols), dtype=np.int8)], axis=0
+            )
+            diff = np.diff(padded.astype(np.int16), axis=0)
+
+            run_lengths = []
+            for ci in range(n_cols):
+                d      = diff[:, ci]
+                starts = np.where(d == 1)[0]
+                ends   = np.where(d == -1)[0]
+                if len(starts) == 0 or len(ends) == 0:
+                    continue
+                runs = ends - starts
+                runs = runs[(runs >= 2) & (runs <= h * 0.15)]
+                if len(runs):
+                    run_lengths.append(runs)
+
+            if run_lengths:
+                all_runs  = np.concatenate(run_lengths)
+                thickness = int(np.median(all_runs))
+                print(f"    [WallThickness] Estimated: {thickness} px")
+                return thickness
+        except Exception as exc:
+            print(f"    [WallThickness] Estimation failed ({exc}), fallback={fallback}")
         return fallback
 
-    def _thin_line_filter(self, walls: np.ndarray, min_thickness: int) -> np.ndarray:
+    def _remove_thin_lines(self, walls: np.ndarray,
+                            min_thickness: int) -> np.ndarray:
+        """Exact GeometryAgent._remove_thin_lines — distance transform CC gate."""
         dist       = cv2.distanceTransform(walls, cv2.DIST_L2, 5)
         thick_mask = dist >= (min_thickness / 2)
-        n, labels, _, _ = cv2.connectedComponentsWithStats(walls, connectivity=8)
-        if n <= 1:
+
+        n_lbl, labels, _, _ = cv2.connectedComponentsWithStats(walls, connectivity=8)
+        if n_lbl <= 1:
             return walls
+
         thick_labels = labels[thick_mask]
         if len(thick_labels) == 0:
             return np.zeros_like(walls)
-        has_thick = np.zeros(n, dtype=bool)
+
+        has_thick               = np.zeros(n_lbl, dtype=bool)
         has_thick[thick_labels] = True
-        lut    = has_thick.astype(np.uint8) * 255
-        lut[0] = 0
-        return lut[labels]
+        keep_lut                = has_thick.astype(np.uint8) * 255
+        keep_lut[0]             = 0
+        return keep_lut[labels]
+
+    def _filter_double_lines_and_thick(
+        self,
+        walls: np.ndarray,
+        min_single_dim: int = 20,
+        double_line_gap: int = 60,
+        double_line_search_pct: int = 12,
+    ) -> np.ndarray:
+        """
+        Exact GeometryAgent._filter_double_lines_and_thick.
+        Keeps blobs that either:
+          (a) have min(bbox_w, bbox_h) >= min_single_dim  (proper wall body), OR
+          (b) have a parallel partner blob within double_line_gap px
+              (double-line wall conventions used in CAD drawings).
+        """
+        n_lbl, labels, stats, _ = cv2.connectedComponentsWithStats(walls, connectivity=8)
+        if n_lbl <= 1:
+            return walls
+
+        # Try ximgproc thinning, fall back to morphological skeleton
+        try:
+            skel_full = cv2.ximgproc.thinning(
+                walls, thinningType=cv2.ximgproc.THINNING_ZHANGSUEN
+            )
+        except AttributeError:
+            skel_full = self._morphological_skeleton(walls)
+
+        skel_bin = (skel_full > 0)
+        keep_ids: set = set()
+
+        thin_candidates = []
+        for i in range(1, n_lbl):
+            bw = int(stats[i, cv2.CC_STAT_WIDTH])
+            bh = int(stats[i, cv2.CC_STAT_HEIGHT])
+            if min(bw, bh) >= min_single_dim:
+                keep_ids.add(i)
+            else:
+                thin_candidates.append(i)
+
+        if not thin_candidates:
+            filtered = np.zeros_like(walls)
+            for i in keep_ids:
+                filtered[labels == i] = 255
+            print(f"    [DblLineFilter] Kept {len(keep_ids)}/{n_lbl-1} blobs "
+                  "(all passed size test)")
+            return filtered
+
+        skel_labels  = labels * skel_bin
+        img_h, img_w = labels.shape
+        probe_dists  = np.arange(3, double_line_gap + 1, 3, dtype=np.float32)
+
+        for i in thin_candidates:
+            blob_skel_ys, blob_skel_xs = np.where(skel_labels == i)
+            if len(blob_skel_ys) < 4:
+                continue
+
+            step = max(1, len(blob_skel_ys) // 80)
+            sy   = blob_skel_ys[::step].astype(np.float32)
+            sx   = blob_skel_xs[::step].astype(np.float32)
+            n_s  = len(sy)
+
+            sy_prev = np.roll(sy,  1);  sy_prev[0]  = sy[0]
+            sy_next = np.roll(sy, -1);  sy_next[-1] = sy[-1]
+            sx_prev = np.roll(sx,  1);  sx_prev[0]  = sx[0]
+            sx_next = np.roll(sx, -1);  sx_next[-1] = sx[-1]
+
+            dr   = (sy_next - sy_prev).astype(np.float32)
+            dc   = (sx_next - sx_prev).astype(np.float32)
+            dlen = np.maximum(1.0, np.hypot(dr, dc))
+
+            pr = (-dc / dlen)[:, np.newaxis]
+            pc = ( dr / dlen)[:, np.newaxis]
+
+            for sign in (1.0, -1.0):
+                rr = np.round(sy[:, np.newaxis] + sign * pr * probe_dists).astype(np.int32)
+                cc = np.round(sx[:, np.newaxis] + sign * pc * probe_dists).astype(np.int32)
+
+                valid   = (rr >= 0) & (rr < img_h) & (cc >= 0) & (cc < img_w)
+                safe_rr = np.clip(rr, 0, img_h - 1)
+                safe_cc = np.clip(cc, 0, img_w - 1)
+                lbl_at  = labels[safe_rr, safe_cc]
+
+                partner_mask = valid & (lbl_at > 0) & (lbl_at != i)
+                hit_any      = partner_mask.any(axis=1)
+                hit_rows     = np.where(hit_any)[0]
+                if len(hit_rows) == 0:
+                    continue
+
+                first_hit_col = partner_mask[hit_rows].argmax(axis=1)
+                partner_ids   = lbl_at[hit_rows, first_hit_col]
+                keep_ids.update(partner_ids.tolist())
+
+                if 100.0 * len(hit_rows) / n_s >= double_line_search_pct:
+                    keep_ids.add(i)
+                    break
+
+        if keep_ids:
+            keep_arr           = np.array(sorted(keep_ids), dtype=np.int32)
+            keep_lut           = np.zeros(n_lbl, dtype=np.uint8)
+            keep_lut[keep_arr] = 255
+            filtered           = keep_lut[labels]
+        else:
+            filtered = np.zeros_like(walls)
+
+        print(f"    [DblLineFilter] Kept {len(keep_ids)}/{n_lbl-1} blobs "
+              f"(min_dim>={min_single_dim}px OR double-line partner found)")
+        return filtered
 
     # ══════════════════════════════════════════════════════════════════════════
     # Stage 5b — Remove fixture symbols  (original)