From 027a33233e4cd3623671b31f0db0b88260943bb0 Mon Sep 17 00:00:00 2001 From: cDc Date: Thu, 15 Jan 2026 19:42:55 +0200 Subject: [PATCH 1/8] add image pairs processing script --- README.md | 21 + examples/example_pairs.csv | 2 + examples/export_point_cloud.py | 360 +++++++++++++ examples/generate_pairs.py | 236 +++++++++ examples/load_results.py | 187 +++++++ scripts/batch_match.py | 920 +++++++++++++++++++++++++++++++++ 6 files changed, 1726 insertions(+) create mode 100644 examples/example_pairs.csv create mode 100755 examples/export_point_cloud.py create mode 100644 examples/generate_pairs.py create mode 100644 examples/load_results.py create mode 100644 scripts/batch_match.py diff --git a/README.md b/README.md index d5b695b..24af0eb 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,27 @@ F, mask = cv2.findFundamentalMat( ``` We additionally provide two demos in the [demos folder](demo), which might help in understanding. +## Batch Matching +For efficiently processing large numbers of image pairs, use the batch matching script: + +```bash +# Create a CSV file with image pairs (image_A_path,image_B_path) +# Then run batch matching +python scripts/batch_match.py \ + --pairs-csv pairs.csv \ + --output-dir outputs/matches/ \ + --cache-images \ + --num-samples 5000 + +# Results are saved as NPZ files with dense warps and optional sampled matches +# See examples/README_batch_matching.md for detailed documentation +``` + +Key features: +- **Intelligent caching**: Automatically reorders pairs and caches images to maximize throughput +- **Resume capability**: Skips already processed pairs, safe to interrupt and restart +- **Flexible output**: Dense warp fields + optional sampled keypoint matches +- **Multiple quality settings**: From `turbo` (fastest) to `precise` (highest quality) ## Setup/Install In your python environment (tested on Linux python 3.12), run: diff --git a/examples/example_pairs.csv b/examples/example_pairs.csv new file mode 100644 index 0000000..f2f7eb8 --- /dev/null +++ b/examples/example_pairs.csv @@ -0,0 +1,2 @@ +image_A_path,image_B_path +../assets/toronto_A.jpg,../assets/toronto_B.jpg diff --git a/examples/export_point_cloud.py b/examples/export_point_cloud.py new file mode 100755 index 0000000..438e0cb --- /dev/null +++ b/examples/export_point_cloud.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 +""" +Export RoMaV2 matches as a 3D point cloud visualization. + +This script loads matched image pairs and their dense correspondences, +then creates a 3D point cloud where: +- Points from image A are positioned at z=0 +- Points from image B are positioned at z=1 +- Corresponding points are connected with lines +- Points are colored by their RGB values from the images + +Usage: + python examples/export_point_cloud.py \ + --match outputs/matches/toronto_A-toronto_B.npz \ + --output visualizations/toronto_match.ply \ + --subsample 5000 +""" + +import argparse +import logging +from pathlib import Path + +import numpy as np +from PIL import Image + +logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + + +def load_image(path: str) -> np.ndarray: + """Load image as RGB numpy array. + + Returns: + np.ndarray: (H, W, 3) array with values in [0, 255] + """ + img = Image.open(path).convert('RGB') + return np.array(img) + + +def bilinear_sample(image: np.ndarray, coords: np.ndarray) -> np.ndarray: + """Sample image at floating-point coordinates using bilinear interpolation. + + Args: + image: (H, W, C) image array + coords: (N, 2) array of (x, y) coordinates + + Returns: + (N, C) sampled values + """ + H, W = image.shape[:2] + x = coords[:, 0] + y = coords[:, 1] + + # Clamp to valid range + x = np.clip(x, 0, W - 1) + y = np.clip(y, 0, H - 1) + + # Get integer parts + x0 = np.floor(x).astype(int) + y0 = np.floor(y).astype(int) + x1 = np.minimum(x0 + 1, W - 1) + y1 = np.minimum(y0 + 1, H - 1) + + # Get fractional parts + fx = x - x0 + fy = y - y0 + + # Bilinear interpolation + fx = fx[:, None] + fy = fy[:, None] + + values = ( + image[y0, x0] * (1 - fx) * (1 - fy) + + image[y0, x1] * fx * (1 - fy) + + image[y1, x0] * (1 - fx) * fy + + image[y1, x1] * fx * fy + ) + + return values + + +def warp_to_correspondences( + warp_AB: np.ndarray, + confidence_AB: np.ndarray, + H_A: int, + W_A: int, + H_B: int, + W_B: int, + subsample: int | None = None, + min_confidence: float = 0.5 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Convert warp field to point correspondences. + + Args: + warp_AB: (H, W, 2) warp field in normalized [-1, 1] coordinates + confidence_AB: (H, W, 4) confidence (overlap logit + precision) + H_A, W_A: Original dimensions of image A + H_B, W_B: Original dimensions of image B + subsample: Maximum number of points to sample (None for all) + min_confidence: Minimum overlap confidence threshold + + Returns: + tuple: (points_A, points_B, confidences) + - points_A: (N, 2) pixel coordinates in image A + - points_B: (N, 2) pixel coordinates in image B + - confidences: (N,) overlap confidence values + """ + H_warp, W_warp = warp_AB.shape[:2] + + # Create grid of coordinates in image A + y_grid, x_grid = np.meshgrid( + np.arange(H_warp), + np.arange(W_warp), + indexing='ij' + ) + + # Convert to pixel coordinates in original image A + points_A = np.stack([ + x_grid.flatten() / W_warp * W_A, + y_grid.flatten() / H_warp * H_A + ], axis=1) + + # Get corresponding points in B from warp field + warp_coords = warp_AB.reshape(-1, 2) # (N, 2) in [-1, 1] + + # Convert from [-1, 1] to pixel coordinates in image B + points_B = np.stack([ + (warp_coords[:, 0] + 1) * W_B / 2, + (warp_coords[:, 1] + 1) * H_B / 2 + ], axis=1) + + # Get confidence (sigmoid of overlap logit) + overlap_logits = confidence_AB[..., 0].flatten() + confidences = 1 / (1 + np.exp(-overlap_logits)) + + # Filter by confidence + valid_mask = confidences > min_confidence + points_A = points_A[valid_mask] + points_B = points_B[valid_mask] + confidences = confidences[valid_mask] + + logger.info(f"Found {len(points_A)} valid correspondences (confidence > {min_confidence})") + + # Subsample if requested + if subsample is not None and len(points_A) > subsample: + # Sample with probability proportional to confidence + probs = confidences / confidences.sum() + indices = np.random.choice(len(points_A), size=subsample, replace=False, p=probs) + points_A = points_A[indices] + points_B = points_B[indices] + confidences = confidences[indices] + logger.info(f"Subsampled to {subsample} points") + + return points_A, points_B, confidences + + +def export_ply( + output_path: Path, + points_A: np.ndarray, + points_B: np.ndarray, + colors_A: np.ndarray, + colors_B: np.ndarray, + scale: float = 1.0, + spacing: float = 1.0, + include_lines: bool = True +): + """Export point cloud to PLY format. + + Args: + output_path: Output PLY file path + points_A: (N, 2) pixel coordinates in image A + points_B: (N, 2) pixel coordinates in image B + colors_A: (N, 3) RGB colors for points A + colors_B: (N, 3) RGB colors for points B + scale: Scale factor for coordinates + spacing: Z-spacing between the two image planes + include_lines: Whether to include line elements connecting correspondences + """ + N = len(points_A) + + # Create 3D coordinates + # Image A at z=0, Image B at z=spacing + vertices_A = np.column_stack([ + points_A[:, 0] * scale, + points_A[:, 1] * scale, + np.zeros(N) + ]) + + vertices_B = np.column_stack([ + points_B[:, 0] * scale, + points_B[:, 1] * scale, + np.full(N, spacing) + ]) + + # Combine vertices + vertices = np.vstack([vertices_A, vertices_B]) + colors = np.vstack([colors_A, colors_B]).astype(np.uint8) + + total_vertices = len(vertices) + + # Write PLY file + with open(output_path, 'w') as f: + # Header + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write(f"element vertex {total_vertices}\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + f.write("property uchar red\n") + f.write("property uchar green\n") + f.write("property uchar blue\n") + + if include_lines: + f.write(f"element edge {N}\n") + f.write("property int vertex1\n") + f.write("property int vertex2\n") + + f.write("end_header\n") + + # Vertices + for i in range(total_vertices): + f.write(f"{vertices[i, 0]:.6f} {vertices[i, 1]:.6f} {vertices[i, 2]:.6f} ") + f.write(f"{colors[i, 0]} {colors[i, 1]} {colors[i, 2]}\n") + + # Edges (correspondences) + if include_lines: + for i in range(N): + f.write(f"{i} {i + N}\n") + + logger.info(f"Saved PLY with {total_vertices} vertices and {N if include_lines else 0} edges to {output_path}") + + +def main(): + parser = argparse.ArgumentParser( + description="Export RoMaV2 matches as 3D point cloud", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Basic export + python examples/export_point_cloud.py \\ + --match outputs/matches/toronto_A-toronto_B.npz \\ + --output visualizations/toronto_match.ply + + # High-density point cloud + python examples/export_point_cloud.py \\ + --match outputs/matches/toronto_A-toronto_B.npz \\ + --output visualizations/toronto_dense.ply \\ + --subsample 20000 --min-confidence 0.3 + + # Points only (no correspondence lines) + python examples/export_point_cloud.py \\ + --match outputs/matches/toronto_A-toronto_B.npz \\ + --output visualizations/toronto_points.ply \\ + --no-lines + """ + ) + + parser.add_argument( + '--match', + type=Path, + required=True, + help='Path to NPZ file with matching results' + ) + parser.add_argument( + '--output', + type=Path, + required=True, + help='Output PLY file path' + ) + parser.add_argument( + '--subsample', + type=int, + default=5000, + help='Maximum number of correspondences to export (default: 5000)' + ) + parser.add_argument( + '--min-confidence', + type=float, + default=0.5, + help='Minimum overlap confidence threshold (default: 0.5)' + ) + parser.add_argument( + '--scale', + type=float, + default=0.01, + help='Coordinate scale factor (default: 0.01)' + ) + parser.add_argument( + '--spacing', + type=float, + default=100.0, + help='Z-spacing between image planes (default: 100.0)' + ) + parser.add_argument( + '--no-lines', + action='store_true', + help='Disable correspondence lines (points only)' + ) + + args = parser.parse_args() + + # Load match data + logger.info(f"Loading match data from {args.match}") + data = np.load(args.match, allow_pickle=True) + + warp_AB = data['warp_AB'] + confidence_AB = data['confidence_AB'] + H_A, W_A = data['image_A_shape'] + H_B, W_B = data['image_B_shape'] + path_A = str(data['image_A_path']) + path_B = str(data['image_B_path']) + + logger.info(f"Image A: {path_A} ({W_A}x{H_A})") + logger.info(f"Image B: {path_B} ({W_B}x{H_B})") + logger.info(f"Warp resolution: {warp_AB.shape[1]}x{warp_AB.shape[0]}") + + # Load images + logger.info("Loading images...") + img_A = load_image(path_A) + img_B = load_image(path_B) + + # Convert warp to correspondences + logger.info("Extracting correspondences from warp field...") + points_A, points_B, confidences = warp_to_correspondences( + warp_AB=warp_AB, + confidence_AB=confidence_AB, + H_A=H_A, + W_A=W_A, + H_B=H_B, + W_B=W_B, + subsample=args.subsample, + min_confidence=args.min_confidence + ) + + # Sample colors from images + logger.info("Sampling colors from images...") + colors_A = bilinear_sample(img_A, points_A) + colors_B = bilinear_sample(img_B, points_B) + + # Export to PLY + logger.info("Exporting point cloud...") + args.output.parent.mkdir(parents=True, exist_ok=True) + export_ply( + output_path=args.output, + points_A=points_A, + points_B=points_B, + colors_A=colors_A, + colors_B=colors_B, + scale=args.scale, + spacing=args.spacing, + include_lines=not args.no_lines + ) + + logger.info("Done!") + return 0 + + +if __name__ == "__main__": + exit(main()) diff --git a/examples/generate_pairs.py b/examples/generate_pairs.py new file mode 100644 index 0000000..2e8500b --- /dev/null +++ b/examples/generate_pairs.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +""" +Helper script to generate a pairs CSV file from image directories. + +This script helps create the input CSV file needed by batch_match.py. +""" + +import argparse +from pathlib import Path +from itertools import combinations +from typing import Optional + + +def generate_all_pairs(image_dir: Path, extensions=None) -> list[tuple[Path, Path]]: + """Generate all possible pairs from images in a directory. + + Args: + image_dir: Directory containing images + extensions: List of valid extensions (default: ['.jpg', '.jpeg', '.png']) + + Returns: + List of (path_A, path_B) tuples + """ + if extensions is None: + extensions = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'] + + # Find all images + images = [] + for ext in extensions: + images.extend(image_dir.glob(f'**/*{ext}')) + + images = sorted(set(images)) + print(f"Found {len(images)} images in {image_dir}") + + # Generate all pairs + pairs = list(combinations(images, 2)) + print(f"Generated {len(pairs)} pairs") + + return pairs + + +def generate_sequential_pairs(image_dir: Path, extensions=None, overlap=1) -> list[tuple[Path, Path]]: + """Generate sequential pairs (e.g., for video frames or ordered images). + + Args: + image_dir: Directory containing images + extensions: List of valid extensions + overlap: Number of following images to pair with each image (default: 1) + + Returns: + List of (path_A, path_B) tuples + """ + if extensions is None: + extensions = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'] + + # Find all images + images = [] + for ext in extensions: + images.extend(image_dir.glob(f'**/*{ext}')) + + images = sorted(set(images)) + print(f"Found {len(images)} images in {image_dir}") + + # Generate sequential pairs + pairs = [] + for i in range(len(images)): + for j in range(1, overlap + 1): + if i + j < len(images): + pairs.append((images[i], images[i + j])) + + print(f"Generated {len(pairs)} sequential pairs (overlap={overlap})") + + return pairs + + +def generate_pairs_from_list(list_file: Path, base_dir: Optional[Path] = None) -> list[tuple[Path, Path]]: + """Generate pairs from a text file listing image paths (one per line). + + Creates all combinations of listed images. + + Args: + list_file: Text file with one image path per line + base_dir: Optional base directory to resolve relative paths + + Returns: + List of (path_A, path_B) tuples + """ + with open(list_file, 'r') as f: + lines = [line.strip() for line in f if line.strip()] + + images = [] + for line in lines: + if base_dir: + img_path = base_dir / line + else: + img_path = Path(line) + + if img_path.exists(): + images.append(img_path) + else: + print(f"Warning: Image not found: {img_path}") + + print(f"Found {len(images)} valid images from {list_file}") + + # Generate all pairs + pairs = list(combinations(images, 2)) + print(f"Generated {len(pairs)} pairs") + + return pairs + + +def save_pairs_csv(pairs: list[tuple[Path, Path]], output_file: Path, relative_to: Optional[Path] = None): + """Save pairs to CSV file. + + Args: + pairs: List of (path_A, path_B) tuples + output_file: Output CSV file path + relative_to: Optional directory to make paths relative to + """ + with open(output_file, 'w') as f: + # Write header + f.write("image_A_path,image_B_path\n") + + # Write pairs + for path_A, path_B in pairs: + if relative_to: + try: + path_A = path_A.relative_to(relative_to) + path_B = path_B.relative_to(relative_to) + except ValueError: + pass # Keep absolute if can't make relative + + f.write(f"{path_A},{path_B}\n") + + print(f"\nSaved {len(pairs)} pairs to {output_file}") + + +def main(): + parser = argparse.ArgumentParser( + description="Generate pairs CSV for batch matching", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" + Examples: + # All pairs from a directory + python examples/generate_pairs.py --image-dir data/images/ --output pairs.csv --mode all + + # Sequential pairs (for video frames) + python examples/generate_pairs.py --image-dir data/frames/ --output pairs.csv --mode sequential --overlap 5 + + # Pairs from a list file + python examples/generate_pairs.py --list-file images.txt --output pairs.csv --mode list + """ + ) + + parser.add_argument( + '--image-dir', + type=Path, + help='Directory containing images' + ) + parser.add_argument( + '--list-file', + type=Path, + help='Text file with image paths (one per line)' + ) + parser.add_argument( + '--output', + type=Path, + required=True, + help='Output CSV file path' + ) + parser.add_argument( + '--mode', + choices=['all', 'sequential', 'list'], + default='all', + help='Pairing mode (default: all)' + ) + parser.add_argument( + '--overlap', + type=int, + default=1, + help='For sequential mode: number of following images to pair with each image (default: 1)' + ) + parser.add_argument( + '--extensions', + nargs='+', + default=['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'], + help='Valid image extensions (default: .jpg .jpeg .png)' + ) + parser.add_argument( + '--relative', + action='store_true', + help='Use relative paths in output CSV' + ) + parser.add_argument( + '--max-pairs', + type=int, + default=None, + help='Maximum number of pairs to generate (useful for large directories)' + ) + + args = parser.parse_args() + + # Generate pairs based on mode + if args.mode == 'list': + if not args.list_file: + parser.error("--list-file required for mode 'list'") + base_dir = args.image_dir if args.image_dir else Path.cwd() + pairs = generate_pairs_from_list(args.list_file, base_dir) + elif args.mode == 'sequential': + if not args.image_dir: + parser.error("--image-dir required for mode 'sequential'") + pairs = generate_sequential_pairs(args.image_dir, args.extensions, args.overlap) + else: # mode == 'all' + if not args.image_dir: + parser.error("--image-dir required for mode 'all'") + pairs = generate_all_pairs(args.image_dir, args.extensions) + + # Limit pairs if requested + if args.max_pairs and len(pairs) > args.max_pairs: + print(f"Limiting to {args.max_pairs} pairs (from {len(pairs)})") + pairs = pairs[:args.max_pairs] + + # Determine relative path base + relative_to = None + if args.relative: + if args.output.parent.exists(): + relative_to = args.output.parent + else: + relative_to = Path.cwd() + + # Save to CSV + save_pairs_csv(pairs, args.output, relative_to) + + +if __name__ == "__main__": + main() diff --git a/examples/load_results.py b/examples/load_results.py new file mode 100644 index 0000000..1cf524d --- /dev/null +++ b/examples/load_results.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +""" +Example script demonstrating how to load and use batch matching results. + +This script shows how to: +1. Load NPZ output files from batch_match.py +2. Extract dense warp fields and confidence +3. Use sampled keypoint matches +4. Visualize results +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + + +def load_results(npz_path: Path) -> dict: + """Load batch matching results from NPZ file.""" + data = np.load(npz_path, allow_pickle=True) + + print(f"\nLoaded results from: {npz_path.name}") + print(f" Pair ID: {data['pair_id']}") + print(f" Setting: {data['setting']}") + print(f" Output shape: {data['output_shape']}") + print(f" Image A: {data['image_A_path']} (original size: {data['image_A_shape']})") + print(f" Image B: {data['image_B_path']} (original size: {data['image_B_shape']})") + + # Extract warp and confidence + warp_AB = data['warp_AB'] # (H, W, 2) in normalized coords [-1, 1] + confidence_AB = data['confidence_AB'] # (H, W, 4) + + # Compute overlap probability from confidence + overlap_logit = confidence_AB[..., 0] + overlap_prob = 1.0 / (1.0 + np.exp(-overlap_logit)) + + print(f"\nWarp field shape: {warp_AB.shape}") + print(f"Overlap range: [{overlap_prob.min():.3f}, {overlap_prob.max():.3f}]") + print(f"Mean overlap: {overlap_prob.mean():.3f}") + + # Check for sampled matches + if 'keypoints_A' in data: + kpts_A = data['keypoints_A'] + kpts_B = data['keypoints_B'] + matches_conf = data['matches_confidence'] + print(f"\nSampled matches: {len(kpts_A)}") + print(f"Match confidence range: [{matches_conf.min():.3f}, {matches_conf.max():.3f}]") + + return { + 'warp_AB': warp_AB, + 'confidence_AB': confidence_AB, + 'overlap_prob': overlap_prob, + 'image_A_path': str(data['image_A_path']), + 'image_B_path': str(data['image_B_path']), + 'image_A_shape': data['image_A_shape'], + 'image_B_shape': data['image_B_shape'], + 'output_shape': data['output_shape'], + 'keypoints_A': data.get('keypoints_A'), + 'keypoints_B': data.get('keypoints_B'), + 'matches_confidence': data.get('matches_confidence'), + } + + +def warp_image(results: dict, save_path: Path): + """Warp image B to image A and save visualization.""" + # Load images + img_A = Image.open(results['image_A_path']).convert('RGB') + img_B = Image.open(results['image_B_path']).convert('RGB') + + # Get output resolution + H_out, W_out = results['output_shape'] + + # Resize images to output resolution + img_A_resized = img_A.resize((W_out, H_out)) + img_B_resized = img_B.resize((W_out, H_out)) + + # Convert to tensors + img_A_tensor = torch.from_numpy(np.array(img_A_resized)).float() / 255.0 + img_B_tensor = torch.from_numpy(np.array(img_B_resized)).float() / 255.0 + + # Permute to (C, H, W) + img_A_tensor = img_A_tensor.permute(2, 0, 1) + img_B_tensor = img_B_tensor.permute(2, 0, 1) + + # Convert warp to tensor + warp_AB = torch.from_numpy(results['warp_AB']).float() + + # Warp image B to A + warped_B = F.grid_sample( + img_B_tensor[None], + warp_AB[None], + mode='bilinear', + align_corners=False + )[0] + + # Get overlap mask + overlap = torch.from_numpy(results['overlap_prob']).float() + + # Create visualization: overlap-masked warp + # White background for non-overlapping regions + white = torch.ones_like(warped_B) + vis = overlap[..., None] * warped_B.permute(1, 2, 0) + (1 - overlap[..., None]) * white.permute(1, 2, 0) + + # Convert to PIL and save + vis_np = (vis.numpy() * 255).astype(np.uint8) + vis_img = Image.fromarray(vis_np) + + # Create side-by-side comparison + combined = Image.new('RGB', (3 * W_out, H_out)) + combined.paste(img_A_resized, (0, 0)) + combined.paste(vis_img, (W_out, 0)) + combined.paste(img_B_resized, (2 * W_out, 0)) + + combined.save(save_path) + print(f"\nSaved warp visualization to: {save_path}") + print(" Layout: [Image A | Warped B (masked) | Image B]") + + +def print_match_stats(results: dict): + """Print statistics about sampled matches.""" + if results['keypoints_A'] is None: + print("\nNo sampled matches available (use --num-samples when running batch_match.py)") + return + + kpts_A = results['keypoints_A'] + kpts_B = results['keypoints_B'] + conf = results['matches_confidence'] + + print(f"\nSampled Match Statistics:") + print(f" Total matches: {len(kpts_A)}") + print(f" Confidence: min={conf.min():.3f}, max={conf.max():.3f}, mean={conf.mean():.3f}") + + # Compute match distances (normalized) + H_A, W_A = results['image_A_shape'] + H_B, W_B = results['image_B_shape'] + + # Normalize to [0, 1] + kpts_A_norm = kpts_A / np.array([W_A, H_A]) + kpts_B_norm = kpts_B / np.array([W_B, H_B]) + + distances = np.linalg.norm(kpts_A_norm - kpts_B_norm, axis=1) + print(f" Match distances (normalized): mean={distances.mean():.3f}, median={np.median(distances):.3f}") + + # Filter by confidence threshold + for threshold in [0.5, 0.7, 0.9]: + high_conf = conf > threshold + print(f" Matches with conf > {threshold}: {high_conf.sum()} ({100*high_conf.sum()/len(conf):.1f}%)") + + +def main(): + parser = argparse.ArgumentParser( + description="Load and analyze batch matching results", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + parser.add_argument( + 'npz_file', + type=Path, + help='Path to NPZ result file from batch_match.py' + ) + parser.add_argument( + '--visualize', + type=Path, + default=None, + help='Save warp visualization to this path' + ) + + args = parser.parse_args() + + # Load results + results = load_results(args.npz_file) + + # Print match statistics + print_match_stats(results) + + # Visualize if requested + if args.visualize: + warp_image(results, args.visualize) + + print("\n" + "="*60) + + +if __name__ == "__main__": + main() diff --git a/scripts/batch_match.py b/scripts/batch_match.py new file mode 100644 index 0000000..e73b8bf --- /dev/null +++ b/scripts/batch_match.py @@ -0,0 +1,920 @@ +#!/usr/bin/env python3 +""" +Batch matching script for RoMaV2. + +Efficiently matches a list of image pairs from a CSV file and saves results to disk. +Supports intelligent caching and pair ordering to maximize throughput. + +Usage: + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ + +CSV Format: + image_A_path,image_B_path + path/to/img1.jpg,path/to/img2.jpg + ... +""" + +import argparse +import csv +import logging +import time +from collections import OrderedDict, defaultdict +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +from tqdm import tqdm + +from romav2 import RoMaV2 + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + datefmt='%H:%M:%S' +) +logger = logging.getLogger(__name__) + + +class LRUImageCache: + """LRU cache for loaded images with original dimension tracking.""" + + def __init__(self, max_size: int = 100): + self.max_size = max_size + self.cache: OrderedDict[str, tuple[torch.Tensor, int, int]] = OrderedDict() + self.hits = 0 + self.misses = 0 + + def get(self, path: str, loader_fn) -> tuple[torch.Tensor, int, int]: + """Get image from cache or load it. + + Returns: + tuple: (image_tensor, H_original, W_original) + """ + if path in self.cache: + self.hits += 1 + # Move to end (most recently used) + self.cache.move_to_end(path) + return self.cache[path] + + self.misses += 1 + # Load image + img_tensor = loader_fn(path) + # Extract original dimensions (before any model resizing) + # img_tensor shape: (1, 3, H, W) or (3, H, W) + if img_tensor.dim() == 4: + _, _, H_orig, W_orig = img_tensor.shape + else: + _, H_orig, W_orig = img_tensor.shape + + # Add to cache + self.cache[path] = (img_tensor, H_orig, W_orig) + + # Evict oldest if over capacity + if len(self.cache) > self.max_size: + self.cache.popitem(last=False) + + return self.cache[path] + + def get_stats(self) -> dict: + """Get cache statistics.""" + total = self.hits + self.misses + hit_rate = self.hits / total if total > 0 else 0.0 + return { + 'hits': self.hits, + 'misses': self.misses, + 'hit_rate': hit_rate, + 'size': len(self.cache) + } + + +def load_pairs_csv(csv_path: Path) -> list[tuple[str, str, str]]: + """Load image pairs from CSV file. + + Args: + csv_path: Path to CSV file with format: image_A_path,image_B_path + + Returns: + List of tuples: (pair_id, path_A, path_B) + """ + pairs = [] + csv_dir = csv_path.parent + + with open(csv_path, 'r') as f: + reader = csv.reader(f) + for i, row in enumerate(reader, 1): + if len(row) < 2: + logger.warning(f"Line {i}: Expected 2 columns, got {len(row)}. Skipping.") + continue + + path_A, path_B, *_ = row + path_A = path_A.strip() + path_B = path_B.strip() + + # Skip empty lines or headers + if not path_A or not path_B or path_A.lower() == 'image_a_path' or path_A.lower() == 'imagea': + continue + + # Resolve relative paths + if not Path(path_A).is_absolute(): + path_A = str(csv_dir / path_A) + if not Path(path_B).is_absolute(): + path_B = str(csv_dir / path_B) + + # Generate pair_id from filenames + stem_A = Path(path_A).stem + stem_B = Path(path_B).stem + pair_id = f"{stem_A}-{stem_B}" + + pairs.append((pair_id, path_A, path_B)) + + logger.info(f"Loaded {len(pairs)} image pairs from {csv_path}") + return pairs + + +def optimize_pair_order(pairs: list[tuple[str, str, str]]) -> list[tuple[str, str, str]]: + """Reorder pairs to maximize cache hits by clustering shared images. + + Uses a simple greedy algorithm to group pairs that share images. + + Args: + pairs: List of (pair_id, path_A, path_B) + + Returns: + Reordered list of pairs + """ + if len(pairs) <= 1: + return pairs + + # Build adjacency graph: pairs that share at least one image + pair_indices = list(range(len(pairs))) + adjacency = defaultdict(list) + + for i in range(len(pairs)): + paths_i = {pairs[i][1], pairs[i][2]} + for j in range(i + 1, len(pairs)): + paths_j = {pairs[j][1], pairs[j][2]} + if paths_i & paths_j: # Share at least one image + adjacency[i].append(j) + adjacency[j].append(i) + + # Greedy traversal: start from pair with most connections + visited = set() + ordered = [] + + # Find starting point (highest degree node) + if adjacency: + start = max(adjacency.keys(), key=lambda k: len(adjacency[k])) + else: + # No connections, return original order + return pairs + + # BFS-like traversal favoring connected pairs + queue = [start] + + while queue: + current = queue.pop(0) + if current in visited: + continue + + visited.add(current) + ordered.append(pairs[current]) + + # Add neighbors, sorted by degree (prefer high-degree nodes) + neighbors = [n for n in adjacency[current] if n not in visited] + neighbors.sort(key=lambda n: len(adjacency[n]), reverse=True) + queue.extend(neighbors) + + # Add any remaining unvisited pairs + for i in pair_indices: + if i not in visited: + ordered.append(pairs[i]) + + logger.info(f"Optimized pair order: {len(ordered)} pairs reordered for better caching") + return ordered + + +def save_results( + output_path: Path, + pair_id: str, + warp_AB: torch.Tensor, + overlap_AB: torch.Tensor, + precision_AB: torch.Tensor, + path_A: str, + path_B: str, + warp_BA: Optional[torch.Tensor] = None, + overlap_BA: Optional[torch.Tensor] = None, + precision_BA: Optional[torch.Tensor] = None, + matches: Optional[torch.Tensor] = None, + matches_confidence: Optional[torch.Tensor] = None, + keypoints_A: Optional[torch.Tensor] = None, + keypoints_B: Optional[torch.Tensor] = None, + calibration: Optional[dict] = None, +): + """Save matching results to NPZ file. + + Args: + output_path: Output file path + pair_id: Pair identifier + warp_AB: (H, W, 2) warp field in normalized coordinates + overlap_AB: (H, W, 1) overlap logit + precision_AB: (H, W, 3) precision parameters [sigma_x, sigma_y, correlation] + path_A, path_B: Image paths + warp_BA, overlap_BA, precision_BA: Reverse direction (optional) + matches: Optional (N, 4) sampled matches in normalized coords + matches_confidence: Optional (N,) confidence for sampled matches + keypoints_A: Optional (N, 2) keypoints in image A (pixel coords) + keypoints_B: Optional (N, 2) keypoints in image B (pixel coords) + calibration: Optional calibration dict. If provided, converts precision to normalized weights. + """ + # Convert to numpy + data = { + 'pair_id': pair_id, + 'warp_AB': warp_AB.cpu().numpy(), + 'overlap_AB': overlap_AB.cpu().numpy(), + 'image_A_path': path_A, + 'image_B_path': path_B, + } + + # Process precision_AB + if calibration is not None: + # Convert to normalized weights and remove raw precision + precision_AB_np = precision_AB.cpu().numpy() + precision_weight_AB = precision_to_weight(precision_AB_np, calibration) + data['precision_weight_AB'] = precision_weight_AB + else: + # Store raw precision + data['precision_AB'] = precision_AB.cpu().numpy() + + # Add reverse direction if available + if warp_BA is not None: + data['warp_BA'] = warp_BA.cpu().numpy() + if overlap_BA is not None: + data['overlap_BA'] = overlap_BA.cpu().numpy() + + if precision_BA is not None: + if calibration is not None: + precision_BA_np = precision_BA.cpu().numpy() + precision_weight_BA = precision_to_weight(precision_BA_np, calibration) + data['precision_weight_BA'] = precision_weight_BA + else: + data['precision_BA'] = precision_BA.cpu().numpy() + + # Add optional sampled matches + if matches is not None: + data['matches'] = matches.cpu().numpy() + if matches_confidence is not None: + data['matches_confidence'] = matches_confidence.cpu().numpy() + if keypoints_A is not None: + data['keypoints_A'] = keypoints_A.cpu().numpy() + if keypoints_B is not None: + data['keypoints_B'] = keypoints_B.cpu().numpy() + + # Ensure output directory exists + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Save + np.savez_compressed(output_path, **data) + + +def calibrate_precision_params(precision_samples: np.ndarray) -> dict: + """Calibrate precision parameters from collected samples. + + Computes threshold and temperature for sigmoid normalization of precision + eigenvalues to [0, 1] weights. + + Args: + precision_samples: Array of shape (num_samples, 3) containing + [sigma_x, sigma_y, correlation] representing 2x2 covariance matrices + + Returns: + Dictionary with calibration parameters: + - 'threshold': eigenvalue threshold for sigmoid + - 'temperature': temperature parameter for sigmoid smoothness + - 'mean_eigenvalue': mean of minimum eigenvalues (for reference) + - 'num_samples': number of samples used + """ + logger.info(f"Calibrating precision parameters from {len(precision_samples)} samples...") + + # Convert list to numpy array if needed + if isinstance(precision_samples, list): + precision_samples = np.array(precision_samples) + + # precision_samples shape: (num_samples, 3) with [sigma_x, sigma_y, correlation] + sigma_x = precision_samples[:, 0] # (num_samples,) + sigma_y = precision_samples[:, 1] # (num_samples,) + corr = precision_samples[:, 2] # (num_samples,) + + # Compute eigenvalues of 2x2 symmetric matrix + # For [[a, b], [b, c]]: eigenvalues = (a+c)/2 ± sqrt((a-c)^2/4 + b^2) + a = sigma_x ** 2 + c = sigma_y ** 2 + b = corr * sigma_x * sigma_y + + trace = a + c + det = a * c - b ** 2 + discriminant = (trace / 2) ** 2 - det + discriminant = np.maximum(discriminant, 0) # Ensure non-negative + + sqrt_disc = np.sqrt(discriminant) + lambda_min = trace / 2 - sqrt_disc + + # Use minimum eigenvalue (smaller = more precise) + all_eigenvalues = lambda_min + all_eigenvalues = all_eigenvalues[np.isfinite(all_eigenvalues)] # Remove NaNs/Infs + + # Compute threshold as median eigenvalue + threshold = float(np.median(all_eigenvalues)) + + # Compute temperature to make sigmoid transition smooth + # Use standard deviation scaled appropriately + temperature = float(np.std(all_eigenvalues) / 2.0) + + # Ensure temperature is positive and reasonable + temperature = max(temperature, 1e-6) + + calibration = { + 'threshold': threshold, + 'temperature': temperature, + 'mean_eigenvalue': float(np.mean(all_eigenvalues)), + 'median_eigenvalue': threshold, + 'std_eigenvalue': float(np.std(all_eigenvalues)), + 'num_samples': len(precision_samples), + } + + logger.info("Calibration parameters:") + logger.info(f" threshold: {calibration['threshold']:.6f}") + logger.info(f" temperature: {calibration['temperature']:.6f}") + logger.info(f" mean eigenvalue: {calibration['mean_eigenvalue']:.6f}") + logger.info(f" std eigenvalue: {calibration['std_eigenvalue']:.6f}") + + return calibration + + +def precision_matrix_to_params(precision_matrix: np.ndarray) -> np.ndarray: + """Convert 2×2 precision matrix to 3-parameter representation. + + Args: + precision_matrix: (H, W, 2, 2) array of 2×2 covariance matrices + + Returns: + (H, W, 3) array with [sigma_x, sigma_y, correlation] + """ + # Extract matrix elements + # Matrix is [[sigma_x^2, rho*sigma_x*sigma_y], [rho*sigma_x*sigma_y, sigma_y^2]] + a = precision_matrix[..., 0, 0] # sigma_x^2 + b = precision_matrix[..., 0, 1] # rho*sigma_x*sigma_y + c = precision_matrix[..., 1, 1] # sigma_y^2 + + # Compute parameters + sigma_x = np.sqrt(np.maximum(a, 1e-10)) + sigma_y = np.sqrt(np.maximum(c, 1e-10)) + corr = b / (sigma_x * sigma_y + 1e-10) + corr = np.clip(corr, -1.0, 1.0) + + # Stack into (H, W, 3) + params = np.stack([sigma_x, sigma_y, corr], axis=-1) + return params.astype(np.float32) + + +def precision_to_weight(precision_params: np.ndarray, calibration: dict) -> np.ndarray: + """Convert 3-channel precision parameters to normalized weight in [0, 1]. + + Computes minimum eigenvalue of 2x2 covariance matrix and applies sigmoid + normalization with learned threshold and temperature. + + Args: + precision_params: (H, W, 3) or (H, W, 2, 2) array with precision parameters + calibration: Calibration dictionary from calibrate_precision_params() + + Returns: + (H, W, 1) array with normalized weights in [0, 1] + """ + # Convert from 2×2 matrix format if needed + if precision_params.ndim == 4 and precision_params.shape[2:] == (2, 2): + precision_params = precision_matrix_to_params(precision_params) + + sigma_x = precision_params[..., 0] + sigma_y = precision_params[..., 1] + corr = precision_params[..., 2] + + # Reconstruct 2x2 covariance matrices + a = sigma_x ** 2 + c = sigma_y ** 2 + b = corr * sigma_x * sigma_y + + # Compute eigenvalues + trace = a + c + det = a * c - b ** 2 + discriminant = (trace / 2) ** 2 - det + discriminant = np.maximum(discriminant, 0) + + sqrt_disc = np.sqrt(discriminant) + lambda_min = trace / 2 - sqrt_disc + + # Apply sigmoid normalization + threshold = calibration['threshold'] + temperature = calibration['temperature'] + + z = (lambda_min - threshold) / temperature + z = np.clip(z, -100, 100) + weight = 1.0 / (1.0 + np.exp(-z)) + + return weight[..., np.newaxis].astype(np.float32) + + +def sample_precision_from_center(precision: np.ndarray, num_samples: int = 1000) -> np.ndarray: + """Sample precision values from center region of image. + + Randomly samples points from the center 70% of the image to avoid edge artifacts. + + Args: + precision: (H, W, 3) or (H, W, 2, 2) precision array + num_samples: Number of samples to draw + + Returns: + (num_samples, 3) array of sampled precision values + """ + # Convert from 2×2 matrix format if needed + if precision.ndim == 4 and precision.shape[2:] == (2, 2): + precision = precision_matrix_to_params(precision) + + # Validate shape + if precision.ndim != 3 or precision.shape[2] != 3: + raise ValueError(f"Expected precision shape (H, W, 3), got {precision.shape}") + + H, W, _ = precision.shape + + # Define center region: 70% of image (15% margin on each side) + h_start = int(H * 0.15) + h_end = int(H * 0.85) + w_start = int(W * 0.15) + w_end = int(W * 0.85) + + center_region = precision[h_start:h_end, w_start:w_end, :] + center_h, center_w, _ = center_region.shape + + # Randomly sample points + num_pixels = center_h * center_w + num_samples = min(num_samples, num_pixels) + + # Flatten and sample + flat_precision = center_region.reshape(-1, 3) + indices = np.random.choice(flat_precision.shape[0], size=num_samples, replace=False) + sampled = flat_precision[indices, :] + + return sampled + + +def update_npz_with_calibrated_precision(npz_path: Path, calibration: dict): + """Update an NPZ file to add calibrated precision weights and remove raw precision. + + Args: + npz_path: Path to NPZ file to update + calibration: Calibration parameters + """ + # Load existing data + data = np.load(npz_path, allow_pickle=True) + updated_data = dict(data) + + # Process precision_AB + if 'precision_AB' in data: + precision_AB = data['precision_AB'] # (H, W, 3) or (H, W, 2, 2) + precision_weight = precision_to_weight(precision_AB, calibration) # (H, W, 1) + updated_data['precision_weight_AB'] = precision_weight + # Remove raw precision to save space + del updated_data['precision_AB'] + + # Process precision_BA if present + if 'precision_BA' in data: + precision_BA = data['precision_BA'] # (H, W, 3) or (H, W, 2, 2) + precision_weight = precision_to_weight(precision_BA, calibration) # (H, W, 1) + updated_data['precision_weight_BA'] = precision_weight + # Remove raw precision to save space + del updated_data['precision_BA'] + + # Save updated data + np.savez_compressed(npz_path, **updated_data) + + +def process_batch( + model: RoMaV2, + batch_pairs: list[tuple[str, str, str]], + output_dir: Path, + cache: Optional[LRUImageCache], + num_samples: Optional[int], + stats: dict, + precision_samples: Optional[np.ndarray] = None, + max_calibration_samples: int = 0, + calibration_sample_size: int = 500, + calibration: Optional[dict] = None, +) -> tuple[list[Path], Optional[np.ndarray]]: + """Process a batch of image pairs. + + Args: + model: RoMaV2 model + batch_pairs: List of (pair_id, path_A, path_B) + output_dir: Output directory + cache: Image cache (or None if caching disabled) + num_samples: Number of matches to sample (or None) + stats: Statistics dictionary to update + precision_samples: Optional list to collect precision parameters for calibration + max_calibration_samples: Maximum number of samples to collect for calibration + calibration_sample_size: Number of points to sample per pair for calibration + calibration: Optional calibration dict. If provided, saves with normalized precision. + + Returns: + List of output file paths that were successfully created + """ + output_paths = [] + + for pair_id, path_A, path_B in batch_pairs: + output_path = output_dir / f"{pair_id}.npz" + + # Skip if already processed + if output_path.exists(): + stats['skipped'] += 1 + continue + + try: + # Load images + if cache is not None: + img_A, H_A, W_A = cache.get(path_A, model._load_image) + img_B, H_B, W_B = cache.get(path_B, model._load_image) + else: + img_A = model._load_image(path_A) + img_B = model._load_image(path_B) + # Extract original dimensions + if img_A.dim() == 4: + _, _, H_A, W_A = img_A.shape + _, _, H_B, W_B = img_B.shape + else: + _, H_A, W_A = img_A.shape + _, H_B, W_B = img_B.shape + + # Match + preds = model.match(img_A, img_B) + + # Get outputs (remove batch dimension) + warp_AB = preds['warp_AB'][0] # (H, W, 2) + overlap_AB = preds['overlap_AB'][0] # (H, W, 1) + precision_AB = preds['precision_AB'][0] # (H, W, 2, 2) + warp_BA = preds['warp_BA'][0] if preds['warp_BA'] is not None else None + overlap_BA = preds['overlap_BA'][0] if preds['overlap_BA'] is not None else None + precision_BA = preds['precision_BA'][0] if preds['precision_BA'] is not None else None + + # Collect precision samples for calibration (sample from center) + if (precision_samples is not None and + len(precision_samples) < max_calibration_samples): + # Sample from center region + precision_AB_np = precision_AB.cpu().numpy() + sampled_AB = sample_precision_from_center(precision_AB_np, calibration_sample_size) + precision_samples = np.vstack([precision_samples, sampled_AB]) + if precision_BA is not None: + precision_BA_np = precision_BA.cpu().numpy() + sampled_BA = sample_precision_from_center(precision_BA_np, calibration_sample_size) + precision_samples = np.vstack([precision_samples, sampled_BA]) + + # Optionally sample matches + matches = None + matches_confidence = None + keypoints_A = None + keypoints_B = None + if num_samples is not None and num_samples > 0: + try: + preds_for_sampling = { + 'warp_AB': preds['warp_AB'], + 'overlap_AB': preds['overlap_AB'], + 'precision_AB': preds['precision_AB'], + } + if preds['warp_BA'] is not None: + preds_for_sampling.update({ + 'warp_BA': preds['warp_BA'], + 'overlap_BA': preds['overlap_BA'], + 'precision_BA': preds['precision_BA'], + }) + + matches, matches_confidence, _, _ = model.sample( + preds_for_sampling, num_samples + ) + + # Convert to pixel coordinates + keypoints_A, keypoints_B = model.to_pixel_coordinates( + matches, H_A, W_A, H_B, W_B + ) + except Exception as e: + logger.warning(f"Failed to sample matches for {pair_id}: {e}") + + # Save results + save_results( + output_path=output_path, + pair_id=pair_id, + warp_AB=warp_AB, + overlap_AB=overlap_AB, + precision_AB=precision_AB, + path_A=path_A, + path_B=path_B, + warp_BA=warp_BA, + overlap_BA=overlap_BA, + precision_BA=precision_BA, + matches=matches, + matches_confidence=matches_confidence, + keypoints_A=keypoints_A, + keypoints_B=keypoints_B, + calibration=calibration, + ) + + stats['processed'] += 1 + output_paths.append(output_path) + + except FileNotFoundError as e: + logger.warning(f"Image not found for {pair_id}: {e}") + stats['failed'] += 1 + except RuntimeError as e: + if "out of memory" in str(e).lower(): + logger.error(f"OOM error on {pair_id}. Try reducing batch size or image resolution.") + stats['failed'] += 1 + # Clear cache to free memory + torch.cuda.empty_cache() if torch.cuda.is_available() else None + else: + logger.error(f"Runtime error on {pair_id}: {e}") + stats['failed'] += 1 + except Exception as e: + import traceback + logger.error(f"Unexpected error on {pair_id}: {e}") + logger.error(f"Traceback: {traceback.format_exc()}") + stats['failed'] += 1 + + return output_paths, precision_samples + + +def main(): + parser = argparse.ArgumentParser( + description="Batch matching for RoMaV2", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" + Examples: + # Basic usage + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ + + # With image caching and match sampling + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + --cache-images --cache-size 200 --num-samples 5000 + + # Fast processing with lower quality + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + --setting fast --batch-size 4 + + # With precision calibration (100 pairs × 1000 samples per pair = 100k total) + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + --calibration-pairs 100 --calibration-sample-size 1000 + """ + ) + + parser.add_argument( + '--pairs-csv', + type=Path, + required=True, + help='Path to CSV file with image pairs (format: image_A_path,image_B_path)' + ) + parser.add_argument( + '--output-dir', + type=Path, + required=True, + help='Output directory for NPZ files' + ) + parser.add_argument( + '--batch-size', + type=int, + default=1, + help='Batch size for processing (default: 1)' + ) + parser.add_argument( + '--setting', + type=str, + choices=['turbo', 'fast', 'base', 'precise'], + default='base', + help='Model setting (default: base)' + ) + parser.add_argument( + '--num-samples', + type=int, + default=None, + help='Number of matches to sample per pair (optional)' + ) + parser.add_argument( + '--cache-images', + action='store_true', + help='Enable LRU image caching' + ) + parser.add_argument( + '--cache-size', + type=int, + default=100, + help='Maximum number of images to cache (default: 100)' + ) + parser.add_argument( + '--no-optimize-order', + action='store_true', + help='Disable pair order optimization (process in original order)' + ) + parser.add_argument( + '--no-bidirectional', + dest='bidirectional', + action='store_false', + default=True, + help='Disable bidirectional matching (only A->B; default is bidirectional)' + ) + parser.add_argument( + '--calibration-pairs', + type=int, + default=15, + help='Number of pairs to use for precision calibration (optional). ' + 'Collects precision statistics from first N pairs, then ' + 'updates all NPZ files with normalized precision weights.' + ) + parser.add_argument( + '--calibration-sample-size', + type=int, + default=1000, + help='Number of random points to sample from center of each pair for calibration (default: 1000). ' + 'Total calibration samples = calibration-pairs × calibration-sample-size' + ) + + args = parser.parse_args() + + # Validate inputs + if not args.pairs_csv.exists(): + logger.error(f"Pairs CSV file not found: {args.pairs_csv}") + return 1 + + # Create output directory + args.output_dir.mkdir(parents=True, exist_ok=True) + + # Load pairs + pairs = load_pairs_csv(args.pairs_csv) + if not pairs: + logger.error("No valid pairs found in CSV file") + return 1 + + # Optimize pair order for caching + if args.cache_images and not args.no_optimize_order: + pairs = optimize_pair_order(pairs) + + # Initialize model + logger.info(f"Initializing RoMaV2 with setting: {args.setting}") + model = RoMaV2() + model.apply_setting(args.setting) + model.bidirectional = args.bidirectional + model.eval() + logger.info(f"Bidirectional matching is {'enabled' if model.bidirectional else 'disabled'}") + + # Initialize cache if enabled + cache = LRUImageCache(max_size=args.cache_size) if args.cache_images else None + if cache: + logger.info(f"Image caching enabled (max size: {args.cache_size})") + + # Initialize calibration if requested + precision_samples = np.empty((0, 3), dtype=np.float32) if args.calibration_pairs is not None else None + calibration = None + max_calibration_samples = 0 + + if precision_samples is not None: + max_calibration_samples = args.calibration_pairs * args.calibration_sample_size + if model.bidirectional: + max_calibration_samples *= 2 # Account for both directions per pair + logger.info(f"Precision calibration enabled:") + logger.info(f" Pairs to calibrate: {args.calibration_pairs}") + logger.info(f" Sample size per pair: {args.calibration_sample_size}") + logger.info(f" Total calibration samples: {max_calibration_samples}") + + # Statistics + stats = { + 'processed': 0, + 'skipped': 0, + 'failed': 0, + 'start_time': time.time(), + } + + # Process pairs + logger.info(f"Processing {len(pairs)} pairs...") + + # Note: Batching is per-pair for now since images can have different resolutions + # Could be extended to batch pairs with same resolution + if args.batch_size > 1: + logger.warning("Batch size > 1 not fully implemented. Processing one pair at a time.") + + processed_files = [] + + # Determine effective batch size for first batch (limited by calibration_pairs if enabled) + first_batch_size = args.batch_size + if precision_samples is not None and args.calibration_pairs is not None: + first_batch_size = min(args.batch_size, args.calibration_pairs) + if first_batch_size < args.batch_size: + logger.info(f"First batch limited to {first_batch_size} pairs (calibration requirement)") + + calibration_complete = False + + with tqdm(total=len(pairs), desc="Matching pairs", unit="pair") as pbar: + idx = 0 + while idx < len(pairs): + # Determine batch size + current_batch_size = args.batch_size if calibration_complete else first_batch_size + batch_end = min(idx + current_batch_size, len(pairs)) + batch = pairs[idx:batch_end] + + for pair_id, path_A, path_B in batch: + # Check if we need to calibrate (only before calibration_complete) + if (not calibration_complete and + precision_samples is not None and + calibration is None and + len(precision_samples) >= max_calibration_samples): + logger.info(f"\n{'='*60}") + logger.info(f"Reached {len(precision_samples)} calibration samples") + logger.info(f"{'='*60}") + + # Calibrate precision parameters + calibration = calibrate_precision_params(precision_samples) + + # Update all previously processed files + logger.info(f"Updating {len(processed_files)} NPZ files with calibrated precision...") + for npz_path in tqdm(processed_files, desc="Updating NPZ files", unit="file"): + update_npz_with_calibrated_precision(npz_path, calibration) + + # Clear samples to free memory + precision_samples = None + calibration_complete = True + logger.info("Calibration complete. Continuing with remaining pairs...\n") + + # Process pair and get output paths (and updated precision samples) + new_outputs, precision_samples = process_batch( + model=model, + batch_pairs=[(pair_id, path_A, path_B)], + output_dir=args.output_dir, + cache=cache, + num_samples=args.num_samples, + stats=stats, + precision_samples=precision_samples, + max_calibration_samples=max_calibration_samples, + calibration_sample_size=args.calibration_sample_size, + calibration=calibration, + ) + + # Track processed files for calibration update (only before calibration) + if not calibration_complete and precision_samples is not None: + processed_files.extend(new_outputs) + + pbar.update(1) + + # Update progress bar description + if stats['processed'] > 0: + elapsed = time.time() - stats['start_time'] + rate = stats['processed'] / elapsed + postfix = { + 'rate': f'{rate:.2f} pairs/s', + 'skipped': stats['skipped'], + 'failed': stats['failed'] + } + if precision_samples is not None: + postfix['cal_samples'] = len(precision_samples) + postfix['cal_target'] = max_calibration_samples + pbar.set_postfix(postfix) + + idx = batch_end + + # Final statistics + elapsed = time.time() - stats['start_time'] + logger.info("\n" + "="*60) + logger.info("BATCH MATCHING COMPLETE") + logger.info("="*60) + logger.info(f"Total pairs: {len(pairs)}") + logger.info(f"Processed: {stats['processed']}") + logger.info(f"Skipped: {stats['skipped']}") + logger.info(f"Failed: {stats['failed']}") + logger.info(f"Elapsed time: {elapsed:.1f}s") + logger.info(f"Throughput: {stats['processed']/elapsed:.2f} pairs/s") + + if calibration is not None: + logger.info(f"\nPrecision Calibration:") + logger.info(f" Calibration pairs: {args.calibration_pairs}") + logger.info(f" Sample size per pair: {args.calibration_sample_size}") + logger.info(f" Total samples collected: {max_calibration_samples}") + logger.info(f" Applied to all {len(pairs)} pairs") + + if cache: + cache_stats = cache.get_stats() + logger.info(f"\nCache Statistics:") + logger.info(f" Hit rate: {cache_stats['hit_rate']:.1%}") + logger.info(f" Hits: {cache_stats['hits']}") + logger.info(f" Misses: {cache_stats['misses']}") + logger.info(f" Final size: {cache_stats['size']}") + + logger.info(f"\nResults saved to: {args.output_dir}") + logger.info("="*60) + + return 0 + + +if __name__ == "__main__": + exit(main()) From 44cd84444562526d1cb2592dd4a40564f64a5f8b Mon Sep 17 00:00:00 2001 From: Mark McCurry Date: Tue, 10 Feb 2026 21:45:52 +0000 Subject: [PATCH 2/8] Optimize Roma2 Inference Gain roughly an 80x speedup in roma2 processing (assuming only the coarse non-bidirectional match is needed) - Introduces Docker container - Add feature caching to romav2.py - Add batch coarse matching inference (using cached features) to romav2.py - Introduce batch_match_min script calculating the minimal coarse matching using the new romav2.py API - Add profiling statements to romav2.py - Add profiling statements to batch_match.py --- build-docker.sh | 2 + cuda-docker.dockerfile | 15 ++ run-docker.sh | 7 + scripts/batch_match.py | 119 ++++++---- scripts/batch_match_min.py | 445 +++++++++++++++++++++++++++++++++++++ src/romav2/romav2.py | 169 ++++++++++++-- 6 files changed, 692 insertions(+), 65 deletions(-) create mode 100644 build-docker.sh create mode 100644 cuda-docker.dockerfile create mode 100644 run-docker.sh create mode 100644 scripts/batch_match_min.py diff --git a/build-docker.sh b/build-docker.sh new file mode 100644 index 0000000..f407367 --- /dev/null +++ b/build-docker.sh @@ -0,0 +1,2 @@ +#docker rmi ubuntu-romav2 +docker build -t ubuntu-romav2 -f cuda-docker.dockerfile . diff --git a/cuda-docker.dockerfile b/cuda-docker.dockerfile new file mode 100644 index 0000000..ee108c9 --- /dev/null +++ b/cuda-docker.dockerfile @@ -0,0 +1,15 @@ +FROM nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04 + +RUN apt-get update +RUN DEBIAN_FRONTEND=noninteractive apt-get install -y tzdata +RUN apt-get install -y python3 +RUN apt-get install -y curl +RUN apt-get install -y git +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.local/bin/:$PATH" +RUN apt-get install -y libgl1 +RUN apt-get install -y libglib2.0-0 libsm6 libxrender1 libxext6 +RUN mkdir /root/roma +WORKDIR /root/roma +COPY . /root/roma/ +RUN uv sync diff --git a/run-docker.sh b/run-docker.sh new file mode 100644 index 0000000..bcb981e --- /dev/null +++ b/run-docker.sh @@ -0,0 +1,7 @@ + +docker run -it --gpus all \ + -v /home/ubuntu/virginia/home/jgibson/data/hypersim/:/data/hypersim:ro \ + -v /home/ubuntu/scratch/:/scratch/ \ + --ipc=host \ + --net=host \ + ubuntu-romav2 diff --git a/scripts/batch_match.py b/scripts/batch_match.py index e73b8bf..a2f620f 100644 --- a/scripts/batch_match.py +++ b/scripts/batch_match.py @@ -88,7 +88,7 @@ def get_stats(self) -> dict: 'size': len(self.cache) } - +@torch.profiler.record_function("load_pairs_csv") def load_pairs_csv(csv_path: Path) -> list[tuple[str, str, str]]: """Load image pairs from CSV file. @@ -133,6 +133,7 @@ def load_pairs_csv(csv_path: Path) -> list[tuple[str, str, str]]: return pairs +@torch.profiler.record_function("optimize_pair_order") def optimize_pair_order(pairs: list[tuple[str, str, str]]) -> list[tuple[str, str, str]]: """Reorder pairs to maximize cache hits by clustering shared images. @@ -195,6 +196,7 @@ def optimize_pair_order(pairs: list[tuple[str, str, str]]) -> list[tuple[str, st return ordered +@torch.profiler.record_function("save_results") def save_results( output_path: Path, pair_id: str, @@ -262,22 +264,28 @@ def save_results( data['precision_BA'] = precision_BA.cpu().numpy() # Add optional sampled matches - if matches is not None: - data['matches'] = matches.cpu().numpy() - if matches_confidence is not None: - data['matches_confidence'] = matches_confidence.cpu().numpy() - if keypoints_A is not None: - data['keypoints_A'] = keypoints_A.cpu().numpy() - if keypoints_B is not None: - data['keypoints_B'] = keypoints_B.cpu().numpy() + with torch.profiler.record_function("convert_to_numpy"): + if matches is not None: + data['matches'] = matches.cpu().numpy() + with torch.profiler.record_function("convert_to_numpy"): + if matches_confidence is not None: + data['matches_confidence'] = matches_confidence.cpu().numpy() + with torch.profiler.record_function("convert_to_numpy"): + if keypoints_A is not None: + data['keypoints_A'] = keypoints_A.cpu().numpy() + with torch.profiler.record_function("convert_to_numpy"): + if keypoints_B is not None: + data['keypoints_B'] = keypoints_B.cpu().numpy() # Ensure output directory exists output_path.parent.mkdir(parents=True, exist_ok=True) # Save - np.savez_compressed(output_path, **data) + with torch.profiler.record_function("save_to_disk"): + np.savez_compressed(output_path, **data) +@torch.profiler.record_function("calibrate_precision_params") def calibrate_precision_params(precision_samples: np.ndarray) -> dict: """Calibrate precision parameters from collected samples. @@ -352,6 +360,7 @@ def calibrate_precision_params(precision_samples: np.ndarray) -> dict: return calibration +@torch.profiler.record_function("precision_matrix_to_params") def precision_matrix_to_params(precision_matrix: np.ndarray) -> np.ndarray: """Convert 2×2 precision matrix to 3-parameter representation. @@ -378,6 +387,7 @@ def precision_matrix_to_params(precision_matrix: np.ndarray) -> np.ndarray: return params.astype(np.float32) +@torch.profiler.record_function("precision_to_weight") def precision_to_weight(precision_params: np.ndarray, calibration: dict) -> np.ndarray: """Convert 3-channel precision parameters to normalized weight in [0, 1]. @@ -424,6 +434,7 @@ def precision_to_weight(precision_params: np.ndarray, calibration: dict) -> np.n return weight[..., np.newaxis].astype(np.float32) +@torch.profiler.record_function("sample_precision_from_center") def sample_precision_from_center(precision: np.ndarray, num_samples: int = 1000) -> np.ndarray: """Sample precision values from center region of image. @@ -467,6 +478,7 @@ def sample_precision_from_center(precision: np.ndarray, num_samples: int = 1000) return sampled +@torch.profiler.record_function("update_npz") def update_npz_with_calibrated_precision(npz_path: Path, calibration: dict): """Update an NPZ file to add calibrated precision weights and remove raw precision. @@ -498,6 +510,7 @@ def update_npz_with_calibrated_precision(npz_path: Path, calibration: dict): np.savez_compressed(npz_path, **updated_data) +@torch.profiler.record_function("process_batch") def process_batch( model: RoMaV2, batch_pairs: list[tuple[str, str, str]], @@ -539,22 +552,24 @@ def process_batch( try: # Load images - if cache is not None: - img_A, H_A, W_A = cache.get(path_A, model._load_image) - img_B, H_B, W_B = cache.get(path_B, model._load_image) - else: - img_A = model._load_image(path_A) - img_B = model._load_image(path_B) - # Extract original dimensions - if img_A.dim() == 4: - _, _, H_A, W_A = img_A.shape - _, _, H_B, W_B = img_B.shape + with torch.profiler.record_function("load_data"): + if cache is not None: + img_A, H_A, W_A = cache.get(path_A, model._load_image) + img_B, H_B, W_B = cache.get(path_B, model._load_image) else: - _, H_A, W_A = img_A.shape - _, H_B, W_B = img_B.shape + img_A = model._load_image(path_A) + img_B = model._load_image(path_B) + # Extract original dimensions + if img_A.dim() == 4: + _, _, H_A, W_A = img_A.shape + _, _, H_B, W_B = img_B.shape + else: + _, H_A, W_A = img_A.shape + _, H_B, W_B = img_B.shape # Match - preds = model.match(img_A, img_B) + with torch.profiler.record_function("DNN_exec"): + preds = model.match(img_A, img_B) # Get outputs (remove batch dimension) warp_AB = preds['warp_AB'][0] # (H, W, 2) @@ -565,7 +580,8 @@ def process_batch( precision_BA = preds['precision_BA'][0] if preds['precision_BA'] is not None else None # Collect precision samples for calibration (sample from center) - if (precision_samples is not None and + with torch.profiler.record_function("prec_gather"): + if (precision_samples is not None and len(precision_samples) < max_calibration_samples): # Sample from center region precision_AB_np = precision_AB.cpu().numpy() @@ -581,7 +597,8 @@ def process_batch( matches_confidence = None keypoints_A = None keypoints_B = None - if num_samples is not None and num_samples > 0: + with torch.profiler.record_function("prec_sample"): + if num_samples is not None and num_samples > 0: try: preds_for_sampling = { 'warp_AB': preds['warp_AB'], @@ -607,7 +624,8 @@ def process_batch( logger.warning(f"Failed to sample matches for {pair_id}: {e}") # Save results - save_results( + with torch.profiler.record_function("save_warps"): + save_results( output_path=output_path, pair_id=pair_id, warp_AB=warp_AB, @@ -764,6 +782,8 @@ def main(): # Initialize model logger.info(f"Initializing RoMaV2 with setting: {args.setting}") + #Use compile=False for full tracing + #model = RoMaV2(RoMaV2.Cfg(compile=False)) model = RoMaV2() model.apply_setting(args.setting) model.bidirectional = args.bidirectional @@ -816,7 +836,11 @@ def main(): calibration_complete = False - with tqdm(total=len(pairs), desc="Matching pairs", unit="pair") as pbar: + activities = [torch.profiler.ProfilerActivity.CPU] + activities += [torch.profiler.ProfilerActivity.CUDA] + with torch.profiler.profile(activities=activities, record_shapes=True) as prof: + with torch.profiler.record_function("matching_pairs"): + with tqdm(total=len(pairs), desc="Matching pairs", unit="pair") as pbar: idx = 0 while idx < len(pairs): # Determine batch size @@ -835,12 +859,14 @@ def main(): logger.info(f"{'='*60}") # Calibrate precision parameters - calibration = calibrate_precision_params(precision_samples) + with torch.profiler.record_function("calibration"): + calibration = calibrate_precision_params(precision_samples) # Update all previously processed files - logger.info(f"Updating {len(processed_files)} NPZ files with calibrated precision...") - for npz_path in tqdm(processed_files, desc="Updating NPZ files", unit="file"): - update_npz_with_calibrated_precision(npz_path, calibration) + with torch.profiler.record_function("calibration_global"): + logger.info(f"Updating {len(processed_files)} NPZ files with calibrated precision...") + for npz_path in tqdm(processed_files, desc="Updating NPZ files", unit="file"): + update_npz_with_calibrated_precision(npz_path, calibration) # Clear samples to free memory precision_samples = None @@ -848,22 +874,24 @@ def main(): logger.info("Calibration complete. Continuing with remaining pairs...\n") # Process pair and get output paths (and updated precision samples) - new_outputs, precision_samples = process_batch( - model=model, - batch_pairs=[(pair_id, path_A, path_B)], - output_dir=args.output_dir, - cache=cache, - num_samples=args.num_samples, - stats=stats, - precision_samples=precision_samples, - max_calibration_samples=max_calibration_samples, - calibration_sample_size=args.calibration_sample_size, - calibration=calibration, - ) + with torch.profiler.record_function("process_batch"): + new_outputs, precision_samples = process_batch( + model=model, + batch_pairs=[(pair_id, path_A, path_B)], + output_dir=args.output_dir, + cache=cache, + num_samples=args.num_samples, + stats=stats, + precision_samples=precision_samples, + max_calibration_samples=max_calibration_samples, + calibration_sample_size=args.calibration_sample_size, + calibration=calibration, + ) # Track processed files for calibration update (only before calibration) - if not calibration_complete and precision_samples is not None: - processed_files.extend(new_outputs) + with torch.profiler.record_function("update_output"): + if not calibration_complete and precision_samples is not None: + processed_files.extend(new_outputs) pbar.update(1) @@ -883,6 +911,7 @@ def main(): idx = batch_end + prof.export_chrome_trace("trace.json") # Final statistics elapsed = time.time() - stats['start_time'] logger.info("\n" + "="*60) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py new file mode 100644 index 0000000..624531b --- /dev/null +++ b/scripts/batch_match_min.py @@ -0,0 +1,445 @@ +#!/usr/bin/env python3 +""" +Batch matching script for RoMaV2. + +Efficiently matches a list of image pairs from a CSV file and saves results to disk. +Supports intelligent caching and pair ordering to maximize throughput. + +Usage: + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ + +CSV Format: + image_A_path,image_B_path + path/to/img1.jpg,path/to/img2.jpg + ... +""" + +import argparse +import csv +import logging +import time +from collections import OrderedDict, defaultdict +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm +from PIL import Image + +from romav2 import RoMaV2 + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + datefmt='%H:%M:%S' +) +logger = logging.getLogger(__name__) + +def load_pairs_csv(csv_path: Path) -> list[tuple[str, str, str]]: + """Load image pairs from CSV file. + + Args: + csv_path: Path to CSV file with format: image_A_path,image_B_path + + Returns: + List of tuples: (pair_id, path_A, path_B) + """ + pairs = [] + csv_dir = csv_path.parent + + with open(csv_path, 'r') as f: + reader = csv.reader(f) + for i, row in enumerate(reader, 1): + if len(row) < 2: + logger.warning(f"Line {i}: Expected 2 columns, got {len(row)}. Skipping.") + continue + + path_A, path_B, *_ = row + path_A = path_A.strip() + path_B = path_B.strip() + + # Skip empty lines or headers + if not path_A or not path_B or path_A.lower() == 'image_a_path' or path_A.lower() == 'imagea': + continue + + # Resolve relative paths + if not Path(path_A).is_absolute(): + path_A = str(csv_dir / path_A) + if not Path(path_B).is_absolute(): + path_B = str(csv_dir / path_B) + + # Generate pair_id from filenames + stem_A = Path(path_A).stem + stem_B = Path(path_B).stem + pair_id = f"{stem_A}-{stem_B}" + + pairs.append((pair_id, path_A, path_B)) + + logger.info(f"Loaded {len(pairs)} image pairs from {csv_path}") + return pairs + +def save_results( + output_path: Path, + pair_id: str, + warp_AB: torch.Tensor, + overlap_AB: torch.Tensor, + path_A: str, + path_B: str, +): + """Save matching results to pt file. + + Args: + output_path: Output file path + pair_id: Pair identifier + warp_AB: (H, W, 2) warp field in normalized coordinates + overlap_AB: (H, W, 1) overlap logit + path_A, path_B: Image paths + """ + # Convert to dict + data = { + 'pair_id': pair_id, + 'warp_AB': warp_AB, + 'overlap_AB': overlap_AB, + 'image_A_path': path_A, + 'image_B_path': path_B, + } + + # Ensure output directory exists + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Save + with torch.profiler.record_function("save_to_disk"): + torch.save(data, output_path) + +@torch.profiler.record_function("load_im") +def load_im(path): + img_pil = Image.open(path) + img_pil = img_pil.convert("RGB") + img = torch.from_numpy(np.array(img_pil)).permute(2, 0, 1) + + if img.dtype == torch.uint8: + img = img.float() / 255.0 + return img + +class ImageDataset(Dataset): + def __init__(self, paths): + self.paths = paths + + def __len__(self): + return len(self.paths) + + def __getitem__(self, idx): + return idx, load_im(self.paths[idx]) + +@torch.profiler.record_function("load_images") +def load_images(pairs): + iset = set() + for idx in range(len(pairs)): + iset.add(pairs[idx][1]) + iset.add(pairs[idx][2]) + iset = list(iset) + print(f"Loading {len(iset)} images...") + dataloader = DataLoader(ImageDataset(iset), + batch_size=1, + prefetch_factor=2, + pin_memory=True, + num_workers=20) + icache = {} + for idx, im in dataloader: + icache[iset[idx]] = im + return icache + + +enroll_queue = [] +enroll_batch = 16 +@torch.profiler.record_function("enroll") +def enroll( + cache, + model, + image_path:str, + icache=None + ): + if(image_path in cache): + return + + if(icache): + img = icache[image_path].to("cuda:0") + else: + img = model._load_image(image_path) + + global enroll_queue + enroll_queue.append([image_path,img]) + if(len(enroll_queue) >= enroll_batch): + enroll_flush(cache, model) + +@torch.profiler.record_function("enroll_flush") +def enroll_flush(cache, model): + global enroll_queue + if(len(enroll_queue) != 0): + bt = model.cache_batch_features(enroll_queue) + for x in bt: + cache[x["path"]] = x + enroll_queue = [] + pass + +def streaming_cache(model, pairs): + iset = set() + for idx in range(len(pairs)): + iset.add(pairs[idx][1]) + iset.add(pairs[idx][2]) + iset = list(iset) + print(f"Loading {len(iset)} images...") + dataloader = DataLoader(ImageDataset(iset), + batch_size=1, + prefetch_factor=2, + pin_memory=True, + num_workers=20) + icache = {} + cache = {} + for idx, im in dataloader: + icache[iset[idx]] = im + enroll(cache, model, iset[idx], icache) + enroll_flush(cache, model) + + return cache + +def process_real_batch( + model: RoMaV2, + batch_pairs: list[tuple[str, str, str]], + output_dir: Path, + stats: dict, + cache, +) -> tuple[list[Path], Optional[np.ndarray]]: + + output_paths = [] + feat1_A = [] + feat1_B = [] + feat2_A = [] + feat2_B = [] + img_A = [] + img_B = [] + + #Assemble batch + for pair_id, path_A, path_B in batch_pairs: + output_path = output_dir / f"{pair_id}.pt" + output_paths.append(output_path) + frame_A = cache[path_A] + frame_B = cache[path_B] + feat1_A.append(frame_A["features"][0]) + feat1_B.append(frame_B["features"][0]) + feat2_A.append(frame_A["features"][1]) + feat2_B.append(frame_B["features"][1]) + img_A.append(frame_A["rescaled"]) + img_B.append(frame_B["rescaled"]) + + batch_A = {"features": [torch.cat(feat1_A), + torch.cat(feat2_A)], + "rescaled": torch.cat(img_A)} + batch_B = {"features": [torch.cat(feat1_B), + torch.cat(feat2_B)], + "rescaled": torch.cat(img_B)} + + #Perform Inference + with torch.profiler.record_function("DNN_exec"): + preds = model.coarse_cached_match(batch_A, batch_B) + + #Dissassemble batch + with torch.profiler.record_function("save_batch"): + if(False): + idx = 0 + for pair_id, path_A, path_B in batch_pairs: + # Get outputs (remove batch dimension) + warp_AB = preds['warp_AB'][idx] # (H, W, 2) + overlap_AB = preds['overlap_AB'][idx] # (H, W, 1) + + # Save results + with torch.profiler.record_function("save_warps"): + save_results( + output_path=output_paths[idx], + pair_id=pair_id, + warp_AB=warp_AB, + overlap_AB=overlap_AB, + path_A=path_A, + path_B=path_B, + ) + idx += 1 + stats['processed'] += 1 + else: + stats['processed'] += len(output_paths) + ids = [] + pA = [] + pB = [] + for pair_id, path_A, path_B in batch_pairs: + ids.append(pair_id) + pA.append(path_A) + pB.append(path_B) + # Convert to dict + data = { + 'pair_id': ids, + 'warp_AB': preds["warp_AB"], + 'overlap_AB': preds["overlap_AB"], + 'image_A_path': pA, + 'image_B_path': pB, + } + + # Ensure output directory exists + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Save + with torch.profiler.record_function("save_to_disk"): + torch.save(data, output_path) + +def main(): + parser = argparse.ArgumentParser( + description="Batch matching for RoMaV2", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" + Examples: + # Basic usage + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ + + # With image caching and match sampling + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + + # Fast processing with lower quality + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + --setting fast --batch-size 4 + + # With precision calibration (100 pairs × 1000 samples per pair = 100k total) + python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + --calibration-pairs 100 --calibration-sample-size 1000 + """ + ) + + parser.add_argument( + '--pairs-csv', + type=Path, + required=True, + help='Path to CSV file with image pairs (format: image_A_path,image_B_path)' + ) + parser.add_argument( + '--output-dir', + type=Path, + required=True, + help='Output directory for NPZ files' + ) + parser.add_argument( + '--batch-size', + type=int, + default=1, + help='Batch size for processing (default: 1)' + ) + parser.add_argument( + '--setting', + type=str, + choices=['turbo', 'fast', 'base', 'precise'], + default='base', + help='Model setting (default: base)' + ) + + args = parser.parse_args() + + # Validate inputs + if not args.pairs_csv.exists(): + logger.error(f"Pairs CSV file not found: {args.pairs_csv}") + return 1 + + # Create output directory + args.output_dir.mkdir(parents=True, exist_ok=True) + + # Load pairs + pairs = load_pairs_csv(args.pairs_csv) + if not pairs: + logger.error("No valid pairs found in CSV file") + return 1 + + # Initialize model + logger.info(f"Initializing RoMaV2 with setting: {args.setting}") + model = RoMaV2(RoMaV2.Cfg(compile=False)) + model.apply_setting(args.setting) + model.eval() + + # Statistics + stats = { + 'processed': 0, + 'skipped': 0, + 'failed': 0, + 'start_time': time.time(), + } + + # Process pairs + logger.info(f"Processing {len(pairs)} pairs...") + + # Note: Batching is per-pair for now since images can have different resolutions + # Could be extended to batch pairs with same resolution + if args.batch_size > 1: + logger.warning("Batch size > 1 not fully implemented. Processing one pair at a time.") + + processed_files = [] + + activities = [torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA] + #if(True): + with torch.profiler.profile(activities=activities, + record_shapes=True, + profile_memory=False,) as prof: + with torch.profiler.record_function("streaming_feature_cache"): + cache = streaming_cache(model, pairs) + + with torch.profiler.record_function("matching_pairs"): + with tqdm(total=len(pairs), desc="Matching pairs", unit="pair") as pbar: + idx = 0 + while idx < len(pairs): + # Determine batch size + current_batch_size = args.batch_size + batch_end = min(idx + current_batch_size, len(pairs)) + batch = pairs[idx:batch_end] + + with torch.profiler.record_function("process_batch"): + new_outputs = process_real_batch( + model=model, + batch_pairs=batch, + output_dir=args.output_dir, + stats=stats, + cache=cache, + ) + + pbar.update(len(batch)) + + # Update progress bar description + if stats['processed'] > 0: + elapsed = time.time() - stats['start_time'] + rate = stats['processed'] / elapsed + postfix = { + 'rate': f'{rate:.2f} pairs/s', + 'skipped': stats['skipped'], + 'failed': stats['failed'] + } + pbar.set_postfix(postfix) + + idx = batch_end + + prof.export_chrome_trace("trace.json") + # Final statistics + elapsed = time.time() - stats['start_time'] + logger.info("\n" + "="*60) + logger.info("BATCH MATCHING COMPLETE") + logger.info("="*60) + logger.info(f"Total pairs: {len(pairs)}") + logger.info(f"Processed: {stats['processed']}") + logger.info(f"Skipped: {stats['skipped']}") + logger.info(f"Failed: {stats['failed']}") + logger.info(f"Elapsed time: {elapsed:.1f}s") + logger.info(f"Throughput: {stats['processed']/elapsed:.2f} pairs/s") + + logger.info(f"\nResults saved to: {args.output_dir}") + logger.info("="*60) + + return 0 + + +if __name__ == "__main__": + exit(main()) diff --git a/src/romav2/romav2.py b/src/romav2/romav2.py index ad80a0b..19480a8 100644 --- a/src/romav2/romav2.py +++ b/src/romav2/romav2.py @@ -29,6 +29,7 @@ logger = logging.getLogger(__name__) +@torch.profiler.record_function("interpolate_warp_conf") def _interpolate_warp_and_confidence( *, warp: torch.Tensor, @@ -176,12 +177,14 @@ def forward( # init preds predictions = OrderedDict() # extract feats - f_A = self.f(img_A_lr) - f_B = self.f(img_B_lr) + with torch.profiler.record_function("DINO_x2"): + f_A = self.f(img_A_lr) + f_B = self.f(img_B_lr) # match feats - matcher_output = self.matcher( - f_A, f_B, img_A=img_A_lr, img_B=img_B_lr, bidirectional=self.bidirectional - ) + with torch.profiler.record_function("matcher"): + matcher_output = self.matcher( + f_A, f_B, img_A=img_A_lr, img_B=img_B_lr, bidirectional=self.bidirectional + ) # return matcher_output predictions["matcher"] = matcher_output warp_AB, confidence_AB = ( @@ -206,8 +209,9 @@ def forward( scale_factor = torch.tensor( (W / self.anchor_width, H / self.anchor_height), device=device ) - refiner_features_A = self.refiner_features(img_A) - refiner_features_B = self.refiner_features(img_B) + with torch.profiler.record_function("refiner_features"): + refiner_features_A = self.refiner_features(img_A) + refiner_features_B = self.refiner_features(img_B) for patch_size_str, refiner in self.refiners.items(): patch_size = int(patch_size_str) zero_out_precision = ( @@ -233,21 +237,23 @@ def forward( f_patch_A = refiner_features_A[patch_size] f_patch_B = refiner_features_B[patch_size] - refiner_output_AB = refiner( - f_A=f_patch_A, - f_B=f_patch_B, - prev_warp=warp_AB, - prev_confidence=confidence_AB, - scale_factor=scale_factor, - ) - if self.bidirectional: - refiner_output_BA = refiner( - f_A=f_patch_B, - f_B=f_patch_A, - prev_warp=warp_BA, - prev_confidence=confidence_BA, + with torch.profiler.record_function("refiner"): + refiner_output_AB = refiner( + f_A=f_patch_A, + f_B=f_patch_B, + prev_warp=warp_AB, + prev_confidence=confidence_AB, scale_factor=scale_factor, ) + if self.bidirectional: + with torch.profiler.record_function("refiner"): + refiner_output_BA = refiner( + f_A=f_patch_B, + f_B=f_patch_A, + prev_warp=warp_BA, + prev_confidence=confidence_BA, + scale_factor=scale_factor, + ) else: refiner_output_BA = None predictions[f"refiner_{patch_size}_AB"] = refiner_output_AB @@ -298,6 +304,128 @@ def _load_image(self, img_like: ImageLike) -> torch.Tensor: img = img[None] return img + @torch.profiler.record_function("cache_features") + @torch.inference_mode() + def cache_features(self, + img_like:ImageLike): + img = self._load_image(img_like) + + img_lr = F.interpolate( + img, + size=(self.H_lr, self.W_lr), + mode="bicubic", + align_corners=False, + antialias=True, + ) + feat = self.f(img_lr) + #XXX move features to CPU? + return {"img": img, + "rescaled": img_lr, + "features": feat, + } + + @torch.profiler.record_function("cache_features") + @torch.inference_mode() + def cache_batch_features(self, inputs): + print(f"Processing a batch of {len(inputs)}") + ibatch = [] + for i in inputs: + img_like = i[1] + img = self._load_image(img_like) + + img_lr = F.interpolate( + img, + size=(self.H_lr, self.W_lr), + mode="bicubic", + align_corners=False, + antialias=True, + ) + ibatch.append(img_lr) + ibatch = torch.cat(ibatch) + feat = self.f(ibatch) + pred = [] + for idx in range(feat[0].shape[0]): + nfeat = [] + #XXX maybe a clone here + for j in range(len(feat)): + nfeat.append(feat[j][idx:idx+1]) + pred.append({"img": inputs[idx][1], + "rescaled": ibatch[idx], + "features": nfeat,#feat[idx], + "path": inputs[idx][0] + }) + return pred + + @torch.inference_mode() + def coarse_cached_match( + self, + frame_A, + frame_B, + ) -> dict[str, torch.Tensor]: + self.eval() + with torch.profiler.record_function("matcher"): + matcher_output = self.matcher( + frame_A["features"], + frame_B["features"], + img_A=frame_A["rescaled"], + img_B=frame_B["rescaled"], + bidirectional=False) + overlap_AB = matcher_output["confidence_AB"][..., :1].sigmoid() + preds = { + "warp_AB": matcher_output["warp_AB"].clone(), + "confidence_AB": matcher_output["confidence_AB"].clone(), + "overlap_AB": overlap_AB.clone(), + "precision_AB": None, + "warp_BA": None, + "confidence_BA": None, + "overlap_BA": None, + "precision_BA": None, + } + return preds + + @torch.inference_mode() + def coarse_match( + self, + img_like_A: ImageLike, + img_like_B: ImageLike, + ) -> dict[str, torch.Tensor]: + self.eval() + img_A = self._load_image(img_like_A) + img_B = self._load_image(img_like_B) + + img_A_lr = F.interpolate( + img_A, + size=(self.H_lr, self.W_lr), + mode="bicubic", + align_corners=False, + antialias=True, + ) + img_B_lr = F.interpolate( + img_B, + size=(self.H_lr, self.W_lr), + mode="bicubic", + align_corners=False, + antialias=True, + ) + + preds = self.coarse_forward(img_A_lr, img_B_lr) + + warp_AB = preds["warp_AB"] + confidence_AB = preds["confidence_AB"] + overlap_AB = confidence_AB[..., :1].sigmoid() + + preds = { + "warp_AB": warp_AB.clone(), + "confidence_AB": confidence_AB.clone(), + "overlap_AB": overlap_AB.clone(), + "precision_AB": None, + "warp_BA": None, + "confidence_BA": None, + "overlap_BA": None, + "precision_BA": None, + } + return preds + @torch.inference_mode() def match( self, @@ -370,6 +498,7 @@ def match( } return preds + @torch.profiler.record_function("sample") def sample(self, preds: dict[str, torch.Tensor], num_corresp: int): warp = preds["warp_AB"] confidence_AB = preds["overlap_AB"] From 5639c33fbe946a58fb00c342f9ec31d9df243611 Mon Sep 17 00:00:00 2001 From: cDc Date: Thu, 26 Feb 2026 07:55:47 +0000 Subject: [PATCH 3/8] fix script --- scripts/batch_match_min.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py index 624531b..0000572 100644 --- a/scripts/batch_match_min.py +++ b/scripts/batch_match_min.py @@ -6,7 +6,7 @@ Supports intelligent caching and pair ordering to maximize throughput. Usage: - python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ + python scripts/batch_match_min.py --pairs-csv pairs.csv --output-dir outputs/ CSV Format: image_A_path,image_B_path @@ -407,7 +407,7 @@ def main(): cache=cache, ) - pbar.update(len(batch)) + pbar.update(len(batch)) # Update progress bar description if stats['processed'] > 0: From 29cd6f9b86f4b609574a2495742d6ae31f998a7c Mon Sep 17 00:00:00 2001 From: cDc Date: Thu, 26 Feb 2026 11:46:09 +0000 Subject: [PATCH 4/8] add npz format --- scripts/batch_match_min.py | 101 +++++++++++++++++++++++++++---------- 1 file changed, 74 insertions(+), 27 deletions(-) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py index 0000572..570aeaa 100644 --- a/scripts/batch_match_min.py +++ b/scripts/batch_match_min.py @@ -88,8 +88,9 @@ def save_results( overlap_AB: torch.Tensor, path_A: str, path_B: str, + save_format: str = 'pt', ): - """Save matching results to pt file. + """Save matching results to pt or npy file. Args: output_path: Output file path @@ -97,22 +98,33 @@ def save_results( warp_AB: (H, W, 2) warp field in normalized coordinates overlap_AB: (H, W, 1) overlap logit path_A, path_B: Image paths + save_format: Format to save ('pt' or 'npy') """ - # Convert to dict - data = { - 'pair_id': pair_id, - 'warp_AB': warp_AB, - 'overlap_AB': overlap_AB, - 'image_A_path': path_A, - 'image_B_path': path_B, - } - # Ensure output directory exists output_path.parent.mkdir(parents=True, exist_ok=True) # Save with torch.profiler.record_function("save_to_disk"): - torch.save(data, output_path) + if save_format == 'npy': + # Convert tensors to numpy and save as npz + np.savez( + output_path, + pair_id=pair_id, + warp_AB=warp_AB.cpu().numpy(), + overlap_AB=overlap_AB.cpu().numpy(), + image_A_path=path_A, + image_B_path=path_B, + ) + else: + # Convert to dict + data = { + 'pair_id': pair_id, + 'warp_AB': warp_AB, + 'overlap_AB': overlap_AB, + 'image_A_path': path_A, + 'image_B_path': path_B, + } + torch.save(data, output_path) @torch.profiler.record_function("load_im") def load_im(path): @@ -212,7 +224,8 @@ def process_real_batch( output_dir: Path, stats: dict, cache, -) -> tuple[list[Path], Optional[np.ndarray]]: + save_format: str = 'pt', +) -> None: output_paths = [] feat1_A = [] @@ -222,9 +235,12 @@ def process_real_batch( img_A = [] img_B = [] + # Determine file extension based on save format + file_ext = '.npz' if save_format == 'npy' else '.pt' + #Assemble batch for pair_id, path_A, path_B in batch_pairs: - output_path = output_dir / f"{pair_id}.pt" + output_path = output_dir / f"{pair_id}{file_ext}" output_paths.append(output_path) frame_A = cache[path_A] frame_B = cache[path_B] @@ -264,6 +280,7 @@ def process_real_batch( overlap_AB=overlap_AB, path_A=path_A, path_B=path_B, + save_format=save_format, ) idx += 1 stats['processed'] += 1 @@ -276,21 +293,35 @@ def process_real_batch( ids.append(pair_id) pA.append(path_A) pB.append(path_B) - # Convert to dict - data = { - 'pair_id': ids, - 'warp_AB': preds["warp_AB"], - 'overlap_AB': preds["overlap_AB"], - 'image_A_path': pA, - 'image_B_path': pB, - } - + + # Use first output path for batch saving + output_path = output_paths[0] + # Ensure output directory exists output_path.parent.mkdir(parents=True, exist_ok=True) # Save with torch.profiler.record_function("save_to_disk"): - torch.save(data, output_path) + if save_format == 'npy': + # Save as npz with arrays + np.savez( + output_path, + pair_id=ids, + warp_AB=preds["warp_AB"].cpu().numpy(), + overlap_AB=preds["overlap_AB"].cpu().numpy(), + image_A_path=pA, + image_B_path=pB, + ) + else: + # Convert to dict + data = { + 'pair_id': ids, + 'warp_AB': preds["warp_AB"], + 'overlap_AB': preds["overlap_AB"], + 'image_A_path': pA, + 'image_B_path': pB, + } + torch.save(data, output_path) def main(): parser = argparse.ArgumentParser( @@ -299,17 +330,21 @@ def main(): epilog=""" Examples: # Basic usage - python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ + python scripts/batch_match_min.py --pairs-csv pairs.csv --output-dir outputs/ + + # Save results in NumPy format + python scripts/batch_match_min.py --pairs-csv pairs.csv --output-dir outputs/ \\ + --save-format npy # With image caching and match sampling - python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + python scripts/batch_match_min.py --pairs-csv pairs.csv --output-dir outputs/ \\ # Fast processing with lower quality - python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + python scripts/batch_match_min.py --pairs-csv pairs.csv --output-dir outputs/ \\ --setting fast --batch-size 4 # With precision calibration (100 pairs × 1000 samples per pair = 100k total) - python scripts/batch_match.py --pairs-csv pairs.csv --output-dir outputs/ \\ + python scripts/batch_match_min.py --pairs-csv pairs.csv --output-dir outputs/ \\ --calibration-pairs 100 --calibration-sample-size 1000 """ ) @@ -339,6 +374,13 @@ def main(): default='base', help='Model setting (default: base)' ) + parser.add_argument( + '--save-format', + type=str, + choices=['pt', 'npy'], + default='pt', + help='Save format for output files: pt (PyTorch) or npy (NumPy) (default: pt)' + ) args = parser.parse_args() @@ -380,6 +422,10 @@ def main(): processed_files = [] + # Adjust file extensions based on save format + file_ext = '.npz' if args.save_format == 'npy' else '.pt' + logger.info(f"Saving results in {args.save_format.upper()} format (extension: {file_ext})") + activities = [torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] #if(True): @@ -405,6 +451,7 @@ def main(): output_dir=args.output_dir, stats=stats, cache=cache, + save_format=args.save_format, ) pbar.update(len(batch)) From 04dca48d2c2ea411bb6494a5abca8a51f0bffbcc Mon Sep 17 00:00:00 2001 From: cDc Date: Thu, 26 Feb 2026 11:47:11 +0000 Subject: [PATCH 5/8] make profiling optional --- scripts/batch_match_min.py | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py index 570aeaa..c1e8719 100644 --- a/scripts/batch_match_min.py +++ b/scripts/batch_match_min.py @@ -381,6 +381,11 @@ def main(): default='pt', help='Save format for output files: pt (PyTorch) or npy (NumPy) (default: pt)' ) + parser.add_argument( + '--save-trace', + action='store_true', + help='Save PyTorch profiler trace to trace.json (default: off)' + ) args = parser.parse_args() @@ -426,12 +431,21 @@ def main(): file_ext = '.npz' if args.save_format == 'npy' else '.pt' logger.info(f"Saving results in {args.save_format.upper()} format (extension: {file_ext})") - activities = [torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA] - #if(True): - with torch.profiler.profile(activities=activities, - record_shapes=True, - profile_memory=False,) as prof: + # Decide whether to use profiler + if args.save_trace: + logger.info("Profiler trace saving enabled") + profiler_context = torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA], + record_shapes=True, + profile_memory=False, + ) + else: + # Use a no-op context manager when profiling is disabled + from contextlib import nullcontext + profiler_context = nullcontext() + + with profiler_context as prof: with torch.profiler.record_function("streaming_feature_cache"): cache = streaming_cache(model, pairs) @@ -469,7 +483,9 @@ def main(): idx = batch_end - prof.export_chrome_trace("trace.json") + if args.save_trace and prof is not None: + prof.export_chrome_trace("trace.json") + # Final statistics elapsed = time.time() - stats['start_time'] logger.info("\n" + "="*60) From 8840ac8f346ada044c54d5c955704fc12ffe5fe6 Mon Sep 17 00:00:00 2001 From: cDc Date: Thu, 26 Feb 2026 12:29:35 +0000 Subject: [PATCH 6/8] export without 1 at batch size --- scripts/batch_match_min.py | 37 +++++++++++++++++++++++++------------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py index c1e8719..536c28b 100644 --- a/scripts/batch_match_min.py +++ b/scripts/batch_match_min.py @@ -106,21 +106,29 @@ def save_results( # Save with torch.profiler.record_function("save_to_disk"): if save_format == 'npy': - # Convert tensors to numpy and save as npz + # Convert tensors to numpy and save as npz (squeeze batch dimension if present) + assert warp_AB.dim() > 3 and warp_AB.shape[0] == 1 or warp_AB.dim() == 3, \ + f"Expected batch dimension of 1, got shape {warp_AB.shape}" + assert overlap_AB.dim() > 3 and overlap_AB.shape[0] == 1 or overlap_AB.dim() == 3, \ + f"Expected batch dimension of 1, got shape {overlap_AB.shape}" np.savez( output_path, pair_id=pair_id, - warp_AB=warp_AB.cpu().numpy(), - overlap_AB=overlap_AB.cpu().numpy(), + warp_AB=warp_AB.squeeze(0).cpu().numpy() if warp_AB.dim() > 3 else warp_AB.cpu().numpy(), + overlap_AB=overlap_AB.squeeze(0).cpu().numpy() if overlap_AB.dim() > 3 else overlap_AB.cpu().numpy(), image_A_path=path_A, image_B_path=path_B, ) else: - # Convert to dict + # Convert to dict (squeeze batch dimension if present) + assert warp_AB.dim() > 3 and warp_AB.shape[0] == 1 or warp_AB.dim() == 3, \ + f"Expected batch dimension of 1, got shape {warp_AB.shape}" + assert overlap_AB.dim() > 3 and overlap_AB.shape[0] == 1 or overlap_AB.dim() == 3, \ + f"Expected batch dimension of 1, got shape {overlap_AB.shape}" data = { 'pair_id': pair_id, - 'warp_AB': warp_AB, - 'overlap_AB': overlap_AB, + 'warp_AB': warp_AB.squeeze(0) if warp_AB.dim() > 3 else warp_AB, + 'overlap_AB': overlap_AB.squeeze(0) if overlap_AB.dim() > 3 else overlap_AB, 'image_A_path': path_A, 'image_B_path': path_B, } @@ -302,22 +310,27 @@ def process_real_batch( # Save with torch.profiler.record_function("save_to_disk"): + assert preds["warp_AB"].shape[0] == 1, \ + f"Expected batch size 1, got {preds['warp_AB'].shape[0]}" + assert preds["overlap_AB"].shape[0] == 1, \ + f"Expected batch size 1, got {preds['overlap_AB'].shape[0]}" + if save_format == 'npy': - # Save as npz with arrays + # Save as npz with arrays (squeeze batch dimension) np.savez( output_path, pair_id=ids, - warp_AB=preds["warp_AB"].cpu().numpy(), - overlap_AB=preds["overlap_AB"].cpu().numpy(), + warp_AB=preds["warp_AB"].squeeze(0).cpu().numpy(), + overlap_AB=preds["overlap_AB"].squeeze(0).cpu().numpy(), image_A_path=pA, image_B_path=pB, ) else: - # Convert to dict + # Convert to dict (squeeze batch dimension) data = { 'pair_id': ids, - 'warp_AB': preds["warp_AB"], - 'overlap_AB': preds["overlap_AB"], + 'warp_AB': preds["warp_AB"].squeeze(0), + 'overlap_AB': preds["overlap_AB"].squeeze(0), 'image_A_path': pA, 'image_B_path': pB, } From 854b6859d5aba84455223c767bc86d5902f68138 Mon Sep 17 00:00:00 2001 From: cDc Date: Mon, 2 Mar 2026 09:48:51 +0000 Subject: [PATCH 7/8] add full-resolution refinement option to cached matching --- scripts/batch_match_min.py | 114 ++++++++++++++++++++++++++++++++++--- 1 file changed, 107 insertions(+), 7 deletions(-) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py index 536c28b..bb56136 100644 --- a/scripts/batch_match_min.py +++ b/scripts/batch_match_min.py @@ -226,6 +226,91 @@ def streaming_cache(model, pairs): return cache +def _refined_cached_match( + model: RoMaV2, + frame_A, + frame_B, +) -> dict[str, torch.Tensor]: + """Perform cached matching with full refinement pipeline. + + This produces full-resolution warps by applying the refinement stages + after the coarse matcher output, using pre-cached features to avoid + redundant feature extraction. + """ + from romav2.romav2 import _interpolate_warp_and_confidence + + model.eval() + + # Get coarse predictions from matcher using cached features + with torch.profiler.record_function("matcher"): + matcher_output = model.matcher( + frame_A["features"], + frame_B["features"], + img_A=frame_A["rescaled"], + img_B=frame_B["rescaled"], + bidirectional=model.bidirectional + ) + + warp_AB = matcher_output["warp_AB"] + confidence_AB = matcher_output["confidence_AB"] + + # Apply refinement stages using cached rescaled images + img_A_lr = frame_A["rescaled"] + img_B_lr = frame_B["rescaled"] + + B, C, H, W = img_A_lr.shape + device = warp_AB.device + scale_factor = torch.tensor( + (W / model.anchor_width, H / model.anchor_height), device=device + ) + + # Extract refiner features + with torch.profiler.record_function("refiner_features"): + refiner_features_A = model.refiner_features(img_A_lr) + refiner_features_B = model.refiner_features(img_B_lr) + + # Refine through patch sizes + for patch_size_str, refiner in model.refiners.items(): + patch_size = int(patch_size_str) + + warp_AB, confidence_AB = _interpolate_warp_and_confidence( + warp=warp_AB, + confidence=confidence_AB, + H=H, + W=W, + patch_size=patch_size, + zero_out_precision=False, + ) + + f_patch_A = refiner_features_A[patch_size] + f_patch_B = refiner_features_B[patch_size] + + with torch.profiler.record_function("refiner"): + refiner_output_AB = refiner( + f_A=f_patch_A, + f_B=f_patch_B, + prev_warp=warp_AB, + prev_confidence=confidence_AB, + scale_factor=scale_factor, + ) + + warp_AB = refiner_output_AB["warp"] + confidence_AB = refiner_output_AB["confidence"] + + overlap_AB = confidence_AB[..., :1].sigmoid() + + preds = { + "warp_AB": warp_AB.clone(), + "confidence_AB": confidence_AB.clone(), + "overlap_AB": overlap_AB.clone(), + "precision_AB": None, + "warp_BA": None, + "confidence_BA": None, + "overlap_BA": None, + "precision_BA": None, + } + return preds + def process_real_batch( model: RoMaV2, batch_pairs: list[tuple[str, str, str]], @@ -233,6 +318,7 @@ def process_real_batch( stats: dict, cache, save_format: str = 'pt', + refine: bool = False, ) -> None: output_paths = [] @@ -261,14 +347,17 @@ def process_real_batch( batch_A = {"features": [torch.cat(feat1_A), torch.cat(feat2_A)], - "rescaled": torch.cat(img_A)} + "rescaled": torch.stack(img_A) if img_A[0].dim() == 3 else torch.cat(img_A)} batch_B = {"features": [torch.cat(feat1_B), torch.cat(feat2_B)], - "rescaled": torch.cat(img_B)} + "rescaled": torch.stack(img_B) if img_B[0].dim() == 3 else torch.cat(img_B)} #Perform Inference with torch.profiler.record_function("DNN_exec"): - preds = model.coarse_cached_match(batch_A, batch_B) + if refine: + preds = _refined_cached_match(model, batch_A, batch_B) + else: + preds = model.coarse_cached_match(batch_A, batch_B) #Dissassemble batch with torch.profiler.record_function("save_batch"): @@ -320,8 +409,8 @@ def process_real_batch( np.savez( output_path, pair_id=ids, - warp_AB=preds["warp_AB"].squeeze(0).cpu().numpy(), - overlap_AB=preds["overlap_AB"].squeeze(0).cpu().numpy(), + warp_AB=preds["warp_AB"].squeeze(0).detach().cpu().numpy(), + overlap_AB=preds["overlap_AB"].squeeze(0).detach().cpu().numpy(), image_A_path=pA, image_B_path=pB, ) @@ -329,8 +418,8 @@ def process_real_batch( # Convert to dict (squeeze batch dimension) data = { 'pair_id': ids, - 'warp_AB': preds["warp_AB"].squeeze(0), - 'overlap_AB': preds["overlap_AB"].squeeze(0), + 'warp_AB': preds["warp_AB"].squeeze(0).detach(), + 'overlap_AB': preds["overlap_AB"].squeeze(0).detach(), 'image_A_path': pA, 'image_B_path': pB, } @@ -399,6 +488,11 @@ def main(): action='store_true', help='Save PyTorch profiler trace to trace.json (default: off)' ) + parser.add_argument( + '--refine', + action='store_true', + help='Enable full-resolution refinement (slower but higher quality, default: off)' + ) args = parser.parse_args() @@ -422,6 +516,11 @@ def main(): model.apply_setting(args.setting) model.eval() + if args.refine: + logger.info("Refinement enabled: outputs will be at full resolution (slower)") + else: + logger.info("Refinement disabled: outputs will be at coarse resolution (4x downsampled)") + # Statistics stats = { 'processed': 0, @@ -479,6 +578,7 @@ def main(): stats=stats, cache=cache, save_format=args.save_format, + refine=args.refine, ) pbar.update(len(batch)) From e1059b8a92dbda2d437b5135ae9d6ca073b3efc2 Mon Sep 17 00:00:00 2001 From: cDc Date: Mon, 2 Mar 2026 10:23:52 +0000 Subject: [PATCH 8/8] add custom resolution option to main function --- scripts/batch_match_min.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/scripts/batch_match_min.py b/scripts/batch_match_min.py index bb56136..a833901 100644 --- a/scripts/batch_match_min.py +++ b/scripts/batch_match_min.py @@ -493,6 +493,12 @@ def main(): action='store_true', help='Enable full-resolution refinement (slower but higher quality, default: off)' ) + parser.add_argument( + '--resolution', + type=str, + default=None, + help='Custom resolution to override setting (e.g., "640" for 640x640 or "800" for 800x600)' + ) args = parser.parse_args() @@ -514,6 +520,17 @@ def main(): logger.info(f"Initializing RoMaV2 with setting: {args.setting}") model = RoMaV2(RoMaV2.Cfg(compile=False)) model.apply_setting(args.setting) + + # Override resolution if specified + if args.resolution: + # Square resolution + res = int(args.resolution) + model.H_lr = res + model.W_lr = res + logger.info(f"Custom resolution: {res}x{res}") + else: + logger.info(f"Using preset resolution: {model.W_lr}x{model.H_lr}") + model.eval() if args.refine: