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..a833901 --- /dev/null +++ b/scripts/batch_match_min.py @@ -0,0 +1,638 @@ +#!/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_min.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_format: str = 'pt', +): + """Save matching results to pt or npy 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 + save_format: Format to save ('pt' or 'npy') + """ + # Ensure output directory exists + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Save + with torch.profiler.record_function("save_to_disk"): + if save_format == 'npy': + # 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.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 (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.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, + } + 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 _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]], + output_dir: Path, + stats: dict, + cache, + save_format: str = 'pt', + refine: bool = False, +) -> None: + + output_paths = [] + feat1_A = [] + feat1_B = [] + feat2_A = [] + feat2_B = [] + 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}{file_ext}" + 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.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.stack(img_B) if img_B[0].dim() == 3 else torch.cat(img_B)} + + #Perform Inference + with torch.profiler.record_function("DNN_exec"): + 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"): + 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, + save_format=save_format, + ) + 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) + + # 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"): + 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 (squeeze batch dimension) + np.savez( + output_path, + pair_id=ids, + 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, + ) + else: + # Convert to dict (squeeze batch dimension) + data = { + 'pair_id': ids, + 'warp_AB': preds["warp_AB"].squeeze(0).detach(), + 'overlap_AB': preds["overlap_AB"].squeeze(0).detach(), + 'image_A_path': pA, + 'image_B_path': pB, + } + 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_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_min.py --pairs-csv pairs.csv --output-dir outputs/ \\ + + # Fast processing with lower quality + 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_min.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( + '--save-format', + type=str, + choices=['pt', 'npy'], + 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)' + ) + parser.add_argument( + '--refine', + 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() + + # 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) + + # 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: + 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, + '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 = [] + + # 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})") + + # 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) + + 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, + save_format=args.save_format, + refine=args.refine, + ) + + 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 + + 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) + 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 06aa2cc..1a8963e 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, @@ -174,12 +175,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 = ( @@ -204,8 +207,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 = ( @@ -231,21 +235,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 @@ -296,6 +302,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, @@ -368,6 +496,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"]