OBJ Sinkにアトラス化を付与したobj_atlas sinkの作成

nusamai/src/sink/obj_atlas/material.rs

@@ -1,6 +1,6 @@
 //! Material mangement
 
-use std::{hash::Hash, path::Path, time::Instant};
+use std::hash::Hash;
 
 use serde::{Deserialize, Serialize};
 use url::Url;
@@ -30,34 +30,3 @@ pub struct Texture {
 pub struct Image {
     pub uri: Url,
 }
-
-// NOTE: temporary implementation
-pub fn load_image(path: &Path) -> std::io::Result<Vec<u8>> {
-    if let Some(ext) = path.extension() {
-        match ext.to_ascii_lowercase().to_str() {
-            Some("tif" | "tiff" | "png") => {
-                let image = image::open(path)
-                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
-
-                let t = Instant::now();
-                let mut writer = std::io::Cursor::new(Vec::new());
-                let encoder = image::codecs::png::PngEncoder::new(&mut writer);
-                image
-                    .write_with_encoder(encoder)
-                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
-                log::debug!("Image encoding took {:?}", t.elapsed());
-
-                Ok(writer.into_inner())
-            }
-            Some("jpg" | "jpeg") => Ok(std::fs::read(path)?),
-            _ => {
-                let err = format!("Unsupported image format: {:?}", path);
-                Err(std::io::Error::new(std::io::ErrorKind::InvalidData, err))
-            }
-        }
-    } else {
-        let err = format!("Unsupported image format: {:?}", path);
-        log::error!("{}", err);
-        Err(std::io::Error::new(std::io::ErrorKind::InvalidData, err))
-    }
-}

nusamai/src/sink/obj_atlas/mod.rs

@@ -5,23 +5,30 @@
 use std::{f64::consts::FRAC_PI_2, path::PathBuf, sync::Mutex};
 
 use ahash::{HashMap, HashMapExt};
+use atlas_packer::{
+    export::{AtlasExporter as _, PngAtlasExporter},
+    pack::TexturePacker,
+    place::{GuillotineTexturePlacer, TexturePlacerConfig},
+    texture::{DownsampleFactor, TextureCache},
+};
 use earcut::{utils3d::project3d_to_2d, Earcut};
 use flatgeom::MultiPolygon;
 use glam::{DMat4, DVec3, DVec4};
 use indexmap::IndexSet;
 use itertools::Itertools;
 use material::{Material, Texture};
+use obj_writer::write;
+use rayon::iter::{IntoParallelIterator, ParallelBridge, ParallelIterator};
+use serde::{Deserialize, Serialize};
+use url::Url;
+
 use nusamai_citygml::{
     object::{ObjectStereotype, Value},
     schema::Schema,
     GeometryType,
 };
 use nusamai_plateau::appearance;
 use nusamai_projection::cartesian::geodetic_to_geocentric;
-use obj_writer::write;
-use rayon::iter::{IntoParallelIterator, ParallelBridge, ParallelIterator};
-use serde::{Deserialize, Serialize};
-use url::Url;
 
 use crate::{
     get_parameter_value,
@@ -368,6 +375,18 @@
             .try_for_each(|(typename, mut features)| {
                 feedback.ensure_not_canceled()?;
 
+                // Texture cache
+                let texture_cache = TextureCache::new(1_000_000_000);
+
+                // file output destination
+                let mut folder_path = self.output_path.clone();
+                let base_folder_name = typename.replace(':', "_").to_string();
+                folder_path.push(&base_folder_name);
+
+                let texture_folder_name = "textures";
+                let atlas_dir = folder_path.join(texture_folder_name);
+                std::fs::create_dir_all(&atlas_dir)?;
+
                 // Triangulation
                 let mut earcutter = Earcut::new();
                 let mut buf3d: Vec<[f64; 3]> = Vec::new();
@@ -377,6 +396,17 @@
                 let mut all_meshes = ObjInfo::new();
                 let mut all_materials = ObjMaterials::new();
 
+                // initialize texture packer
+                let config = TexturePlacerConfig {
+                    width: 4096,
+                    height: 4096,
+                    padding: 0,
+                };
+                let placer = GuillotineTexturePlacer::new(config.clone());
+                let exporter = PngAtlasExporter::default();
+                let ext = exporter.clone().get_extension().to_string();
+                let mut packer = TexturePacker::new(placer, exporter);
+
                 for feature in features.features.iter_mut() {
                     feedback.ensure_not_canceled()?;
 
@@ -395,20 +425,113 @@
                             [v_enu[0], v_enu[1], v_enu[2], u, v]
                         });
 
-                    for (poly, &orig_mat_id) in feature
+                    for (poly_count, (mut poly, &orig_mat_id)) in feature
                         .polygons
                         .iter()
                         .zip_eq(feature.polygon_material_ids.iter())
+                        .enumerate()
                     {
+                        let mat = &feature.materials[orig_mat_id as usize];
+                        let t = mat.base_texture.clone();
+                        let mut new_mat = mat.clone();
+                        if let Some(base_texture) = t {
+                            // textured
+                            let original_vertices = poly
+                                .raw_coords()
+                                .iter()
+                                .map(|[x, y, z, u, v]| (*x, *y, *z, *u, *v))
+                                .collect::<Vec<(f64, f64, f64, f64, f64)>>();
+
+                            let texture = texture_cache.get_or_insert(
+                                &original_vertices
+                                    .iter()
+                                    .map(|(_, _, _, u, v)| (*u, *v))
+                                    .collect::<Vec<_>>(),
+                                &base_texture.uri.to_file_path().unwrap(),
+                                &DownsampleFactor::new(&1.0).value(),
+                            );
+
+                            // Unique id required for placement in atlas
+                            let texture_id = format!(
+                                "{}_{}_{}",
+                                base_folder_name, feature.feature_id, poly_count
+                            );
+                            let info = packer.add_texture(texture_id, texture);
+
+                            let atlas_placed_uv_coords = info
+                                .placed_uv_coords
+                                .iter()
+                                .map(|(u, v)| ({ *u }, { *v }))
+                                .collect::<Vec<(f64, f64)>>();
+                            let updated_vertices = original_vertices
+                                .iter()
+                                .zip(atlas_placed_uv_coords.iter())
+                                .map(|((x, y, z, _, _), (u, v))| (*x, *y, *z, *u, *v))
+                                .collect::<Vec<(f64, f64, f64, f64, f64)>>();
+
+                            // Apply the UV coordinates placed in the atlas to the original polygon
+                            poly.transform_inplace(|&[x, y, z, _, _]| {
+                                let (u, v) = updated_vertices
+                                    .iter()
+                                    .find(|(x_, y_, z_, _, _)| {
+                                        (*x_ - x).abs() < 1e-6
+                                            && (*y_ - y).abs() < 1e-6
+                                            && (*z_ - z).abs() < 1e-6
+                                    })
+                                    .map(|(_, _, _, u, v)| (*u, *v))
+                                    .unwrap();
+                                [x, y, z, u, v]
+                            });
+
+                            let atlas_file_name = info.atlas_id.to_string();
+
+                            let atlas_uri =
+                                atlas_dir.join(atlas_file_name).with_extension(ext.clone());
+
+                            // update material
+                            new_mat = material::Material {
+                                base_color: mat.base_color,
+                                base_texture: Some(material::Texture {
+                                    uri: Url::from_file_path(atlas_uri).unwrap(),
+                                }),
+                            };
+                        }
+
                         let num_outer = match poly.hole_indices().first() {
                             Some(&v) => v as usize,
                             None => poly.raw_coords().len(),
                         };
 
-                        let poly_material = &feature.materials[orig_mat_id as usize];
+                        let poly_material = new_mat;
                         let poly_color = poly_material.base_color;
                         let poly_texture = poly_material.base_texture.as_ref();
-                        let poly_material_key = format!("{}_{}", feature.feature_id, orig_mat_id);
+                        let texture_name = poly_texture.map_or_else(
+                            || "".to_string(),
+                            |t| {
+                                t.uri
+                                    .to_file_path()
+                                    .unwrap()
+                                    .file_stem()
+                                    .unwrap()
+                                    .to_str()
+                                    .unwrap()
+                                    .to_string()
+                            },
+                        );
+                        let poly_material_key = poly_material.base_texture.as_ref().map_or_else(
+                            || {
+                                format!(
+                                    "material_{}_{}_{}",
+                                    poly_color[0], poly_color[1], poly_color[2]
+                                )
+                            },
+                            |_| {
+                                format!(
+                                    "{}_{}_{}",
+                                    base_folder_name, texture_folder_name, texture_name
+                                )
+                            },
+                        );
 
                         all_materials.insert(
                             poly_material_key.clone(),
@@ -442,15 +565,12 @@
                     all_meshes.insert(feature.feature_id.clone(), feature_mesh);
                 }
 
+                packer.finalize();
+                packer.export(&atlas_dir, &texture_cache, config.width, config.height);
+
                 feedback.ensure_not_canceled()?;
 
                 // Write OBJ file
-                let mut folder_path = self.output_path.clone();
-                let file_name = typename.replace(':', "_").to_string();
-                folder_path.push(&file_name);
-
-                std::fs::create_dir_all(&folder_path)?;
-
                 write(
                     all_meshes,
                     all_materials,

nusamai/src/sink/obj_atlas/obj_writer.rs

@@ -3,9 +3,8 @@
 use std::path::PathBuf;
 use std::{collections::HashMap, path::Path};
 
-use super::{material, ObjInfo, ObjMaterials};
+use super::{ObjInfo, ObjMaterials};
 use crate::pipeline::PipelineError;
-use material::load_image;
 
 pub fn write(
     meshes: ObjInfo,
@@ -16,20 +15,13 @@
     let mut material_cache: HashMap<String, String> = HashMap::new();
 
     write_mtl(&materials, &mut material_cache, &folder_path)?;
-    write_obj(
-        &meshes,
-        &materials,
-        &mut material_cache,
-        &folder_path,
-        is_split,
-    )?;
+    write_obj(&meshes, &mut material_cache, &folder_path, is_split)?;
 
     Ok(())
 }
 
 fn write_obj(
     meshes: &ObjInfo,
-    materials: &ObjMaterials,
     material_cache: &mut HashMap<String, String>,
     folder_path: &Path,
     is_split: bool,
@@ -54,13 +46,14 @@
         for tex_coord in &mesh.uvs {
             writeln!(obj_writer, "vt {} {}", tex_coord[0], tex_coord[1])?;
         }
+
         for (material_key, indices) in &mesh.primitives {
             if material_cache.contains_key(material_key) {
                 writeln!(obj_writer, "usemtl {}", material_key)?;
             } else {
-                let m = materials.get(material_key).unwrap();
-                let (r, g, b) = (m.base_color[0], m.base_color[1], m.base_color[2]);
-                writeln!(obj_writer, "usemtl material_{}_{}_{}", r, g, b)?;
+                // todo: Add error handling
+                println!("Material not found: {}", material_key);
+                continue;
             }
             for index in indices.chunks(3) {
                 writeln!(
@@ -103,25 +96,19 @@
         if let Some(uri) = &material.texture_uri {
             if let Ok(path) = uri.to_file_path() {
                 writeln!(mtl_writer, "newmtl {}", material_key)?;
-                let content = load_image(&path)?;
-
-                let textures_dir = folder_path.join("textures");
-                std::fs::create_dir_all(&textures_dir)?;
 
-                let image_file_name = format!("{}.jpg", material_key);
-                let image_path = textures_dir.join(&image_file_name);
-                std::fs::write(&image_path, content)?;
+                let texture_name = path.file_name().unwrap().to_str().unwrap();
+                writeln!(mtl_writer, "map_Kd .\\textures\\{}", texture_name)?;
 
-                writeln!(mtl_writer, "map_Kd .\\textures\\{}", image_file_name)?;
-                material_cache.insert(material_key.to_string(), image_file_name);
+                material_cache.insert(material_key.to_string(), path.to_str().unwrap().to_string());
             }
         } else {
             let (r, g, b) = (
                 material.base_color[0],
                 material.base_color[1],
                 material.base_color[2],
             );
-            let color_key = format!("{:.6}_{:.6}_{:.6}", r, g, b);
+            let color_key = format!("{}_{}_{}", r, g, b);
             let material_key = format!("material_{}_{}_{}", r, g, b);
             if material_cache.contains_key(&material_key) {
                 continue;