Add two-video drone 3DGS pipeline with Apple Silicon fixes
- main.py: extract frames from two videos, run COLMAP feature extraction - match_features.py: Python-based within-video SIFT matching via OpenCV (replaces colmap exhaustive_matcher which segfaults on ARM64 in COLMAP 4.x) - match_crossvideo.py: exhaustive cross-video matching (v1×v2) to stitch two flights into a single COLMAP model - run.sh: entry point for frame extraction + feature extraction - train_splat.sh: ns-process-data → splatfacto → .ply export, with correct PATH for Homebrew ffmpeg and MPS device flags for Apple Silicon - .gitignore: exclude videos, generated scene data, venv, logs - README.md: full pipeline walkthrough, all known issues and fixes Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Jon
155
main.py
Normal file
155
main.py
Normal file
@@ -0,0 +1,155 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Two-video drone footage -> 3DGS pipeline for Apple Silicon Macs.
|
||||
|
||||
Usage:
|
||||
python drone_3dgs_pipeline.py \
|
||||
--video1 path/to/flight1.mp4 \
|
||||
--video2 path/to/flight2.mp4 \
|
||||
--output_dir ./my_scene \
|
||||
--fps 2
|
||||
|
||||
What it does:
|
||||
1. Extracts frames from both videos at the given fps (default 2 fps).
|
||||
2. Pools them into one folder with non-colliding names.
|
||||
3. Runs COLMAP feature extraction, matching, and sparse reconstruction.
|
||||
4. Hands you a Nerfstudio-ready folder structure to train splatfacto.
|
||||
|
||||
Then run:
|
||||
ns-train splatfacto --data ./my_scene
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def run(cmd, check=True):
|
||||
"""Run a shell command, streaming output."""
|
||||
print(f"\n>>> {' '.join(str(c) for c in cmd)}\n")
|
||||
result = subprocess.run(cmd, check=check)
|
||||
return result
|
||||
|
||||
|
||||
def check_dependencies():
|
||||
"""Verify ffmpeg and colmap are installed."""
|
||||
for tool in ("ffmpeg", "colmap"):
|
||||
if shutil.which(tool) is None:
|
||||
sys.exit(f"ERROR: {tool} not found. Install with: brew install {tool}")
|
||||
|
||||
|
||||
def extract_frames(video_path: Path, out_dir: Path, fps: float, prefix: str):
|
||||
"""Extract frames from a video using ffmpeg at the given fps."""
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
pattern = str(out_dir / f"{prefix}_%05d.jpg")
|
||||
run([
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-r", str(fps),
|
||||
"-q:v", "2", # JPEG quality (2 = high)
|
||||
"-y", # overwrite
|
||||
pattern,
|
||||
])
|
||||
count = len(list(out_dir.glob(f"{prefix}_*.jpg")))
|
||||
print(f"Extracted {count} frames from {video_path.name} -> {out_dir}")
|
||||
return count
|
||||
|
||||
|
||||
def run_colmap(workspace: Path, images_dir: Path, use_gpu: bool = False):
|
||||
"""Run the COLMAP SfM pipeline: feature extraction, matching, mapping."""
|
||||
sparse_dir = workspace / "sparse"
|
||||
sparse_dir.mkdir(parents=True, exist_ok=True)
|
||||
db_path = workspace / "database.db"
|
||||
|
||||
# Step 1: feature extraction (COLMAP 4.x renamed SiftExtraction → FeatureExtraction)
|
||||
run([
|
||||
"colmap", "feature_extractor",
|
||||
"--database_path", str(db_path),
|
||||
"--image_path", str(images_dir),
|
||||
"--ImageReader.single_camera", "1",
|
||||
"--ImageReader.camera_model", "OPENCV",
|
||||
"--FeatureExtraction.use_gpu", "1" if use_gpu else "0",
|
||||
])
|
||||
|
||||
# Step 2: Python matcher (COLMAP 4.x exhaustive_matcher segfaults on Apple Silicon ARM64;
|
||||
# match_features.py reads SIFT descriptors from the DB and matches via OpenCV BFMatcher)
|
||||
script = Path(__file__).parent / "match_features.py"
|
||||
run([sys.executable, str(script), "--db", str(db_path)])
|
||||
|
||||
# Step 3: sparse reconstruction (this is the slow one)
|
||||
run([
|
||||
"colmap", "mapper",
|
||||
"--database_path", str(db_path),
|
||||
"--image_path", str(images_dir),
|
||||
"--output_path", str(sparse_dir),
|
||||
])
|
||||
|
||||
# COLMAP writes to sparse/0/ by default
|
||||
print(f"\nCOLMAP done. Reconstruction in {sparse_dir}/0/")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--video1", type=Path, required=True)
|
||||
parser.add_argument("--video2", type=Path, required=True)
|
||||
parser.add_argument("--output_dir", type=Path, required=True)
|
||||
parser.add_argument("--fps", type=float, default=2.0,
|
||||
help="Frames per second to extract (default: 2). "
|
||||
"Higher = more frames, slower training, better quality.")
|
||||
parser.add_argument("--use_gpu", action="store_true",
|
||||
help="Try GPU SIFT in COLMAP (often unreliable on M1; default off).")
|
||||
args = parser.parse_args()
|
||||
|
||||
check_dependencies()
|
||||
|
||||
if not args.video1.exists() or not args.video2.exists():
|
||||
sys.exit("ERROR: one or both video files not found.")
|
||||
|
||||
workspace = args.output_dir
|
||||
images_dir = workspace / "images"
|
||||
workspace.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Extract frames from both videos into the SAME images folder, with prefixes
|
||||
# so filenames don't collide. COLMAP treats them as one set automatically.
|
||||
print(f"\n=== Extracting frames from {args.video1.name} ===")
|
||||
n1 = extract_frames(args.video1, images_dir, args.fps, prefix="v1")
|
||||
print(f"\n=== Extracting frames from {args.video2.name} ===")
|
||||
n2 = extract_frames(args.video2, images_dir, args.fps, prefix="v2")
|
||||
total = n1 + n2
|
||||
|
||||
print(f"\n=== Total frames: {total} ===")
|
||||
if total > 800:
|
||||
print("WARNING: lots of frames. Consider lowering --fps. "
|
||||
"Exhaustive matching will be slow; switch to sequential_matcher if needed.")
|
||||
|
||||
# Run COLMAP
|
||||
print("\n=== Running COLMAP (this is the slow part, get a coffee) ===")
|
||||
run_colmap(workspace, images_dir, use_gpu=args.use_gpu)
|
||||
|
||||
# Print next steps
|
||||
print(f"""
|
||||
========================================================================
|
||||
DONE with SfM. Your scene is at: {workspace}
|
||||
|
||||
Next, train the splat. Two options on Mac:
|
||||
|
||||
OPTION 1: Nerfstudio (Python, scriptable)
|
||||
pip install nerfstudio
|
||||
ns-process-data images --data {images_dir} --output-dir {workspace}/ns_data \\
|
||||
--skip-colmap --colmap-model-path {workspace}/sparse/0
|
||||
ns-train splatfacto --data {workspace}/ns_data
|
||||
|
||||
OPTION 2: Brush (Rust binary, faster on Mac)
|
||||
Download from https://github.com/ArthurBrussee/brush
|
||||
brush --source {workspace}
|
||||
|
||||
View result:
|
||||
- Online: https://playcanvas.com/supersplat/editor (drag .ply in)
|
||||
- Local: install SuperSplat or use Nerfstudio's viewer
|
||||
========================================================================
|
||||
""")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user