SHA256
Funktioning rewrite, that is able to use gpu on a more variety of
systems
This commit is contained in:
+777
@@ -0,0 +1,777 @@
|
||||
//! Entry point: parse CLI, run the requested stage(s), stream progress.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::process::ExitCode;
|
||||
|
||||
use clap::{Args, Parser, Subcommand};
|
||||
|
||||
use pipeline_360::error::{CancelFlag, PipelineError};
|
||||
use pipeline_360::gpu::GpuContext;
|
||||
use pipeline_360::logging::{self, Logger};
|
||||
use pipeline_360::project::ProjectLayout;
|
||||
use pipeline_360::stages::cubemap;
|
||||
use pipeline_360::stages::extract::{self, BadFrameAction, ExtractConfig, MotionMode, Selection};
|
||||
use pipeline_360::stages::masks::{self, MasksConfig, Projection};
|
||||
use pipeline_360::stages::rig::{self, Layout, LinkMode};
|
||||
|
||||
/// CLI parsed into a typed struct. `gui_log_output` is read from either the
|
||||
/// parent or subcommand position.
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(
|
||||
name = "pipeline_360",
|
||||
version,
|
||||
about = "360° -> COLMAP cubemap rig pipeline"
|
||||
)]
|
||||
struct Cli {
|
||||
/// Whether to emit `[progress] ` lines for the GUI (accepted in any position).
|
||||
#[arg(long, hide = true)]
|
||||
gui_log_output: bool,
|
||||
|
||||
/// Launch the native GUI (requires the `gui` feature). Also launched when
|
||||
/// the binary is run with no subcommand.
|
||||
#[arg(long, hide = true)]
|
||||
gui: bool,
|
||||
|
||||
#[command(subcommand)]
|
||||
command: Option<Command>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Subcommand)]
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
enum Command {
|
||||
/// Run the full pipeline: extract -> masks -> cubemap -> rig.
|
||||
All(AllArgs),
|
||||
/// Stage 1 only: extract sharpest/best frames from video.
|
||||
Extract(ExtractArgs),
|
||||
/// Stage 2 only: generate YOLO SfM masks (planned).
|
||||
Masks(MasksArgs),
|
||||
/// Stage 3 only: project equirect frames to cube faces.
|
||||
Cubemap(CubemapArgs),
|
||||
/// Stage 4 only: generate COLMAP rig config + nested tree (or --self-test).
|
||||
Rig(RigArgs),
|
||||
}
|
||||
|
||||
#[derive(Debug, Args, Clone)]
|
||||
struct ExtractOpts {
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 30,
|
||||
help = "Select one frame per this many frames (chunk mode)"
|
||||
)]
|
||||
chunk_size: u32,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 1920,
|
||||
help = "Resize width for analysis; 0 = native"
|
||||
)]
|
||||
scale_width: u32,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 4,
|
||||
help = "Worker count (parity; GPU is the parallelism)"
|
||||
)]
|
||||
workers: u32,
|
||||
#[arg(
|
||||
long,
|
||||
default_value = "output_frame_%05d.png",
|
||||
help = "printf-style output pattern"
|
||||
)]
|
||||
output_pattern: String,
|
||||
#[arg(long, default_value = "png", help = "png | jpg")]
|
||||
output_format: String,
|
||||
#[arg(long, default_value_t = 95)]
|
||||
jpeg_quality: u32,
|
||||
#[arg(
|
||||
long = "extract-regenerate-metadata",
|
||||
help = "Ignore existing metadata CSV and regenerate"
|
||||
)]
|
||||
reuse_metadata_inv: bool,
|
||||
|
||||
// --- Quality thresholds ("throw away bad frames"). ---
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 0.0,
|
||||
help = "Drop frames whose variance-of-Laplacian is below this"
|
||||
)]
|
||||
min_sharpness: f64,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 0.0,
|
||||
help = "Drop frames whose mean luma is below this (0..255)"
|
||||
)]
|
||||
min_exposure: f64,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 255.0,
|
||||
help = "Drop frames whose mean luma is above this (0..255)"
|
||||
)]
|
||||
max_exposure: f64,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 0.0,
|
||||
help = "Drop frames whose luma std is below this"
|
||||
)]
|
||||
min_contrast: f64,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 0.0,
|
||||
help = "Min inter-frame motion (0..255) — parallax baseline"
|
||||
)]
|
||||
min_motion: f64,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 1.0e12,
|
||||
help = "Max inter-frame motion (0..255) — reject jitter/blur"
|
||||
)]
|
||||
max_motion: f64,
|
||||
|
||||
// --- Selection ---
|
||||
#[arg(long, default_value = "chunk", help = "chunk | spread")]
|
||||
selection: String,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 0,
|
||||
help = "Target frame count (spread mode); 0 = all survivors / chunk-based"
|
||||
)]
|
||||
target_count: u32,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 0,
|
||||
help = "Min frame-index spacing between selected frames"
|
||||
)]
|
||||
min_gap: u32,
|
||||
#[arg(
|
||||
long,
|
||||
default_value = "keep",
|
||||
help = "skip | keep | warn — chunk with no surviving frame"
|
||||
)]
|
||||
bad_frame_action: String,
|
||||
#[arg(
|
||||
long,
|
||||
default_value = "frame_diff",
|
||||
help = "off | frame_diff — inter-frame motion proxy"
|
||||
)]
|
||||
motion_mode: String,
|
||||
#[arg(long, default_value_t = 8, help = "Frames per GPU batch (throughput)")]
|
||||
batch: u32,
|
||||
}
|
||||
|
||||
impl ExtractOpts {
|
||||
fn to_config(&self) -> Result<ExtractConfig, PipelineError> {
|
||||
Ok(ExtractConfig {
|
||||
chunk_size: self.chunk_size,
|
||||
scale_width: self.scale_width,
|
||||
workers: self.workers,
|
||||
output_pattern: self.output_pattern.clone(),
|
||||
output_format: self.output_format.clone(),
|
||||
jpeg_quality: self.jpeg_quality,
|
||||
reuse_metadata: !self.reuse_metadata_inv,
|
||||
min_sharpness: self.min_sharpness,
|
||||
min_exposure: self.min_exposure,
|
||||
max_exposure: self.max_exposure,
|
||||
min_contrast: self.min_contrast,
|
||||
min_motion: self.min_motion,
|
||||
max_motion: self.max_motion,
|
||||
selection: Selection::parse(&self.selection)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
target_count: self.target_count,
|
||||
min_gap: self.min_gap,
|
||||
bad_frame_action: BadFrameAction::parse(&self.bad_frame_action)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
motion_mode: MotionMode::parse(&self.motion_mode)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
batch: self.batch,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Shared extract flags for the `all` subcommand, with `--extract-*` prefixes
|
||||
/// (mirrors the Python `all` CLI the GUI drives).
|
||||
#[derive(Debug, Args, Clone)]
|
||||
struct ExtractOptsAll {
|
||||
#[arg(long = "extract-chunk-size", default_value_t = 30)]
|
||||
chunk_size: u32,
|
||||
#[arg(long = "extract-scale-width", default_value_t = 1920)]
|
||||
scale_width: u32,
|
||||
#[arg(long = "extract-workers", default_value_t = 4)]
|
||||
workers: u32,
|
||||
#[arg(
|
||||
long = "extract-output-pattern",
|
||||
default_value = "output_frame_%05d.png"
|
||||
)]
|
||||
output_pattern: String,
|
||||
#[arg(long = "extract-output-format", default_value = "png")]
|
||||
output_format: String,
|
||||
#[arg(long = "extract-jpeg-quality", default_value_t = 95)]
|
||||
jpeg_quality: u32,
|
||||
#[arg(long = "extract-regenerate-metadata")]
|
||||
reuse_metadata_inv: bool,
|
||||
#[arg(long = "extract-min-sharpness", default_value_t = 0.0)]
|
||||
min_sharpness: f64,
|
||||
#[arg(long = "extract-min-exposure", default_value_t = 0.0)]
|
||||
min_exposure: f64,
|
||||
#[arg(long = "extract-max-exposure", default_value_t = 255.0)]
|
||||
max_exposure: f64,
|
||||
#[arg(long = "extract-min-contrast", default_value_t = 0.0)]
|
||||
min_contrast: f64,
|
||||
#[arg(long = "extract-min-motion", default_value_t = 0.0)]
|
||||
min_motion: f64,
|
||||
#[arg(long = "extract-max-motion", default_value_t = 1.0e12)]
|
||||
max_motion: f64,
|
||||
#[arg(long = "extract-selection", default_value = "chunk")]
|
||||
selection: String,
|
||||
#[arg(long = "extract-target-count", default_value_t = 0)]
|
||||
target_count: u32,
|
||||
#[arg(long = "extract-min-gap", default_value_t = 0)]
|
||||
min_gap: u32,
|
||||
#[arg(long = "extract-bad-frame-action", default_value = "keep")]
|
||||
bad_frame_action: String,
|
||||
#[arg(long = "extract-motion-mode", default_value = "frame_diff")]
|
||||
motion_mode: String,
|
||||
#[arg(long = "extract-batch", default_value_t = 8)]
|
||||
batch: u32,
|
||||
}
|
||||
|
||||
impl ExtractOptsAll {
|
||||
fn to_config(&self) -> Result<ExtractConfig, PipelineError> {
|
||||
Ok(ExtractConfig {
|
||||
chunk_size: self.chunk_size,
|
||||
scale_width: self.scale_width,
|
||||
workers: self.workers,
|
||||
output_pattern: self.output_pattern.clone(),
|
||||
output_format: self.output_format.clone(),
|
||||
jpeg_quality: self.jpeg_quality,
|
||||
reuse_metadata: !self.reuse_metadata_inv,
|
||||
min_sharpness: self.min_sharpness,
|
||||
min_exposure: self.min_exposure,
|
||||
max_exposure: self.max_exposure,
|
||||
min_contrast: self.min_contrast,
|
||||
min_motion: self.min_motion,
|
||||
max_motion: self.max_motion,
|
||||
selection: Selection::parse(&self.selection)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
target_count: self.target_count,
|
||||
min_gap: self.min_gap,
|
||||
bad_frame_action: BadFrameAction::parse(&self.bad_frame_action)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
motion_mode: MotionMode::parse(&self.motion_mode)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
batch: self.batch,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
struct AllArgs {
|
||||
#[arg(long)]
|
||||
project_dir: PathBuf,
|
||||
#[arg(long, num_args = 1..)]
|
||||
video: Vec<PathBuf>,
|
||||
#[arg(long, default_value = "extract")]
|
||||
start_stage: String,
|
||||
#[arg(long)]
|
||||
skip_masks: bool,
|
||||
#[arg(long, hide = true)]
|
||||
gui_log_output: bool,
|
||||
#[command(flatten)]
|
||||
extract: ExtractOptsAll,
|
||||
#[command(flatten)]
|
||||
mask: MasksOptsAll,
|
||||
#[arg(long = "cubemap-face-size", default_value_t = 1600)]
|
||||
cubemap_face_size: u32,
|
||||
#[arg(long = "cubemap-skip-faces", default_value = "")]
|
||||
cubemap_skip_faces: String,
|
||||
#[arg(long = "cubemap-jpeg-quality", default_value_t = 95)]
|
||||
cubemap_jpeg_quality: u32,
|
||||
#[arg(long = "cubemap-batch", default_value_t = 4)]
|
||||
cubemap_batch: u32,
|
||||
#[arg(long = "rig-ref-face", default_value = "F")]
|
||||
rig_ref_face: String,
|
||||
#[arg(long = "rig-intrinsics-convention", default_value = "linspace")]
|
||||
rig_intrinsics_convention: String,
|
||||
#[arg(long = "rig-link-mode", default_value = "symlink")]
|
||||
rig_link_mode: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
struct ExtractArgs {
|
||||
#[arg(long, num_args = 1..)]
|
||||
video: Vec<PathBuf>,
|
||||
#[arg(long)]
|
||||
output_dir: PathBuf,
|
||||
#[arg(long, hide = true)]
|
||||
gui_log_output: bool,
|
||||
#[command(flatten)]
|
||||
extract: ExtractOpts,
|
||||
}
|
||||
|
||||
// Shared masks-stage flags for the `masks` subcommand (unprefixed).
|
||||
#[derive(Debug, Args, Clone)]
|
||||
struct MasksOpts {
|
||||
#[arg(long, default_value = "person,car,bus,truck,bicycle,motorcycle")]
|
||||
classes: String,
|
||||
#[arg(
|
||||
long,
|
||||
default_value = "yolo11x-seg.onnx",
|
||||
help = "ONNX segmentation model path"
|
||||
)]
|
||||
model: String,
|
||||
#[arg(
|
||||
long,
|
||||
default_value = "cpu",
|
||||
help = "ort execution device (cpu|cuda|rocm)"
|
||||
)]
|
||||
device: String,
|
||||
#[arg(long, default_value_t = 0.35)]
|
||||
conf: f32,
|
||||
#[arg(long, default_value_t = 1280)]
|
||||
imgsz: u32,
|
||||
#[arg(long, default_value_t = 6)]
|
||||
dilate: u32,
|
||||
#[arg(long, default_value = "cubemap", help = "cubemap | equirect")]
|
||||
projection: String,
|
||||
#[arg(long, default_value_t = 1024)]
|
||||
face_size: u32,
|
||||
#[arg(long, default_value = "")]
|
||||
skip_faces: String,
|
||||
#[arg(long, default_value_t = 8)]
|
||||
mask_batch: u32,
|
||||
#[arg(long, default_value_t = 8)]
|
||||
tile_batch: u32,
|
||||
#[arg(long = "mask-empty-if-none")]
|
||||
empty_if_none: bool,
|
||||
#[arg(long, default_value_t = 0.45, help = "NMS IoU threshold")]
|
||||
iou: f32,
|
||||
#[arg(
|
||||
long,
|
||||
help = "Directory for the MIGraphX compiled-model cache (GPU path). First \
|
||||
run compiles+saves; later runs load. Pair with --device rocm. \
|
||||
Defaults to ~/.cache/pipeline_360/migraphx."
|
||||
)]
|
||||
ep_cache: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl MasksOpts {
|
||||
fn to_config(&self) -> Result<MasksConfig, PipelineError> {
|
||||
let skip = cubemap::parse_skip_faces(&self.skip_faces)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
Ok(MasksConfig {
|
||||
classes: masks::resolve_classes(&self.classes)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
model: PathBuf::from(&self.model),
|
||||
device: self.device.clone(),
|
||||
conf: self.conf,
|
||||
imgsz: self.imgsz,
|
||||
dilate: self.dilate,
|
||||
projection: Projection::parse(&self.projection)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
face_size: self.face_size,
|
||||
skip_faces: skip,
|
||||
batch: self.mask_batch,
|
||||
tile_batch: self.tile_batch,
|
||||
empty_if_none: self.empty_if_none,
|
||||
iou: self.iou,
|
||||
ep_cache: self.ep_cache.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Shared masks-stage flags for `all` (with `--mask-*` prefixes, matching Python).
|
||||
#[derive(Debug, Args, Clone)]
|
||||
struct MasksOptsAll {
|
||||
#[arg(
|
||||
long = "mask-classes",
|
||||
default_value = "person,car,bus,truck,bicycle,motorcycle"
|
||||
)]
|
||||
classes: String,
|
||||
#[arg(long = "mask-model", default_value = "yolo11x-seg.onnx")]
|
||||
model: String,
|
||||
#[arg(long = "mask-device", default_value = "cpu")]
|
||||
device: String,
|
||||
#[arg(long = "mask-conf", default_value_t = 0.35)]
|
||||
conf: f32,
|
||||
#[arg(long = "mask-imgsz", default_value_t = 1280)]
|
||||
imgsz: u32,
|
||||
#[arg(long = "mask-dilate", default_value_t = 6)]
|
||||
dilate: u32,
|
||||
#[arg(long = "mask-projection", default_value = "cubemap")]
|
||||
projection: String,
|
||||
#[arg(long = "mask-face-size", default_value_t = 1024)]
|
||||
face_size: u32,
|
||||
#[arg(long = "mask-skip-faces", default_value = "")]
|
||||
skip_faces: String,
|
||||
#[arg(long = "mask-batch", default_value_t = 8)]
|
||||
mask_batch: u32,
|
||||
#[arg(long = "mask-tile-batch", default_value_t = 8)]
|
||||
mask_tile_batch: u32,
|
||||
#[arg(long = "mask-empty-if-none")]
|
||||
empty_if_none: bool,
|
||||
#[arg(long = "mask-iou", default_value_t = 0.45)]
|
||||
iou: f32,
|
||||
#[arg(
|
||||
long = "mask-ep-cache",
|
||||
help = "Directory for the MIGraphX compiled-model cache (GPU path). First \
|
||||
run compiles+saves; later runs load. Pair with --mask-device rocm. \
|
||||
Defaults to ~/.cache/pipeline_360/migraphx."
|
||||
)]
|
||||
ep_cache: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl MasksOptsAll {
|
||||
fn to_config(&self) -> Result<MasksConfig, PipelineError> {
|
||||
let skip = cubemap::parse_skip_faces(&self.skip_faces)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
Ok(MasksConfig {
|
||||
classes: masks::resolve_classes(&self.classes)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
model: PathBuf::from(&self.model),
|
||||
device: self.device.clone(),
|
||||
conf: self.conf,
|
||||
imgsz: self.imgsz,
|
||||
dilate: self.dilate,
|
||||
projection: Projection::parse(&self.projection)
|
||||
.map_err(|e| PipelineError(e.to_string()))?,
|
||||
face_size: self.face_size,
|
||||
skip_faces: skip,
|
||||
batch: self.mask_batch,
|
||||
tile_batch: self.mask_tile_batch,
|
||||
empty_if_none: self.empty_if_none,
|
||||
iou: self.iou,
|
||||
ep_cache: self.ep_cache.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
struct MasksArgs {
|
||||
#[arg(long)]
|
||||
input_dir: PathBuf,
|
||||
#[arg(long)]
|
||||
output_dir: PathBuf,
|
||||
#[arg(long, hide = true)]
|
||||
gui_log_output: bool,
|
||||
#[command(flatten)]
|
||||
opts: MasksOpts,
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
struct CubemapArgs {
|
||||
#[arg(long)]
|
||||
frames_dir: PathBuf,
|
||||
#[arg(long)]
|
||||
masks_dir: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
output_frames_dir: PathBuf,
|
||||
#[arg(long)]
|
||||
output_masks_dir: Option<PathBuf>,
|
||||
#[arg(long, default_value_t = 1600)]
|
||||
face_size: u32,
|
||||
#[arg(long, default_value = "")]
|
||||
skip_faces: String,
|
||||
#[arg(long, default_value_t = 95)]
|
||||
jpeg_quality: u32,
|
||||
#[arg(long)]
|
||||
resume: bool,
|
||||
#[arg(long, default_value_t = 4, help = "Frames per GPU batch (throughput)")]
|
||||
cubemap_batch: u32,
|
||||
#[arg(long, hide = true)]
|
||||
gui_log_output: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
struct RigArgs {
|
||||
#[arg(long)]
|
||||
images_dir: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
masks_dir: Option<PathBuf>,
|
||||
#[arg(long, default_value = "rig_config.json")]
|
||||
output: PathBuf,
|
||||
#[arg(long)]
|
||||
face_size: Option<u32>,
|
||||
#[arg(long, default_value = "auto")]
|
||||
layout: String,
|
||||
#[arg(long)]
|
||||
link_tree: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
masks_link_tree: Option<PathBuf>,
|
||||
#[arg(long, default_value = "symlink")]
|
||||
link_mode: String,
|
||||
#[arg(long, default_value = "F")]
|
||||
ref_face: String,
|
||||
#[arg(long, default_value = "linspace")]
|
||||
intrinsics_convention: String,
|
||||
#[arg(long)]
|
||||
self_test: bool,
|
||||
#[arg(long, hide = true)]
|
||||
gui_log_output: bool,
|
||||
}
|
||||
|
||||
fn run() -> Result<(), PipelineError> {
|
||||
let cli = Cli::parse();
|
||||
let sub_gui = match &cli.command {
|
||||
Some(Command::All(a)) => a.gui_log_output,
|
||||
Some(Command::Extract(a)) => a.gui_log_output,
|
||||
Some(Command::Masks(a)) => a.gui_log_output,
|
||||
Some(Command::Cubemap(a)) => a.gui_log_output,
|
||||
Some(Command::Rig(a)) => a.gui_log_output,
|
||||
None => false,
|
||||
};
|
||||
let gui_log = cli.gui_log_output || sub_gui;
|
||||
let logger: Logger = if gui_log {
|
||||
Some(&|m| println!("{m}"))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let cancel = CancelFlag::new();
|
||||
|
||||
// No subcommand (or explicit --gui): launch the native GUI if available.
|
||||
if cli.command.is_none() || cli.gui {
|
||||
#[cfg(feature = "gui")]
|
||||
{
|
||||
pipeline_360::gui::launch().map_err(|e| PipelineError(format!("GUI: {e}")))?;
|
||||
return Ok(());
|
||||
}
|
||||
#[cfg(not(feature = "gui"))]
|
||||
{
|
||||
use clap::CommandFactory;
|
||||
Cli::command()
|
||||
.print_help()
|
||||
.map_err(|e| PipelineError(format!("print_help: {e}")))?;
|
||||
println!("\n\n(no subcommand given; rebuild with `--features gui` to launch the GUI)");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
match cli.command {
|
||||
Some(Command::Rig(RigArgs {
|
||||
self_test: true, ..
|
||||
})) => {
|
||||
let ok = pipeline_360::selftest::run_selftest().map_err(PipelineError)?;
|
||||
logging::emit(&format!("self-test: {ok} check groups passed."), logger);
|
||||
}
|
||||
Some(Command::Rig(a)) => {
|
||||
let images_dir = a.images_dir.as_deref().ok_or_else(|| {
|
||||
PipelineError("--images-dir is required (or use --self-test)".into())
|
||||
})?;
|
||||
if a.masks_link_tree.is_some() && a.masks_dir.is_none() {
|
||||
return Err(PipelineError(
|
||||
"--masks-link-tree requires --masks-dir".into(),
|
||||
));
|
||||
}
|
||||
run_rig(images_dir, &a, logger)?;
|
||||
}
|
||||
Some(Command::All(a)) => run_all(&a, logger, &cancel)?,
|
||||
Some(Command::Extract(a)) => {
|
||||
let cfg = a.extract.to_config()?;
|
||||
let gpu = pollster::block_on(GpuContext::new("extract"))
|
||||
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
|
||||
extract::run_extract_stage(&a.video, &a.output_dir, &cfg, &gpu, &cancel, logger)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
}
|
||||
Some(Command::Masks(a)) => {
|
||||
let cfg = a.opts.to_config()?;
|
||||
let gpu = pollster::block_on(GpuContext::new("masks"))
|
||||
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
|
||||
masks::run_masks_stage(&a.input_dir, &a.output_dir, &cfg, &gpu, &cancel, logger)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
}
|
||||
Some(Command::Cubemap(a)) => run_cubemap(a, &cancel, logger)?,
|
||||
None => unreachable!(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run_cubemap(a: CubemapArgs, cancel: &CancelFlag, logger: Logger) -> Result<(), PipelineError> {
|
||||
if a.masks_dir.is_some() && a.output_masks_dir.is_none() {
|
||||
return Err(PipelineError(
|
||||
"--masks-dir given but --output-masks-dir missing".into(),
|
||||
));
|
||||
}
|
||||
let gpu = pollster::block_on(GpuContext::new("cubemap"))
|
||||
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
|
||||
let skip =
|
||||
cubemap::parse_skip_faces(&a.skip_faces).map_err(|e| PipelineError(e.to_string()))?;
|
||||
let (fout, mout) = cubemap::run_cubemap_stage(
|
||||
&a.frames_dir,
|
||||
&a.output_frames_dir,
|
||||
a.masks_dir.as_deref(),
|
||||
a.output_masks_dir.as_deref(),
|
||||
a.face_size,
|
||||
&skip,
|
||||
a.jpeg_quality,
|
||||
a.resume,
|
||||
a.cubemap_batch,
|
||||
&gpu,
|
||||
cancel,
|
||||
logger,
|
||||
)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
logging::emit(&format!("Cubemap output: {}", fout.display()), logger);
|
||||
if let Some(m) = mout {
|
||||
logging::emit(&format!("Cubemap masks: {}", m.display()), logger);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run_rig(images_dir: &std::path::Path, a: &RigArgs, logger: Logger) -> Result<(), PipelineError> {
|
||||
let layout = match a.layout.as_str() {
|
||||
"flat" => Layout::Flat,
|
||||
"nested" => Layout::Nested,
|
||||
_ => Layout::Auto,
|
||||
};
|
||||
let mode = LinkMode::parse(&a.link_mode)?;
|
||||
let out = rig::run_rig_stage(
|
||||
images_dir,
|
||||
&a.output,
|
||||
a.masks_dir.as_deref(),
|
||||
a.face_size,
|
||||
layout,
|
||||
&a.ref_face,
|
||||
&a.intrinsics_convention,
|
||||
a.link_tree.as_deref(),
|
||||
a.masks_link_tree.as_deref(),
|
||||
mode,
|
||||
)?;
|
||||
for msg in &out.messages {
|
||||
logging::emit(msg, logger);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run_all(a: &AllArgs, logger: Logger, cancel: &CancelFlag) -> Result<(), PipelineError> {
|
||||
let project = ProjectLayout::new(&a.project_dir);
|
||||
project
|
||||
.create_all()
|
||||
.map_err(|e| PipelineError(format!("mkdir: {e}")))?;
|
||||
|
||||
let stage_idx = |s: &str| match s {
|
||||
"extract" => 0,
|
||||
"cubemap" => 1,
|
||||
"masks" => 2,
|
||||
"rig" => 3,
|
||||
_ => 0,
|
||||
};
|
||||
let from = stage_idx(&a.start_stage);
|
||||
// NOTE: order is extract -> cubemap -> masks -> rig. Masks run on the
|
||||
// cube-face images the cubemap stage produces, so cubemap must come first
|
||||
// (masking equirect frames is useless, and round-tripping e2c->YOLO->c2e
|
||||
// then e2c again was wasteful). `02_masks` is no longer written by `all`.
|
||||
let stages = ["extract", "cubemap", "masks", "rig"];
|
||||
|
||||
if a.video.is_empty() && from == 0 {
|
||||
return Err(PipelineError(
|
||||
"--video is required when starting from the extract stage.".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// One GPU device shared by every GPU stage in the run.
|
||||
let gpu = pollster::block_on(GpuContext::new("pipeline"))
|
||||
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
|
||||
|
||||
for (i, stage) in stages.iter().enumerate() {
|
||||
if i < from {
|
||||
continue;
|
||||
}
|
||||
cancel
|
||||
.check()
|
||||
.map_err(|_| PipelineError("Cancelled.".into()))?;
|
||||
logging::stage_header(i + 1, 4, stage_title(stage), logger);
|
||||
match *stage {
|
||||
"extract" => {
|
||||
let cfg = a.extract.to_config()?;
|
||||
extract::run_extract_stage(&a.video, &project.frames, &cfg, &gpu, cancel, logger)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
}
|
||||
"masks" => {
|
||||
if a.skip_masks {
|
||||
logging::emit("Skipping masks stage (--skip-masks).", logger);
|
||||
} else {
|
||||
let cfg = a.mask.to_config()?;
|
||||
// Masks run on the cube-face images the cubemap stage
|
||||
// produced, and write cube-face masks next to them.
|
||||
masks::run_masks_stage(
|
||||
&project.cubemap_images,
|
||||
&project.cubemap_masks,
|
||||
&cfg,
|
||||
&gpu,
|
||||
cancel,
|
||||
logger,
|
||||
)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
}
|
||||
}
|
||||
"cubemap" => {
|
||||
let skip = cubemap::parse_skip_faces(&a.cubemap_skip_faces)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
// In the `all` flow the cubemap stage only projects images; the
|
||||
// masks stage runs afterwards on those cube faces (no equirect
|
||||
// mask round-trip, no double e2c projection).
|
||||
cubemap::run_cubemap_stage(
|
||||
&project.frames,
|
||||
&project.cubemap_images,
|
||||
None,
|
||||
None,
|
||||
a.cubemap_face_size,
|
||||
&skip,
|
||||
a.cubemap_jpeg_quality,
|
||||
false,
|
||||
a.cubemap_batch,
|
||||
&gpu,
|
||||
cancel,
|
||||
logger,
|
||||
)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
}
|
||||
"rig" => {
|
||||
let mode = LinkMode::parse(&a.rig_link_mode)?;
|
||||
let out = rig::run_rig_stage(
|
||||
&project.cubemap_images,
|
||||
&project.rig_config,
|
||||
if a.skip_masks {
|
||||
None
|
||||
} else {
|
||||
Some(project.cubemap_masks.as_path())
|
||||
},
|
||||
None,
|
||||
Layout::Auto,
|
||||
&a.rig_ref_face,
|
||||
&a.rig_intrinsics_convention,
|
||||
Some(&project.colmap_images),
|
||||
if a.skip_masks {
|
||||
None
|
||||
} else {
|
||||
Some(&project.colmap_masks)
|
||||
},
|
||||
mode,
|
||||
)
|
||||
.map_err(|e| PipelineError(e.to_string()))?;
|
||||
for msg in &out.messages {
|
||||
logging::emit(msg, logger);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
logging::emit("Pipeline complete.", logger);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn stage_title(stage: &str) -> &'static str {
|
||||
match stage {
|
||||
"extract" => "Extract Sharpest Frames",
|
||||
"masks" => "Generate SfM Masks",
|
||||
"cubemap" => "Project to Cubemap",
|
||||
"rig" => "Generate Rig Config",
|
||||
_ => "Unknown",
|
||||
}
|
||||
}
|
||||
|
||||
fn main() -> ExitCode {
|
||||
match run() {
|
||||
Ok(()) => ExitCode::SUCCESS,
|
||||
Err(e) => {
|
||||
eprintln!("Error: {e}");
|
||||
ExitCode::FAILURE
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user