oxihuman_export/
weight_map_export.rs

1// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3
4//! Export bone/skin weight maps as images or data files.
5
6/// A single bone weight assignment for a vertex.
7#[allow(dead_code)]
8#[derive(Clone, Debug)]
9pub struct BoneWeight {
10    pub bone_index: u32,
11    pub weight: f32,
12}
13
14/// Configuration for weight map generation.
15#[allow(dead_code)]
16pub struct WeightMapConfig {
17    pub width: u32,
18    pub height: u32,
19    pub bone_index: u32,
20    pub gamma: f32,
21}
22
23/// Pixel buffer storing grayscale weight values.
24#[allow(dead_code)]
25pub struct WeightMapBuffer {
26    pub pixels: Vec<f32>,
27    pub width: u32,
28    pub height: u32,
29}
30
31/// Statistics for a weight map buffer.
32#[allow(dead_code)]
33pub struct WeightMapStats {
34    pub min: f32,
35    pub max: f32,
36    pub avg: f32,
37}
38
39/// Type alias for weight-per-vertex data (vertex_index, weights).
40#[allow(dead_code)]
41pub type VertexWeights = Vec<Vec<BoneWeight>>;
42
43#[allow(dead_code)]
44pub fn default_weight_map_config(width: u32, height: u32) -> WeightMapConfig {
45    WeightMapConfig {
46        width,
47        height,
48        bone_index: 0,
49        gamma: 1.0,
50    }
51}
52
53#[allow(dead_code)]
54pub fn new_weight_map_buffer(width: u32, height: u32) -> WeightMapBuffer {
55    let count = (width * height) as usize;
56    WeightMapBuffer {
57        pixels: vec![0.0; count],
58        width,
59        height,
60    }
61}
62
63#[allow(dead_code)]
64pub fn set_weight_pixel(buffer: &mut WeightMapBuffer, x: u32, y: u32, value: f32) {
65    if x < buffer.width && y < buffer.height {
66        let idx = (y * buffer.width + x) as usize;
67        buffer.pixels[idx] = value;
68    }
69}
70
71#[allow(dead_code)]
72pub fn get_weight_pixel(buffer: &WeightMapBuffer, x: u32, y: u32) -> f32 {
73    if x < buffer.width && y < buffer.height {
74        let idx = (y * buffer.width + x) as usize;
75        buffer.pixels[idx]
76    } else {
77        0.0
78    }
79}
80
81/// Rasterize per-vertex weights to UV space for a specific bone.
82#[allow(dead_code)]
83pub fn weight_map_from_vertices(
84    vertex_weights: &[Vec<BoneWeight>],
85    uvs: &[[f32; 2]],
86    bone_index: u32,
87    width: u32,
88    height: u32,
89) -> WeightMapBuffer {
90    let mut buffer = new_weight_map_buffer(width, height);
91    let count = vertex_weights.len().min(uvs.len());
92    for i in 0..count {
93        let w = vertex_weights[i]
94            .iter()
95            .find(|bw| bw.bone_index == bone_index)
96            .map_or(0.0, |bw| bw.weight);
97        let u = uvs[i][0].clamp(0.0, 1.0);
98        let v = uvs[i][1].clamp(0.0, 1.0);
99        let px = (u * (width as f32 - 1.0)).round() as u32;
100        let py = (v * (height as f32 - 1.0)).round() as u32;
101        set_weight_pixel(&mut buffer, px, py, w);
102    }
103    buffer
104}
105
106/// Encode weight map as grayscale PPM (P5 PGM).
107#[allow(dead_code)]
108pub fn encode_weight_map_ppm(buffer: &WeightMapBuffer) -> Vec<u8> {
109    let header = format!("P5\n{} {}\n255\n", buffer.width, buffer.height);
110    let mut out = header.into_bytes();
111    for &val in &buffer.pixels {
112        let byte = (val.clamp(0.0, 1.0) * 255.0).round() as u8;
113        out.push(byte);
114    }
115    out
116}
117
118/// Normalize weights so that per-vertex weights sum to 1.0.
119#[allow(dead_code)]
120pub fn normalize_weights(vertex_weights: &mut [Vec<BoneWeight>]) {
121    for weights in vertex_weights.iter_mut() {
122        let sum: f32 = weights.iter().map(|bw| bw.weight).sum();
123        if sum > 1e-8 {
124            for bw in weights.iter_mut() {
125                bw.weight /= sum;
126            }
127        }
128    }
129}
130
131/// Get the top N highest weights per vertex.
132#[allow(dead_code)]
133pub fn top_n_weights(weights: &[BoneWeight], n: usize) -> Vec<BoneWeight> {
134    let mut sorted: Vec<BoneWeight> = weights.to_vec();
135    sorted.sort_by(|a, b| {
136        b.weight
137            .partial_cmp(&a.weight)
138            .unwrap_or(std::cmp::Ordering::Equal)
139    });
140    sorted.truncate(n);
141    sorted
142}
143
144/// Export weight map buffer to CSV string.
145#[allow(dead_code)]
146pub fn weight_map_to_csv(buffer: &WeightMapBuffer) -> String {
147    let mut out = String::from("x,y,weight\n");
148    for y in 0..buffer.height {
149        for x in 0..buffer.width {
150            let val = get_weight_pixel(buffer, x, y);
151            out.push_str(&format!("{},{},{:.6}\n", x, y, val));
152        }
153    }
154    out
155}
156
157/// Blend two weight maps with factor t (0 = all a, 1 = all b).
158#[allow(dead_code)]
159pub fn blend_weight_maps(a: &WeightMapBuffer, b: &WeightMapBuffer, t: f32) -> WeightMapBuffer {
160    let width = a.width;
161    let height = a.height;
162    let count = (width * height) as usize;
163    let t = t.clamp(0.0, 1.0);
164    let mut pixels = Vec::with_capacity(count);
165    for i in 0..count {
166        let va = a.pixels[i];
167        let vb = if i < b.pixels.len() { b.pixels[i] } else { 0.0 };
168        pixels.push(va * (1.0 - t) + vb * t);
169    }
170    WeightMapBuffer {
171        pixels,
172        width,
173        height,
174    }
175}
176
177/// Return total pixel count of the weight map.
178#[allow(dead_code)]
179pub fn weight_map_pixel_count(buffer: &WeightMapBuffer) -> usize {
180    buffer.pixels.len()
181}
182
183/// Compute min, max, and average weight values.
184#[allow(dead_code)]
185pub fn weight_map_stats(buffer: &WeightMapBuffer) -> WeightMapStats {
186    if buffer.pixels.is_empty() {
187        return WeightMapStats {
188            min: 0.0,
189            max: 0.0,
190            avg: 0.0,
191        };
192    }
193    let mut min = f32::MAX;
194    let mut max = f32::MIN;
195    let mut sum = 0.0f64;
196    for &v in &buffer.pixels {
197        if v < min {
198            min = v;
199        }
200        if v > max {
201            max = v;
202        }
203        sum += v as f64;
204    }
205    let avg = (sum / buffer.pixels.len() as f64) as f32;
206    WeightMapStats { min, max, avg }
207}
208
209/// Clear all pixels to zero.
210#[allow(dead_code)]
211pub fn clear_weight_map(buffer: &mut WeightMapBuffer) {
212    for px in buffer.pixels.iter_mut() {
213        *px = 0.0;
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn test_default_config() {
223        let cfg = default_weight_map_config(256, 256);
224        assert_eq!(cfg.width, 256);
225        assert_eq!(cfg.height, 256);
226        assert_eq!(cfg.bone_index, 0);
227        assert!((cfg.gamma - 1.0).abs() < f32::EPSILON);
228    }
229
230    #[test]
231    fn test_new_buffer() {
232        let buf = new_weight_map_buffer(8, 8);
233        assert_eq!(buf.width, 8);
234        assert_eq!(buf.height, 8);
235        assert_eq!(buf.pixels.len(), 64);
236    }
237
238    #[test]
239    fn test_set_get_pixel() {
240        let mut buf = new_weight_map_buffer(4, 4);
241        set_weight_pixel(&mut buf, 2, 3, 0.75);
242        let v = get_weight_pixel(&buf, 2, 3);
243        assert!((v - 0.75).abs() < f32::EPSILON);
244    }
245
246    #[test]
247    fn test_get_pixel_out_of_bounds() {
248        let buf = new_weight_map_buffer(4, 4);
249        let v = get_weight_pixel(&buf, 10, 10);
250        assert!((v - 0.0).abs() < f32::EPSILON);
251    }
252
253    #[test]
254    fn test_set_pixel_out_of_bounds() {
255        let mut buf = new_weight_map_buffer(4, 4);
256        set_weight_pixel(&mut buf, 10, 10, 1.0);
257        // No panic; all pixels stay 0
258        for &px in &buf.pixels {
259            assert!((px - 0.0).abs() < f32::EPSILON);
260        }
261    }
262
263    #[test]
264    fn test_weight_map_from_vertices() {
265        let vw = vec![
266            vec![BoneWeight {
267                bone_index: 0,
268                weight: 1.0,
269            }],
270            vec![BoneWeight {
271                bone_index: 1,
272                weight: 0.5,
273            }],
274        ];
275        let uvs = [[0.0, 0.0], [0.5, 0.5]];
276        let buf = weight_map_from_vertices(&vw, &uvs, 0, 8, 8);
277        assert_eq!(buf.width, 8);
278        // The first vertex at (0,0) should have weight 1.0
279        assert!((get_weight_pixel(&buf, 0, 0) - 1.0).abs() < f32::EPSILON);
280    }
281
282    #[test]
283    fn test_encode_ppm_starts_with_p5() {
284        let buf = new_weight_map_buffer(2, 2);
285        let ppm = encode_weight_map_ppm(&buf);
286        assert!(ppm.starts_with(b"P5"));
287    }
288
289    #[test]
290    fn test_encode_ppm_size() {
291        let buf = new_weight_map_buffer(4, 4);
292        let ppm = encode_weight_map_ppm(&buf);
293        let header = "P5\n4 4\n255\n".to_string();
294        assert_eq!(ppm.len(), header.len() + 16);
295    }
296
297    #[test]
298    fn test_normalize_weights() {
299        let mut vw = vec![vec![
300            BoneWeight {
301                bone_index: 0,
302                weight: 2.0,
303            },
304            BoneWeight {
305                bone_index: 1,
306                weight: 3.0,
307            },
308        ]];
309        normalize_weights(&mut vw);
310        let sum: f32 = vw[0].iter().map(|bw| bw.weight).sum();
311        assert!((sum - 1.0).abs() < 1e-6);
312    }
313
314    #[test]
315    fn test_normalize_weights_zero_sum() {
316        let mut vw = vec![vec![BoneWeight {
317            bone_index: 0,
318            weight: 0.0,
319        }]];
320        normalize_weights(&mut vw);
321        // Should not panic, weight stays 0
322        assert!((vw[0][0].weight - 0.0).abs() < f32::EPSILON);
323    }
324
325    #[test]
326    fn test_top_n_weights() {
327        let weights = vec![
328            BoneWeight {
329                bone_index: 0,
330                weight: 0.1,
331            },
332            BoneWeight {
333                bone_index: 1,
334                weight: 0.5,
335            },
336            BoneWeight {
337                bone_index: 2,
338                weight: 0.3,
339            },
340            BoneWeight {
341                bone_index: 3,
342                weight: 0.1,
343            },
344        ];
345        let top = top_n_weights(&weights, 2);
346        assert_eq!(top.len(), 2);
347        assert!((top[0].weight - 0.5).abs() < f32::EPSILON);
348        assert!((top[1].weight - 0.3).abs() < f32::EPSILON);
349    }
350
351    #[test]
352    fn test_weight_map_to_csv() {
353        let mut buf = new_weight_map_buffer(2, 2);
354        set_weight_pixel(&mut buf, 0, 0, 1.0);
355        let csv = weight_map_to_csv(&buf);
356        assert!(csv.starts_with("x,y,weight\n"));
357        assert!(csv.contains("0,0,1.000000"));
358    }
359
360    #[test]
361    fn test_blend_weight_maps() {
362        let mut a = new_weight_map_buffer(2, 2);
363        let mut b = new_weight_map_buffer(2, 2);
364        for px in a.pixels.iter_mut() {
365            *px = 0.0;
366        }
367        for px in b.pixels.iter_mut() {
368            *px = 1.0;
369        }
370        let blended = blend_weight_maps(&a, &b, 0.5);
371        for &px in &blended.pixels {
372            assert!((px - 0.5).abs() < f32::EPSILON);
373        }
374    }
375
376    #[test]
377    fn test_weight_map_pixel_count() {
378        let buf = new_weight_map_buffer(8, 4);
379        assert_eq!(weight_map_pixel_count(&buf), 32);
380    }
381
382    #[test]
383    fn test_weight_map_stats() {
384        let mut buf = new_weight_map_buffer(4, 4);
385        set_weight_pixel(&mut buf, 0, 0, 0.2);
386        set_weight_pixel(&mut buf, 1, 0, 0.8);
387        let stats = weight_map_stats(&buf);
388        assert!((stats.min - 0.0).abs() < f32::EPSILON);
389        assert!((stats.max - 0.8).abs() < f32::EPSILON);
390        assert!(stats.avg > 0.0);
391    }
392
393    #[test]
394    fn test_weight_map_stats_empty() {
395        let buf = WeightMapBuffer {
396            pixels: vec![],
397            width: 0,
398            height: 0,
399        };
400        let stats = weight_map_stats(&buf);
401        assert!((stats.min - 0.0).abs() < f32::EPSILON);
402        assert!((stats.max - 0.0).abs() < f32::EPSILON);
403    }
404
405    #[test]
406    fn test_clear_weight_map() {
407        let mut buf = new_weight_map_buffer(4, 4);
408        set_weight_pixel(&mut buf, 1, 1, 0.9);
409        clear_weight_map(&mut buf);
410        for &px in &buf.pixels {
411            assert!((px - 0.0).abs() < f32::EPSILON);
412        }
413    }
414}
oxihuman_export/weight_map_export.rs

oxihuman_export/
weight_map_export.rs
