reconstruct3d 0.1.0__py3-none-any.whl

reconstruct3d/track_sfm.py

@@ -0,0 +1,253 @@
+import os
+import sys
+import argparse
+import cv2
+import numpy as np
+from reconstruct3d.core import (
+    SfMDatabase, Features, create_frontend,
+    filter_triangulation,
+    DB_FILE, MAP_STATE_FILE, TRACKED_CLOUD_FILE,
+)
+from reconstruct3d.init_sfm import export_ply
+from reconstruct3d.bundle_adjust import run_local_ba, fuse_map_points
+
+sys.modules['__main__'].Features = Features
+
+def triangulate_new_points(k, ref, map3d, db, cam, frontend, min_angle=1.5):
+    """
+    Triangula nuevas correspondencias entre el frame k y un frame de referencia.
+    Impone restricciones epipolares, finitud espacial y quiralidad (Z > 0).
+    """
+    fk, f_ref = db.features[k], db.features[ref]
+    matches = frontend.match(fk, f_ref)
+    if not matches: return 0
+
+    pts_k = fk.pts[[m[0] for m in matches]]
+    pts_ref = f_ref.pts[[m[1] for m in matches]]
+
+    E, mask = cv2.findEssentialMat(pts_k, pts_ref, cam.K, cv2.RANSAC, 0.999, 1.0)
+    if E is None: return 0
+
+    valid_m = np.where(mask.ravel() > 0)[0]
+    new_matches = []
+
+    # Filtrar: Solo triangular si NINGUNO de los keypoints existe ya en el mapa 3D
+    for idx in valid_m:
+        k_idx, ref_idx = matches[idx]
+        if k_idx not in map3d['obs'].get(k, {}) and ref_idx not in map3d['obs'].get(ref, {}):
+            new_matches.append((k_idx, ref_idx))
+
+    if not new_matches: return 0
+
+    pts_k_new = np.float32([fk.pts[m[0]] for m in new_matches])
+    pts_ref_new = np.float32([f_ref.pts[m[1]] for m in new_matches])
+    colors_new = np.array([fk.colors[m[0]] for m in new_matches])
+
+    # Matrices de proyección P = K[R|t]
+    Pk = cam.K @ map3d['poses'][k][:3, :4]
+    Pref = cam.K @ map3d['poses'][ref][:3, :4]
+
+    pts4d = cv2.triangulatePoints(Pk, Pref, pts_k_new.T, pts_ref_new.T)
+    pts3d = (pts4d[:3, :] / pts4d[3, :]).T
+
+    valid_mask = filter_triangulation(
+        pts3d, map3d['poses'][k], map3d['poses'][ref],
+        pts_k_new, pts_ref_new, cam.K,
+        max_reproj_err=6.0, min_angle_deg=min_angle,
+    )
+
+    valid_pts = pts3d[valid_mask]
+    valid_colors = colors_new[valid_mask]
+    valid_matches = [new_matches[i] for i, v in enumerate(valid_mask) if v]
+
+    # Inyectar al mapa global
+    if len(valid_pts) > 0:
+        start_idx = len(map3d['points'])
+        map3d['points'] = np.vstack((map3d['points'], valid_pts))
+        map3d['colors'] = np.vstack((map3d['colors'], valid_colors))
+
+        for i, (k_idx, ref_idx) in enumerate(valid_matches):
+            pt_idx = start_idx + i
+            map3d['obs'][k][k_idx] = pt_idx
+            map3d['obs'][ref][ref_idx] = pt_idx
+
+    return len(valid_pts)
+
+def track_frames(out_dir, m_window=3, min_angle=1.5, local_ba=True, ba_every=1, fuse=True):
+    db = SfMDatabase()
+    db.load(os.path.join(out_dir, DB_FILE))
+    map_path = os.path.join(out_dir, MAP_STATE_FILE)
+    map3d = np.load(map_path, allow_pickle=True).item()
+    # Si el map_state es viejo y no trae 'frontend'/'camera', heredarlo de la DB
+    if 'frontend' not in map3d:
+        map3d['frontend'] = db.frontend_name
+
+    cam = db.get_camera()  # intrínsecos persistidos (consistentes con extract/init)
+    if 'camera' not in map3d:
+        map3d['camera'] = cam.to_dict()
+    print(f"[*] Frontend activo: '{db.frontend_name}'")
+    frontend = create_frontend(db.frontend_name)
+
+    frames_all = sorted(db.features.keys())
+    registered = set(map3d['poses'].keys())
+
+    # Espaciado real entre frames extraídos (= k_skip usado en core.py). Se
+    # deriva de los datos para no tener que mantenerlo sincronizado a mano: la
+    # ventana de vecinos para PnP es ±m_window frames extraídos.
+    diffs = np.diff(frames_all)
+    step = int(np.median(diffs)) if len(diffs) else 1
+    window = step * m_window
+    print(f"[*] Espaciado detectado: {step} frames  |  ventana PnP: ±{window} frames ({m_window} vecinos)")
+
+    ba_msg = f"BA local cada {ba_every} reg." if local_ba else "OFF"
+    print(f"[*] Iniciando tracking incremental. Camaras base: {len(registered)} | Puntos: {len(map3d['points'])} | {ba_msg}")
+    pts_of = lambda f: db.features[f].pts
+    reg_count = 0
+
+    while True:
+        candidates = [f for f in frames_all if f not in registered and any(abs(f - r) <= window for r in registered)]
+        if not candidates:
+            print("[!] No hay más candidatos viables en la ventana. Fin del tracking.")
+            break
+
+        # Registrar primero el candidato MÁS CERCANO a una cámara ya registrada:
+        # baseline corto + mucho solape = PnP bien condicionado y mínimo arrastre
+        # de error. Procesarlos por índice de frame registraba saltos largos
+        # primero (geometría débil) y propagaba el error frame a frame.
+        candidates.sort(key=lambda f: min(abs(f - r) for r in registered))
+
+        progress_made = False
+
+        for k in candidates:
+            print(f"  -> Evaluando frame {k:04d}...", end="\r")
+            fk = db.features[k]
+            obs_dict = {}
+
+            ref_frames = [r for r in registered if abs(k - r) <= window]
+            for ref in ref_frames:
+                f_ref = db.features[ref]
+                matches = frontend.match(fk, f_ref)
+                if not matches: continue
+
+                pts_k = fk.pts[[m[0] for m in matches]]
+                pts_ref = f_ref.pts[[m[1] for m in matches]]
+                E, mask = cv2.findEssentialMat(pts_k, pts_ref, cam.K, cv2.RANSAC, 0.999, 1.0)
+                if E is None: continue
+
+                valid_m = np.where(mask.ravel() > 0)[0]
+                for idx in valid_m:
+                    k_idx, ref_idx = matches[idx]
+                    if ref_idx in map3d['obs'].get(ref, {}):
+                        obs_dict[k_idx] = map3d['obs'][ref][ref_idx]
+
+            if len(obs_dict) < 8:
+                continue
+
+            k_indices = list(obs_dict.keys())
+            p2d_candidate = np.array([fk.pts[idx] for idx in k_indices], dtype=np.float32)
+            p3d_candidate = np.array([map3d['points'][obs_dict[idx]] for idx in k_indices], dtype=np.float32)
+
+            # Semilla: pose de la cámara registrada más cercana. Como ahora
+            # registramos de cerca a lejos, es un guess muy bueno y estabiliza PnP.
+            r_near = min(registered, key=lambda r: abs(k - r))
+            rvec0, _ = cv2.Rodrigues(map3d['poses'][r_near][:3, :3])
+            tvec0 = map3d['poses'][r_near][:3, 3].reshape(3, 1).astype(np.float64)
+
+            succ, rvec, tvec, inliers = cv2.solvePnPRansac(
+                p3d_candidate, p2d_candidate, cam.K, None,
+                rvec=rvec0.copy(), tvec=tvec0.copy(), useExtrinsicGuess=True,
+                reprojectionError=3.0, iterationsCount=100, confidence=0.999,
+            )
+
+            if succ and inliers is not None and len(inliers) > 8:
+                # Refinar la pose con TODOS los inliers (LM): RANSAC solo ajusta
+                # sobre la muestra mínima, esto minimiza la reproyección global.
+                inl = inliers.ravel()
+                rvec, tvec = cv2.solvePnPRefineLM(
+                    p3d_candidate[inl], p2d_candidate[inl], cam.K, None, rvec, tvec)
+                R, _ = cv2.Rodrigues(rvec)
+                map3d['poses'][k] = np.vstack((np.hstack((R, tvec)), [0,0,0,1]))
+                registered.add(k)
+
+                if k not in map3d['obs']: map3d['obs'][k] = {}
+                for i_inlier in inliers.ravel():
+                    idx_2d = k_indices[i_inlier]
+                    map3d['obs'][k][idx_2d] = obs_dict[idx_2d]
+
+                new_pts_total = 0
+                for ref in ref_frames:
+                    new_pts_total += triangulate_new_points(k, ref, map3d, db, cam, frontend, min_angle=min_angle)
+
+                # BA LOCAL: re-optimiza ventana (k + vecinos) + sus puntos justo
+                # tras insertarlos, antes de que el error se propague al siguiente
+                # frame. Fija ancla + cámaras externas (ataca #2 arrastre y #4 lazo).
+                reg_count += 1
+                if local_ba and reg_count % ba_every == 0:
+                    win = [r for r in registered if abs(k - r) <= window]
+                    run_local_ba(map3d, pts_of, cam.K, win, max_nfev=20)
+
+                print(f"[+] Frame {k:04d} registrado | PnP Inliers: {len(inliers):03d} | Nuevos Puntos 3D: {new_pts_total}")
+
+                np.save(map_path, map3d)
+
+                # --- FASE DE EXPORTACIÓN CON CENTRADO ESPACIAL ---
+                # Calcular la mediana ignorando potenciales (aunque improbables) NaNs residuales
+                valid_global = np.isfinite(map3d['points']).all(axis=1)
+                if np.sum(valid_global) > 0:
+                    median_pt = np.median(map3d['points'][valid_global], axis=0)
+                else:
+                    median_pt = np.zeros(3)
+
+                centered_points = map3d['points'][valid_global] - median_pt
+                valid_colors = map3d['colors'][valid_global]
+
+                cam_centers = [-pose[:3, :3].T @ pose[:3, 3:] for pose in map3d['poses'].values()]
+                centered_cams = [c.ravel() - median_pt for c in cam_centers]
+
+                export_ply(os.path.join(out_dir, TRACKED_CLOUD_FILE), centered_points, valid_colors, centered_cams)
+
+                progress_made = True
+                break
+
+        if not progress_made:
+            print(f"\n[!] Bucle estancado. No se lograron registrar más cámaras de los {len(candidates)} candidatos.")
+            break
+
+    # FUSIÓN DE PUNTOS (#1 asociación de datos): une duplicados del mismo punto
+    # físico triangulado por separado en frames distintos (tracks disjuntos),
+    # acortando la nube y reforzando restricciones. Conservador (vóxel fino).
+    if fuse:
+        valid = np.isfinite(map3d['points']).all(axis=1)
+        if valid.sum() > 0:
+            scale = float(np.median(np.linalg.norm(
+                map3d['points'][valid] - np.median(map3d['points'][valid], axis=0), axis=1)))
+            voxel = scale / 800.0
+            n_fused = fuse_map_points(map3d, voxel)
+            print(f"[*] Fusión de puntos: {n_fused} duplicados unidos (vóxel={voxel:.4f}).")
+            np.save(map_path, map3d)
+            valid = np.isfinite(map3d['points']).all(axis=1)
+            median_pt = np.median(map3d['points'][valid], axis=0)
+            cam_centers = [-(p[:3, :3].T @ p[:3, 3:]).ravel() - median_pt for p in map3d['poses'].values()]
+            export_ply(os.path.join(out_dir, TRACKED_CLOUD_FILE),
+                       map3d['points'][valid] - median_pt, map3d['colors'][valid], cam_centers)
+
+    print(f"\nTracking finalizado. Cámaras: {len(map3d['poses'])} | Puntos totales: {len(map3d['points'])}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="SfM track: incremental camera registration")
+    parser.add_argument("--out", required=True, help="Output dir from core/init runs.")
+    parser.add_argument("--m-window", type=int, default=3,
+                        help="Ventana PnP: registrar usando ±m frames extraídos como referencia "
+                             "(default: 3). El espaciado real (k_skip) se autodetecta de los datos.")
+    parser.add_argument("--min-angle", type=float, default=1.5,
+                        help="Ángulo mínimo de triangulación en grados (default: 1.5). Más alto "
+                             "descarta puntos de baseline corta que generan ruido de escala.")
+    parser.add_argument("--no-local-ba", action="store_true",
+                        help="Desactiva el Bundle Adjustment local por ventana deslizante.")
+    parser.add_argument("--ba-every", type=int, default=1,
+                        help="Ejecutar BA local cada N cámaras registradas (default: 1 = siempre).")
+    parser.add_argument("--no-fuse", action="store_true",
+                        help="Desactiva la fusión de puntos duplicados al final del tracking.")
+    args = parser.parse_args()
+    track_frames(args.out, m_window=args.m_window, min_angle=args.min_angle,
+                 local_ba=not args.no_local_ba, ba_every=args.ba_every, fuse=not args.no_fuse)

reconstruct3d/viewer.py

@@ -0,0 +1,108 @@
+import os
+import sys
+import argparse
+import cv2
+import numpy as np
+from reconstruct3d.core import CameraConfig, MAP_STATE_FILE
+
+def main(video_path, out_dir):
+    map3d = np.load(os.path.join(out_dir, MAP_STATE_FILE), allow_pickle=True).item()
+    cam = CameraConfig.from_dict(map3d.get('camera'))  # intrínsecos del map_state
+    
+    points3d = map3d['points']
+    colors3d = map3d['colors']
+    
+    cap = cv2.VideoCapture(video_path)
+    registered_frames = sorted(map3d['poses'].keys())
+    num_frames = len(registered_frames)
+    
+    window_name = "POV Proyeccion 3D"
+    cv2.namedWindow(window_name)
+    
+    # Trackbars para Opacidad y Línea de tiempo
+    cv2.createTrackbar("Fondo %", window_name, 50, 100, lambda x: None)
+    cv2.createTrackbar("Frame", window_name, 0, max(0, num_frames - 1), lambda x: None)
+
+    # En Windows la ventana y sus trackbars no quedan registradas hasta que el
+    # event loop de HighGUI hace un ciclo; sin este "pump" inicial el primer
+    # getTrackbarPos lanza "NULL window".
+    cv2.imshow(window_name, np.zeros((10, 10, 3), dtype=np.uint8))
+    cv2.waitKey(1)
+
+    # Variables de estado para evitar re-renderizados innecesarios
+    last_idx = -1
+    last_alpha = -1
+    
+    while True:
+        # 1. Leer estado actual de los sliders
+        idx = cv2.getTrackbarPos("Frame", window_name)
+        alpha_val = cv2.getTrackbarPos("Fondo %", window_name)
+        
+        # 2. Solo procesar si el usuario movió algún control
+        if idx != last_idx or alpha_val != last_alpha:
+            f_idx = registered_frames[idx]
+            alpha = alpha_val / 100.0
+            
+            cap.set(cv2.CAP_PROP_POS_FRAMES, f_idx)
+            ret, frame = cap.read()
+            if not ret: 
+                break
+            
+            frame = cam.preprocess(frame)
+            h, w = frame.shape[:2]
+            
+            pose = map3d['poses'][f_idx]
+            R, t = pose[:3, :3], pose[:3, 3]
+            
+            # Transformación isométrica al espacio de cámara local
+            pts_c = (R @ points3d.T).T + t
+            
+            # Restricción de quiralidad (Z > 0)
+            mask_z = pts_c[:, 2] > 0
+            pts_c_valid = pts_c[mask_z]
+            colors_valid = colors3d[mask_z]
+            
+            # Composición de fondo
+            black_bg = np.zeros_like(frame)
+            view = cv2.addWeighted(frame, alpha, black_bg, 1 - alpha, 0)
+            
+            if len(pts_c_valid) > 0:
+                # Proyección en plano imagen
+                fx, fy = cam.K[0, 0], cam.K[1, 1]
+                cx, cy = cam.K[0, 2], cam.K[1, 2]
+                
+                u = (fx * (pts_c_valid[:, 0] / pts_c_valid[:, 2]) + cx).astype(int)
+                v = (fy * (pts_c_valid[:, 1] / pts_c_valid[:, 2]) + cy).astype(int)
+                
+                # Filtrado de límites del sensor
+                valid_img = (u >= 0) & (u < w) & (v >= 0) & (v < h)
+                u, v = u[valid_img], v[valid_img]
+                final_colors = colors_valid[valid_img]
+                
+                # Renderizado de inliers con color original
+                for pt_u, pt_v, color in zip(u, v, final_colors):
+                    c = (int(color[0]), int(color[1]), int(color[2]))
+                    cv2.circle(view, (pt_u, pt_v), 2, c, -1)
+                    
+            # Actualizar título de la ventana con el frame real
+            cv2.setWindowTitle(window_name, f"POV Proyeccion 3D | Frame de Video: {f_idx}")
+            cv2.imshow(window_name, view)
+            
+            # Guardar estado
+            last_idx = idx
+            last_alpha = alpha_val
+        
+        # Mantener el bucle vivo y escuchar la tecla 'q' o 'ESC' para salir
+        key = cv2.waitKey(30) & 0xFF
+        if key == ord('q') or key == 27:
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="POV viewer over reconstructed map")
+    parser.add_argument("video", help="Path to input video (mp4)")
+    parser.add_argument("--out", required=True, help="Output dir containing map_state.npy (from track_sfm.py).")
+    args = parser.parse_args()
+    main(args.video, args.out)