Merge branch 'crabman/batch_map_generator' into 'master'

Batch map generation example

See merge request veloren/veloren!4334

Cargo.lock

@@ -1131,6 +1131,19 @@ name = "conrod_winit"
 version = "0.63.0"
 source = "git+https://gitlab.com/veloren/conrod.git?branch=copypasta_0.7#59fddc617696e68d28a75c2137a08c2572efb986"
 
+[[package]]
+name = "console"
+version = "0.15.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0e1f83fc076bd6dd27517eacdf25fef6c4dfe5f1d7448bafaaf3a26f13b5e4eb"
+dependencies = [
+ "encode_unicode",
+ "lazy_static",
+ "libc",
+ "unicode-width",
+ "windows-sys 0.52.0",
+]
+
 [[package]]
 name = "constant_time_eq"
 version = "0.3.0"
@@ -1937,6 +1950,12 @@ dependencies = [
  "bytemuck",
 ]
 
+[[package]]
+name = "encode_unicode"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f"
+
 [[package]]
 name = "encoding_rs"
 version = "0.8.33"
@@ -3138,6 +3157,19 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "indicatif"
+version = "0.17.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "763a5a8f45087d6bcea4222e7b72c291a054edf80e4ef6efd2a4979878c7bea3"
+dependencies = [
+ "console",
+ "instant",
+ "number_prefix",
+ "portable-atomic",
+ "unicode-width",
+]
+
 [[package]]
 name = "indoc"
 version = "2.0.4"
@@ -4329,6 +4361,12 @@ dependencies = [
  "syn 2.0.48",
 ]
 
+[[package]]
+name = "number_prefix"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
+
 [[package]]
 name = "objc"
 version = "0.2.7"
@@ -4746,6 +4784,12 @@ dependencies = [
  "miniz_oxide",
 ]
 
+[[package]]
+name = "portable-atomic"
+version = "1.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0"
+
 [[package]]
 name = "portpicker"
 version = "0.1.0"
@@ -7401,6 +7445,7 @@ dependencies = [
  "fxhash",
  "hashbrown 0.13.2",
  "image",
+ "indicatif",
  "itertools 0.10.5",
  "kiddo",
  "lazy_static",
@@ -7418,6 +7463,7 @@ dependencies = [
  "rstar",
  "rusqlite",
  "serde",
+ "signal-hook",
  "strum 0.24.1",
  "svg_fmt",
  "tracing",
@@ -8117,7 +8163,7 @@ dependencies = [
  "js-sys",
  "log",
  "naga",
- "parking_lot 0.11.2",
+ "parking_lot 0.12.1",
  "profiling",
  "raw-window-handle 0.5.2",
  "serde",
@@ -8142,7 +8188,7 @@ dependencies = [
  "codespan-reporting",
  "log",
  "naga",
- "parking_lot 0.11.2",
+ "parking_lot 0.12.1",
  "profiling",
  "raw-window-handle 0.5.2",
  "ron",
@@ -8177,13 +8223,13 @@ dependencies = [
  "js-sys",
  "khronos-egl",
  "libc",
- "libloading 0.7.4",
+ "libloading 0.8.1",
  "log",
  "metal",
  "naga",
  "objc",
  "once_cell",
- "parking_lot 0.11.2",
+ "parking_lot 0.12.1",
  "profiling",
  "range-alloc",
  "raw-window-handle 0.5.2",

world/Cargo.toml

@@ -8,7 +8,8 @@ edition = "2021"
 use-dyn-lib = ["common-dynlib"]
 be-dyn-lib = []
 simd = ["vek/platform_intrinsics", "packed_simd"]
-bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "clap", "rstar"]
+bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "rstar", "cli"]
+cli = ["clap", "signal-hook", "indicatif"]
 
 default = ["simd"]
 
@@ -51,6 +52,8 @@ num-traits = { workspace = true, optional = true }
 fallible-iterator = { version = "0.2.0", optional = true }
 rstar = { version = "0.10", optional = true }
 clap = { workspace = true, optional = true }
+signal-hook = { version = "0.3.6", optional = true }
+indicatif = { version = "0.17.8", optional = true }
 
 
 [dev-dependencies]
@@ -80,3 +83,7 @@ required-features = ["bin_compression"]
 [[example]]
 name = "heightmap_visualization"
 required-features = ["bin_compression"]
+
+[[example]]
+name = "batch_generate"
+required-features = ["cli"]

world/examples/batch_generate.rs

@@ -0,0 +1,378 @@
+use std::{
+    fs::{self, create_dir_all, File},
+    io::Write,
+    ops::RangeInclusive,
+    path::{Path, PathBuf},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+
+use clap::{Parser, Subcommand};
+use common::{
+    resources::MapKind,
+    terrain::{
+        map::{MapConfig, MapSample},
+        uniform_idx_as_vec2,
+    },
+};
+use image::{codecs::png::PngEncoder, ColorType, DynamicImage, GenericImage, ImageEncoder};
+use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
+use rand::{thread_rng, Rng};
+use rayon::ThreadPool;
+use serde::{Deserialize, Serialize};
+use tracing::{debug, error, info, info_span, Level, Span};
+use tracing_subscriber::EnvFilter;
+use vek::{Aabr, Rgb, Vec2};
+use veloren_world::{
+    sim::{get_horizon_map, sample_pos, sample_wpos, FileOpts, GenOpts, WorldOpts, WorldSimStage},
+    IndexOwned, World, WorldGenerateStage, CONFIG,
+};
+
+#[derive(Parser)]
+struct Cli {
+    #[command(subcommand)]
+    subcommand: Action,
+    /// Whether .bin files should be saved for maps
+    #[arg(short, long)]
+    save_bin: bool,
+    /// Hide progress bars
+    #[arg(short, long)]
+    no_progress: bool,
+    /// Path to where maps are saved
+    #[arg(long)]
+    maps_path: Option<PathBuf>,
+}
+
+#[derive(Subcommand)]
+enum Action {
+    /// Generate maps in a loop using the provided configuration
+    Batch {
+        /// Configuration to use for map generation
+        config: String,
+        /// How many maps will be generated in parallel
+        #[arg(short, long)]
+        threads: Option<usize>,
+    },
+    /// Generate a map from the .ron file emitted by the batch command
+    Regenerate {
+        config: String,
+        /// Override erosion quality
+        #[arg(long)]
+        erosion_quality: Option<f32>,
+    },
+}
+
+#[derive(Debug, Clone, Deserialize)]
+struct BatchGenerateConfig {
+    scale: RangeInclusive<f64>,
+    size: (u32, u32),
+    kind: MapKind,
+    erosion_quality: RangeInclusive<f32>,
+}
+
+impl BatchGenerateConfig {
+    fn gen_rand(&self) -> GenOpts {
+        GenOpts {
+            x_lg: self.size.0,
+            y_lg: self.size.1,
+            scale: thread_rng().gen_range(self.scale.clone()),
+            map_kind: self.kind,
+            erosion_quality: thread_rng().gen_range(self.erosion_quality.clone()),
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct MapGenConfig {
+    seed: u32,
+    gen_opts: GenOpts,
+}
+
+fn main() {
+    tracing_subscriber::fmt()
+        .with_max_level(Level::WARN)
+        .with_env_filter(EnvFilter::from_default_env())
+        .init();
+
+    let command = Cli::parse();
+
+    let maps_path = command.maps_path.unwrap_or(PathBuf::from("maps"));
+
+    match command.subcommand {
+        Action::Batch { config, threads } => do_batch_generate(
+            config,
+            command.save_bin,
+            threads,
+            command.no_progress,
+            maps_path,
+        ),
+        Action::Regenerate {
+            config,
+            erosion_quality,
+        } => do_regenerate(
+            config,
+            maps_path,
+            erosion_quality,
+            command.no_progress,
+            command.save_bin,
+        ),
+    }
+}
+
+fn generate_one(
+    seed: u32,
+    base_path: &Path,
+    gen_opts: GenOpts,
+    (save_bin, save_image, save_metadata): (bool, bool, bool),
+    span: &Span,
+    threadpool: &ThreadPool,
+    progress: Option<ProgressBar>,
+) -> (World, IndexOwned) {
+    if let Some(progress) = &progress {
+        progress.set_message(seed.to_string());
+    }
+
+    let (world, index) = World::generate(
+        seed,
+        WorldOpts {
+            seed_elements: false,
+            world_file: if save_bin {
+                FileOpts::Save(base_path.with_extension("bin"), gen_opts.clone())
+            } else {
+                FileOpts::Generate(gen_opts.clone())
+            },
+            calendar: None,
+        },
+        threadpool,
+        &|stage| {
+            if let WorldGenerateStage::WorldSimGenerate(WorldSimStage::Erosion(percentage)) = stage
+            {
+                if let Some(progress) = &progress {
+                    progress.set_position(percentage as u64);
+                }
+
+                span.in_scope(|| {
+                    info!("Erosion progress: {percentage:02.0}%");
+                })
+            }
+        },
+    );
+
+    if save_image {
+        let index_ref = index.as_index_ref();
+        let sampler = world.sim();
+        let map_size_lg = sampler.map_size_lg();
+
+        let horizons = get_horizon_map(
+            map_size_lg,
+            Aabr {
+                min: Vec2::zero(),
+                max: map_size_lg.chunks().map(|e| e as i32),
+            },
+            CONFIG.sea_level,
+            CONFIG.sea_level + sampler.max_height,
+            |posi| {
+                let sample = sampler.get(uniform_idx_as_vec2(map_size_lg, posi)).unwrap();
+
+                sample.basement.max(sample.water_alt)
+            },
+            |a| a,
+            |h| h,
+        )
+        .ok();
+
+        let mut map_config = MapConfig::orthographic(map_size_lg, 0.0..=sampler.max_height);
+        map_config.horizons = horizons.as_ref();
+        map_config.is_shaded = true;
+        map_config.is_stylized_topo = true;
+        let map = sampler.get_map(index_ref, None);
+
+        let mut image = DynamicImage::new(
+            map_size_lg.chunks().x as u32,
+            map_size_lg.chunks().y as u32,
+            image::ColorType::Rgba8,
+        );
+
+        map_config.generate(
+            |pos| {
+                let default_sample = sample_pos(&map_config, sampler, index_ref, None, pos);
+                let [r, g, b, _a] = map.rgba[pos].to_le_bytes();
+
+                MapSample {
+                    rgb: Rgb::new(r, g, b),
+                    ..default_sample
+                }
+            },
+            |wpos| sample_wpos(&map_config, sampler, wpos),
+            |pos, (r, g, b, a)| {
+                image.put_pixel(
+                    pos.x as u32,
+                    map_size_lg.chunks().y as u32 - pos.y as u32 - 1,
+                    [r, g, b, a].into(),
+                )
+            },
+        );
+
+        let mut image_file =
+            File::create_new(base_path.with_extension("png")).expect("Could not create map file");
+
+        if let Err(error) = PngEncoder::new(&mut image_file).write_image(
+            image.as_bytes(),
+            map_size_lg.chunks().x as u32,
+            map_size_lg.chunks().y as u32,
+            ColorType::Rgba8,
+        ) {
+            error!(?error, "Could not write image data");
+        }
+
+        let _ = image_file.flush();
+    }
+
+    if save_metadata {
+        // Write config
+        if let Err(error) = fs::write(
+            base_path.with_extension("ron"),
+            ron::ser::to_string_pretty(&MapGenConfig { seed, gen_opts }, Default::default())
+                .unwrap(),
+        ) {
+            error!(?error, "Colud not write map configuration file");
+        }
+    }
+
+    info!("Finished writing map to: {}", base_path.display());
+    if let Some(progress) = progress {
+        progress.finish()
+    }
+
+    (world, index)
+}
+
+fn do_regenerate(
+    config: String,
+    maps_path: PathBuf,
+    erosion_quality: Option<f32>,
+    no_progress: bool,
+    save_bin: bool,
+) {
+    let mut config: MapGenConfig =
+        ron::from_str(&fs::read_to_string(config).expect("Failed to read generation file"))
+            .expect("Could not parse generation file");
+
+    let base_path = if let Some(erosion_quality) = erosion_quality {
+        config.gen_opts.erosion_quality = erosion_quality;
+        maps_path.join(format!("{}_{:03}", config.seed, erosion_quality * 100.0))
+    } else {
+        maps_path.join(config.seed.to_string())
+    };
+
+    let span = info_span!("Generating map", map = ?config);
+    let pool = rayon::ThreadPoolBuilder::new().build().unwrap();
+
+    generate_one(
+        config.seed,
+        &base_path,
+        config.gen_opts,
+        (save_bin, true, true),
+        &span,
+        &pool,
+        (!no_progress).then(progress_bar),
+    );
+}
+
+fn do_batch_generate(
+    file: String,
+    save_bin: bool,
+    threads: Option<usize>,
+    no_progress: bool,
+    maps_path: PathBuf,
+) {
+    let config: BatchGenerateConfig =
+        ron::from_str(&fs::read_to_string(file).expect("Failed to read generator config file"))
+            .expect("Could not parse generator config");
+
+    #[cfg(debug_assertions)]
+    tracing::warn!("For best performance, run this in release mode");
+
+    let threads = threads.unwrap_or(1);
+
+    let mut handles = vec![];
+
+    let map_i = Arc::new(AtomicUsize::new(0));
+    let shutdown_started = Arc::new(std::sync::atomic::AtomicBool::new(false));
+
+    debug!("Registering shutdown signal");
+    use signal_hook::consts::signal::*;
+    let _ = signal_hook::flag::register_conditional_default(SIGINT, Arc::clone(&shutdown_started));
+    let _ = signal_hook::flag::register(SIGINT, Arc::clone(&shutdown_started));
+
+    create_dir_all(&maps_path).unwrap();
+
+    let progress_bars = (!no_progress).then(MultiProgress::new);
+
+    for thread_id in 0..threads {
+        info!(?thread_id, "Starting thread");
+        let config = config.clone();
+        let map_i = Arc::clone(&map_i);
+        let shutdown_started = Arc::clone(&shutdown_started);
+        let maps_path = maps_path.clone();
+        let progress_bars = progress_bars.clone();
+
+        let h = std::thread::spawn::<_, ()>(move || {
+            loop {
+                let progress = progress_bars.as_ref().map(|bars| {
+                    let progress = progress_bar();
+                    bars.add(progress.clone());
+                    progress
+                });
+
+                if shutdown_started.load(Ordering::Relaxed) {
+                    info!(?thread_id, "Shutting down thread");
+                    break;
+                }
+
+                let map_i = map_i.fetch_add(1, Ordering::SeqCst);
+
+                if let Some(progress) = &progress {
+                    progress.set_prefix(format!("Map {}", map_i));
+                }
+
+                let seed = thread_rng().gen::<u32>();
+                let span = info_span!("generate", map_i, thread_id);
+                let _guard = span.enter();
+                let gen_opts = config.gen_rand();
+                let base_path = maps_path.join(seed.to_string());
+
+                let threadpool = rayon::ThreadPoolBuilder::new().build().unwrap();
+
+                info!("Starting world generation");
+                generate_one(
+                    seed,
+                    &base_path,
+                    gen_opts,
+                    (save_bin, true, true),
+                    &span,
+                    &threadpool,
+                    progress,
+                );
+            }
+        });
+
+        handles.push(h);
+    }
+
+    for handle in handles {
+        let _ = handle.join();
+    }
+}
+
+fn progress_bar() -> ProgressBar {
+    ProgressBar::new(100).with_style(
+        ProgressStyle::with_template(
+            "[{elapsed_precise}] [{eta:6}] {prefix:8} {msg:15} [{wide_bar:.red/cyan}] {percent:3}%",
+        )
+        .unwrap()
+        .progress_chars("#>~"),
+    )
+}