reconstruct3d 0.1.0__py3-none-any.whl

reconstruct3d/chunked.py

@@ -0,0 +1,335 @@
+"""Reconstrucción por fragmentos (chunks) con solape + fusión de nubes.
+
+Para videos largos, reconstruir todo de una vez es lento y acumula deriva. Esta
+estrategia parte el video en fragmentos que se SOLAPAN, reconstruye cada uno de
+forma INDEPENDIENTE (y en PARALELO, ya que no comparten estado), y luego alinea
+las sub-nubes usando los frames compartidos en el solape.
+
+Idea de la alineación (Sim(3) por Umeyama):
+  Dos chunks contiguos extraen exactamente los mismos frames absolutos en su zona
+  de solape (todos muestrean la MISMA rejilla global de k_skip). Cada chunk estima
+  la pose de esos frames en su propio sistema de coordenadas y a su propia escala
+  arbitraria. Los CENTROS DE CÁMARA de los frames compartidos son entonces
+  correspondencias 3D-3D directas entre el sistema del chunk y el sistema global
+  -> resolvemos una similaridad (escala s, rotación R, traslación t) con el método
+  de Umeyama y la aplicamos a los puntos y cámaras del chunk para llevarlos al
+  marco global. Encadenando chunk a chunk, todo queda en un único marco.
+
+Salida: `merged_cloud.ply` (nube global) y `merged_state.npy` (poses globales +
+puntos + cámara), compatibles con `dense` y `view`.
+
+Uso:
+    uv run python chunked.py data/video.mp4 --chunk 80 --overlap 20 --chunk-jobs 4
+    uv run python pipeline.py chunked data/video.mp4 --chunk 80 --overlap 20
+"""
+import os
+import argparse
+from concurrent.futures import ProcessPoolExecutor
+
+import cv2
+import numpy as np
+
+from reconstruct3d.core import MAP_STATE_FILE, CameraConfig, DEFAULT_FRONTEND, resolve_workers
+from reconstruct3d.init_sfm import export_ply
+
+MERGED_CLOUD_FILE = "merged_cloud.ply"
+MERGED_STATE_FILE = "merged_state.npy"
+
+
+# --------------------------------------------------------------------------- #
+# Alineación de similaridad (Umeyama)                                          #
+# --------------------------------------------------------------------------- #
+def umeyama(src, dst, with_scale=True):
+    """Similaridad que mapea `src` -> `dst`: dst ≈ s·R·src + t.
+
+    src, dst: (N,3) correspondencias. Devuelve (s, R(3x3), t(3,)).
+    Método de Umeyama (1991), robusto a reflexiones vía corrección del
+    determinante.
+    """
+    src = np.asarray(src, dtype=np.float64)
+    dst = np.asarray(dst, dtype=np.float64)
+    n = src.shape[0]
+    mu_s, mu_d = src.mean(0), dst.mean(0)
+    sc, dc = src - mu_s, dst - mu_d
+    cov = (dc.T @ sc) / n
+    U, D, Vt = np.linalg.svd(cov)
+    S = np.eye(3)
+    if np.linalg.det(U) * np.linalg.det(Vt) < 0:
+        S[2, 2] = -1.0
+    R = U @ S @ Vt
+    if with_scale:
+        var_s = (sc ** 2).sum() / n
+        s = float((D * np.diag(S)).sum() / var_s) if var_s > 1e-12 else 1.0
+    else:
+        s = 1.0
+    t = mu_d - s * R @ mu_s
+    return s, R, t
+
+
+def apply_sim3(pts, s, R, t):
+    """Aplica dst = s·R·pts + t a (N,3)."""
+    return (s * (R @ np.asarray(pts, dtype=np.float64).T).T) + t
+
+
+def _camera_centers(poses):
+    """{frame: pose 4x4 world->cam} -> {frame: centro de cámara (3,) en mundo}."""
+    return {f: (-(P[:3, :3].T @ P[:3, 3])).ravel() for f, P in poses.items()}
+
+
+# --------------------------------------------------------------------------- #
+# Particionado en chunks                                                       #
+# --------------------------------------------------------------------------- #
+def plan_chunks(n_frames, k_skip, chunk, overlap):
+    """Parte la rejilla global de frames muestreados en chunks solapados.
+
+    Todos los chunks muestrean la MISMA rejilla (múltiplos de k_skip desde 0),
+    así que los frames del solape tienen índices absolutos idénticos entre chunks
+    contiguos (requisito para alinear por centros de cámara compartidos).
+
+    Devuelve lista de (start_frame, end_frame) en frames absolutos.
+    """
+    if overlap >= chunk:
+        raise ValueError(f"overlap ({overlap}) debe ser menor que chunk ({chunk}).")
+    grid = list(range(0, n_frames, k_skip))
+    step = chunk - overlap
+    spans = []
+    a = 0
+    while a < len(grid):
+        b = min(a + chunk, len(grid))
+        start_frame = grid[a]
+        end_frame = grid[b - 1] + 1  # exclusivo; +1 para incluir el último
+        spans.append((start_frame, end_frame))
+        if b >= len(grid):
+            break
+        a += step
+    return spans
+
+
+# --------------------------------------------------------------------------- #
+# Worker de chunk (corre en su propio proceso)                                 #
+# --------------------------------------------------------------------------- #
+def _run_chunk_worker(cfg):
+    """Reconstruye un chunk completo (extract->init->track->[ba]) en su out_dir.
+    Debe ser top-level y con `cfg` picklable para ProcessPoolExecutor (spawn)."""
+    from reconstruct3d.core import process_core, CameraConfig
+    from reconstruct3d.init_sfm import init_reconstruction
+    from reconstruct3d.track_sfm import track_frames
+
+    out = cfg["out"]
+    os.makedirs(out, exist_ok=True)
+    camera = CameraConfig.from_dict(cfg["camera"]) if cfg["camera"] else None
+    try:
+        process_core(
+            cfg["video"], out, frontend_name=cfg["frontend"], k_skip=cfg["k_skip"],
+            m_window=cfg["m_window"], start_frame=cfg["start_frame"], end_frame=cfg["end_frame"],
+            camera=camera, reuse_db=not cfg["force"], jobs=cfg["inner_jobs"],
+        )
+        init_reconstruction(out)
+        track_frames(out, m_window=cfg["m_window"], min_angle=cfg["min_angle"],
+                     local_ba=cfg["local_ba"], ba_every=cfg["ba_every"], fuse=cfg["fuse"])
+        if cfg["ba"]:
+            from reconstruct3d.bundle_adjust import run_bundle_adjustment
+            run_bundle_adjustment(out, max_nfev=cfg["ba_max_nfev"])
+        return {"chunk": cfg["chunk"], "out": out, "ok": True}
+    except Exception as e:  # noqa: BLE001 - reportamos y seguimos con otros chunks
+        return {"chunk": cfg["chunk"], "out": out, "ok": False, "error": repr(e)}
+
+
+# --------------------------------------------------------------------------- #
+# Orquestación                                                                 #
+# --------------------------------------------------------------------------- #
+def run_chunked(video_path, out_dir, frontend_name=DEFAULT_FRONTEND, k_skip=5, m_window=3,
+                chunk=80, overlap=20, min_angle=1.5, local_ba=True, ba_every=1, fuse=True,
+                ba=False, ba_max_nfev=200, chunk_jobs=0, inner_jobs=1, camera=None,
+                min_overlap=4, voxel=0.0):
+    os.makedirs(out_dir, exist_ok=True)
+    cam = camera or CameraConfig()
+
+    cap = cv2.VideoCapture(video_path)
+    n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+    if n_frames <= 0:
+        raise RuntimeError(f"No se pudo leer el conteo de frames de {video_path}")
+
+    spans = plan_chunks(n_frames, k_skip, chunk, overlap)
+    chunk_jobs = resolve_workers(chunk_jobs, frontend_name)
+    chunk_jobs = min(chunk_jobs, len(spans))
+    print(f"[chunked] {len(spans)} chunks (chunk={chunk}, overlap={overlap} frames muestreados) "
+          f"| procesos paralelos={chunk_jobs} | hilos internos/chunk={inner_jobs}")
+    for c, (s, e) in enumerate(spans):
+        print(f"   chunk {c:02d}: frames absolutos [{s}, {e})")
+
+    cfgs = []
+    for c, (s, e) in enumerate(spans):
+        cfgs.append(dict(
+            chunk=c, out=os.path.join(out_dir, f"chunk_{c:02d}"), video=video_path,
+            frontend=frontend_name, k_skip=k_skip, m_window=m_window,
+            start_frame=s, end_frame=e, camera=cam.to_dict(), force=True,
+            inner_jobs=inner_jobs, min_angle=min_angle, local_ba=local_ba,
+            ba_every=ba_every, fuse=fuse, ba=ba, ba_max_nfev=ba_max_nfev,
+        ))
+
+    # --- Fase 1: reconstruir cada chunk (en paralelo) ---
+    if chunk_jobs == 1:
+        results = [_run_chunk_worker(cfg) for cfg in cfgs]
+    else:
+        with ProcessPoolExecutor(max_workers=chunk_jobs) as ex:
+            results = list(ex.map(_run_chunk_worker, cfgs))
+
+    for r in sorted(results, key=lambda r: r["chunk"]):
+        status = "OK" if r["ok"] else f"FALLÓ ({r.get('error', '')})"
+        print(f"[chunked] chunk {r['chunk']:02d}: {status}")
+
+    # --- Fase 2: alinear y fusionar ---
+    return merge_chunks(out_dir, [r["out"] for r in sorted(results, key=lambda r: r["chunk"]) if r["ok"]],
+                        cam, min_overlap=min_overlap, voxel=voxel)
+
+
+def merge_chunks(out_dir, chunk_dirs, cam, min_overlap=4, voxel=0.0):
+    """Alinea las sub-nubes por solape (Umeyama) y escribe la nube fusionada."""
+    chunks = []
+    for d in chunk_dirs:
+        p = os.path.join(d, MAP_STATE_FILE)
+        if not os.path.exists(p):
+            print(f"[merge] {d}: sin {MAP_STATE_FILE}, se omite.")
+            continue
+        m = np.load(p, allow_pickle=True).item()
+        if len(m.get("poses", {})) < 2:
+            print(f"[merge] {d}: <2 poses, se omite.")
+            continue
+        chunks.append((d, m))
+
+    if not chunks:
+        print("[merge] No hay chunks reconstruidos para fusionar.")
+        return None
+
+    # El primer chunk define el marco global.
+    global_centers = {}     # frame_abs -> centro en marco global
+    merged_pts, merged_cols = [], []
+    merged_poses = {}       # frame_abs -> pose global 4x4 world->cam
+    n_used = 0
+
+    for i, (d, m) in enumerate(chunks):
+        poses = m["poses"]
+        centers = _camera_centers(poses)
+        pts = np.asarray(m["points"], dtype=np.float64)
+        cols = np.asarray(m["colors"])
+        finite = np.isfinite(pts).all(axis=1)
+        pts, cols = pts[finite], cols[finite]
+
+        if i == 0:
+            s, R, t = 1.0, np.eye(3), np.zeros(3)
+        else:
+            shared = [f for f in centers if f in global_centers]
+            if len(shared) < min_overlap:
+                print(f"[merge] chunk dir '{os.path.basename(d)}': solape insuficiente "
+                      f"({len(shared)} < {min_overlap} frames), se omite del merge.")
+                continue
+            src = np.array([centers[f] for f in shared])       # marco local del chunk
+            dst = np.array([global_centers[f] for f in shared])  # marco global
+            s, R, t = umeyama(src, dst, with_scale=True)
+            # Diagnóstico: residual RMS de la alineación
+            res = np.linalg.norm(apply_sim3(src, s, R, t) - dst, axis=1)
+            print(f"[merge] chunk dir '{os.path.basename(d)}': {len(shared)} frames de solape, "
+                  f"escala={s:.3f}, RMS alineación={res.mean():.4f}")
+
+        merged_pts.append(apply_sim3(pts, s, R, t))
+        merged_cols.append(cols)
+        # Registrar/actualizar centros y poses globales de este chunk
+        for f, c in centers.items():
+            gc = apply_sim3(c[None], s, R, t)[0]
+            global_centers.setdefault(f, gc)
+            if f not in merged_poses:
+                merged_poses[f] = _transform_pose(poses[f], s, R, t)
+        n_used += 1
+
+    if not merged_pts:
+        print("[merge] Ningún chunk pudo alinearse.")
+        return None
+
+    pts = np.concatenate(merged_pts)
+    cols = np.concatenate(merged_cols)
+
+    # Dedup del solape por vóxel (opcional). Escala automática si voxel<=0.
+    if voxel <= 0:
+        center = pts.mean(0)
+        scale = float(np.median(np.linalg.norm(pts - center, axis=1)))
+        voxel = scale / 600.0 if scale > 0 else 0.0
+    if voxel > 0:
+        keys = np.floor(pts / voxel).astype(np.int64)
+        _, idx = np.unique(keys, axis=0, return_index=True)
+        pts, cols = pts[idx], cols[idx]
+
+    median_pt = np.median(pts, axis=0)
+    cam_centers = [c - median_pt for c in global_centers.values()]
+    out_ply = os.path.join(out_dir, MERGED_CLOUD_FILE)
+    export_ply(out_ply, pts - median_pt, cols, cam_centers)
+
+    merged_state = {
+        "points": pts, "colors": cols, "poses": merged_poses,
+        "camera": cam.to_dict(), "frontend": "", "obs": {},
+    }
+    np.save(os.path.join(out_dir, MERGED_STATE_FILE), merged_state)
+    # También como map_state.npy para que `dense`/`view --out <out_dir>` funcionen
+    # directo sobre la nube fusionada (ambos leen map_state.npy).
+    np.save(os.path.join(out_dir, MAP_STATE_FILE), merged_state)
+    print(f"[merge] {n_used} chunks fusionados | {len(pts):,} puntos | "
+          f"{len(merged_poses)} cámaras -> {out_ply}")
+    print(f"[merge] estado global -> {MERGED_STATE_FILE} (+ copia como {MAP_STATE_FILE} "
+          f"para dense/view con --out {out_dir})")
+    return out_ply
+
+
+def _transform_pose(P, s, R, t):
+    """Lleva una pose world->cam del marco local al global bajo la similaridad
+    (s,R,t) que mapea puntos locales a globales (X_g = s R X_l + t).
+
+    Para puntos: X_g = sR X_l + t. La cámara cumple x_cam = R_wc X_l + t_wc.
+    En el marco global: x_cam = R_wc' X_g + t_wc' con
+        R_wc' = R_wc R^T,  t_wc' = t_wc - R_wc' t,  y la traslación de cámara
+    se reescala por s (la escala del mundo). Se mantiene la orientación de la
+    cámara; solo cambia su posición/escala.
+    """
+    R_wc, t_wc = P[:3, :3], P[:3, 3]
+    R_new = R_wc @ R.T
+    t_new = s * t_wc - R_new @ t
+    out = np.eye(4)
+    out[:3, :3] = R_new
+    out[:3, 3] = t_new
+    return out
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser(description="Reconstrucción por chunks con solape + fusión de nubes.")
+    p.add_argument("video", help="Video de entrada (mp4/MOV).")
+    p.add_argument("--out", default=None, help="Directorio de salida (default outputs/<frontend>/).")
+    p.add_argument("--frontend", default=DEFAULT_FRONTEND, help="sift|orb|spglue (default sift).")
+    p.add_argument("--camera", default=None, help="JSON de intrínsecos (ver camera.example.json).")
+    p.add_argument("--k-skip", type=int, default=5, help="1 de cada k frames (default 5).")
+    p.add_argument("--m-window", type=int, default=3, help="Ventana de pares/PnP (default 3).")
+    p.add_argument("--chunk", type=int, default=80, help="Frames muestreados por chunk (default 80).")
+    p.add_argument("--overlap", type=int, default=20, help="Frames muestreados de solape (default 20).")
+    p.add_argument("--min-angle", type=float, default=1.5, help="Ángulo mín. de triangulación.")
+    p.add_argument("--no-local-ba", action="store_true", help="Desactiva BA local en cada chunk.")
+    p.add_argument("--ba-every", type=int, default=1)
+    p.add_argument("--no-fuse", action="store_true")
+    p.add_argument("--ba", action="store_true", help="BA global por chunk antes de fusionar.")
+    p.add_argument("--ba-max-nfev", type=int, default=200)
+    p.add_argument("--chunk-jobs", type=int, default=0,
+                   help="Chunks reconstruidos en paralelo (procesos; 0=auto, 1=secuencial).")
+    p.add_argument("--inner-jobs", type=int, default=1,
+                   help="Hilos de extracción/matching dentro de cada chunk (default 1).")
+    p.add_argument("--min-overlap", type=int, default=4,
+                   help="Frames de solape mínimos para alinear un chunk (default 4).")
+    p.add_argument("--voxel", type=float, default=0.0, help="Vóxel de dedup del merge (0=auto).")
+    args = p.parse_args()
+
+    out = args.out or os.path.join("outputs", args.frontend)
+    camera = CameraConfig.from_json(args.camera) if args.camera else None
+    run_chunked(
+        args.video, out, frontend_name=args.frontend, k_skip=args.k_skip, m_window=args.m_window,
+        chunk=args.chunk, overlap=args.overlap, min_angle=args.min_angle,
+        local_ba=not args.no_local_ba, ba_every=args.ba_every, fuse=not args.no_fuse,
+        ba=args.ba, ba_max_nfev=args.ba_max_nfev, chunk_jobs=args.chunk_jobs,
+        inner_jobs=args.inner_jobs, camera=camera, min_overlap=args.min_overlap, voxel=args.voxel,
+    )