Update app.py

app.py

@@ -821,11 +821,14 @@ def apply_user_lines_to_walls(
 
 
 def segment_rooms_flood(walls: np.ndarray) -> np.ndarray:
-    h, w = walls.shape
-    walls[:5, :] = 255; walls[-5:, :] = 255
-    walls[:, :5] = 255; walls[:, -5:] = 255
-
-    filled = walls.copy()
+    """Flood-fill room segmentation. Does NOT modify the input."""
+    h, w  = walls.shape
+    work  = walls.copy()          # never mutate the caller's array
+    # seal border so exterior flood fill can reach everywhere outside walls
+    work[:5, :] = 255;  work[-5:, :] = 255
+    work[:, :5] = 255;  work[:, -5:]  = 255
+
+    filled = work.copy()
     mask   = np.zeros((h+2, w+2), np.uint8)
     for sx, sy in [(0,0),(w-1,0),(0,h-1),(w-1,h-1),
                    (w//2,0),(w//2,h-1),(0,h//2),(w-1,h//2)]:
@@ -833,6 +836,7 @@ def segment_rooms_flood(walls: np.ndarray) -> np.ndarray:
             cv2.floodFill(filled, mask, (sx, sy), 255)
 
     rooms = cv2.bitwise_not(filled)
+    # remove wall pixels from room mask (use original walls, not border-sealed)
     rooms = cv2.bitwise_and(rooms, cv2.bitwise_not(walls))
     rooms = cv2.morphologyEx(rooms, cv2.MORPH_OPEN, np.ones((2,2), np.uint8))
     return rooms
@@ -888,21 +892,25 @@ def generate_prompts(
     inv  = cv2.bitwise_not(walls_mask)
     n, labels, stats, centroids = cv2.connectedComponentsWithStats(inv, connectivity=8)
 
+    # minimum room area = 0.0001 of image — much lower than SAM_CLOSET_THR=300 was
+    min_prompt_area = max(200, int(h * w * 0.0001))
+
     pts, lbls = [], []
     for i in range(1, n):
         area = int(stats[i, cv2.CC_STAT_AREA])
-        if area < SAM_CLOSET_THR:
+        if area < min_prompt_area:
             continue
         bx = int(stats[i, cv2.CC_STAT_LEFT]); by = int(stats[i, cv2.CC_STAT_TOP])
         bw = int(stats[i, cv2.CC_STAT_WIDTH]); bh = int(stats[i, cv2.CC_STAT_HEIGHT])
-        if bx <= 5 and by <= 5 and bx+bw >= w-5 and by+bh >= h-5:
+        # skip only if it covers virtually the entire image (background region)
+        if bx <= 2 and by <= 2 and bx+bw >= w-2 and by+bh >= h-2:
             continue
         cx = int(np.clip(centroids[i][0], 0, w-1))
         cy = int(np.clip(centroids[i][1], 0, h-1))
         if walls_mask[cy, cx] > 0:
             found = False
-            for dy in range(-10, 11, 2):
-                for dx in range(-10, 11, 2):
+            for dy in range(-15, 16, 2):
+                for dx in range(-15, 16, 2):
                     ny2, nx2 = cy+dy, cx+dx
                     if 0<=ny2<h and 0<=nx2<w and walls_mask[ny2,nx2]==0:
                         cx, cy = nx2, ny2; found = True; break
@@ -913,6 +921,8 @@ def generate_prompts(
     for pt in _find_thick_wall_neg_prompts(walls_mask):
         pts.append(list(pt)); lbls.append(0)
 
+    print(f"    [prompts] {sum(l==1 for l in lbls)} positive + "
+          f"{sum(l==0 for l in lbls)} negative prompts")
     return np.array(pts, dtype=np.float32), np.array(lbls, dtype=np.int32)
 
 
@@ -932,35 +942,57 @@ def segment_with_sam(
     img_rgb: np.ndarray,
     walls: np.ndarray,
     sam_ckpt: str,
-) -> List[Dict]:
-    """Returns list of room dicts with keys: mask, score, prompt."""
-    rooms_flood = segment_rooms_flood(walls.copy())
+    rooms_flood: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, List[Dict]]:
+    """
+    Exact port of GeometryAgent.segment_with_sam().
+    Returns (rooms_mask, sam_room_masks_list).
+    rooms_mask is a single combined binary mask — same as GeometryAgent output.
+    """
+    if rooms_flood is None:
+        rooms_flood = segment_rooms_flood(walls.copy())
+
+    sam_room_masks: List[Dict] = []   # mirrors self._sam_room_masks
 
     try:
         import torch
         from segment_anything import sam_model_registry, SamPredictor
+
+        if not Path(sam_ckpt).exists():
+            print("    [SAM] Model not found — using flood-fill")
+            return rooms_flood, []
+
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        sam    = sam_model_registry["vit_h"](checkpoint=sam_ckpt)
+        print(f"    [SAM] Loading vit_h on {device}")
+        sam = sam_model_registry["vit_h"](checkpoint=sam_ckpt)
         sam.to(device); sam.eval()
         predictor = SamPredictor(sam)
+
     except Exception as e:
-        print(f"[SAM] Load failed: {e} — fallback to flood-fill")
-        return _flood_fill_rooms(rooms_flood)
+        print(f"    [SAM] Load failed ({e}) — using flood-fill")
+        return rooms_flood, []
+
+    all_points, all_labels = generate_prompts(walls, rooms_flood)
+    if len(all_points) == 0:
+        print("    [SAM] No prompt centroids — using flood-fill")
+        return rooms_flood, []
 
-    pts, lbls = generate_prompts(walls, rooms_flood)
-    if len(pts) == 0:
-        return _flood_fill_rooms(rooms_flood)
+    pos_pts = [(p, l) for p, l in zip(all_points, all_labels) if l == 1]
+    neg_pts = [p      for p, l in zip(all_points, all_labels) if l == 0]
+    print(f"    [SAM] {len(pos_pts)} room prompts + {len(neg_pts)} wall-negative prompts")
 
     predictor.set_image(img_rgb)
-    pos_pts = [(tuple(p), int(l)) for p, l in zip(pts, lbls) if l == 1]
-    neg_pts = [tuple(p) for p, l in zip(pts, lbls) if l == 0]
+
+    h, w     = walls.shape
+    sam_mask = np.zeros((h, w), dtype=np.uint8)
+    accepted = 0
 
     neg_coords = np.array(neg_pts, dtype=np.float32) if neg_pts else None
     neg_lbls   = np.zeros(len(neg_pts), dtype=np.int32) if neg_pts else None
-    denoise_k  = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
+    denoise_k  = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
 
-    results = []
     for (px, py), lbl in pos_pts:
+        px, py = int(px), int(py)
         if neg_coords is not None:
             pt_c = np.vstack([[[px, py]], neg_coords])
             pt_l = np.concatenate([[lbl], neg_lbls])
@@ -968,80 +1000,229 @@ def segment_with_sam(
             pt_c = np.array([[px, py]], dtype=np.float32)
             pt_l = np.array([lbl],      dtype=np.int32)
 
-        masks, scores, _ = predictor.predict(
-            point_coords=pt_c, point_labels=pt_l, multimask_output=True
-        )
-        best = int(np.argmax(scores))
-        if float(scores[best]) < SAM_MIN_SCORE: continue
+        try:
+            masks, scores, _ = predictor.predict(
+                point_coords=pt_c, point_labels=pt_l, multimask_output=True
+            )
+        except Exception as e:
+            print(f"    [SAM] predict failed ({e})")
+            continue
 
-        m = (masks[best] > 0).astype(np.uint8) * 255
-        m = cv2.bitwise_and(m, rooms_flood)
-        m = cv2.morphologyEx(m, cv2.MORPH_OPEN, denoise_k)
-        if not np.any(m): continue
+        best_idx   = int(np.argmax(scores))
+        best_score = float(scores[best_idx])
+        if best_score < SAM_MIN_SCORE:
+            continue
 
-        results.append({"mask": m, "score": float(scores[best]), "prompt": (px, py)})
+        best_mask = (masks[best_idx] > 0).astype(np.uint8) * 255
+        # AND with flood-fill constraint — exact same as GeometryAgent
+        best_mask = cv2.bitwise_and(best_mask, rooms_flood)
+        best_mask = cv2.morphologyEx(best_mask, cv2.MORPH_OPEN, denoise_k, iterations=1)
 
-    if not results:
-        return _flood_fill_rooms(rooms_flood)
-    return results
+        if not np.any(best_mask):
+            continue
 
+        sam_room_masks.append({
+            "mask"  : best_mask.copy(),
+            "score" : best_score,
+            "prompt": (px, py),
+        })
+        sam_mask = cv2.bitwise_or(sam_mask, best_mask)
+        accepted += 1
 
-def _flood_fill_rooms(rooms_flood: np.ndarray) -> List[Dict]:
-    contours, _ = cv2.findContours(
-        rooms_flood, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
-    )
-    results = []
-    for cnt in contours:
-        m = np.zeros_like(rooms_flood)
-        cv2.drawContours(m, [cnt], -1, 255, -1)
-        M  = cv2.moments(cnt)
-        cx = int(M["m10"]/M["m00"]) if M["m00"] else 0
-        cy = int(M["m01"]/M["m00"]) if M["m00"] else 0
-        results.append({"mask": m, "score": 1.0, "prompt": (cx, cy)})
-    return results
+    print(f"    [SAM] Accepted {accepted}/{len(pos_pts)} masks")
 
+    if accepted == 0:
+        print("    [SAM] No masks accepted — using flood-fill")
+        return rooms_flood, []
 
-def filter_room_masks(
-    room_masks: List[Dict], img_shape: Tuple
-) -> List[Dict]:
+    return sam_mask, sam_room_masks
+
+
+def _flood_fill_fallback(rooms_flood: np.ndarray) -> Tuple[np.ndarray, List]:
+    """Mirrors GeometryAgent flood-fill path — returns (rooms_mask, [])."""
+    return rooms_flood, []
+
+
+# ════════════════════════════════════════════════════════════════════════════
+#  FILTER ROOM REGIONS  — exact port of GeometryAgent.filter_room_regions()
+# ════════════════════════════════════════════════════════════════════════════
+
+def filter_room_regions(
+    rooms_mask: np.ndarray, img_shape: Tuple
+) -> Tuple[np.ndarray, List]:
+    """
+    Exact port of GeometryAgent.filter_room_regions().
+    Finds contours on the COMBINED rooms_mask, filters them.
+    Returns (valid_mask, valid_rooms) where valid_rooms is a list of contours.
+    """
     h, w     = img_shape[:2]
     img_area = float(h * w)
+
     min_area = img_area * MIN_ROOM_AREA_FRAC
     max_area = img_area * MAX_ROOM_AREA_FRAC
     min_dim  = w        * MIN_ROOM_DIM_FRAC
-    margin   = max(5.0, w * BORDER_MARGIN_FRAC)
-
-    valid = []
-    for entry in room_masks:
-        m = entry["mask"]
-        cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        if not cnts: continue
-        cnt  = max(cnts, key=cv2.contourArea)
-        area = cv2.contourArea(cnt)
-        if not (min_area <= area <= max_area): continue
-        bx, by, bw, bh = cv2.boundingRect(cnt)
-        if bx < margin or by < margin or bx+bw > w-margin or by+bh > h-margin:
-            continue
-        if bw < min_dim and bh < min_dim: continue
-        asp = max(bw, bh) / (min(bw, bh) + 1e-6)
-        if asp > MAX_ASPECT_RATIO: continue
-        if (area / (bw*bh+1e-6)) < MIN_EXTENT: continue
+    margin   = max(5.0,  w * BORDER_MARGIN_FRAC)
+
+    print(f"    [filter] min_area={min_area:.0f}  max_area={max_area:.0f}  "
+          f"min_dim={min_dim:.0f}  margin={margin:.0f}")
+
+    contours, _ = cv2.findContours(rooms_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        print("    [filter] no contours found in rooms_mask")
+        return np.zeros_like(rooms_mask), []
+
+    bboxes = np.array([cv2.boundingRect(c) for c in contours], dtype=np.float32)
+    areas  = np.array([cv2.contourArea(c)   for c in contours], dtype=np.float32)
+    bx, by, bw_arr, bh_arr = bboxes[:,0], bboxes[:,1], bboxes[:,2], bboxes[:,3]
+
+    area_ok   = (areas >= min_area) & (areas <= max_area)
+    border_ok = (bx >= margin) & (by >= margin) & \
+                (bx + bw_arr <= w - margin) & (by + bh_arr <= h - margin)
+    dim_ok    = (bw_arr >= min_dim) | (bh_arr >= min_dim)
+    aspect    = np.maximum(bw_arr, bh_arr) / (np.minimum(bw_arr, bh_arr) + 1e-6)
+    aspect_ok = aspect <= MAX_ASPECT_RATIO
+    extent_ok = (areas / (bw_arr * bh_arr + 1e-6)) >= MIN_EXTENT
+
+    cheap_pass = np.where(area_ok & border_ok & dim_ok & aspect_ok & extent_ok)[0]
+
+    valid_mask  = np.zeros_like(rooms_mask)
+    valid_rooms = []
+    for i in cheap_pass:
+        cnt  = contours[i]
         hull = cv2.convexHull(cnt)
         ha   = cv2.contourArea(hull)
-        if ha > 0 and (area / ha) < MIN_SOLIDITY: continue
-
-        entry = dict(entry)
-        entry["contour"] = cnt
-        entry["area_px"] = area
-        valid.append(entry)
+        if ha > 0 and (areas[i] / ha) >= MIN_SOLIDITY:
+            cv2.drawContours(valid_mask, [cnt], -1, 255, -1)
+            valid_rooms.append(cnt)
 
-    return valid
+    print(f"    [filter] {len(valid_rooms)} valid rooms from {len(contours)} contours  "
+          f"(area_ok={int(area_ok.sum())}  border_ok={int((area_ok&border_ok).sum())}  "
+          f"dim_ok={int((area_ok&border_ok&dim_ok).sum())})")
+    return valid_mask, valid_rooms
 
 
 def pixel_area_to_m2(area_px: float) -> float:
     return area_px * (2.54 / DPI) ** 2 * (SCALE_FACTOR ** 2) / 10000
 
 
+# ════════════════════════════════════════════════════════════════════════════
+#  SAM MASK → CONTOUR  — exact port of GeometryAgent._match_sam_mask_to_contour
+# ════════════════════════════════════════════════════════════════════════════
+
+def _mask_to_contour_flat(mask: np.ndarray) -> List[float]:
+    """Exact port of GeometryAgent._mask_to_contour_flat()."""
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+    if not contours:
+        return []
+    largest = max(contours, key=cv2.contourArea)
+    pts = largest[:, 0, :].tolist()
+    return [v for pt in pts for v in pt]
+
+
+def _match_sam_mask_to_contour(
+    contour: np.ndarray,
+    sam_room_masks: List[Dict],
+) -> Tuple[Dict, List[float], float]:
+    """
+    Exact port of GeometryAgent._match_sam_mask_to_contour().
+    Returns (rle_dict, sam_contour_flat, score).
+    """
+    if not sam_room_masks:
+        return _contour_to_rle_and_flat(contour)
+
+    sam_h, sam_w = sam_room_masks[0]["mask"].shape
+    contour_mask = np.zeros((sam_h, sam_w), dtype=np.uint8)
+    cv2.drawContours(contour_mask, [contour], -1, 255, thickness=-1)
+
+    best_iou   = 0.0
+    best_entry = None
+
+    for entry in sam_room_masks:
+        m = entry["mask"]
+        if m.shape != contour_mask.shape:
+            continue
+        inter = np.count_nonzero(cv2.bitwise_and(m, contour_mask))
+        if inter == 0:
+            continue
+        union = np.count_nonzero(cv2.bitwise_or(m, contour_mask))
+        iou   = inter / (union + 1e-6)
+        if iou > best_iou:
+            best_iou   = iou
+            best_entry = entry
+
+    if best_entry is None or best_iou < 0.05:
+        return _contour_to_rle_and_flat(contour)
+
+    sam_contour_flat = _mask_to_contour_flat(best_entry["mask"])
+    if not sam_contour_flat:
+        raw_pts = contour[:, 0, :].tolist()
+        sam_contour_flat = [v for pt in raw_pts for v in pt]
+
+    return mask_to_rle(best_entry["mask"]), sam_contour_flat, best_entry["score"]
+
+
+def _contour_to_rle_and_flat(contour: np.ndarray) -> Tuple[Dict, List[float], float]:
+    """Fallback when no SAM mask matches — exact port of GeometryAgent._contour_to_rle_and_flat()."""
+    x, y, rw, rh = cv2.boundingRect(contour)
+    canvas_h = rh + y + 20
+    canvas_w = rw + x + 20
+    canvas   = np.zeros((canvas_h, canvas_w), dtype=np.uint8)
+    cv2.drawContours(canvas, [contour], -1, 255, thickness=-1)
+    raw_pts  = contour[:, 0, :].tolist()
+    flat_pts = [v for pt in raw_pts for v in pt]
+    return mask_to_rle(canvas), flat_pts, 1.0
+
+
+# ════════════════════════════════════════════════════════════════════════════
+#  MEASURE AND LABEL ROOMS — exact port of GeometryAgent.measure_and_label_rooms
+# ════════════════════════════════════════════════════════════════════════════
+
+def measure_and_label_rooms(
+    img: np.ndarray,
+    valid_rooms: List,
+    sam_room_masks: List[Dict],
+) -> List[Dict]:
+    """Exact port of GeometryAgent.measure_and_label_rooms()."""
+    room_data = []
+
+    for idx, contour in enumerate(valid_rooms, 1):
+        x, y, rw, rh = cv2.boundingRect(contour)
+
+        # OCR — try to find a room label
+        label = run_ocr_on_room(img, contour)
+        if not label or not validate_label(label):
+            # GeometryAgent skips rooms with no valid label
+            # but here we keep them with a fallback so users can see them
+            label = f"ROOM {idx}"
+
+        area_px = cv2.contourArea(contour)
+        M  = cv2.moments(contour)
+        cx = int(M["m10"] / M["m00"]) if M["m00"] else x + rw // 2
+        cy = int(M["m01"] / M["m00"]) if M["m00"] else y + rh // 2
+
+        _, raw_seg_flat, sam_score = _match_sam_mask_to_contour(contour, sam_room_masks)
+
+        room_data.append({
+            "id"        : len(room_data) + 1,
+            "label"     : label,
+            "contour"   : contour,          # kept for annotation
+            "segmentation"    : [raw_seg_flat],
+            "raw_segmentation": [raw_seg_flat],
+            "sam_score" : round(sam_score, 4),
+            "score"     : round(sam_score, 4),
+            "area"      : area_px,
+            "area_px"   : area_px,
+            "area_m2"   : round(pixel_area_to_m2(area_px), 2),
+            "bbox"      : [x, y, rw, rh],
+            "centroid"  : [cx, cy],
+            "confidence": 0.95,
+            "isAi"      : True,
+        })
+
+    print(f"    [label] {len(room_data)} rooms labeled")
+    return room_data
+
+
 def run_ocr_on_room(img_bgr: np.ndarray, contour: np.ndarray) -> Optional[str]:
     try:
         import easyocr
@@ -1357,59 +1538,52 @@ def cb_undo_door_line(state):
 def cb_run_sam(state):
     walls = state.get("walls")
     img   = state.get("img_cropped")
+    img_clean = state.get("img_clean")
     if walls is None or img is None:
         return None, None, state, "Run preprocessing first."
 
-    print("[SAM] Downloading checkpoint if needed…")
-    ckpt = download_sam_if_needed()
-    if ckpt is None:
-        return None, None, state, "❌ SAM checkpoint download failed."
+    print(f"[SAM] walls={walls.shape}  wall_px={np.count_nonzero(walls)}")
 
-    print("[SAM] Segmenting rooms…")
     img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    raw     = segment_with_sam(img_rgb, walls.copy(), ckpt)
+    ckpt    = download_sam_if_needed()
 
-    print("[SAM] Filtering rooms…")
-    filtered = filter_room_masks(raw, img.shape)
+    # ── Segment (SAM or flood-fill fallback) — mirrors GeometryAgent.run() ──
+    sam_enabled = ckpt is not None and Path(ckpt).exists()
 
-    print("[SAM] Running OCR…")
-    rooms = []
-    for idx, entry in enumerate(filtered, 1):
-        cnt   = entry["contour"]
-        label = run_ocr_on_room(img, cnt)
-        if not label or not validate_label(label):
-            label = f"ROOM {idx}"
-
-        M  = cv2.moments(cnt)
-        cx = int(M["m10"]/M["m00"]) if M["m00"] else 0
-        cy = int(M["m01"]/M["m00"]) if M["m00"] else 0
-        area_px  = entry["area_px"]
-        area_m2  = pixel_area_to_m2(area_px)
-        bx,by,bw,bh = cv2.boundingRect(cnt)
-
-        rooms.append({
-            "id":         idx,
-            "label":      label,
-            "contour":    cnt,
-            "mask":       entry["mask"],
-            "score":      entry["score"],
-            "area_px":    round(area_px, 1),
-            "area_m2":    round(area_m2, 2),
-            "bbox":       [bx, by, bw, bh],
-            "centroid":   [cx, cy],
-            "confidence": 0.95,
-        })
+    if sam_enabled:
+        rooms_mask, sam_room_masks = segment_with_sam(img_rgb, walls.copy(), ckpt)
+    else:
+        print("    [SAM] SAM disabled — using flood-fill")
+        rooms_mask = segment_rooms_flood(walls.copy())
+        sam_room_masks = []
+
+    # store sam_room_masks for _match_sam_mask_to_contour
+    state["_sam_room_masks"] = sam_room_masks
+
+    flood_px = int(np.count_nonzero(rooms_mask))
+    print(f"[SAM] rooms_mask px={flood_px}")
+    if flood_px == 0:
+        return None, None, state, "⚠ rooms_mask is empty — walls may be blocking everything."
+
+    # ── filter_room_regions ��� exact GeometryAgent method ────────────────────
+    valid_mask, valid_rooms = filter_room_regions(rooms_mask, img.shape)
+    if not valid_rooms:
+        return None, None, state, "⚠ No valid room regions detected after filtering."
+
+    # ── measure_and_label_rooms ──────────────────────────────────────────────
+    src = img_clean if img_clean is not None else img
+    rooms = measure_and_label_rooms(src, valid_rooms, sam_room_masks)
+    if not rooms:
+        return None, None, state, "⚠ No rooms passed labeling / OCR."
 
-    state["rooms"]       = rooms
+    state["rooms"]        = rooms
     state["selected_ids"] = []
 
-    print("[SAM] Rendering…")
     annotated = build_annotated_image(img, rooms)
     state["annotated"] = annotated
 
     table = [[r["id"], r["label"], f"{r['area_m2']} m²", f"{r['score']:.2f}"]
              for r in rooms]
-
     ann_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
     return ann_rgb, table, state, f"✅ {len(rooms)} rooms detected."