diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9d2cb6d --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +* +!.gitignore +!pipeline_360.py +!pipeline_360_gui.py +!requirements.txt diff --git a/pipeline_360.py b/pipeline_360.py new file mode 100644 index 0000000..8368056 --- /dev/null +++ b/pipeline_360.py @@ -0,0 +1,1585 @@ +#!/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=.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 ".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/.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.+)_(?P[FRBLUD])$") +MASK_FLAT_RE = re.compile(r"^(?P.+)_(?P[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()) diff --git a/pipeline_360_gui.py b/pipeline_360_gui.py new file mode 100644 index 0000000..c7ec1e3 --- /dev/null +++ b/pipeline_360_gui.py @@ -0,0 +1,542 @@ +#!/usr/bin/env python3 +""" +pipeline_360_gui.py -- single-form GUI for pipeline_360.py. + +One form, all settings, one Run button that runs the whole pipeline +(extract -> masks -> cubemap -> rig) in sequence, streaming progress into +a log pane. Mirrors extract_sharpest_frame_gui.py's approach of shelling +out to the CLI script and parsing "[progress] "-prefixed lines for +in-place progress updates -- except this wraps all four stages instead of +one, via `pipeline_360.py all --gui-log-output ...`. + +Differences from the original GUI this is modeled on: + - English-only. The original's EN/JA translation table added a fair + amount of boilerplate for a tool this much larger; happy to add + localisation back if it's wanted. + - Cross-platform Stop button. The original only force-killed via + Windows' `taskkill`, which silently wouldn't have worked on Linux. + This one kills the whole process group via SIGTERM/SIGKILL on POSIX + (and still uses taskkill on Windows), so Stop actually stops the + ProcessPoolExecutor workers the extract/cubemap stages spawn, not + just the parent CLI process. + - No PyInstaller-frozen-exe branch -- always invokes + `sys.executable pipeline_360.py all ...`. Add that branch back if you + package this as a standalone binary later. + +Requires: tkinter (usually stdlib; on some Linux distros it's a separate +package, e.g. `sudo pacman -S tk` on Arch/CachyOS). +""" + +import os +import queue +import signal +import subprocess +import sys +import threading +import tkinter as tk +from pathlib import Path +from tkinter import filedialog, messagebox, ttk +from typing import Optional + +from pipeline_360 import PROGRESS_PREFIX, COCO_MOVABLE_DEFAULTS, FACES + +APP_TITLE = "360 Photogrammetry Pipeline" + + +class ToolTip: + def __init__(self, widget, text: str = "") -> None: + self.widget = widget + self.text = text + self.tip_window = None + widget.bind("", self.show_tip) + widget.bind("", self.hide_tip) + + def show_tip(self, _event=None) -> None: + if self.tip_window or not self.text: + return + x = self.widget.winfo_rootx() + 20 + y = self.widget.winfo_rooty() + self.widget.winfo_height() + 8 + self.tip_window = tw = tk.Toplevel(self.widget) + tw.wm_overrideredirect(True) + tw.wm_geometry(f"+{x}+{y}") + tk.Label(tw, text=self.text, justify=tk.LEFT, background="#ffffe0", + relief=tk.SOLID, borderwidth=1, font=("Segoe UI", 9), wraplength=360).pack(ipadx=4) + + def hide_tip(self, _event=None) -> None: + if self.tip_window: + self.tip_window.destroy() + self.tip_window = None + + +# ------------------------------------------------------------------------ +# Field spec: one entry per pipeline_360.py 'all' flag exposed in the form. +# kind: "int" | "float" | "str" | "choice" | "bool" +# ------------------------------------------------------------------------ + +class Field: + def __init__(self, flag, label, kind, default, choices=None, tip=""): + self.flag = flag # CLI flag, e.g. "--extract-chunk-size" + self.label = label + self.kind = kind + self.default = default + self.choices = choices or [] + self.tip = tip + + +SECTIONS = [ + ("Extract Sharpest Frames", [ + Field("--extract-chunk-size", "Chunk size", "int", 30, + tip="Pick one sharpest frame out of every N video frames."), + Field("--extract-scale-width", "Scale width (px)", "int", 1920, + tip="Resize width used only for sharpness scoring. 0 = full resolution."), + Field("--extract-workers", "Workers", "int", 4, + tip="Parallel worker processes for the sharpness pass."), + Field("--extract-output-format", "Output format", "choice", "png", choices=["png", "jpg"]), + Field("--extract-jpeg-quality", "JPEG quality", "int", 95, + tip="Only used when output format is jpg."), + Field("--extract-regenerate-metadata", "Regenerate sharpness cache", "bool", False, + tip="Ignore any existing _sharpness_metadata.csv in 01_frames and recompute it."), + ]), + ("Generate SfM Masks (YOLO)", [ + Field("--mask-classes", "Classes to mask", "str", ",".join(COCO_MOVABLE_DEFAULTS), + tip="Comma-separated COCO class names to exclude from reconstruction."), + Field("--mask-model", "Model", "str", "yolo11x-seg.pt", + tip="Ultralytics seg model name or path."), + Field("--mask-device", "Device", "str", "cuda", + tip="cuda, cuda:0, cpu, or mps. ROCm PyTorch keeps the 'cuda' device string."), + Field("--mask-conf", "Confidence threshold", "float", 0.35), + Field("--mask-imgsz", "Inference size", "int", 1280), + Field("--mask-dilate", "Dilate (px)", "int", 6, + tip="Grows mask edges to cover segmentation-boundary error / motion blur."), + Field("--mask-projection", "Projection", "choice", "cubemap", choices=["cubemap", "equirect"], + tip="cubemap fixes YOLO missing people near the poles; equirect runs on the raw frame."), + Field("--mask-face-size", "Detection face size (px)", "int", 1024), + Field("--mask-skip-faces", "Skip faces", "str", "", + tip="Comma list of F/R/B/L/U/D to skip detection on, e.g. U,D."), + Field("--mask-batch", "Frame batch", "int", 8), + Field("--mask-tile-batch", "GPU tile batch", "int", 8, + tip="Bounds VRAM use directly for cubemap projection. Lower this first on OOM."), + Field("--mask-empty-if-none", "Write empty masks for clean frames", "bool", False, + tip="Off by default to save disk space; a missing mask already means 'unmasked'."), + ]), + ("Project to Cubemap", [ + Field("--cubemap-face-size", "Face size (px)", "int", 1600, + tip="Final training image resolution per cube face."), + Field("--cubemap-skip-faces", "Skip faces", "str", "", + tip="Comma list of F/R/B/L/U/D to exclude from the exported dataset, e.g. U,D."), + Field("--cubemap-jpeg-quality", "JPEG quality", "int", 95), + Field("--cubemap-workers", "Workers", "int", 0, tip="0 = use all CPU cores."), + Field("--cubemap-resume", "Resume (skip already-written frames)", "bool", False), + ]), + ("Generate Rig Config", [ + Field("--rig-ref-face", "Reference face", "choice", "F", choices=list(FACES), + tip="Defines the rig's coordinate frame. F (front) is the sane default."), + Field("--rig-intrinsics-convention", "Intrinsics convention", "choice", "linspace", + choices=["linspace", "half"], + tip="linspace is exact for py360convert output; only change this if COLMAP rejects the schema."), + Field("--rig-link-mode", "Link mode", "choice", "symlink", choices=["symlink", "hardlink", "copy"], + tip="How the flat cubemap output is relinked into COLMAP's required nested tree."), + ]), +] + + +class PipelineGui(tk.Tk): + def __init__(self) -> None: + super().__init__() + self.title(APP_TITLE) + self.geometry("920x760") + self.minsize(760, 560) + + self.project_dir = tk.StringVar() + self.video_path = tk.StringVar() + self.start_stage = tk.StringVar(value="extract") + self.skip_masks = tk.BooleanVar(value=False) + self.status_text = tk.StringVar(value="Ready") + + self.field_vars: dict[str, tk.Variable] = {} + self.field_widgets: dict[str, tk.Widget] = {} + + self.worker_process: Optional[subprocess.Popen] = None + self.log_thread: Optional[threading.Thread] = None + self.log_queue: "queue.Queue" = queue.Queue() + self.stop_requested = False + self.is_closing = False + self.progress_line_active = False + + self._build_widgets() + self.protocol("WM_DELETE_WINDOW", self._on_close_requested) + self.after(100, self._drain_log_queue) + + # -- layout ----------------------------------------------------------- + + def _build_widgets(self) -> None: + outer = ttk.Frame(self, padding=12) + outer.pack(fill="both", expand=True) + outer.rowconfigure(1, weight=3) + outer.rowconfigure(3, weight=2) + outer.columnconfigure(0, weight=1) + + top = ttk.Frame(outer) + top.grid(row=0, column=0, sticky="ew", pady=(0, 8)) + top.columnconfigure(1, weight=1) + + ttk.Label(top, text="Project directory").grid(row=0, column=0, sticky="w", padx=(0, 8), pady=4) + entry = ttk.Entry(top, textvariable=self.project_dir) + entry.grid(row=0, column=1, sticky="ew", pady=4) + ttk.Button(top, text="Browse...", command=self._browse_project_dir).grid(row=0, column=2, padx=(8, 0), pady=4) + + ttk.Label(top, text="Video file").grid(row=1, column=0, sticky="w", padx=(0, 8), pady=4) + self.video_entry = ttk.Entry(top, textvariable=self.video_path) + self.video_entry.grid(row=1, column=1, sticky="ew", pady=4) + self.video_button = ttk.Button(top, text="Browse...", command=self._browse_video) + self.video_button.grid(row=1, column=2, padx=(8, 0), pady=4) + + ttk.Label(top, text="Start from stage").grid(row=2, column=0, sticky="w", padx=(0, 8), pady=4) + stage_combo = ttk.Combobox(top, state="readonly", values=["extract", "masks", "cubemap", "rig"], + textvariable=self.start_stage, width=14) + stage_combo.grid(row=2, column=1, sticky="w", pady=4) + stage_combo.bind("<>", lambda _e: self._sync_video_requirement()) + ToolTip(stage_combo, "Resuming from 'masks' or later assumes 01_frames (and 02_masks) already exist " + "in the project directory from a prior run.") + + self.skip_masks_check = ttk.Checkbutton(top, text="Skip masking entirely", variable=self.skip_masks) + self.skip_masks_check.grid(row=2, column=2, sticky="w", padx=(8, 0), pady=4) + + # Scrollable settings area ----------------------------------------- + settings_outer = ttk.Frame(outer) + settings_outer.grid(row=1, column=0, sticky="nsew") + settings_outer.rowconfigure(0, weight=1) + settings_outer.columnconfigure(0, weight=1) + + canvas = tk.Canvas(settings_outer, highlightthickness=0) + vscroll = ttk.Scrollbar(settings_outer, orient="vertical", command=canvas.yview) + canvas.configure(yscrollcommand=vscroll.set) + canvas.grid(row=0, column=0, sticky="nsew") + vscroll.grid(row=0, column=1, sticky="ns") + + inner = ttk.Frame(canvas) + inner_id = canvas.create_window((0, 0), window=inner, anchor="nw") + + def _on_inner_configure(_event): + canvas.configure(scrollregion=canvas.bbox("all")) + + def _on_canvas_configure(event): + canvas.itemconfigure(inner_id, width=event.width) + + inner.bind("", _on_inner_configure) + canvas.bind("", _on_canvas_configure) + + def _on_mousewheel(event): + delta = 1 if getattr(event, "num", None) == 5 or event.delta < 0 else -1 + canvas.yview_scroll(delta, "units") + + canvas.bind_all("", _on_mousewheel) # Windows / macOS + canvas.bind_all("", _on_mousewheel) # Linux X11 scroll up + canvas.bind_all("", _on_mousewheel) # Linux X11 scroll down + + for title, fields in SECTIONS: + self._build_section(inner, title, fields) + + # Buttons ------------------------------------------------------- + actions = ttk.Frame(outer) + actions.grid(row=2, column=0, sticky="ew", pady=(8, 4)) + self.run_button = ttk.Button(actions, text="Run pipeline", command=self._start_run) + self.run_button.pack(side="right") + self.stop_button = ttk.Button(actions, text="Stop", command=self._stop_run, state="disabled") + self.stop_button.pack(side="right", padx=(0, 8)) + ttk.Button(actions, text="Clear log", command=self._clear_log).pack(side="right", padx=(0, 8)) + ttk.Label(actions, textvariable=self.status_text).pack(side="left") + + # Log ------------------------------------------------------------- + log_frame = ttk.LabelFrame(outer, text="Log") + log_frame.grid(row=3, column=0, sticky="nsew") + log_frame.rowconfigure(0, weight=1) + log_frame.columnconfigure(0, weight=1) + + self.log_text = tk.Text(log_frame, wrap="word", state="disabled", height=14) + self.log_text.grid(row=0, column=0, sticky="nsew") + log_scroll = ttk.Scrollbar(log_frame, orient="vertical", command=self.log_text.yview) + log_scroll.grid(row=0, column=1, sticky="ns") + self.log_text.configure(yscrollcommand=log_scroll.set) + + self._sync_video_requirement() + + def _build_section(self, parent, title: str, fields) -> None: + frame = ttk.LabelFrame(parent, text=title, padding=8) + frame.pack(fill="x", expand=True, padx=4, pady=6) + frame.columnconfigure(1, weight=1) + frame.columnconfigure(3, weight=1) + + for i, field in enumerate(fields): + row, col_base = divmod(i, 2) + col = col_base * 2 + + if field.kind == "bool": + var = tk.BooleanVar(value=field.default) + widget = ttk.Checkbutton(frame, text=field.label, variable=var) + widget.grid(row=row, column=col, columnspan=2, sticky="w", padx=(0, 16), pady=3) + else: + ttk.Label(frame, text=field.label).grid(row=row, column=col, sticky="w", padx=(0, 6), pady=3) + if field.kind == "choice": + var = tk.StringVar(value=str(field.default)) + widget = ttk.Combobox(frame, state="readonly", values=field.choices, textvariable=var, width=14) + else: + var = tk.StringVar(value=str(field.default)) + widget = ttk.Entry(frame, textvariable=var, width=18) + widget.grid(row=row, column=col + 1, sticky="w", padx=(0, 16), pady=3) + + if field.tip: + ToolTip(widget, field.tip) + + self.field_vars[field.flag] = var + self.field_widgets[field.flag] = widget + + # -- state sync --------------------------------------------------------- + + def _sync_video_requirement(self) -> None: + needs_video = self.start_stage.get() == "extract" + state = "normal" if needs_video else "disabled" + self.video_entry.configure(state=state) + self.video_button.configure(state=state) + + # -- browse dialogs ------------------------------------------------- + + def _browse_project_dir(self) -> None: + selected = filedialog.askdirectory(title="Select project directory") + if selected: + self.project_dir.set(selected) + + def _browse_video(self) -> None: + selected = filedialog.askopenfilename(title="Select video file") + if selected: + self.video_path.set(selected) + + # -- log helpers (mirrors extract_sharpest_frame_gui.py's in-place + # progress-line replacement so tqdm-style updates don't spam the log) -- + + def _clear_log(self) -> None: + self.log_text.configure(state="normal") + self.log_text.delete("1.0", tk.END) + if self.progress_line_active: + self.log_text.mark_unset("progress_start") + self.log_text.mark_unset("progress_end") + self.log_text.configure(state="disabled") + self.progress_line_active = False + + def _append_log(self, message: str, replace: bool = False) -> None: + self.log_text.configure(state="normal") + if replace: + if self.progress_line_active: + self.log_text.delete("progress_start", "progress_end") + else: + self.log_text.mark_set("progress_start", "end-1c") + self.log_text.mark_gravity("progress_start", tk.LEFT) + self.log_text.insert("progress_start", message + "\n") + self.log_text.mark_set("progress_end", "progress_start +1 lines") + self.progress_line_active = True + else: + if self.progress_line_active: + self.log_text.mark_unset("progress_start") + self.log_text.mark_unset("progress_end") + self.progress_line_active = False + self.log_text.insert(tk.END, message + "\n") + self.log_text.see(tk.END) + self.log_text.configure(state="disabled") + + # -- command building / validation ----------------------------------- + + def _validate_and_build_command(self) -> Optional[list]: + project_dir = self.project_dir.get().strip() + if not project_dir: + messagebox.showerror(APP_TITLE, "Select a project directory.") + return None + + start_stage = self.start_stage.get() + video = self.video_path.get().strip() + if start_stage == "extract" and not video: + messagebox.showerror(APP_TITLE, "Select a video file, or choose a later start stage.") + return None + + base_dir = Path(__file__).parent + command = [sys.executable, str(base_dir / "pipeline_360.py"), "all", "--gui-log-output", + "--project-dir", project_dir, "--start-stage", start_stage] + if start_stage == "extract": + command += ["--video", video] + if self.skip_masks.get(): + command += ["--skip-masks"] + + for _title, fields in SECTIONS: + for field in fields: + var = self.field_vars[field.flag] + if field.kind == "bool": + if var.get(): + command.append(field.flag) + continue + + raw = var.get().strip() if isinstance(var, tk.StringVar) else str(var.get()) + if field.kind == "int": + try: + int(raw) + except ValueError: + messagebox.showerror(APP_TITLE, f"'{field.label}' must be a whole number (got {raw!r}).") + return None + elif field.kind == "float": + try: + float(raw) + except ValueError: + messagebox.showerror(APP_TITLE, f"'{field.label}' must be a number (got {raw!r}).") + return None + command += [field.flag, raw] + + return command + + # -- process lifecycle ------------------------------------------------ + + def _is_worker_running(self) -> bool: + return self.worker_process is not None and self.worker_process.poll() is None + + def _spawn(self, command: list) -> subprocess.Popen: + kwargs = {} + if sys.platform == "win32": + kwargs["creationflags"] = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) + else: + kwargs["start_new_session"] = True # own process group, so we can kill it (and its children) together + return subprocess.Popen( + command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, encoding="utf-8", errors="replace", bufsize=1, **kwargs, + ) + + def _force_kill_process_tree(self, process: subprocess.Popen) -> None: + if sys.platform == "win32": + subprocess.run(["taskkill", "/PID", str(process.pid), "/T", "/F"], capture_output=True, text=True, check=False) + try: + process.wait(timeout=2.0) + except subprocess.TimeoutExpired: + process.kill() + else: + try: + pgid = os.getpgid(process.pid) + os.killpg(pgid, signal.SIGTERM) + process.wait(timeout=2.0) + except (ProcessLookupError, subprocess.TimeoutExpired): + try: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + except PermissionError: + pass + try: + process.wait(timeout=1.0) + except subprocess.TimeoutExpired: + pass + + def _stream_process_output(self, process: subprocess.Popen) -> None: + assert process.stdout is not None + for line in process.stdout: + self.log_queue.put(("log", line.rstrip("\r\n"))) + process.stdout.close() + return_code = process.wait() + + if self.stop_requested: + self.log_queue.put(("cancelled", None)) + elif return_code == 0: + self.log_queue.put(("done", None)) + else: + self.log_queue.put(("error", f"Pipeline exited with code {return_code}. See log above for details.")) + + def _set_running_state(self, running: bool) -> None: + state = "disabled" if running else "normal" + self.run_button.configure(state=state) + for widget in self.field_widgets.values(): + widget.configure(state=("disabled" if running else ("readonly" if isinstance(widget, ttk.Combobox) else "normal"))) + self.skip_masks_check.configure(state=state) + self.stop_button.configure(state="normal" if running else "disabled") + self.status_text.set("Running..." if running else "Ready") + if not running: + self._sync_video_requirement() + + def _start_run(self) -> None: + if self._is_worker_running(): + return + command = self._validate_and_build_command() + if command is None: + return + + self.stop_requested = False + self._set_running_state(True) + self._append_log(f"Running: {' '.join(command)}") + + try: + self.worker_process = self._spawn(command) + except OSError as exc: + self._set_running_state(False) + messagebox.showerror(APP_TITLE, f"Could not start pipeline_360.py: {exc}") + return + + self.log_thread = threading.Thread(target=self._stream_process_output, args=(self.worker_process,), daemon=True) + self.log_thread.start() + + def _stop_run(self) -> None: + if not self._is_worker_running(): + return + self.status_text.set("Stopping...") + self._append_log("Stopping...") + self.stop_button.configure(state="disabled") + self.stop_requested = True + self._force_kill_process_tree(self.worker_process) + self.worker_process = None + self.log_thread = None + self._set_running_state(False) + + def _on_close_requested(self) -> None: + if self._is_worker_running(): + if not messagebox.askyesno(APP_TITLE, "The pipeline is still running. Force stop it and close the window?"): + return + self.status_text.set("Stopping...") + self._append_log("Stopping...") + self._force_kill_process_tree(self.worker_process) + self.is_closing = True + self.destroy() + + def _drain_log_queue(self) -> None: + if self.is_closing: + return + while True: + try: + kind, payload = self.log_queue.get_nowait() + except queue.Empty: + break + + if kind == "log": + if payload.startswith(PROGRESS_PREFIX): + self._append_log(payload[len(PROGRESS_PREFIX):], replace=True) + else: + self._append_log(payload) + elif kind == "done": + self._append_log("Done: pipeline completed successfully.") + self.worker_process = None + self.log_thread = None + self._set_running_state(False) + messagebox.showinfo(APP_TITLE, "Pipeline completed successfully.") + elif kind == "cancelled": + self._append_log("Cancelled.") + self.worker_process = None + self.log_thread = None + self._set_running_state(False) + elif kind == "error": + self._append_log(f"Error: {payload}") + self.worker_process = None + self.log_thread = None + self._set_running_state(False) + messagebox.showerror(APP_TITLE, payload) + + self.after(100, self._drain_log_queue) + + +def main() -> None: + app = PipelineGui() + app.mainloop() + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..577a0cc --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +# Core pipeline (extract, cubemap, rig stages) -- always needed +opencv-python>=4.8 +numpy>=1.24 +tqdm>=4.65 +py360convert>=1.0 + +# Masks stage only (YOLO segmentation) -- skip these if you always run +# with --skip-masks +ultralytics>=8.3 + +# NOTE on ultralytics/torch + ROCm: +# `pip install ultralytics` will pull in a CUDA-only torch wheel as a +# dependency if torch isn't already installed. On your RX 9070 XT setup, +# install the ROCm build of torch FIRST (from https://pytorch.org, ROCm +# variant), THEN install ultralytics, so it detects torch is already +# satisfied and doesn't overwrite it with the CUDA build. Once ROCm torch +# is in place, --mask-device cuda keeps working as-is (ROCm torch keeps +# the "cuda" device string). + +# tkinter (needed for pipeline_360_gui.py) is not pip-installable -- it's +# a system package. On CachyOS/Arch: sudo pacman -S tk