phasm_core/stego/armor/
selection.rs

1// Copyright (c) 2026 Christoph Gaffga
2// SPDX-License-Identifier: GPL-3.0-only
3// https://github.com/cgaffga/phasmcore
4
5//! Frequency-based coefficient selection for Armor embedding.
6//!
7//! Selects DCT coefficient positions suitable for robust STDM embedding
8//! based on zigzag frequency index. Low-to-mid frequency AC positions
9//! (zigzag 1..=MAX_ARMOR_ZIGZAG) are selected; DC and high-frequency
10//! positions are excluded.
11
12use crate::codec::jpeg::dct::DctGrid;
13use crate::codec::jpeg::zigzag::NATURAL_TO_ZIGZAG;
14use crate::stego::cost::CostMap;
15use super::embedding::MAX_ARMOR_ZIGZAG;
16
17/// Cost assigned to stable positions (low, uniform cost for permutation compatibility).
18const STABLE_COST: f32 = 1.0;
19
20/// Compute a stability map for the given DCT grid.
21///
22/// Includes all AC coefficient positions with zigzag index 1..=MAX_ARMOR_ZIGZAG.
23/// Returns a `CostMap` where selected positions have `STABLE_COST` and
24/// DC/high-frequency positions have `WET_COST`.
25pub fn compute_stability_map(grid: &DctGrid, _qt: &crate::codec::jpeg::dct::QuantTable) -> CostMap {
26    compute_stability_map_freq_only(grid)
27}
28
29/// Frequency-only selection: include all zigzag 1..=MAX_ARMOR_ZIGZAG positions.
30fn compute_stability_map_freq_only(grid: &DctGrid) -> CostMap {
31    let bw = grid.blocks_wide();
32    let bt = grid.blocks_tall();
33    let mut cost_map = CostMap::new(bw, bt);
34
35    for br in 0..bt {
36        for bc in 0..bw {
37            for i in 0..8 {
38                for j in 0..8 {
39                    if i == 0 && j == 0 { continue; }
40                    let freq_idx = i * 8 + j;
41                    let zz = NATURAL_TO_ZIGZAG[freq_idx];
42                    if (1..=MAX_ARMOR_ZIGZAG).contains(&zz) {
43                        cost_map.set(br, bc, i, j, STABLE_COST);
44                    }
45                }
46            }
47        }
48    }
49
50    cost_map
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use crate::codec::jpeg::dct::{DctGrid, QuantTable};
57    use crate::stego::cost::WET_COST;
58
59    #[test]
60    fn stability_map_dc_is_wet_lowfreq_ac_is_stable() {
61        // Small grid (2x2 = 4 blocks).
62        let mut grid = DctGrid::new(2, 2);
63        let qt = QuantTable::new([8; 64]);
64
65        grid.set(0, 0, 0, 0, 100);
66
67        let cost_map = compute_stability_map(&grid, &qt);
68        for br in 0..2 {
69            for bc in 0..2 {
70                // DC should be WET
71                assert_eq!(
72                    cost_map.get(br, bc, 0, 0),
73                    WET_COST,
74                    "DC ({br},{bc},0,0) should be WET"
75                );
76                // Low-freq AC (0,1) = zigzag 1 -> STABLE (frequency-only mode)
77                assert_eq!(
78                    cost_map.get(br, bc, 0, 1),
79                    STABLE_COST,
80                    "AC ({br},{bc},0,1) should be STABLE (zigzag 1)"
81                );
82                // (1,0) = zigzag 2 -> STABLE
83                assert_eq!(
84                    cost_map.get(br, bc, 1, 0),
85                    STABLE_COST,
86                    "AC ({br},{bc},1,0) should be STABLE (zigzag 2)"
87                );
88            }
89        }
90    }
91
92    #[test]
93    fn stability_map_excludes_high_freq() {
94        // Small grid (1x1 = 1 block).
95        let mut grid = DctGrid::new(1, 1);
96        let qt = QuantTable::new([8; 64]);
97
98        grid.set(0, 0, 1, 0, 50);
99        grid.set(0, 0, 0, 1, 1);
100
101        let cost_map = compute_stability_map(&grid, &qt);
102
103        // Low-freq AC positions should be STABLE (frequency-only selection)
104        assert_eq!(cost_map.get(0, 0, 1, 0), STABLE_COST); // zigzag 2
105        assert_eq!(cost_map.get(0, 0, 0, 1), STABLE_COST); // zigzag 1
106        // Mid-freq AC should also be STABLE (zigzag 11, 12 within 1..=15)
107        assert_eq!(cost_map.get(0, 0, 3, 1), STABLE_COST, "zigzag 11 should be STABLE");
108        assert_eq!(cost_map.get(0, 0, 2, 2), STABLE_COST, "zigzag 12 should be STABLE");
109        // High-freq AC (7,7) = zigzag 63 -> WET (excluded, beyond MAX_ARMOR_ZIGZAG)
110        assert_eq!(cost_map.get(0, 0, 7, 7), WET_COST);
111        // DC should be WET
112        assert_eq!(cost_map.get(0, 0, 0, 0), WET_COST);
113    }
114}
phasm_core/stego/armor/selection.rs

phasm_core/stego/armor/
selection.rs
