modify/
modify.rs

1use clap::Parser;
2use glam::*;
3
4use wgpu_3dgs_editor::{self as gs};
5
6/// The command line arguments.
7#[derive(Parser, Debug)]
8#[command(
9    version,
10    about,
11    long_about = "\
12    A 3D Gaussian splatting editor to apply basic modifier to all Gaussians in a model.
13    "
14)]
15struct Args {
16    /// Path to the .ply file.
17    #[arg(short, long)]
18    model: String,
19
20    /// The output path for the modified .ply file.
21    #[arg(short, long, default_value = "target/output.ply")]
22    output: String,
23
24    /// Whether to override the RGB color of the selected Gaussians.
25    #[arg(long)]
26    override_rgb: bool,
27
28    /// If [`Args::override_rgb`], then it is used to override the RGB color,
29    /// otherwise it is used to apply HSV modifications.
30    ///
31    /// Normally hue (H) is in [0, 1], saturation (S) and value (V) are in [0, 2].
32    /// This function adds the hue and multiplies saturation and value.
33    #[arg(
34        long,
35        allow_hyphen_values = true,
36        num_args = 3,
37        value_delimiter = ',',
38        default_value = "0.0,1.0,1.0"
39    )]
40    rgb_or_hsv: Vec<f32>,
41
42    /// Alpha is multiplied with the original alpha.
43    #[arg(long, allow_hyphen_values = true, default_value = "1.0")]
44    alpha: f32,
45
46    /// Contrast is applied to the RGB color.
47    ///
48    /// Normally the range is [-1, 1].
49    #[arg(long, allow_hyphen_values = true, default_value = "0.0")]
50    contrast: f32,
51
52    /// Exposure is applied to the RGB color.
53    ///
54    /// Normally the range is [-5, 5].
55    #[arg(long, allow_hyphen_values = true, default_value = "0.0")]
56    exposure: f32,
57
58    /// Gamma is applied to the RGB color.
59    ///
60    /// Normally the range is [0, 5].
61    #[arg(long, allow_hyphen_values = true, default_value = "1.0")]
62    gamma: f32,
63}
64
65type GaussianPod = gs::core::GaussianPodWithShSingleCov3dRotScaleConfigs;
66
67#[tokio::main]
68async fn main() {
69    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
70
71    let args = Args::parse();
72    let model_path = &args.model;
73
74    log::debug!("Creating wgpu instance");
75    let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
76
77    log::debug!("Requesting adapter");
78    let adapter = instance
79        .request_adapter(&wgpu::RequestAdapterOptions::default())
80        .await
81        .expect("adapter");
82
83    log::debug!("Requesting device");
84    let (device, queue) = adapter
85        .request_device(&wgpu::DeviceDescriptor {
86            label: Some("Device"),
87            required_features: wgpu::Features::empty(),
88            required_limits: adapter.limits(),
89            memory_hints: wgpu::MemoryHints::default(),
90            trace: wgpu::Trace::Off,
91        })
92        .await
93        .expect("device");
94
95    log::debug!("Creating gaussians");
96    let f = std::fs::File::open(model_path).expect("ply file");
97    let mut reader = std::io::BufReader::new(f);
98    let gaussians = gs::core::Gaussians::read_ply(&mut reader).expect("gaussians");
99
100    log::debug!("Creating editor");
101    let editor = gs::Editor::<GaussianPod>::new(&device, &gaussians);
102
103    log::debug!("Creating basic modifier");
104    let basic_modifier = gs::BasicModifier::<GaussianPod>::new(
105        &device,
106        &editor.gaussians_buffer,
107        &editor.model_transform_buffer,
108        &editor.gaussian_transform_buffer,
109    );
110
111    log::debug!("Configuring modifiers");
112    match args.override_rgb {
113        true => basic_modifier
114            .basic_color_modifiers_buffer
115            .update_with_override_rgb(
116                &queue,
117                Vec3::from_slice(&args.rgb_or_hsv),
118                args.alpha,
119                args.contrast,
120                args.exposure,
121                args.gamma,
122            ),
123        false => basic_modifier
124            .basic_color_modifiers_buffer
125            .update_with_hsv_modifiers(
126                &queue,
127                Vec3::from_slice(&args.rgb_or_hsv),
128                args.alpha,
129                args.contrast,
130                args.exposure,
131                args.gamma,
132            ),
133    }
134
135    log::info!("Starting editing process");
136    let time = std::time::Instant::now();
137
138    log::debug!("Editing Gaussians");
139    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
140        label: Some("Edit Encoder"),
141    });
142
143    editor.apply(
144        &device,
145        &mut encoder,
146        [&basic_modifier as &dyn gs::Modifier<GaussianPod>],
147    );
148
149    queue.submit(Some(encoder.finish()));
150
151    #[allow(unused_must_use)]
152    device.poll(wgpu::PollType::Wait);
153
154    log::info!("Editing process completed in {:?}", time.elapsed());
155
156    log::debug!("Downloading Gaussians");
157    let modified_gaussians = gs::core::Gaussians {
158        gaussians: editor
159            .gaussians_buffer
160            .download_gaussians(&device, &queue)
161            .await
162            .expect("gaussians download"),
163    };
164
165    log::debug!("Writing modified Gaussians to output file");
166    let output_file = std::fs::File::create(&args.output).expect("output file");
167    let mut writer = std::io::BufWriter::new(output_file);
168    modified_gaussians
169        .write_ply(&mut writer)
170        .expect("write modified Gaussians to output file");
171
172    log::info!("Modified Gaussians written to {}", args.output);
173}
modify/modify.rs

modify/
modify.rs
