SHA256
1586 lines
69 KiB
Python
1586 lines
69 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
pipeline_360.py -- End-to-end 360 photogrammetry prep pipeline for COLMAP.
|
|
|
|
Merges what used to be four separate scripts (extract_sharpest_frame.py,
|
|
generate_sfm_masks.py, project_to_cubemap.py, generate_rig_config.py) into
|
|
one tool with one project directory layout. The old multi-session merge
|
|
step (merge_frames_masks.py) is gone -- this pipeline is single-session.
|
|
|
|
Also gone: the Metashape/RealityCapture mask "convention" options that used
|
|
to be threaded through every stage. Everything downstream commits to
|
|
COLMAP's rig-based reconstruction now, so there was never a real reason for
|
|
a mask written for stage 2 to end up in a different format than stage 3
|
|
expects. One convention (COLMAP's: black=excluded, filename=<image>.png)
|
|
everywhere.
|
|
|
|
PIPELINE
|
|
--------
|
|
video.mp4
|
|
| [1: extract] sharpest frame per --extract-chunk-size window
|
|
v
|
|
01_frames/ equirect frames (output_frame_00001.png, ...)
|
|
| [2: masks] YOLO-seg, cubemap-projected detection to fix
|
|
| pole distortion, reprojected back to equirect
|
|
v
|
|
02_masks/ equirect masks (output_frame_00001.png.png, ...)
|
|
| [3: cubemap] each equirect frame+mask -> 6 perspective faces
|
|
v
|
|
03_cubemap/
|
|
images/ output_frame_00001_F.png, ..._R.png, ...
|
|
masks/ output_frame_00001_F.png.png, ...
|
|
| [4: rig] COLMAP rig config + nested {FACE}/{stem} tree
|
|
v
|
|
04_colmap/
|
|
images/F/output_frame_00001.png, images/R/..., ...
|
|
masks/F/output_frame_00001.png.png, ...
|
|
rig_config.json
|
|
|
|
Run the whole thing:
|
|
python pipeline_360.py all --project-dir ./project --video ./raw.mp4
|
|
|
|
Resume from a later stage (e.g. you already have frames+masks and just
|
|
want to redo the cubemap/rig stages with a different --cubemap-face-size):
|
|
python pipeline_360.py all --project-dir ./project --start-stage cubemap
|
|
|
|
Run stages individually (mirrors the original standalone scripts, useful
|
|
for debugging a single stage in isolation):
|
|
python pipeline_360.py extract --video raw.mp4 --output-dir 01_frames
|
|
python pipeline_360.py masks --input-dir 01_frames --output-dir 02_masks
|
|
python pipeline_360.py cubemap --frames-dir 01_frames --masks-dir 02_masks \\
|
|
--output-frames-dir 03_cubemap/images --output-masks-dir 03_cubemap/masks
|
|
python pipeline_360.py rig --images-dir 03_cubemap/images --masks-dir 03_cubemap/masks \\
|
|
--output 04_colmap/rig_config.json --link-tree 04_colmap/images \\
|
|
--masks-link-tree 04_colmap/masks
|
|
|
|
Self-test for the rig-config geometry/schema logic (no video/GPU needed):
|
|
python pipeline_360.py rig --self-test
|
|
|
|
Requirements:
|
|
pip install opencv-python numpy tqdm py360convert ultralytics --break-system-packages
|
|
(ultralytics is only needed for the masks stage; the other three stages
|
|
work without it.)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import math
|
|
import multiprocessing
|
|
import re
|
|
import struct
|
|
import sys
|
|
import zlib
|
|
from collections import defaultdict
|
|
from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, as_completed, wait
|
|
from math import ceil
|
|
from pathlib import Path
|
|
from time import monotonic
|
|
from typing import Callable, Dict, List, Optional, Tuple
|
|
|
|
try:
|
|
import cv2
|
|
except ImportError:
|
|
cv2 = None
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
from tqdm import tqdm
|
|
except ImportError:
|
|
tqdm = None
|
|
|
|
|
|
# ==========================================================================
|
|
# Shared: logging / progress / cancellation
|
|
#
|
|
# One shared plumbing layer used by every stage, so the GUI only has to
|
|
# understand one line protocol (plain lines = log, "[progress] "-prefixed
|
|
# lines = replace the last line) regardless of which stage is running.
|
|
# ==========================================================================
|
|
|
|
Logger = Optional[Callable[[str], None]]
|
|
CancelChecker = Optional[Callable[[], bool]]
|
|
PROGRESS_PREFIX = "[progress] "
|
|
SPINNER_FRAMES = "|/-\\"
|
|
MULTIPROCESS_PROGRESS_INTERVAL = 0.2
|
|
|
|
IMAGE_EXTENSIONS = frozenset({".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"})
|
|
FACES: Tuple[str, ...] = ("F", "R", "B", "L", "U", "D")
|
|
FACE_NAMES = {"F": "front", "R": "right", "B": "back", "L": "left", "U": "up", "D": "down"}
|
|
|
|
|
|
class PipelineError(Exception):
|
|
pass
|
|
|
|
|
|
class PipelineCancelled(PipelineError):
|
|
pass
|
|
|
|
|
|
class GuiTqdm(tqdm if tqdm is not None else object):
|
|
def __init__(self, *args, logger: Logger = None, **kwargs) -> None:
|
|
self.logger = logger
|
|
self._last_displayed_message = ""
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def update(self, n=1):
|
|
displayed = super().update(n)
|
|
if self.logger is not None:
|
|
self.display()
|
|
return displayed
|
|
|
|
def display(self, msg=None, pos=None) -> None:
|
|
if self.logger is None:
|
|
super().display(msg=msg, pos=pos)
|
|
return
|
|
progress_message = msg if msg is not None else str(self)
|
|
cleaned = progress_message.strip()
|
|
if cleaned and cleaned != self._last_displayed_message:
|
|
self._last_displayed_message = cleaned
|
|
emit_progress(cleaned, self.logger)
|
|
|
|
|
|
def emit(message: str, logger: Logger = None) -> None:
|
|
if logger is None:
|
|
print(message)
|
|
else:
|
|
logger(message)
|
|
|
|
|
|
def emit_progress(message: str, logger: Logger = None) -> None:
|
|
if logger is not None:
|
|
logger(f"{PROGRESS_PREFIX}{message}")
|
|
|
|
|
|
def print_stage_header(index: int, total: int, title: str, logger: Logger = None) -> None:
|
|
emit("", logger)
|
|
emit(f"==== STAGE {index}/{total}: {title} ====", logger)
|
|
|
|
|
|
def create_progress(total: Optional[int], desc: str, logger: Logger = None):
|
|
ensure_tqdm_available()
|
|
kwargs = {
|
|
"total": total, "desc": desc, "unit": "frame",
|
|
"dynamic_ncols": True, "leave": False,
|
|
"bar_format": "{l_bar}{bar:20}{r_bar}",
|
|
}
|
|
if logger is None:
|
|
return tqdm(**kwargs)
|
|
kwargs.update(mininterval=0, miniters=1, ascii=False)
|
|
return GuiTqdm(logger=logger, **kwargs)
|
|
|
|
|
|
def ensure_opencv_available() -> None:
|
|
if cv2 is None:
|
|
raise PipelineError("opencv-python is required. Install with: pip install opencv-python --break-system-packages")
|
|
|
|
|
|
def ensure_tqdm_available() -> None:
|
|
if tqdm is None:
|
|
raise PipelineError("tqdm is required. Install with: pip install tqdm --break-system-packages")
|
|
|
|
|
|
def ensure_py360convert_available():
|
|
try:
|
|
import py360convert
|
|
except ImportError:
|
|
raise PipelineError("py360convert is required. Install with: pip install py360convert --break-system-packages")
|
|
return py360convert
|
|
|
|
|
|
def raise_if_cancelled(should_cancel: CancelChecker = None) -> None:
|
|
if should_cancel is not None and should_cancel():
|
|
raise PipelineCancelled("Processing was cancelled.")
|
|
|
|
|
|
def gather_images(input_dir: Path) -> List[Path]:
|
|
frames = sorted(f for f in input_dir.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS)
|
|
if not frames:
|
|
raise PipelineError(f"No image files found in {input_dir}")
|
|
return frames
|
|
|
|
|
|
# ==========================================================================
|
|
# STAGE 1: extract sharpest frame per chunk from a video
|
|
# ==========================================================================
|
|
#
|
|
# Picks the sharpest (least motion-blurred) frame out of every
|
|
# --extract-chunk-size window, using variance-of-Laplacian as the sharpness
|
|
# metric. Two-pass: analyze every frame's sharpness to a CSV (parallelized
|
|
# across --extract-workers, since this is the slow part), then extract just
|
|
# the winners.
|
|
|
|
def get_frame_count(capture) -> Optional[int]:
|
|
total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
return total if total > 0 else None
|
|
|
|
|
|
def compute_sharpness(frame, scale_width: int) -> float:
|
|
height, width = frame.shape[:2]
|
|
if scale_width > 0 and width > scale_width:
|
|
ratio = scale_width / float(width)
|
|
frame = cv2.resize(frame, (scale_width, max(1, int(height * ratio))), interpolation=cv2.INTER_AREA)
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
return float(cv2.Laplacian(gray, cv2.CV_64F).var())
|
|
|
|
|
|
def build_analysis_ranges(total_frames: int, workers: int) -> List[Tuple[int, int]]:
|
|
task_count = max(workers * 8, workers)
|
|
range_size = max(1, ceil(total_frames / task_count))
|
|
ranges, start = [], 0
|
|
while start < total_frames:
|
|
end = min(total_frames, start + range_size)
|
|
ranges.append((start, end))
|
|
start = end
|
|
return ranges
|
|
|
|
|
|
def analyze_video_range(video_path: str, start_frame: int, end_frame: int, scale_width: int) -> List[Tuple[int, float]]:
|
|
ensure_opencv_available()
|
|
capture = cv2.VideoCapture(video_path)
|
|
if not capture.isOpened():
|
|
raise PipelineError(f"Failed to open video: {video_path}")
|
|
capture.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
|
|
results, frame_number = [], start_frame
|
|
while frame_number < end_frame:
|
|
ok, frame = capture.read()
|
|
if not ok:
|
|
break
|
|
results.append((frame_number, compute_sharpness(frame, scale_width)))
|
|
frame_number += 1
|
|
capture.release()
|
|
return results
|
|
|
|
|
|
def analyze_video_single_process(video_file, metadata_path, scale_width, total_frames, logger=None, should_cancel=None) -> int:
|
|
capture = cv2.VideoCapture(str(video_file))
|
|
if not capture.isOpened():
|
|
raise PipelineError(f"Failed to open video: {video_file}")
|
|
frame_number = 0
|
|
try:
|
|
with metadata_path.open("w", newline="", encoding="utf-8") as fh:
|
|
writer = csv.writer(fh)
|
|
writer.writerow(["frame", "sharpness"])
|
|
with create_progress(total=total_frames, desc="Analyze", logger=logger) as bar:
|
|
while True:
|
|
raise_if_cancelled(should_cancel)
|
|
ok, frame = capture.read()
|
|
if not ok:
|
|
break
|
|
sharpness = compute_sharpness(frame, scale_width)
|
|
writer.writerow([frame_number, f"{sharpness:.10f}"])
|
|
frame_number += 1
|
|
bar.update(1)
|
|
except PipelineCancelled:
|
|
capture.release()
|
|
if metadata_path.exists():
|
|
metadata_path.unlink()
|
|
raise
|
|
capture.release()
|
|
return frame_number
|
|
|
|
|
|
def analyze_video_multi_process(video_file, metadata_path, scale_width, total_frames, workers, logger=None, should_cancel=None) -> int:
|
|
ranges = build_analysis_ranges(total_frames, workers)
|
|
completed: Dict[int, List[Tuple[int, float]]] = {}
|
|
next_index, written, spinner, last_refresh = 0, 0, 0, 0.0
|
|
|
|
try:
|
|
with metadata_path.open("w", newline="", encoding="utf-8") as fh:
|
|
writer = csv.writer(fh)
|
|
writer.writerow(["frame", "sharpness"])
|
|
with create_progress(total=total_frames, desc="Analyze", logger=logger) as bar:
|
|
with ProcessPoolExecutor(max_workers=workers) as executor:
|
|
future_map = {
|
|
executor.submit(analyze_video_range, str(video_file), s, e, scale_width): i
|
|
for i, (s, e) in enumerate(ranges)
|
|
}
|
|
bar.set_postfix_str(f"{SPINNER_FRAMES[spinner]} {len(future_map)} ranges")
|
|
bar.display()
|
|
last_refresh = monotonic()
|
|
|
|
while future_map:
|
|
raise_if_cancelled(should_cancel)
|
|
done, _ = wait(future_map, timeout=0.1, return_when=FIRST_COMPLETED)
|
|
if not done:
|
|
now = monotonic()
|
|
if now - last_refresh >= MULTIPROCESS_PROGRESS_INTERVAL:
|
|
spinner = (spinner + 1) % len(SPINNER_FRAMES)
|
|
bar.set_postfix_str(f"{SPINNER_FRAMES[spinner]} {len(future_map)} ranges")
|
|
bar.display()
|
|
last_refresh = now
|
|
continue
|
|
for fut in done:
|
|
idx = future_map.pop(fut)
|
|
completed[idx] = fut.result()
|
|
bar.update(len(completed[idx]))
|
|
spinner = (spinner + 1) % len(SPINNER_FRAMES)
|
|
bar.set_postfix_str(f"{SPINNER_FRAMES[spinner]} {len(future_map)} ranges")
|
|
bar.display()
|
|
last_refresh = monotonic()
|
|
while next_index in completed:
|
|
for frame_number, sharpness in completed.pop(next_index):
|
|
writer.writerow([frame_number, f"{sharpness:.10f}"])
|
|
written += 1
|
|
next_index += 1
|
|
bar.set_postfix_str("done")
|
|
bar.display()
|
|
except PipelineCancelled:
|
|
if metadata_path.exists():
|
|
metadata_path.unlink()
|
|
raise
|
|
return written
|
|
|
|
|
|
def analyze_video(video_file, metadata_path, scale_width, workers=4, logger=None, should_cancel=None) -> None:
|
|
ensure_opencv_available()
|
|
ensure_tqdm_available()
|
|
metadata_path.parent.mkdir(parents=True, exist_ok=True)
|
|
emit("[1/3] Analyzing sharpness...", logger)
|
|
|
|
probe = cv2.VideoCapture(str(video_file))
|
|
if not probe.isOpened():
|
|
raise PipelineError(f"Failed to open video: {video_file}")
|
|
total_frames = get_frame_count(probe)
|
|
probe.release()
|
|
|
|
if workers <= 1 or total_frames is None or total_frames <= 1:
|
|
n = analyze_video_single_process(video_file, metadata_path, scale_width, total_frames, logger, should_cancel)
|
|
else:
|
|
n = analyze_video_multi_process(video_file, metadata_path, scale_width, total_frames, workers, logger, should_cancel)
|
|
|
|
if n == 0:
|
|
raise PipelineError("No frames could be read from the video.")
|
|
|
|
|
|
def parse_best_frames(metadata_path: Path, chunk_size: int, logger: Logger = None, should_cancel: CancelChecker = None) -> List[int]:
|
|
ensure_tqdm_available()
|
|
frame_data: List[Tuple[int, float]] = []
|
|
try:
|
|
with metadata_path.open("r", encoding="utf-8") as fh:
|
|
total_lines = max(0, sum(1 for _ in fh) - 1)
|
|
with metadata_path.open("r", newline="", encoding="utf-8") as fh:
|
|
reader = csv.DictReader(fh)
|
|
with create_progress(total=total_lines or None, desc="Parse", logger=logger) as bar:
|
|
for row in reader:
|
|
raise_if_cancelled(should_cancel)
|
|
try:
|
|
frame_data.append((int(row["frame"]), float(row["sharpness"])))
|
|
except (KeyError, TypeError, ValueError):
|
|
pass
|
|
finally:
|
|
bar.update(1)
|
|
except FileNotFoundError as exc:
|
|
raise PipelineError(f"Metadata file not found: {metadata_path}") from exc
|
|
|
|
if not frame_data:
|
|
raise PipelineError("No valid frame data was found in metadata.")
|
|
|
|
best, chunk = [], []
|
|
for item in frame_data:
|
|
chunk.append(item)
|
|
if len(chunk) == chunk_size:
|
|
best.append(max(chunk, key=lambda x: x[1])[0])
|
|
chunk = []
|
|
if chunk:
|
|
best.append(max(chunk, key=lambda x: x[1])[0])
|
|
return best
|
|
|
|
|
|
def format_output_filename(pattern: str, index: int, output_format: str) -> str:
|
|
try:
|
|
name = pattern % index
|
|
except (TypeError, ValueError):
|
|
if "%d" not in pattern:
|
|
name = pattern
|
|
else:
|
|
raise PipelineError("--extract-output-pattern must be a valid printf-style pattern such as output_frame_%05d.png")
|
|
|
|
path = Path(name)
|
|
target_suffix = f".{output_format}"
|
|
if path.suffix.lower() in {".jpg", ".jpeg", ".png"}:
|
|
return str(path.with_suffix(target_suffix))
|
|
if path.suffix:
|
|
raise PipelineError("--extract-output-pattern must end with .png or .jpg, or omit the extension entirely.")
|
|
return f"{name}{target_suffix}"
|
|
|
|
|
|
def save_frame(frame, output_file: Path, jpeg_quality: int, output_format: str) -> None:
|
|
output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
if output_format == "jpg":
|
|
ok = cv2.imwrite(str(output_file), frame, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
|
|
else:
|
|
ok = cv2.imwrite(str(output_file), frame)
|
|
if not ok:
|
|
raise PipelineError(f"Failed to write image: {output_file}")
|
|
|
|
|
|
def extract_frames(video_file, frame_numbers, output_dir, output_pattern, jpeg_quality, output_format, logger=None, should_cancel=None) -> None:
|
|
ensure_opencv_available()
|
|
ensure_tqdm_available()
|
|
if not frame_numbers:
|
|
emit("No frames selected for extraction.", logger)
|
|
return
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
targets = set(frame_numbers)
|
|
|
|
capture = cv2.VideoCapture(str(video_file))
|
|
if not capture.isOpened():
|
|
raise PipelineError(f"Failed to open video: {video_file}")
|
|
emit(f"[3/3] Extracting frames... ({len(frame_numbers)} frames)", logger)
|
|
total = get_frame_count(capture)
|
|
|
|
frame_number, saved, output_index = 0, 0, 1
|
|
with create_progress(total=total, desc="Extract", logger=logger) as bar:
|
|
while True:
|
|
raise_if_cancelled(should_cancel)
|
|
ok, frame = capture.read()
|
|
if not ok:
|
|
break
|
|
if frame_number in targets:
|
|
name = format_output_filename(output_pattern, output_index, output_format)
|
|
save_frame(frame, output_dir / name, jpeg_quality, output_format)
|
|
output_index += 1
|
|
saved += 1
|
|
if saved == len(targets):
|
|
bar.update(1)
|
|
break
|
|
frame_number += 1
|
|
bar.update(1)
|
|
capture.release()
|
|
|
|
if saved != len(targets):
|
|
raise PipelineError(f"Some selected frames could not be extracted. Expected {len(targets)}, saved {saved}.")
|
|
|
|
|
|
def run_extract_stage(
|
|
video: Path, output_dir: Path, chunk_size: int = 30, scale_width: int = 1920, workers: int = 4,
|
|
output_pattern: str = "output_frame_%05d.png", output_format: str = "png", jpeg_quality: int = 95,
|
|
reuse_metadata: bool = True, logger: Logger = None, should_cancel: CancelChecker = None,
|
|
) -> Path:
|
|
"""Returns output_dir (where the extracted equirect frames now live)."""
|
|
ensure_opencv_available()
|
|
ensure_tqdm_available()
|
|
if not video.exists():
|
|
raise PipelineError(f"Input video was not found: {video}")
|
|
if chunk_size <= 0:
|
|
raise PipelineError("--extract-chunk-size must be 1 or greater.")
|
|
if workers <= 0:
|
|
raise PipelineError("--extract-workers must be 1 or greater.")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
metadata_path = output_dir / "_sharpness_metadata.csv"
|
|
|
|
raise_if_cancelled(should_cancel)
|
|
if metadata_path.exists() and reuse_metadata:
|
|
emit(f"Using existing sharpness metadata: {metadata_path}", logger)
|
|
else:
|
|
emit(f"Generating sharpness metadata: {metadata_path}", logger)
|
|
analyze_video(video, metadata_path, scale_width, workers, logger, should_cancel)
|
|
|
|
raise_if_cancelled(should_cancel)
|
|
emit("[2/3] Parsing metadata...", logger)
|
|
best_frames = parse_best_frames(metadata_path, chunk_size, logger, should_cancel)
|
|
|
|
extract_frames(video, best_frames, output_dir, output_pattern, jpeg_quality, output_format, logger, should_cancel)
|
|
|
|
emit("Done: sharp frames extracted.", logger)
|
|
emit(f"Output directory: {output_dir.resolve()}", logger)
|
|
return output_dir
|
|
|
|
|
|
# ==========================================================================
|
|
# STAGE 2: YOLO segmentation masks (cubemap-projected detection)
|
|
# ==========================================================================
|
|
#
|
|
# Masks excluded/movable objects (people, vehicles, ...) so they don't
|
|
# pollute the SfM point cloud. Detection is run on cube-face reprojections
|
|
# rather than the raw equirect frame by default, since equirect projection
|
|
# stretches anything near the poles enough that YOLO misses it -- then
|
|
# masks are reprojected back to equirect space and merged. Output is always
|
|
# COLMAP convention: black(0)=excluded, filename is "<image filename>.png".
|
|
|
|
COCO_MOVABLE_DEFAULTS = ["person", "car", "bus", "truck", "bicycle", "motorcycle"]
|
|
|
|
|
|
def load_yolo_model(model_name: str, device: str):
|
|
try:
|
|
from ultralytics import YOLO
|
|
except ImportError:
|
|
raise PipelineError(
|
|
"ultralytics is not installed. Run:\n pip install ultralytics --break-system-packages\n"
|
|
"(ROCm users: install the ROCm torch build first so this doesn't pull in a CUDA-only wheel.)"
|
|
)
|
|
model = YOLO(model_name)
|
|
try:
|
|
model.to(device)
|
|
except Exception as exc:
|
|
print(f"[warn] could not move model to device={device} ({exc}); falling back to cpu", file=sys.stderr)
|
|
device = "cpu"
|
|
return model, device
|
|
|
|
|
|
def build_mask(result, shape_hw, target_ids, dilate_px) -> Tuple[np.ndarray, int]:
|
|
h, w = shape_hw
|
|
mask = np.zeros((h, w), dtype=np.uint8)
|
|
if result.masks is None or result.boxes is None or len(result.boxes) == 0:
|
|
return mask, 0
|
|
|
|
cls_ids = result.boxes.cls.cpu().numpy().astype(int)
|
|
seg_masks = result.masks.data.cpu().numpy()
|
|
hits = 0
|
|
for i, cid in enumerate(cls_ids):
|
|
if cid not in target_ids:
|
|
continue
|
|
m = seg_masks[i]
|
|
if m.shape != (h, w):
|
|
m = cv2.resize(m, (w, h), interpolation=cv2.INTER_NEAREST)
|
|
mask[m > 0.5] = 255
|
|
hits += 1
|
|
|
|
if hits and dilate_px > 0:
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_px * 2 + 1, dilate_px * 2 + 1))
|
|
mask = cv2.dilate(mask, kernel)
|
|
return mask, hits
|
|
|
|
|
|
def _chunked(seq, size):
|
|
for i in range(0, len(seq), size):
|
|
yield seq[i:i + size]
|
|
|
|
|
|
def run_predict_chunked(model, images, device, conf, imgsz, tile_batch):
|
|
all_results = []
|
|
for chunk in _chunked(images, max(1, tile_batch)):
|
|
results = model.predict(source=chunk, device=device, conf=conf, imgsz=imgsz, verbose=False, retina_masks=True)
|
|
all_results.extend(results)
|
|
if device != "cpu":
|
|
try:
|
|
import torch
|
|
torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
return all_results
|
|
|
|
|
|
def _process_mask_batch_equirect(batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate_px):
|
|
results = run_predict_chunked(model, [str(p) for p in batch_paths], device, conf, imgsz, tile_batch)
|
|
out = []
|
|
for result in results:
|
|
h, w = result.orig_shape
|
|
raw_mask, hits = build_mask(result, (h, w), target_ids, dilate_px)
|
|
out.append((raw_mask, hits, (h, w)))
|
|
return out
|
|
|
|
|
|
def _process_mask_batch_cubemap(batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate_px, face_size, skip_faces):
|
|
py360convert = ensure_py360convert_available()
|
|
|
|
frame_records, flat_images, flat_index = [], [], []
|
|
for local_idx, frame_path in enumerate(batch_paths):
|
|
img = cv2.imread(str(frame_path))
|
|
if img is None:
|
|
print(f"[warn] could not read {frame_path}, skipping", file=sys.stderr)
|
|
frame_records.append(None)
|
|
continue
|
|
h, w = img.shape[:2]
|
|
faces = py360convert.e2c(img, face_w=face_size, mode="bilinear", cube_format="dict")
|
|
frame_records.append({"h": h, "w": w, "faces": faces})
|
|
for face_name, face_img in faces.items():
|
|
if face_name in skip_faces:
|
|
continue
|
|
flat_images.append(face_img)
|
|
flat_index.append((local_idx, face_name))
|
|
|
|
results = run_predict_chunked(model, flat_images, device, conf, imgsz, tile_batch) if flat_images else []
|
|
|
|
per_frame_face_masks = [dict() if rec is not None else None for rec in frame_records]
|
|
per_frame_hits = [0] * len(frame_records)
|
|
for (local_idx, face_name), result in zip(flat_index, results):
|
|
face_img = frame_records[local_idx]["faces"][face_name]
|
|
fh, fw = face_img.shape[:2]
|
|
raw_mask, hits = build_mask(result, (fh, fw), target_ids, dilate_px)
|
|
per_frame_face_masks[local_idx][face_name] = raw_mask
|
|
per_frame_hits[local_idx] += hits
|
|
|
|
out = []
|
|
for local_idx, rec in enumerate(frame_records):
|
|
if rec is None:
|
|
out.append((None, 0, None))
|
|
continue
|
|
h, w = rec["h"], rec["w"]
|
|
mask_dict_3ch = {}
|
|
for face_name in FACES:
|
|
face_shape = rec["faces"][face_name].shape[:2]
|
|
m = per_frame_face_masks[local_idx].get(face_name)
|
|
if m is None:
|
|
m = np.zeros(face_shape, dtype=np.uint8)
|
|
mask_dict_3ch[face_name] = cv2.merge([m, m, m])
|
|
equirect_mask_3ch = py360convert.c2e(mask_dict_3ch, h=h, w=w, mode="nearest", cube_format="dict")
|
|
raw_equirect_mask = equirect_mask_3ch[:, :, 0]
|
|
_, raw_equirect_mask = cv2.threshold(raw_equirect_mask, 127, 255, cv2.THRESH_BINARY)
|
|
out.append((raw_equirect_mask, per_frame_hits[local_idx], (h, w)))
|
|
return out
|
|
|
|
|
|
def mask_path_for(frame_path: Path, output_dir: Path) -> Path:
|
|
"""COLMAP convention only: mask_path/<image filename>.png"""
|
|
return output_dir / f"{frame_path.name}.png"
|
|
|
|
|
|
def run_masks_stage(
|
|
input_dir: Path, output_dir: Path,
|
|
classes: str = ",".join(COCO_MOVABLE_DEFAULTS), model_name: str = "yolo11x-seg.pt", device: str = "cuda",
|
|
conf: float = 0.35, imgsz: int = 1280, dilate: int = 6,
|
|
projection: str = "cubemap", face_size: int = 1024, skip_faces_raw: str = "",
|
|
batch: int = 8, tile_batch: int = 8, empty_if_none: bool = False,
|
|
logger: Logger = None, should_cancel: CancelChecker = None,
|
|
) -> Path:
|
|
"""Returns output_dir (where equirect-space COLMAP-convention masks now live)."""
|
|
ensure_opencv_available()
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
frames = gather_images(input_dir)
|
|
emit(f"Found {len(frames)} frames in {input_dir}", logger)
|
|
|
|
if projection == "cubemap":
|
|
ensure_py360convert_available()
|
|
|
|
skip_faces = {f.strip().upper() for f in skip_faces_raw.split(",") if f.strip()}
|
|
unknown = skip_faces - set(FACES)
|
|
if unknown:
|
|
raise PipelineError(f"--mask-skip-faces has unrecognised face(s) {sorted(unknown)}; valid: {list(FACES)}")
|
|
|
|
model, device = load_yolo_model(model_name, device)
|
|
|
|
wanted = {c.strip().lower() for c in classes.split(",") if c.strip()}
|
|
name_to_id = {v.lower(): k for k, v in model.names.items()}
|
|
unknown_classes = wanted - set(name_to_id.keys())
|
|
if unknown_classes:
|
|
print(f"[warn] class names not in model label set, ignored: {sorted(unknown_classes)}", file=sys.stderr)
|
|
target_ids = {name_to_id[n] for n in wanted if n in name_to_id}
|
|
if not target_ids:
|
|
raise PipelineError("No valid target classes resolved. Check --mask-classes against the model's class list.")
|
|
emit(f"Masking classes: {sorted(wanted & set(name_to_id.keys()))} -> ids {sorted(target_ids)}", logger)
|
|
emit(f"Projection: {projection}" + (f" (skipping faces: {sorted(skip_faces)})" if skip_faces else ""), logger)
|
|
|
|
written = skipped = total_detections = 0
|
|
with create_progress(total=len(frames), desc="Masks", logger=logger) as bar:
|
|
for batch_start in range(0, len(frames), batch):
|
|
raise_if_cancelled(should_cancel)
|
|
batch_paths = frames[batch_start:batch_start + batch]
|
|
|
|
if projection == "cubemap":
|
|
batch_results = _process_mask_batch_cubemap(
|
|
batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate, face_size, skip_faces
|
|
)
|
|
else:
|
|
batch_results = _process_mask_batch_equirect(batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate)
|
|
|
|
for frame_path, (raw_mask, hit_count, _shape) in zip(batch_paths, batch_results):
|
|
if raw_mask is None:
|
|
bar.update(1)
|
|
continue
|
|
total_detections += hit_count
|
|
out_path = mask_path_for(frame_path, output_dir)
|
|
if hit_count == 0 and not empty_if_none:
|
|
skipped += 1
|
|
bar.update(1)
|
|
continue
|
|
final_mask = 255 - raw_mask # COLMAP: black(0) = excluded
|
|
cv2.imwrite(str(out_path), final_mask)
|
|
written += 1
|
|
bar.update(1)
|
|
|
|
emit(f"Done. {written} masks written, {skipped} frame(s) had no target objects (unmasked), "
|
|
f"{total_detections} total target-object instances masked.", logger)
|
|
emit(f"Masks in: {output_dir}", logger)
|
|
return output_dir
|
|
|
|
|
|
# ==========================================================================
|
|
# STAGE 3: project equirect frames+masks to 6 cube faces
|
|
# ==========================================================================
|
|
#
|
|
# Neither COLMAP nor Brush has an equirectangular camera model, so this
|
|
# splits each equirect frame (and its matching mask, if present) into 6
|
|
# perspective cube-face crops -- the "virtual camera rig" images that
|
|
# stage 4's rig config describes.
|
|
|
|
def _cubemap_input_mask_path(frame_path: Path, masks_dir: Path) -> Path:
|
|
return masks_dir / f"{frame_path.name}.png"
|
|
|
|
|
|
def cubemap_process_one(args_tuple):
|
|
(frame_path_str, masks_dir_str, out_frames_dir_str, out_masks_dir_str,
|
|
face_size, skip_faces, jpeg_quality, resume) = args_tuple
|
|
|
|
import py360convert # per-worker import: required for ProcessPoolExecutor pickling safety
|
|
|
|
frame_path = Path(frame_path_str)
|
|
out_frames_dir = Path(out_frames_dir_str)
|
|
out_masks_dir = Path(out_masks_dir_str) if out_masks_dir_str else None
|
|
masks_dir = Path(masks_dir_str) if masks_dir_str else None
|
|
|
|
active_faces = [f for f in FACES if f not in skip_faces]
|
|
|
|
if resume:
|
|
expected = [out_frames_dir / f"{frame_path.stem}_{f}{frame_path.suffix}" for f in active_faces]
|
|
if all(e.exists() for e in expected):
|
|
return frame_path.name, len(active_faces), 0, "skipped (resume)"
|
|
|
|
img = cv2.imread(str(frame_path))
|
|
if img is None:
|
|
return frame_path.name, 0, 0, "unreadable"
|
|
h, w = img.shape[:2]
|
|
|
|
faces = py360convert.e2c(img, face_w=face_size, mode="bilinear", cube_format="dict")
|
|
|
|
mask_faces = None
|
|
if masks_dir is not None and out_masks_dir is not None:
|
|
mask_path = _cubemap_input_mask_path(frame_path, masks_dir)
|
|
if mask_path.exists():
|
|
mask_img = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
|
|
if mask_img is not None and mask_img.shape[:2] == (h, w):
|
|
mask_3ch = cv2.merge([mask_img, mask_img, mask_img])
|
|
mask_faces_3ch = py360convert.e2c(mask_3ch, face_w=face_size, mode="nearest", cube_format="dict")
|
|
mask_faces = {}
|
|
for fname, m3 in mask_faces_3ch.items():
|
|
m = m3[:, :, 0]
|
|
_, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
|
|
mask_faces[fname] = m
|
|
|
|
frames_written = masks_written = 0
|
|
for face_name in active_faces:
|
|
face_img = faces[face_name]
|
|
out_frame_name = f"{frame_path.stem}_{face_name}{frame_path.suffix}"
|
|
out_path = out_frames_dir / out_frame_name
|
|
if frame_path.suffix.lower() in (".jpg", ".jpeg"):
|
|
cv2.imwrite(str(out_path), face_img, [cv2.IMWRITE_JPEG_QUALITY, jpeg_quality])
|
|
else:
|
|
cv2.imwrite(str(out_path), face_img)
|
|
frames_written += 1
|
|
|
|
if mask_faces is not None:
|
|
cv2.imwrite(str(out_masks_dir / f"{out_frame_name}.png"), mask_faces[face_name])
|
|
masks_written += 1
|
|
|
|
return frame_path.name, frames_written, masks_written, "ok"
|
|
|
|
|
|
def run_cubemap_stage(
|
|
frames_dir: Path, output_frames_dir: Path, masks_dir: Optional[Path] = None, output_masks_dir: Optional[Path] = None,
|
|
face_size: int = 1600, skip_faces_raw: str = "", jpeg_quality: int = 95, workers: int = 0, resume: bool = False,
|
|
logger: Logger = None, should_cancel: CancelChecker = None,
|
|
) -> Tuple[Path, Optional[Path]]:
|
|
"""Returns (output_frames_dir, output_masks_dir_or_None)."""
|
|
ensure_opencv_available()
|
|
ensure_py360convert_available()
|
|
|
|
if masks_dir is not None and output_masks_dir is None:
|
|
raise PipelineError("masks_dir given but output_masks_dir missing")
|
|
|
|
output_frames_dir.mkdir(parents=True, exist_ok=True)
|
|
if output_masks_dir is not None:
|
|
output_masks_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
skip_faces = {f.strip().upper() for f in skip_faces_raw.split(",") if f.strip()}
|
|
unknown = skip_faces - set(FACES)
|
|
if unknown:
|
|
raise PipelineError(f"--cubemap-skip-faces has unrecognised face(s) {sorted(unknown)}; valid: {list(FACES)}")
|
|
|
|
frames = gather_images(frames_dir)
|
|
emit(f"Found {len(frames)} equirect frames in {frames_dir}", logger)
|
|
active_faces = [f for f in FACES if f not in skip_faces]
|
|
emit(f"Active faces: {active_faces}", logger)
|
|
emit(f"Expected output frames: {len(frames) * len(active_faces)}", logger)
|
|
|
|
tasks = [
|
|
(str(f), str(masks_dir) if masks_dir else None, str(output_frames_dir),
|
|
str(output_masks_dir) if output_masks_dir else None, face_size, skip_faces, jpeg_quality, resume)
|
|
for f in frames
|
|
]
|
|
pool_workers = workers if workers > 0 else None
|
|
|
|
total_frames_written = total_masks_written = unreadable = 0
|
|
with create_progress(total=len(tasks), desc="Cubemap", logger=logger) as bar:
|
|
with ProcessPoolExecutor(max_workers=pool_workers) as executor:
|
|
futures = [executor.submit(cubemap_process_one, t) for t in tasks]
|
|
for fut in as_completed(futures):
|
|
raise_if_cancelled(should_cancel)
|
|
name, fw, mw, status = fut.result()
|
|
total_frames_written += fw
|
|
total_masks_written += mw
|
|
if status == "unreadable":
|
|
unreadable += 1
|
|
print(f"[warn] could not read {name}", file=sys.stderr)
|
|
bar.update(1)
|
|
|
|
emit(f"Done. {total_frames_written} face-crop images written to {output_frames_dir}", logger)
|
|
if output_masks_dir:
|
|
emit(f"{total_masks_written} face-crop masks written to {output_masks_dir}", logger)
|
|
if unreadable:
|
|
emit(f"{unreadable} frames were unreadable and skipped", logger)
|
|
return output_frames_dir, (output_masks_dir if output_masks_dir else None)
|
|
|
|
|
|
# ==========================================================================
|
|
# STAGE 4: COLMAP rig config + nested {FACE}/{stem} tree
|
|
# ==========================================================================
|
|
#
|
|
# Declares that the 6 cube faces from one frame share a single optical
|
|
# center with fixed relative rotations, so bundle adjustment solves one
|
|
# 6-DoF rig pose per frame instead of 6 independent poses. Also relinks the
|
|
# flat "{stem}_{FACE}.ext" output from stage 3 into the nested
|
|
# "{FACE}/{stem}.ext" tree COLMAP's rig_configurator requires (it matches
|
|
# rig cameras by filename PREFIX only -- see build_colmap_rig_config()'s
|
|
# docstring for the schema this assumes and why).
|
|
#
|
|
# Conventions verified against py360convert source: xyzcube() puts
|
|
# +X=right, +Y=up, +Z=front in an endpoint-inclusive linspace grid, so
|
|
# fx=fy=(N-1)/2 (not the naive N/2) and F->R is a clean 90deg yaw. See
|
|
# FACE_AXES / compute_intrinsics below for the full derivation.
|
|
|
|
Vec3 = Tuple[float, float, float]
|
|
Mat3 = Tuple[Vec3, Vec3, Vec3]
|
|
|
|
FACE_AXES: Dict[str, Tuple[Vec3, Vec3, Vec3]] = {
|
|
# image right image down view direction
|
|
"F": ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)),
|
|
"R": ((0.0, 0.0, -1.0), (0.0, 1.0, 0.0), (1.0, 0.0, 0.0)),
|
|
"B": ((-1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, -1.0)),
|
|
"L": ((0.0, 0.0, 1.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)),
|
|
"U": ((1.0, 0.0, 0.0), (0.0, 0.0, 1.0), (0.0, -1.0, 0.0)),
|
|
"D": ((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0)),
|
|
}
|
|
|
|
|
|
def face_rotation_matrix(face: str) -> Mat3:
|
|
right, down, forward = FACE_AXES[face]
|
|
return (right, down, forward)
|
|
|
|
|
|
def mat_mul(a: Mat3, b: Mat3) -> Mat3:
|
|
return tuple(tuple(sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3)) for i in range(3)) # type: ignore
|
|
|
|
|
|
def mat_transpose(m: Mat3) -> Mat3:
|
|
return tuple(tuple(m[j][i] for j in range(3)) for i in range(3)) # type: ignore
|
|
|
|
|
|
def mat_det(m: Mat3) -> float:
|
|
return (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
|
|
- m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
|
|
+ m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))
|
|
|
|
|
|
def assert_valid_rotation(m: Mat3, label: str, tol: float = 1e-9) -> None:
|
|
det = mat_det(m)
|
|
if abs(det - 1.0) > tol:
|
|
raise ValueError(f"{label}: det={det!r}, expected +1 (not a proper rotation)")
|
|
ident = mat_mul(m, mat_transpose(m))
|
|
for i in range(3):
|
|
for j in range(3):
|
|
expected = 1.0 if i == j else 0.0
|
|
if abs(ident[i][j] - expected) > tol:
|
|
raise ValueError(f"{label}: R @ R.T is not identity ({ident!r})")
|
|
|
|
|
|
def quaternion_from_matrix(m: Mat3) -> Tuple[float, float, float, float]:
|
|
"""Shepperd's method -> COLMAP's (w, x, y, z) order, sign-canonicalised to w>=0."""
|
|
m00, m01, m02 = m[0]
|
|
m10, m11, m12 = m[1]
|
|
m20, m21, m22 = m[2]
|
|
trace = m00 + m11 + m22
|
|
|
|
if trace > 0.0:
|
|
s = math.sqrt(trace + 1.0) * 2.0
|
|
w, x, y, z = 0.25 * s, (m21 - m12) / s, (m02 - m20) / s, (m10 - m01) / s
|
|
elif m00 > m11 and m00 > m22:
|
|
s = math.sqrt(1.0 + m00 - m11 - m22) * 2.0
|
|
w, x, y, z = (m21 - m12) / s, 0.25 * s, (m01 + m10) / s, (m02 + m20) / s
|
|
elif m11 > m22:
|
|
s = math.sqrt(1.0 + m11 - m00 - m22) * 2.0
|
|
w, x, y, z = (m02 - m20) / s, (m01 + m10) / s, 0.25 * s, (m12 + m21) / s
|
|
else:
|
|
s = math.sqrt(1.0 + m22 - m00 - m11) * 2.0
|
|
w, x, y, z = (m10 - m01) / s, (m02 + m20) / s, (m12 + m21) / s, 0.25 * s
|
|
|
|
norm = math.sqrt(w * w + x * x + y * y + z * z)
|
|
w, x, y, z = w / norm, x / norm, y / norm, z / norm
|
|
if w < 0.0:
|
|
w, x, y, z = -w, -x, -y, -z
|
|
return (w, x, y, z)
|
|
|
|
|
|
def cam_from_rig_rotation(face: str, ref_face: str) -> Tuple[float, float, float, float]:
|
|
r_face = face_rotation_matrix(face)
|
|
r_ref = face_rotation_matrix(ref_face)
|
|
assert_valid_rotation(r_face, f"FACE_AXES[{face}]")
|
|
assert_valid_rotation(r_ref, f"FACE_AXES[{ref_face}]")
|
|
r = mat_mul(r_face, mat_transpose(r_ref))
|
|
assert_valid_rotation(r, f"cam_from_rig[{face}<-{ref_face}]")
|
|
return quaternion_from_matrix(r)
|
|
|
|
|
|
def rotation_angle_deg(q: Tuple[float, float, float, float]) -> float:
|
|
w = min(1.0, max(-1.0, abs(q[0])))
|
|
return math.degrees(2.0 * math.acos(w))
|
|
|
|
|
|
def compute_intrinsics(face_size: int, convention: str) -> Tuple[float, float, float, float]:
|
|
"""'linspace' (default, exact for py360convert): f=(N-1)/2. 'half': f=N/2."""
|
|
if face_size < 2:
|
|
raise ValueError(f"face_size must be >= 2, got {face_size}")
|
|
if convention == "linspace":
|
|
f = (face_size - 1) / 2.0
|
|
elif convention == "half":
|
|
f = face_size / 2.0
|
|
else:
|
|
raise ValueError(f"unknown intrinsics convention: {convention!r}")
|
|
c = face_size / 2.0
|
|
return (f, f, c, c)
|
|
|
|
|
|
FLAT_RE = re.compile(r"^(?P<stem>.+)_(?P<face>[FRBLUD])$")
|
|
MASK_FLAT_RE = re.compile(r"^(?P<stem>.+)_(?P<face>[FRBLUD])\.[A-Za-z0-9]+$")
|
|
|
|
|
|
def _iter_images(root: Path):
|
|
for p in sorted(root.rglob("*")):
|
|
if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS:
|
|
yield p
|
|
|
|
|
|
def scan_cubemap_images(images_dir: Path, layout: str) -> Tuple[Dict[str, Dict[str, Path]], str]:
|
|
"""{stem: {FACE: path}}. flat='{stem}_{FACE}.ext', nested='{FACE}/{stem}.ext', auto detects."""
|
|
if not images_dir.is_dir():
|
|
raise PipelineError(f"--images-dir is not a directory: {images_dir}")
|
|
if layout == "auto":
|
|
subdirs = {p.name.upper() for p in images_dir.iterdir() if p.is_dir()}
|
|
layout = "nested" if subdirs and subdirs <= set(FACES) else "flat"
|
|
|
|
frames: Dict[str, Dict[str, Path]] = defaultdict(dict)
|
|
if layout == "flat":
|
|
for path in _iter_images(images_dir):
|
|
m = FLAT_RE.match(path.stem)
|
|
if m:
|
|
frames[m.group("stem")][m.group("face")] = path
|
|
else:
|
|
for face in FACES:
|
|
face_dir = images_dir / face
|
|
if face_dir.is_dir():
|
|
for path in _iter_images(face_dir):
|
|
frames[path.stem][face] = path
|
|
|
|
if not frames:
|
|
expected = "'{stem}_{FACE}.png'" if layout == "flat" else "'{FACE}/{stem}.png'"
|
|
raise PipelineError(f"No cubemap face images found under {images_dir} (layout={layout}). Expected {expected}.")
|
|
return dict(frames), layout
|
|
|
|
|
|
def scan_cubemap_masks(masks_dir: Path, layout: str) -> Tuple[Dict[str, Dict[str, Path]], str]:
|
|
"""Mirrors scan_cubemap_images(); a frame/face with no mask is expected, not an error."""
|
|
if not masks_dir.is_dir():
|
|
raise PipelineError(f"--masks-dir is not a directory: {masks_dir}")
|
|
if layout == "auto":
|
|
subdirs = {p.name.upper() for p in masks_dir.iterdir() if p.is_dir()}
|
|
layout = "nested" if subdirs and subdirs <= set(FACES) else "flat"
|
|
|
|
frames: Dict[str, Dict[str, Path]] = defaultdict(dict)
|
|
if layout == "flat":
|
|
for p in sorted(masks_dir.rglob("*")):
|
|
if not p.is_file() or p.suffix.lower() != ".png":
|
|
continue
|
|
m = MASK_FLAT_RE.match(p.stem)
|
|
if m:
|
|
frames[m.group("stem")][m.group("face")] = p
|
|
else:
|
|
for face in FACES:
|
|
face_dir = masks_dir / face
|
|
if not face_dir.is_dir():
|
|
continue
|
|
for p in sorted(face_dir.rglob("*")):
|
|
if not p.is_file() or p.suffix.lower() != ".png":
|
|
continue
|
|
frames[Path(p.stem).stem][face] = p # strip outer ".png" mask suffix only
|
|
|
|
return dict(frames), layout
|
|
|
|
|
|
def _png_dimensions(path: Path) -> Optional[Tuple[int, int]]:
|
|
try:
|
|
with path.open("rb") as fh:
|
|
head = fh.read(24)
|
|
except OSError:
|
|
return None
|
|
if len(head) < 24 or head[:8] != b"\x89PNG\r\n\x1a\n" or head[12:16] != b"IHDR":
|
|
return None
|
|
return struct.unpack(">II", head[16:24])
|
|
|
|
|
|
def image_dimensions(path: Path) -> Optional[Tuple[int, int]]:
|
|
dims = _png_dimensions(path)
|
|
if dims is not None:
|
|
return dims
|
|
try:
|
|
from PIL import Image
|
|
except ImportError:
|
|
return None
|
|
try:
|
|
with Image.open(path) as im:
|
|
return im.size
|
|
except OSError:
|
|
return None
|
|
|
|
|
|
def detect_face_size(frames: Dict[str, Dict[str, Path]], sample_limit: int = 24) -> int:
|
|
seen: Dict[Tuple[int, int], Path] = {}
|
|
checked = 0
|
|
for faces in frames.values():
|
|
for path in faces.values():
|
|
dims = image_dimensions(path)
|
|
if dims is None:
|
|
continue
|
|
seen.setdefault(dims, path)
|
|
checked += 1
|
|
if checked >= sample_limit:
|
|
break
|
|
if checked >= sample_limit:
|
|
break
|
|
if not seen:
|
|
raise PipelineError("could not read any image dimensions (non-PNG without Pillow?). Pass --rig-face-size explicitly.")
|
|
if len(seen) > 1:
|
|
detail = ", ".join(f"{w}x{h} (e.g. {p.name})" for (w, h), p in seen.items())
|
|
raise PipelineError(f"inconsistent face dimensions: {detail}. Pass --rig-face-size explicitly.")
|
|
(width, height), example = next(iter(seen.items()))
|
|
if width != height:
|
|
raise PipelineError(f"faces must be square, got {width}x{height} ({example.name}).")
|
|
return width
|
|
|
|
|
|
def build_colmap_rig_config(present_faces: List[str], ref_face: str, intrinsics: Tuple[float, float, float, float],
|
|
camera_model: str = "PINHOLE") -> List[dict]:
|
|
"""
|
|
SCHEMA (COLMAP >= 3.12, `colmap rig_configurator --rig_config_path`)
|
|
Top level: JSON array holding ONE rig definition, reused for every frame
|
|
(rig_configurator groups images by name-after-prefix into per-frame
|
|
Frames -- you do not emit one rig per source panorama).
|
|
|
|
[{"cameras": [
|
|
{"image_prefix": "F/", "ref_sensor": true,
|
|
"camera_model_name": "PINHOLE", "camera_params": [fx,fy,cx,cy]},
|
|
{"image_prefix": "R/",
|
|
"cam_from_rig_rotation": [w,x,y,z], "cam_from_rig_translation": [0,0,0],
|
|
"camera_model_name": "PINHOLE", "camera_params": [fx,fy,cx,cy]}
|
|
]}]
|
|
|
|
Read from colmap/doc/rigs.rst + src/colmap/scene/rig.cc:
|
|
* image_prefix: matched via StringStartsWith(); no suffix field exists,
|
|
which is why the nested {FACE}/{stem} layout is mandatory.
|
|
* exactly one ref_sensor per rig (identity pose, no cam_from_rig_*).
|
|
* cam_from_rig_rotation is W-FIRST (Eigen::Quaterniond(w,x,y,z) ctor).
|
|
* cam_from_rig_translation all-zero: every face shares one optical center.
|
|
Not run against a real COLMAP binary -- verify against your build if
|
|
rig_configurator rejects this.
|
|
"""
|
|
fx, fy, cx, cy = intrinsics
|
|
if ref_face not in present_faces:
|
|
raise ValueError(f"reference face {ref_face!r} not among present faces {present_faces}")
|
|
|
|
cameras = []
|
|
ordered = [ref_face] + [f for f in FACES if f in present_faces and f != ref_face]
|
|
for face in ordered:
|
|
entry: dict = {"image_prefix": f"{face}/"}
|
|
if face == ref_face:
|
|
entry["ref_sensor"] = True
|
|
else:
|
|
qw, qx, qy, qz = cam_from_rig_rotation(face, ref_face)
|
|
entry["cam_from_rig_rotation"] = [qw, qx, qy, qz]
|
|
entry["cam_from_rig_translation"] = [0.0, 0.0, 0.0]
|
|
entry["camera_model_name"] = camera_model
|
|
entry["camera_params"] = [fx, fy, cx, cy]
|
|
cameras.append(entry)
|
|
return [{"cameras": cameras}]
|
|
|
|
|
|
def write_rig_config(config: List[dict], output: Path) -> None:
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
output.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def materialise_nested_tree(frames: Dict[str, Dict[str, Path]], dest: Path, mode: str = "symlink") -> int:
|
|
"""Rebuild flat '{stem}_{FACE}.ext' as '{FACE}/{stem}.ext' under dest (COLMAP matches by prefix only)."""
|
|
count = 0
|
|
for stem, faces in sorted(frames.items()):
|
|
for face, src in sorted(faces.items()):
|
|
face_dir = dest / face
|
|
face_dir.mkdir(parents=True, exist_ok=True)
|
|
target = face_dir / f"{stem}{src.suffix}"
|
|
_link(src, target, mode)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def materialise_mask_tree(mask_frames: Dict[str, Dict[str, Path]], image_frames: Dict[str, Dict[str, Path]],
|
|
dest: Path, mode: str = "symlink") -> Tuple[int, int]:
|
|
"""Relink masks to '{FACE}/{stem}{image_ext}.png', mirroring materialise_nested_tree()'s image tree.
|
|
Returns (linked, orphans) -- an orphan mask has no matching image (stem/face drift)."""
|
|
linked = orphans = 0
|
|
for stem, faces in sorted(mask_frames.items()):
|
|
for face, mask_src in sorted(faces.items()):
|
|
image_src = image_frames.get(stem, {}).get(face)
|
|
mask_name = f"{stem}{image_src.suffix}.png" if image_src is not None else f"{stem}.png.png"
|
|
if image_src is None:
|
|
orphans += 1
|
|
face_dir = dest / face
|
|
face_dir.mkdir(parents=True, exist_ok=True)
|
|
_link(mask_src, face_dir / mask_name, mode)
|
|
linked += 1
|
|
return linked, orphans
|
|
|
|
|
|
def _link(src: Path, target: Path, mode: str) -> None:
|
|
if target.exists() or target.is_symlink():
|
|
target.unlink()
|
|
resolved = src.resolve()
|
|
if mode == "symlink":
|
|
target.symlink_to(resolved)
|
|
elif mode == "hardlink":
|
|
target.hardlink_to(resolved)
|
|
elif mode == "copy":
|
|
target.write_bytes(resolved.read_bytes())
|
|
else:
|
|
raise ValueError(f"unknown link mode: {mode!r}")
|
|
|
|
|
|
def run_rig_stage(
|
|
images_dir: Path, output: Path, masks_dir: Optional[Path] = None,
|
|
face_size: Optional[int] = None, layout: str = "auto", ref_face: str = "F",
|
|
intrinsics_convention: str = "linspace", link_tree: Optional[Path] = None,
|
|
masks_link_tree: Optional[Path] = None, link_mode: str = "symlink",
|
|
logger: Logger = None,
|
|
) -> Tuple[Path, Optional[Path]]:
|
|
"""Returns (link_tree_or_images_dir, masks_link_tree_or_None) -- the paths to hand COLMAP."""
|
|
if masks_link_tree is not None and masks_dir is None:
|
|
raise PipelineError("masks_link_tree requires masks_dir")
|
|
|
|
frames, resolved_layout = scan_cubemap_images(images_dir, layout)
|
|
present_faces = sorted({f for faces in frames.values() for f in faces}, key=FACES.index)
|
|
missing_faces = [f for f in FACES if f not in present_faces]
|
|
|
|
resolved_face_size = face_size or detect_face_size(frames)
|
|
intrinsics = compute_intrinsics(resolved_face_size, intrinsics_convention)
|
|
|
|
if ref_face not in present_faces:
|
|
emit(f"warning: --rig-ref-face {ref_face} not present; using {present_faces[0]} as rig frame.", logger)
|
|
ref_face = present_faces[0]
|
|
|
|
partial = {s: sorted(f) for s, f in frames.items() if set(f) != set(present_faces)}
|
|
|
|
config = build_colmap_rig_config(present_faces, ref_face, intrinsics)
|
|
write_rig_config(config, output)
|
|
|
|
fx, fy, cx, cy = intrinsics
|
|
emit(f"scanned : {images_dir} (layout={resolved_layout})", logger)
|
|
emit(f"frames : {len(frames)}", logger)
|
|
emit(f"faces present: {','.join(present_faces)}" + (f" absent: {','.join(missing_faces)}" if missing_faces else ""), logger)
|
|
emit(f"face size : {resolved_face_size}px" + (" (auto-detected)" if face_size is None else " (override)"), logger)
|
|
emit(f"intrinsics : PINHOLE fx={fx:.6g} fy={fy:.6g} cx={cx:.6g} cy={cy:.6g} [{intrinsics_convention}]", logger)
|
|
emit(f"rig frame : {ref_face} ({FACE_NAMES[ref_face]}) -- 1 rig definition, {len(frames)} frames, {len(present_faces)} cameras", logger)
|
|
if partial:
|
|
shown = list(partial.items())[:3]
|
|
emit(f"partial frames: {len(partial)} frame(s) missing >=1 face, e.g. " + "; ".join(f"{s}=[{','.join(f)}]" for s, f in shown), logger)
|
|
emit(f"wrote : {output}", logger)
|
|
|
|
final_images_path = images_dir
|
|
if link_tree is not None:
|
|
n = materialise_nested_tree(frames, link_tree, link_mode)
|
|
emit(f"linked : {n} images into {link_tree} ({link_mode})", logger)
|
|
final_images_path = link_tree
|
|
elif resolved_layout == "flat":
|
|
emit("", logger)
|
|
emit("IMPORTANT: COLMAP matches rig cameras by filename PREFIX only, so this flat", logger)
|
|
emit(" '{stem}_{FACE}.png' naming cannot be used directly with --rig-link-tree unset.", logger)
|
|
|
|
final_masks_path = None
|
|
if masks_dir is not None:
|
|
mask_frames, mask_layout = scan_cubemap_masks(masks_dir, "auto")
|
|
n_masks = sum(len(f) for f in mask_frames.values())
|
|
n_images = sum(len(f) for f in frames.values())
|
|
emit("", logger)
|
|
emit(f"masks scanned: {masks_dir} (layout={mask_layout})", logger)
|
|
emit(f"masks found : {n_masks} across {len(mask_frames)} frame(s) (of {n_images} face images -- missing ones are unmasked)", logger)
|
|
|
|
unknown_stems = set(mask_frames) - set(frames)
|
|
if unknown_stems:
|
|
emit(f"warning: {len(unknown_stems)} mask frame(s) have no matching image frame (stem drift?), e.g. "
|
|
+ ", ".join(sorted(unknown_stems)[:3]), logger)
|
|
|
|
if masks_link_tree is not None:
|
|
linked, orphans = materialise_mask_tree(mask_frames, frames, masks_link_tree, link_mode)
|
|
emit(f"masks linked : {linked} into {masks_link_tree} ({link_mode})", logger)
|
|
if orphans:
|
|
emit(f"warning: {orphans} mask(s) linked with a fallback name, no matching image found -- verify by hand.", logger)
|
|
final_masks_path = masks_link_tree
|
|
elif mask_layout == "flat":
|
|
emit("", logger)
|
|
emit("IMPORTANT: masks must mirror the nested image tree too. Set --rig-masks-link-tree.", logger)
|
|
final_masks_path = masks_dir
|
|
|
|
return final_images_path, final_masks_path
|
|
|
|
|
|
# ==========================================================================
|
|
# Full pipeline orchestration ("all")
|
|
# ==========================================================================
|
|
|
|
STAGE_ORDER = ["extract", "masks", "cubemap", "rig"]
|
|
|
|
|
|
def run_pipeline(args, logger: Logger = None, should_cancel: CancelChecker = None) -> None:
|
|
project_dir: Path = args.project_dir
|
|
frames_dir = project_dir / "01_frames"
|
|
masks_dir = project_dir / "02_masks"
|
|
cubemap_images_dir = project_dir / "03_cubemap" / "images"
|
|
cubemap_masks_dir = project_dir / "03_cubemap" / "masks"
|
|
colmap_images_dir = project_dir / "04_colmap" / "images"
|
|
colmap_masks_dir = project_dir / "04_colmap" / "masks"
|
|
rig_config_path = project_dir / "04_colmap" / "rig_config.json"
|
|
|
|
start_idx = STAGE_ORDER.index(args.start_stage)
|
|
total_stages = 4
|
|
|
|
if start_idx <= STAGE_ORDER.index("extract"):
|
|
if not args.video:
|
|
raise PipelineError("--video is required when starting from the extract stage.")
|
|
print_stage_header(1, total_stages, "Extract Sharpest Frames", logger)
|
|
run_extract_stage(
|
|
video=args.video, output_dir=frames_dir,
|
|
chunk_size=args.extract_chunk_size, scale_width=args.extract_scale_width, workers=args.extract_workers,
|
|
output_pattern=args.extract_output_pattern, output_format=args.extract_output_format,
|
|
jpeg_quality=args.extract_jpeg_quality, reuse_metadata=args.extract_reuse_metadata,
|
|
logger=logger, should_cancel=should_cancel,
|
|
)
|
|
else:
|
|
emit(f"Skipping extract stage; using existing frames in {frames_dir}", logger)
|
|
raise_if_cancelled(should_cancel)
|
|
|
|
have_masks = False
|
|
if args.skip_masks:
|
|
emit("Skipping masks stage (--skip-masks).", logger)
|
|
elif start_idx <= STAGE_ORDER.index("masks"):
|
|
print_stage_header(2, total_stages, "Generate SfM Masks", logger)
|
|
run_masks_stage(
|
|
input_dir=frames_dir, output_dir=masks_dir,
|
|
classes=args.mask_classes, model_name=args.mask_model, device=args.mask_device,
|
|
conf=args.mask_conf, imgsz=args.mask_imgsz, dilate=args.mask_dilate,
|
|
projection=args.mask_projection, face_size=args.mask_face_size, skip_faces_raw=args.mask_skip_faces,
|
|
batch=args.mask_batch, tile_batch=args.mask_tile_batch, empty_if_none=args.mask_empty_if_none,
|
|
logger=logger, should_cancel=should_cancel,
|
|
)
|
|
have_masks = True
|
|
else:
|
|
have_masks = masks_dir.is_dir() and any(masks_dir.glob("*.png"))
|
|
emit(f"Skipping masks stage; {'using' if have_masks else 'no'} existing masks in {masks_dir}", logger)
|
|
raise_if_cancelled(should_cancel)
|
|
|
|
print_stage_header(3, total_stages, "Project to Cubemap", logger)
|
|
run_cubemap_stage(
|
|
frames_dir=frames_dir, output_frames_dir=cubemap_images_dir,
|
|
masks_dir=(masks_dir if have_masks else None), output_masks_dir=(cubemap_masks_dir if have_masks else None),
|
|
face_size=args.cubemap_face_size, skip_faces_raw=args.cubemap_skip_faces,
|
|
jpeg_quality=args.cubemap_jpeg_quality, workers=args.cubemap_workers, resume=args.cubemap_resume,
|
|
logger=logger, should_cancel=should_cancel,
|
|
)
|
|
raise_if_cancelled(should_cancel)
|
|
|
|
have_cubemap_masks = have_masks and cubemap_masks_dir.is_dir() and any(cubemap_masks_dir.rglob("*.png"))
|
|
|
|
print_stage_header(4, total_stages, "Generate Rig Config", logger)
|
|
run_rig_stage(
|
|
images_dir=cubemap_images_dir, output=rig_config_path,
|
|
masks_dir=(cubemap_masks_dir if have_cubemap_masks else None),
|
|
ref_face=args.rig_ref_face, intrinsics_convention=args.rig_intrinsics_convention,
|
|
link_tree=colmap_images_dir, masks_link_tree=(colmap_masks_dir if have_cubemap_masks else None),
|
|
link_mode=args.rig_link_mode, logger=logger,
|
|
)
|
|
|
|
emit("", logger)
|
|
emit("==== PIPELINE COMPLETE ====", logger)
|
|
emit(f"COLMAP images : {colmap_images_dir}", logger)
|
|
if have_cubemap_masks:
|
|
emit(f"COLMAP masks : {colmap_masks_dir}", logger)
|
|
emit(f"Rig config : {rig_config_path}", logger)
|
|
emit("", logger)
|
|
emit("Next, run COLMAP:", logger)
|
|
mask_flag = f" --ImageReader.mask_path {colmap_masks_dir}" if have_cubemap_masks else ""
|
|
emit(f" colmap feature_extractor --database_path db.db --image_path {colmap_images_dir}"
|
|
f"{mask_flag} --ImageReader.single_camera_per_folder 1", logger)
|
|
emit(f" colmap rig_configurator --database_path db.db --rig_config_path {rig_config_path}", logger)
|
|
emit(" colmap exhaustive_matcher --database_path db.db # or vocab_tree_matcher at scale", logger)
|
|
emit(f" colmap mapper --database_path db.db --image_path {colmap_images_dir} --output_path {project_dir / '05_sparse'}"
|
|
f" --Mapper.ba_refine_sensor_from_rig 0 --Mapper.ba_refine_focal_length 0 --Mapper.ba_refine_extra_params 0", logger)
|
|
|
|
|
|
# ==========================================================================
|
|
# CLI
|
|
# ==========================================================================
|
|
|
|
def _add_extract_args(p: argparse.ArgumentParser, prefix: str = "") -> None:
|
|
d = f"--{prefix}" if prefix else "--"
|
|
p.add_argument(f"{d}chunk-size", dest=f"{prefix}chunk_size".replace("-", "_"), type=int, default=30,
|
|
help="Select one sharpest frame per this many frames (default: 30)")
|
|
p.add_argument(f"{d}scale-width", dest=f"{prefix}scale_width".replace("-", "_"), type=int, default=1920,
|
|
help="Resize width for sharpness analysis; 0 = original size (default: 1920)")
|
|
p.add_argument(f"{d}workers", dest=f"{prefix}workers".replace("-", "_"), type=int, default=4,
|
|
help="Worker processes for sharpness analysis (default: 4)")
|
|
p.add_argument(f"{d}output-pattern", dest=f"{prefix}output_pattern".replace("-", "_"), default="output_frame_%05d.png",
|
|
help="printf-style output filename pattern")
|
|
p.add_argument(f"{d}output-format", dest=f"{prefix}output_format".replace("-", "_"), choices=("png", "jpg"), default="png")
|
|
p.add_argument(f"{d}jpeg-quality", dest=f"{prefix}jpeg_quality".replace("-", "_"), type=int, default=95)
|
|
p.add_argument(f"{d}regenerate-metadata", dest=f"{prefix}reuse_metadata".replace("-", "_"), action="store_false", default=True,
|
|
help="Ignore existing _sharpness_metadata.csv and regenerate it")
|
|
|
|
|
|
def _add_masks_args(p: argparse.ArgumentParser, prefix: str) -> None:
|
|
d = f"--{prefix}"
|
|
p.add_argument(f"{d}classes", dest=f"{prefix.replace('-', '_')}classes", default=",".join(COCO_MOVABLE_DEFAULTS))
|
|
p.add_argument(f"{d}model", dest=f"{prefix.replace('-', '_')}model", default="yolo11x-seg.pt")
|
|
p.add_argument(f"{d}device", dest=f"{prefix.replace('-', '_')}device", default="cuda")
|
|
p.add_argument(f"{d}conf", dest=f"{prefix.replace('-', '_')}conf", type=float, default=0.35)
|
|
p.add_argument(f"{d}imgsz", dest=f"{prefix.replace('-', '_')}imgsz", type=int, default=1280)
|
|
p.add_argument(f"{d}dilate", dest=f"{prefix.replace('-', '_')}dilate", type=int, default=6)
|
|
p.add_argument(f"{d}projection", dest=f"{prefix.replace('-', '_')}projection", choices=("cubemap", "equirect"), default="cubemap")
|
|
p.add_argument(f"{d}face-size", dest=f"{prefix.replace('-', '_')}face_size", type=int, default=1024)
|
|
p.add_argument(f"{d}skip-faces", dest=f"{prefix.replace('-', '_')}skip_faces", default="")
|
|
p.add_argument(f"{d}batch", dest=f"{prefix.replace('-', '_')}batch", type=int, default=8)
|
|
p.add_argument(f"{d}tile-batch", dest=f"{prefix.replace('-', '_')}tile_batch", type=int, default=8)
|
|
p.add_argument(f"{d}empty-if-none", dest=f"{prefix.replace('-', '_')}empty_if_none", action="store_true")
|
|
|
|
|
|
def _add_cubemap_args(p: argparse.ArgumentParser, prefix: str) -> None:
|
|
d = f"--{prefix}"
|
|
r = lambda name: f"{prefix.replace('-', '_')}{name}" # noqa: E731
|
|
p.add_argument(f"{d}face-size", dest=r("face_size"), type=int, default=1600)
|
|
p.add_argument(f"{d}skip-faces", dest=r("skip_faces"), default="")
|
|
p.add_argument(f"{d}jpeg-quality", dest=r("jpeg_quality"), type=int, default=95)
|
|
p.add_argument(f"{d}workers", dest=r("workers"), type=int, default=0, help="0 = all CPU cores")
|
|
p.add_argument(f"{d}resume", dest=r("resume"), action="store_true")
|
|
|
|
|
|
def _add_rig_args(p: argparse.ArgumentParser, prefix: str) -> None:
|
|
d = f"--{prefix}"
|
|
r = lambda name: f"{prefix.replace('-', '_')}{name}" # noqa: E731
|
|
p.add_argument(f"{d}ref-face", dest=r("ref_face"), default="F", choices=list(FACES))
|
|
p.add_argument(f"{d}intrinsics-convention", dest=r("intrinsics_convention"), choices=("linspace", "half"), default="linspace")
|
|
p.add_argument(f"{d}link-mode", dest=r("link_mode"), choices=("symlink", "hardlink", "copy"), default="symlink")
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(prog="pipeline_360.py", description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--gui-log-output", action="store_true", help=argparse.SUPPRESS)
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
|
|
p_all = sub.add_parser("all", help="Run the full pipeline in one project directory.")
|
|
p_all.add_argument("--project-dir", required=True, type=Path)
|
|
p_all.add_argument("--video", type=Path, default=None, help="Required unless --start-stage skips extraction.")
|
|
p_all.add_argument("--start-stage", choices=STAGE_ORDER, default="extract")
|
|
p_all.add_argument("--skip-masks", action="store_true", help="Skip YOLO masking entirely.")
|
|
_add_extract_args(p_all, prefix="extract-")
|
|
_add_masks_args(p_all, prefix="mask-")
|
|
_add_cubemap_args(p_all, prefix="cubemap-")
|
|
_add_rig_args(p_all, prefix="rig-")
|
|
|
|
p_extract = sub.add_parser("extract", help="Stage 1 only: extract sharpest frames from a video.")
|
|
p_extract.add_argument("--video", required=True, type=Path)
|
|
p_extract.add_argument("--output-dir", required=True, type=Path)
|
|
_add_extract_args(p_extract)
|
|
|
|
p_masks = sub.add_parser("masks", help="Stage 2 only: generate YOLO SfM masks.")
|
|
p_masks.add_argument("--input-dir", required=True, type=Path)
|
|
p_masks.add_argument("--output-dir", required=True, type=Path)
|
|
_add_masks_args(p_masks, prefix="")
|
|
|
|
p_cubemap = sub.add_parser("cubemap", help="Stage 3 only: project equirect frames+masks to cube faces.")
|
|
p_cubemap.add_argument("--frames-dir", required=True, type=Path)
|
|
p_cubemap.add_argument("--masks-dir", type=Path, default=None)
|
|
p_cubemap.add_argument("--output-frames-dir", required=True, type=Path)
|
|
p_cubemap.add_argument("--output-masks-dir", type=Path, default=None)
|
|
_add_cubemap_args(p_cubemap, prefix="")
|
|
|
|
p_rig = sub.add_parser("rig", help="Stage 4 only: generate the COLMAP rig config.")
|
|
p_rig.add_argument("--images-dir", type=Path, default=None)
|
|
p_rig.add_argument("--masks-dir", type=Path, default=None)
|
|
p_rig.add_argument("--output", type=Path, default=Path("rig_config.json"))
|
|
p_rig.add_argument("--face-size", type=int, default=None)
|
|
p_rig.add_argument("--layout", choices=("auto", "flat", "nested"), default="auto")
|
|
p_rig.add_argument("--link-tree", type=Path, default=None)
|
|
p_rig.add_argument("--masks-link-tree", type=Path, default=None)
|
|
p_rig.add_argument("--self-test", action="store_true")
|
|
_add_rig_args(p_rig, prefix="")
|
|
|
|
return parser
|
|
|
|
|
|
def _cli_logger(gui_log_output: bool) -> Logger:
|
|
if not gui_log_output:
|
|
return None
|
|
if hasattr(sys.stdout, "reconfigure"):
|
|
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
|
return lambda message: print(message, flush=True)
|
|
|
|
|
|
def main(argv: Optional[List[str]] = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
logger = _cli_logger(args.gui_log_output)
|
|
|
|
try:
|
|
if args.command == "all":
|
|
args.project_dir.mkdir(parents=True, exist_ok=True)
|
|
run_pipeline(args, logger=logger)
|
|
|
|
elif args.command == "extract":
|
|
run_extract_stage(
|
|
video=args.video, output_dir=args.output_dir, chunk_size=args.chunk_size,
|
|
scale_width=args.scale_width, workers=args.workers, output_pattern=args.output_pattern,
|
|
output_format=args.output_format, jpeg_quality=args.jpeg_quality,
|
|
reuse_metadata=args.reuse_metadata, logger=logger,
|
|
)
|
|
|
|
elif args.command == "masks":
|
|
run_masks_stage(
|
|
input_dir=args.input_dir, output_dir=args.output_dir, classes=args.classes, model_name=args.model,
|
|
device=args.device, conf=args.conf, imgsz=args.imgsz, dilate=args.dilate, projection=args.projection,
|
|
face_size=args.face_size, skip_faces_raw=args.skip_faces, batch=args.batch, tile_batch=args.tile_batch,
|
|
empty_if_none=args.empty_if_none, logger=logger,
|
|
)
|
|
|
|
elif args.command == "cubemap":
|
|
if args.masks_dir is not None and args.output_masks_dir is None:
|
|
parser.error("--masks-dir given but --output-masks-dir missing")
|
|
run_cubemap_stage(
|
|
frames_dir=args.frames_dir, output_frames_dir=args.output_frames_dir, masks_dir=args.masks_dir,
|
|
output_masks_dir=args.output_masks_dir, face_size=args.face_size, skip_faces_raw=args.skip_faces,
|
|
jpeg_quality=args.jpeg_quality, workers=args.workers, resume=args.resume, logger=logger,
|
|
)
|
|
|
|
elif args.command == "rig":
|
|
if args.self_test:
|
|
return run_self_test()
|
|
if args.images_dir is None:
|
|
parser.error("--images-dir is required (or use --self-test)")
|
|
if args.masks_link_tree is not None and args.masks_dir is None:
|
|
parser.error("--masks-link-tree requires --masks-dir")
|
|
run_rig_stage(
|
|
images_dir=args.images_dir, output=args.output, masks_dir=args.masks_dir, face_size=args.face_size,
|
|
layout=args.layout, ref_face=args.ref_face, intrinsics_convention=args.intrinsics_convention,
|
|
link_tree=args.link_tree, masks_link_tree=args.masks_link_tree, link_mode=args.link_mode, logger=logger,
|
|
)
|
|
except PipelineCancelled:
|
|
emit("Cancelled.", logger)
|
|
return 130
|
|
except PipelineError as exc:
|
|
emit(f"Error: {exc}", logger)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
# ==========================================================================
|
|
# Self-test (stage 4 geometry/schema only -- no video/GPU required)
|
|
# ==========================================================================
|
|
|
|
def _write_dummy_png(path: Path, width: int, height: int) -> None:
|
|
raw = b"".join(b"\x00" + bytes(width) for _ in range(height))
|
|
|
|
def chunk(tag: bytes, data: bytes) -> bytes:
|
|
body = tag + data
|
|
return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)
|
|
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_bytes(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0))
|
|
+ chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))
|
|
|
|
|
|
def _apply(m: Mat3, v: Vec3) -> Vec3:
|
|
return tuple(sum(m[i][k] * v[k] for k in range(3)) for i in range(3)) # type: ignore
|
|
|
|
|
|
def run_self_test() -> int:
|
|
import tempfile
|
|
n, ok = 64, 0
|
|
|
|
for face in FACES:
|
|
assert_valid_rotation(face_rotation_matrix(face), f"FACE_AXES[{face}]")
|
|
ok += 1
|
|
|
|
q_r = cam_from_rig_rotation("R", "F")
|
|
expected_r = (math.sqrt(2) / 2, 0.0, -math.sqrt(2) / 2, 0.0)
|
|
assert all(abs(a - b) < 1e-12 for a, b in zip(q_r, expected_r)), q_r
|
|
assert abs(rotation_angle_deg(q_r) - 90.0) < 1e-9
|
|
ok += 1
|
|
|
|
r_mat = face_rotation_matrix("R")
|
|
assert max(abs(a - b) for a, b in zip(_apply(mat_transpose(r_mat), (0, 0, 1)), (1, 0, 0))) < 1e-12
|
|
for face, fwd in (("B", (0, 0, -1)), ("L", (-1, 0, 0)), ("U", (0, -1, 0)), ("D", (0, 1, 0))):
|
|
got = _apply(mat_transpose(face_rotation_matrix(face)), (0, 0, 1))
|
|
assert max(abs(a - b) for a, b in zip(got, fwd)) < 1e-12, (face, got)
|
|
ok += 1
|
|
|
|
fx, fy, cx, cy = compute_intrinsics(n, "linspace")
|
|
assert (fx, fy, cx, cy) == ((n - 1) / 2, (n - 1) / 2, n / 2, n / 2)
|
|
ok += 1
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
flat, masks_flat = root / "images", root / "masks"
|
|
|
|
for i in (1, 2):
|
|
for face in FACES:
|
|
_write_dummy_png(flat / f"frame_{i:05d}_{face}.png", n, n)
|
|
for face in ("F", "R", "B", "L"):
|
|
_write_dummy_png(flat / f"frame_{3:05d}_{face}.png", n, n)
|
|
for face in FACES:
|
|
_write_dummy_png(masks_flat / f"frame_{1:05d}_{face}.png.png", n, n)
|
|
for face in ("F", "R"):
|
|
_write_dummy_png(masks_flat / f"frame_{3:05d}_{face}.png.png", n, n)
|
|
|
|
out, link_tree, masks_link_tree = root / "rig.json", root / "nested", root / "masks_nested"
|
|
run_rig_stage(images_dir=flat, output=out, masks_dir=masks_flat, link_tree=link_tree,
|
|
masks_link_tree=masks_link_tree, link_mode="copy")
|
|
config = json.loads(out.read_text())
|
|
cams = config[0]["cameras"]
|
|
assert len(cams) == 6
|
|
refs = [c for c in cams if c.get("ref_sensor")]
|
|
assert len(refs) == 1 and refs[0]["image_prefix"] == "F/"
|
|
by_prefix = {c["image_prefix"]: c for c in cams}
|
|
got_r = by_prefix["R/"]["cam_from_rig_rotation"]
|
|
assert all(abs(a - b) < 1e-12 for a, b in zip(got_r, expected_r))
|
|
assert (link_tree / "F" / "frame_00001.png").is_file()
|
|
assert (masks_link_tree / "F" / "frame_00001.png.png").is_file()
|
|
assert not (masks_link_tree / "F" / "frame_00002.png.png").exists()
|
|
assert (masks_link_tree / "R" / "frame_00003.png.png").is_file()
|
|
ok += 1
|
|
|
|
partial_dir = root / "partial"
|
|
for i in (1, 2):
|
|
for face in ("F", "R", "B", "L"):
|
|
_write_dummy_png(partial_dir / f"frame_{i:05d}_{face}.png", n, n)
|
|
out2 = root / "rig_partial.json"
|
|
run_rig_stage(images_dir=partial_dir, output=out2)
|
|
cfg2 = json.loads(out2.read_text())
|
|
assert len(cfg2[0]["cameras"]) == 4
|
|
ok += 1
|
|
|
|
print(f"\nself-test: {ok} check groups passed.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
multiprocessing.freeze_support()
|
|
raise SystemExit(main())
|