reconstruct3d 0.1.0__py3-none-any.whl

reconstruct3d/core.py

@@ -0,0 +1,515 @@
+import cv2
+import os
+import sys
+import argparse
+import pickle
+import threading
+from abc import ABC, abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+DB_FILE = "sfm_data.pkl"
+MAP_STATE_FILE = "map_state.npy"
+INIT_CLOUD_FILE = "init_cloud.ply"
+TRACKED_CLOUD_FILE = "tracked_cloud.ply"
+DEFAULT_FRONTEND = "sift"
+
+# Intrínsecos por defecto: iPhone calibrado a 540x960 (yaml ORB-SLAM3) con
+# checkerboard (error reproyección 0.40 px). El video crudo viene a 2160x3840
+# (exactamente 4x), se redimensiona a PROC_SIZE antes de extraer y K se usa tal
+# cual. Distorsión ~0 -> no se rectifica.
+DEFAULT_K = np.array([[801.25, 0, 188.67], [0, 801.25, 390.22], [0, 0, 1]], dtype=np.float64)
+DEFAULT_PROC_SIZE = (540, 960)  # (W, H) — debe coincidir con la resolución de calibración
+
+
+class CameraConfig:
+    """Intrínsecos de cámara. **Dependen del dispositivo de grabación**, por eso
+    se pueden cambiar pasando un JSON con `--camera` (ver `camera.example.json`).
+    Si no se especifica, se usan los valores por defecto (iPhone @ 540x960).
+
+    K corresponde a la resolución PROC_SIZE: el frame se redimensiona a PROC_SIZE
+    antes de extraer features, así que K debe estar calibrada para ESA resolución.
+
+    La config elegida en `extract` se persiste en `sfm_data.pkl` y `map_state.npy`,
+    de modo que todas las etapas posteriores usan la misma sin re-especificarla.
+    """
+    PROC_SIZE = DEFAULT_PROC_SIZE  # default de clase (compat. retro)
+
+    def __init__(self, K=None, dist_coeffs=None, proc_size=None):
+        self.K = np.array(DEFAULT_K if K is None else K, dtype=np.float64)
+        self.dist_coeffs = (np.zeros(5, dtype=np.float64) if dist_coeffs is None
+                            else np.asarray(dist_coeffs, dtype=np.float64).ravel())
+        self.PROC_SIZE = tuple(DEFAULT_PROC_SIZE if proc_size is None else proc_size)
+
+    def preprocess(self, frame):
+        """Lleva el frame a la resolución de calibración (y rectifica si hay
+        distorsión). Todos los keypoints viven en este espacio PROC_SIZE."""
+        w, h = self.PROC_SIZE
+        if (frame.shape[1], frame.shape[0]) != (w, h):
+            frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA)
+        if np.any(self.dist_coeffs):
+            frame = cv2.undistort(frame, self.K, self.dist_coeffs)
+        return frame
+
+    def to_dict(self):
+        return {"K": self.K.tolist(), "dist_coeffs": self.dist_coeffs.tolist(),
+                "proc_size": list(self.PROC_SIZE)}
+
+    @classmethod
+    def from_dict(cls, d):
+        """Construye desde un dict. Acepta K explícita 3x3 o el atajo
+        fx/fy/cx/cy, y proc_size o width/height. Dict vacío/None -> defaults."""
+        if not d:
+            return cls()
+        K = d.get("K")
+        if K is None and all(k in d for k in ("fx", "fy", "cx", "cy")):
+            K = [[d["fx"], 0, d["cx"]], [0, d["fy"], d["cy"]], [0, 0, 1]]
+        proc = d.get("proc_size")
+        if proc is None and "width" in d and "height" in d:
+            proc = (d["width"], d["height"])
+        return cls(K=K, dist_coeffs=d.get("dist_coeffs"), proc_size=proc)
+
+    @classmethod
+    def from_json(cls, path):
+        """Carga intrínsecos desde un archivo JSON (ver camera.example.json)."""
+        import json
+        with open(path) as f:
+            return cls.from_dict(json.load(f))
+
+class Features:
+    def __init__(self, pts, descriptors, colors, scores=None, image_size=None):
+        self.pts = pts                  # (N, 2) float32 — pixel coords in processed (resized) frame
+        self.descriptors = descriptors  # frontend-specific: SIFT (N,128) f32, ORB (N,32) u8, SP (N,256) f32
+        self.colors = colors            # (N, 3) BGR uint8
+        self.scores = scores            # (N,) float32 or None
+        self.image_size = image_size    # (W, H) float32 or None — needed by LightGlue
+
+    def __setstate__(self, state):
+        # Backward-compat: old pickles predate scores/image_size
+        state.setdefault('scores', None)
+        state.setdefault('image_size', None)
+        self.__dict__.update(state)
+
+
+class FrontEnd(ABC):
+    """Pluggable feature extractor + matcher. Each preset is atomic:
+    extractor and matcher are bound together because they must agree on
+    descriptor type (binary vs float) and matching strategy (BF vs GNN)."""
+    name: str = ""
+
+    @abstractmethod
+    def extract(self, frame: np.ndarray) -> Features:
+        ...
+
+    @abstractmethod
+    def match(self, feat_i: Features, feat_j: Features) -> list:
+        """Return list of (idx_i, idx_j) tuples."""
+        ...
+
+
+def filter_triangulation(pts3d, pose_a, pose_b, pts2d_a, pts2d_b, K,
+                         max_reproj_err=4.0, min_angle_deg=1.0):
+    """Boolean mask of points that pass: finite, chirality in both cameras,
+    reprojection error ≤ max_reproj_err in both views, and triangulation
+    angle ≥ min_angle_deg. Poses are world→cam (4x4)."""
+    n = len(pts3d)
+    if n == 0:
+        return np.zeros(0, dtype=bool)
+
+    Ra, ta = pose_a[:3, :3], pose_a[:3, 3:]
+    Rb, tb = pose_b[:3, :3], pose_b[:3, 3:]
+    pts_a = (Ra @ pts3d.T + ta).T
+    pts_b = (Rb @ pts3d.T + tb).T
+
+    valid = np.isfinite(pts3d).all(axis=1)
+    valid &= pts_a[:, 2] > 0
+    valid &= pts_b[:, 2] > 0
+
+    with np.errstate(invalid='ignore', divide='ignore'):
+        proj_a = (K @ pts_a.T).T
+        proj_a = proj_a[:, :2] / proj_a[:, 2:3]
+        proj_b = (K @ pts_b.T).T
+        proj_b = proj_b[:, :2] / proj_b[:, 2:3]
+        err_a = np.linalg.norm(proj_a - pts2d_a, axis=1)
+        err_b = np.linalg.norm(proj_b - pts2d_b, axis=1)
+    valid &= np.nan_to_num(err_a, nan=np.inf) <= max_reproj_err
+    valid &= np.nan_to_num(err_b, nan=np.inf) <= max_reproj_err
+
+    ca = (-Ra.T @ ta).ravel()
+    cb = (-Rb.T @ tb).ravel()
+    ray_a = pts3d - ca
+    ray_b = pts3d - cb
+    na = np.linalg.norm(ray_a, axis=1, keepdims=True)
+    nb = np.linalg.norm(ray_b, axis=1, keepdims=True)
+    with np.errstate(invalid='ignore', divide='ignore'):
+        ray_a = ray_a / na
+        ray_b = ray_b / nb
+        cos_ang = np.clip(np.einsum('ij,ij->i', ray_a, ray_b), -1, 1)
+        angles_deg = np.degrees(np.arccos(cos_ang))
+    valid &= np.nan_to_num(angles_deg, nan=0.0) >= min_angle_deg
+
+    return valid
+
+
+def _sample_colors(frame, pts):
+    if len(pts) == 0:
+        return np.empty((0, 3), dtype=np.uint8)
+    h, w = frame.shape[:2]
+    ix = np.clip(pts[:, 0].astype(int), 0, w - 1)
+    iy = np.clip(pts[:, 1].astype(int), 0, h - 1)
+    return frame[iy, ix]
+
+
+class SiftFrontEnd(FrontEnd):
+    name = "sift"
+
+    def __init__(self, ratio=0.75):
+        self.sift = cv2.SIFT_create()
+        self.bf = cv2.BFMatcher(cv2.NORM_L2)
+        self.ratio = ratio
+
+    def extract(self, frame):
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        kp, des = self.sift.detectAndCompute(gray, None)
+        if kp is None or len(kp) == 0:
+            return Features(np.empty((0, 2), dtype=np.float32), None, np.empty((0, 3), dtype=np.uint8))
+        pts = np.float32([k.pt for k in kp])
+        return Features(pts, des, _sample_colors(frame, pts))
+
+    def match(self, feat_i, feat_j):
+        d1, d2 = feat_i.descriptors, feat_j.descriptors
+        if d1 is None or d2 is None or len(d1) < 2 or len(d2) < 2:
+            return []
+        knn = self.bf.knnMatch(d1, d2, k=2)
+        out = []
+        for pair in knn:
+            if len(pair) < 2:
+                continue
+            m, n = pair
+            if m.distance < self.ratio * n.distance:
+                out.append((m.queryIdx, m.trainIdx))
+        return out
+
+
+class OrbFrontEnd(FrontEnd):
+    name = "orb"
+
+    def __init__(self, nfeatures=5000, ratio=0.8):
+        self.orb = cv2.ORB_create(
+            nfeatures=nfeatures, scaleFactor=1.2, nlevels=8,
+            edgeThreshold=15, fastThreshold=10,
+        )
+        self.bf = cv2.BFMatcher(cv2.NORM_HAMMING)
+        self.ratio = ratio
+
+    def extract(self, frame):
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        kp, des = self.orb.detectAndCompute(gray, None)
+        if kp is None or len(kp) == 0:
+            return Features(np.empty((0, 2), dtype=np.float32), None, np.empty((0, 3), dtype=np.uint8))
+        pts = np.float32([k.pt for k in kp])
+        return Features(pts, des, _sample_colors(frame, pts))
+
+    def match(self, feat_i, feat_j):
+        d1, d2 = feat_i.descriptors, feat_j.descriptors
+        if d1 is None or d2 is None or len(d1) < 2 or len(d2) < 2:
+            return []
+        knn = self.bf.knnMatch(d1, d2, k=2)
+        out = []
+        for pair in knn:
+            if len(pair) < 2:
+                continue
+            m, n = pair
+            if m.distance < self.ratio * n.distance:
+                out.append((m.queryIdx, m.trainIdx))
+        return out
+
+
+class SuperPointLightGlueFrontEnd(FrontEnd):
+    name = "spglue"
+
+    def __init__(self, max_keypoints=2048, device="cpu"):
+        # Lazy import — only spglue users need torch installed
+        import torch
+        from lightglue import SuperPoint, LightGlue
+        self._torch = torch
+        self.device = torch.device(device)
+        self.extractor = SuperPoint(max_num_keypoints=max_keypoints).eval().to(self.device)
+        self.matcher = LightGlue(features="superpoint").eval().to(self.device)
+
+    def _to_tensor(self, frame):
+        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        t = self._torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0
+        return t.to(self.device)
+
+    def extract(self, frame):
+        h, w = frame.shape[:2]
+        img_size = np.array([w, h], dtype=np.float32)
+
+        with self._torch.inference_mode():
+            feats = self.extractor.extract(self._to_tensor(frame))
+
+        pts = feats["keypoints"][0].cpu().numpy().astype(np.float32)
+        if len(pts) == 0:
+            return Features(
+                np.empty((0, 2), dtype=np.float32), None, np.empty((0, 3), dtype=np.uint8),
+                scores=None, image_size=img_size,
+            )
+        scores = feats["keypoint_scores"][0].cpu().numpy().astype(np.float32)
+        desc = feats["descriptors"][0].cpu().numpy().astype(np.float32)
+        return Features(pts, desc, _sample_colors(frame, pts), scores=scores, image_size=img_size)
+
+    def match(self, feat_i, feat_j):
+        if feat_i.descriptors is None or feat_j.descriptors is None:
+            return []
+        if len(feat_i.pts) < 2 or len(feat_j.pts) < 2:
+            return []
+
+        torch = self._torch
+
+        def as_input(f):
+            return {
+                "keypoints":       torch.from_numpy(f.pts).float().unsqueeze(0).to(self.device),
+                "keypoint_scores": torch.from_numpy(f.scores).float().unsqueeze(0).to(self.device),
+                "descriptors":     torch.from_numpy(f.descriptors).float().unsqueeze(0).to(self.device),
+                "image_size":      torch.from_numpy(f.image_size).float().unsqueeze(0).to(self.device),
+            }
+
+        with torch.inference_mode():
+            out = self.matcher({"image0": as_input(feat_i), "image1": as_input(feat_j)})
+
+        m = out["matches"][0].cpu().numpy()  # (K, 2)
+        return [(int(a), int(b)) for a, b in m]
+
+
+_FRONTEND_REGISTRY = {
+    SiftFrontEnd.name: SiftFrontEnd,
+    OrbFrontEnd.name: OrbFrontEnd,
+    SuperPointLightGlueFrontEnd.name: SuperPointLightGlueFrontEnd,
+}
+
+
+def create_frontend(name: str) -> FrontEnd:
+    if name not in _FRONTEND_REGISTRY:
+        raise ValueError(f"Frontend '{name}' desconocido. Opciones: {sorted(_FRONTEND_REGISTRY)}")
+    return _FRONTEND_REGISTRY[name]()
+
+
+# Los detectores/matchers de OpenCV (SIFT, ORB, BFMatcher) NO son thread-safe si
+# se comparte la misma instancia entre hilos. Para paralelizar damos a cada hilo
+# su propio frontend vía almacenamiento thread-local (creado perezosamente).
+_thread_frontends = threading.local()
+
+
+def tls_frontend(frontend_name: str) -> FrontEnd:
+    fe = getattr(_thread_frontends, "fe", None)
+    if fe is None or getattr(_thread_frontends, "name", None) != frontend_name:
+        fe = create_frontend(frontend_name)
+        _thread_frontends.fe = fe
+        _thread_frontends.name = frontend_name
+    return fe
+
+
+def resolve_workers(jobs: int, frontend_name: str = "") -> int:
+    """Número de hilos a usar. jobs<=0 -> nº de CPUs; jobs==1 -> secuencial.
+    'spglue' (torch) se fuerza a 1: el modelo no es thread-safe y torch ya
+    paraleliza internamente sobre los núcleos."""
+    n = jobs if jobs and jobs > 0 else (os.cpu_count() or 1)
+    if frontend_name == "spglue":
+        return 1
+    return max(1, n)
+
+
+class SfMDatabase:
+    def __init__(self, frontend_name=DEFAULT_FRONTEND, camera=None):
+        self.features = {}        # idx -> Features
+        self.pairwise = []        # list of dicts
+        self.frontend_name = frontend_name
+        self.camera = camera      # dict de CameraConfig.to_dict() o None (defaults)
+
+    def get_camera(self):
+        """CameraConfig persistida en la DB (o la default si no hay)."""
+        return CameraConfig.from_dict(self.camera)
+
+    def save(self, filepath):
+        with open(filepath, 'wb') as f:
+            pickle.dump({
+                'features': self.features,
+                'pairwise': self.pairwise,
+                'frontend_name': self.frontend_name,
+                'camera': self.camera,
+            }, f)
+
+    def load(self, filepath):
+        with open(filepath, 'rb') as f:
+            data = pickle.load(f)
+            self.features = data['features']
+            self.pairwise = data['pairwise']
+            self.frontend_name = data.get('frontend_name', DEFAULT_FRONTEND)
+            self.camera = data.get('camera')  # backward-compat: pickles viejos
+
+
+def process_core(video_path, out_dir, frontend_name=DEFAULT_FRONTEND, k_skip=5, m_window=3,
+                 max_seconds=None, start_seconds=0.0, start_frame=None, end_frame=None,
+                 camera=None, reuse_db=True, jobs=0):
+    print(f"[*] k_skip={k_skip}  m_window={m_window}")
+    cam = camera or CameraConfig()
+    db = SfMDatabase(frontend_name=frontend_name, camera=cam.to_dict())
+    db_path = os.path.join(out_dir, DB_FILE)
+    if os.path.exists(db_path) and reuse_db:
+        # Reutilizar la DB existente (no recalcular). Pasa reuse_db=False para
+        # forzar la re-extracción. Antes esto era un input() interactivo, lo que
+        # impedía automatizar el pipeline.
+        db.load(db_path)
+        print(f"  [DB cargada] {db_path}  frontend='{db.frontend_name}'  frames={len(db.features)}")
+        print("      (usa --force para re-extraer desde cero)")
+        return db
+
+    workers = resolve_workers(jobs, frontend_name)
+    print(f"[*] Cámara: K=[fx={cam.K[0,0]:.2f} fy={cam.K[1,1]:.2f} "
+          f"cx={cam.K[0,2]:.2f} cy={cam.K[1,2]:.2f}]  proc_size={cam.PROC_SIZE}")
+    print(f"Inicializando frontend '{frontend_name}'...")
+    frontend = create_frontend(frontend_name)
+
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
+    frames_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if max_seconds is not None:
+        limit = int(fps * max_seconds)
+        if limit < frames_count:
+            print(f"[*] Recortando a {max_seconds}s ({limit}/{frames_count} frames @ {fps:.2f} fps)")
+            frames_count = limit
+    # Rango explícito de frames (usado por el procesamiento por chunks). Tiene
+    # prioridad sobre los segundos.
+    if end_frame is not None:
+        frames_count = min(frames_count, int(end_frame))
+    sf = int(start_frame) if start_frame is not None else int(fps * start_seconds)
+    if sf > 0:
+        print(f"[*] Inicio en frame {sf}")
+
+    frame_idxs = list(range(sf, frames_count, k_skip))
+    print(f"Extrayendo features [{frontend_name}] (Salto: {k_skip}, hilos: {workers})...")
+    if workers == 1:
+        for i in tqdm(frame_idxs):
+            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
+            ret, frame = cap.read()
+            if not ret: break
+            db.features[i] = frontend.extract(cam.preprocess(frame))
+    else:
+        # El video se decodifica en serie (VideoCapture no es thread-safe), pero
+        # la extracción de features (lo costoso) se reparte en hilos. Buffer
+        # acotado para no cargar todo el video a memoria; los frames ya van
+        # reducidos a PROC_SIZE.
+        batch = max(8, workers * 4)
+        with ThreadPoolExecutor(max_workers=workers) as ex:
+            pbar = tqdm(total=len(frame_idxs))
+            buf_idx, buf_fr = [], []
+
+            def _extract(frame):
+                return tls_frontend(frontend_name).extract(frame)
+
+            def flush():
+                for idx, feat in zip(buf_idx, ex.map(_extract, buf_fr)):
+                    db.features[idx] = feat
+                pbar.update(len(buf_idx))
+                buf_idx.clear(); buf_fr.clear()
+
+            for i in frame_idxs:
+                cap.set(cv2.CAP_PROP_POS_FRAMES, i)
+                ret, frame = cap.read()
+                if not ret: break
+                buf_idx.append(i); buf_fr.append(cam.preprocess(frame))
+                if len(buf_idx) >= batch:
+                    flush()
+            if buf_idx:
+                flush()
+            pbar.close()
+    cap.release()
+
+    indices = sorted(db.features.keys())
+    # Tareas de pares (ambas direcciones, como antes): cada frame contra sus
+    # ±m_window vecinos extraídos. Son independientes -> se paralelizan.
+    tasks = []
+    for idx_pos in range(len(indices)):
+        start_pos, end_pos = max(0, idx_pos - m_window), min(len(indices), idx_pos + m_window + 1)
+        for j_pos in range(start_pos, end_pos):
+            if j_pos == idx_pos:
+                continue
+            tasks.append((indices[idx_pos], indices[j_pos]))
+
+    def _match_pair(pair):
+        i_frame, j_frame = pair
+        fe = frontend if workers == 1 else tls_frontend(frontend_name)
+        fi, fj = db.features[i_frame], db.features[j_frame]
+        matches = fe.match(fi, fj)
+        if len(matches) < 5:
+            return None
+        pts1 = fi.pts[[m[0] for m in matches]]
+        pts2 = fj.pts[[m[1] for m in matches]]
+        E, mask = cv2.findEssentialMat(pts1, pts2, cam.K, cv2.RANSAC, 0.999, 1.0)
+        if E is None:
+            return None
+        return {"frame_i": i_frame, "frame_j": j_frame,
+                "inliers": int(mask.sum()), "matches": len(matches),
+                "E": E, "mask": mask.ravel()}
+
+    print(f"Calculando M. Esencial (Ventana: +/-{m_window}, hilos: {workers})...")
+    if workers == 1:
+        results = (_match_pair(t) for t in tqdm(tasks))
+        db.pairwise = [r for r in results if r is not None]
+    else:
+        with ThreadPoolExecutor(max_workers=workers) as ex:
+            results = list(tqdm(ex.map(_match_pair, tasks), total=len(tasks)))
+        db.pairwise = [r for r in results if r is not None]
+
+    db.save(db_path)
+    return db
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="SfM core: feature extraction + pairwise essential matrices")
+    parser.add_argument("video", help="Path to input video (mp4)")
+    parser.add_argument(
+        "--frontend", choices=sorted(_FRONTEND_REGISTRY), default=DEFAULT_FRONTEND,
+        help=f"Detector + matcher preset (default: {DEFAULT_FRONTEND}). "
+             "'sift'=SIFT+BF/L2, 'orb'=ORB+BF/Hamming, 'spglue'=SuperPoint+LightGlue (CPU/GPU)."
+    )
+    parser.add_argument(
+        "--out", default=None,
+        help="Output directory for all artifacts (default: outputs/{frontend}/)."
+    )
+    parser.add_argument(
+        "--max-seconds", type=float, default=None,
+        help="Procesar hasta el segundo N del video (límite final, default: video completo)."
+    )
+    parser.add_argument(
+        "--start-seconds", type=float, default=0.0,
+        help="Empezar a procesar desde el segundo N (default: 0). Útil para saltar "
+             "un arranque degenerado de casi-rotación que envenena la escala."
+    )
+    parser.add_argument("--k-skip", type=int, default=5, help="Process 1 out of every k frames (default: 5).")
+    parser.add_argument("--m-window", type=int, default=3,
+                        help="Pairwise window: compare each frame with ±m extracted neighbors "
+                             "(default: 3). Wider = mejor baseline para escenas lentas, pero más costoso.")
+    parser.add_argument("--force", action="store_true",
+                        help="Re-extraer aunque ya exista sfm_data.pkl (por defecto se reutiliza).")
+    parser.add_argument("--camera", default=None,
+                        help="JSON con intrínsecos de la cámara del dispositivo "
+                             "(ver camera.example.json). Default: iPhone @ 540x960.")
+    parser.add_argument("--jobs", type=int, default=0,
+                        help="Hilos para extracción/matching (0=auto=nº de CPUs; 1=secuencial).")
+    args = parser.parse_args()
+
+    out_dir = args.out or os.path.join("outputs", args.frontend)
+    os.makedirs(out_dir, exist_ok=True)
+    print(f"[*] Output dir: {out_dir}")
+
+    camera = CameraConfig.from_json(args.camera) if args.camera else None
+    db = process_core(args.video, out_dir, frontend_name=args.frontend,
+                      k_skip=args.k_skip, m_window=args.m_window, max_seconds=args.max_seconds,
+                      start_seconds=args.start_seconds, camera=camera, reuse_db=not args.force,
+                      jobs=args.jobs)
+    df = pd.DataFrame(db.pairwise).sort_values("inliers", ascending=False)
+    print(f"\nFrontend usado: {db.frontend_name}")
+    print("Top candidatos:\n", df[['frame_i', 'frame_j', 'inliers', 'matches']].head())