scripts/tests/test_smart_pairing.py

#!/usr/bin/env python3
"""
Quick test of smart pairing optimization.
Uses pre-computed poses to skip DA3 inference.
"""

import logging
import sys
from pathlib import Path
import numpy as np

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from ylff.services.ba_validator import BAValidator  # noqa: E402

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def test_smart_pairing():
    """Test smart pairing with a small set of images."""

    # Use existing images from previous run
    image_dir = project_root / "data" / "ba_validation_results_rickroll_v2" / "ba_work" / "images"

    if not image_dir.exists():
        logger.error(f"Image directory not found: {image_dir}")
        logger.info("Please run BA validation first to generate images")
        return

    image_paths = sorted(list(image_dir.glob("*.jpg")))[:10]  # Use first 10 images
    image_paths = [str(p) for p in image_paths]

    logger.info(f"Testing with {len(image_paths)} images")

    # Create dummy poses (simulate DA3 output)
    # For a video, poses should be relatively close together
    np.random.seed(42)
    base_pose = np.eye(4)[:3, :]
    poses = []
    for i in range(len(image_paths)):
        # Small random translation (simulating camera movement)
        pose = base_pose.copy()
        pose[:3, 3] = np.random.randn(3) * 0.1 * i  # Gradually move
        poses.append(pose)
    poses = np.array(poses)

    logger.info(f"Generated {len(poses)} poses")

    # Initialize validator
    validator = BAValidator(
        work_dir=project_root / "data" / "test_smart_pairing",
    )

    # Test pair generation
    logger.info("\n=== Testing Pair Generation ===")

    # Sequential
    pairs_seq = validator._generate_smart_pairs(image_paths, sequential_only=True)
    logger.info(f"Sequential pairs: {len(pairs_seq)}")

    # Spatial
    pairs_spatial = validator._generate_smart_pairs(
        image_paths,
        poses=poses,
        max_baseline=0.3,
        min_baseline=0.05,
        max_pairs_per_image=5,
    )
    logger.info(f"Spatial pairs: {len(pairs_spatial)}")

    # Exhaustive
    pairs_exhaustive = validator._generate_smart_pairs(image_paths)
    logger.info(f"Exhaustive pairs: {len(pairs_exhaustive)}")

    logger.info("\n=== Speedup Analysis ===")
    logger.info(f"Sequential: {len(pairs_exhaustive) / len(pairs_seq):.1f}x fewer pairs")
    logger.info(f"Spatial: {len(pairs_exhaustive) / len(pairs_spatial):.1f}x fewer pairs")

    # Test actual matching (if features exist)
    features_path = (
        project_root / "data" / "ba_validation_results_rickroll_v2" / "ba_work" / "features.h5"
    )

    if features_path.exists():
        logger.info("\n=== Testing Matching with Smart Pairs ===")

        import time
        from hloc import match_features

        # Test with sequential pairs
        pairs_file_seq = validator.work_dir / "pairs_seq.txt"
        with open(pairs_file_seq, "w") as f:
            for img1, img2 in pairs_seq:
                f.write(f"{Path(img1).name} {Path(img2).name}\n")

        matches_file_seq = validator.work_dir / "matches_seq.h5"

        logger.info(f"Matching {len(pairs_seq)} sequential pairs...")
        start = time.time()

        try:
            match_conf = match_features.confs["superpoint+lightglue"]
            match_features.main(
                conf=match_conf,
                pairs=pairs_file_seq,
                features=features_path,
                matches=matches_file_seq,
            )
            elapsed_seq = time.time() - start
            logger.info(f"✓ Sequential matching completed in {elapsed_seq:.2f}s")
        except Exception as e:
            logger.error(f"Matching failed: {e}")
            elapsed_seq = None

        # Test with exhaustive pairs (if we have time)
        if len(pairs_exhaustive) < 50:  # Only if reasonable
            pairs_file_exh = validator.work_dir / "pairs_exh.txt"
            with open(pairs_file_exh, "w") as f:
                for img1, img2 in pairs_exhaustive:
                    f.write(f"{Path(img1).name} {Path(img2).name}\n")

            matches_file_exh = validator.work_dir / "matches_exh.h5"

            logger.info(f"Matching {len(pairs_exhaustive)} exhaustive pairs...")
            start = time.time()

            try:
                match_features.main(
                    conf=match_conf,
                    pairs=pairs_file_exh,
                    features=features_path,
                    matches=matches_file_exh,
                )
                elapsed_exh = time.time() - start
                logger.info(f"✓ Exhaustive matching completed in {elapsed_exh:.2f}s")

                if elapsed_seq:
                    speedup = elapsed_exh / elapsed_seq
                    logger.info(f"\n=== Speedup: {speedup:.1f}x ===")
            except Exception as e:
                logger.error(f"Matching failed: {e}")
    else:
        logger.info(f"\nFeatures not found at {features_path}")
        logger.info("Skipping matching test. Run full BA validation first.")

    logger.info("\n=== Test Complete ===")


if __name__ == "__main__":
    test_smart_pairing()