endToEnd2.py

import cv2
import os
import sys
import shutil
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent
os.chdir(PROJECT_ROOT)
sys.path.insert(0, str(PROJECT_ROOT))

from src.segmentation.predictor import FloorPlanPredictor
from src.segmentation.visualizer import SegmentationVisualizer
from src.geometry.pipeline import GeometryPipeline
from src.reconstruction.pipeline import ReconstructionPipeline

# GAN/Phase 5 fully removed as requested
GAN_AVAILABLE = False

MODEL_PATH    = PROJECT_ROOT / "models" / "best.pt"
OUTPUTS_DIR   = PROJECT_ROOT / "outputs"
GENERATED_DIR = PROJECT_ROOT / "generated_models"

print(f"Project root: {PROJECT_ROOT}")
print(f"Model exists: {MODEL_PATH.exists()}")

def run_pipeline(sample_image: Path):
    stem = sample_image.stem

    print(f"\n{'='*60}")
    print(f"Processing: {sample_image.name}")
    print(f"{'='*60}")

    # ── Phase 2: Segmentation (raw image, CPU) ────────────────────────
    print("\n[Phase 2] Segmentation...")
    best_conf   = None
    best_result = None

    # Use CPU by default to match cvlab behavior
    for conf in [0.35, 0.25, 0.15, 0.10, 0.05]:
        predictor  = FloorPlanPredictor(str(MODEL_PATH), confidence=conf, device="cpu")
        seg_result = predictor.predict(str(sample_image))
        total      = seg_result.summary["total_elements"]
        print(f"  conf={conf}: {total} elements detected")
        if total > 0 and best_result is None:
            best_conf   = conf
            best_result = seg_result

    # Load original color image for drawing
    base_img = cv2.imread(str(sample_image), cv2.IMREAD_COLOR)
    if base_img is not None and len(base_img.shape) == 2:
        base_img = cv2.cvtColor(base_img, cv2.COLOR_GRAY2BGR)

    out_dir  = GENERATED_DIR / stem
    if out_dir.exists():
        shutil.rmtree(out_dir)
    out_dir.mkdir(parents=True)

    ann_path = out_dir / f"{stem}_detections.png"

    if best_result is None:
        print("  ⚠ No elements detected.")
        if base_img is not None:
            cv2.imwrite(str(ann_path), base_img)
        return

    print(f"  ✓ conf={best_conf} → {best_result.summary['total_elements']} elements")
    print(f"  {best_result.summary['by_class']}")

    # Save detection overlay into output folder
    viz       = SegmentationVisualizer()
    annotated = viz.draw(base_img, best_result)
    cv2.imwrite(str(ann_path), annotated)
    print(f"  ✓ Detection overlay → {ann_path.name}")

    # ── Phase 3: Geometry ─────────────────────────────────────────────
    print("\n[Phase 3] Geometry reconstruction...")
    img = cv2.imread(str(sample_image))
    geo_result, _, _ = GeometryPipeline().run_and_visualize(
        best_result, img,
        image_path=str(sample_image),
        output_dir=str(OUTPUTS_DIR / "geometry"),
    )

    # ── Phase 4: 3D Reconstruction ────────────────────────────────────
    print("\n[Phase 4] 3D reconstruction...")

    seg_has_elements = best_result is not None and len(best_result.elements) > 0
    has_polygons = len(geo_result.vectorization.all_polygons) > 0

    if not has_polygons and not seg_has_elements:
        print("  ⚠ No geometry to extrude — skipping 3D.")
        return

    model_3d = ReconstructionPipeline().reconstruct(
        geo_result,
        output_dir=str(out_dir),
        stem=stem,
        floorplan_image_path=str(sample_image),
        render_image_path=None, # Explicitly no GAN render
    )
    print(f"  ✓ {model_3d.summary}")

    print(f"\n  Output → {out_dir}/")
    print(f"  ├── {stem}_detections.png")
    print(f"  ├── {stem}.gltf")
    print(f"  ├── {stem}.obj")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        run_pipeline(Path(sys.argv[1]))
    else:
        exts = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tiff", "*.tif")
        samples = sorted(p for ext in exts for p in (PROJECT_ROOT / "samples").glob(ext))
        for sample in samples:
            run_pipeline(sample)