scan_core/aruco/

dictionary.rs

use std::collections::HashMap;

pub struct Dictionary {
    pub n_bits: usize,
    pub tau: u32,
    pub codes: HashMap<u64, usize>, // code -> id
    pub code_list: Vec<u64>,
}

impl Dictionary {
    pub fn new(name: &str) -> Option<Self> {
        match name {
            "ARUCO_4X4_50" => Some(Self::aruco_4x4_50()),
            "ARUCO_5X5" => Some(Self::aruco_5x5()),
            "ARUCO_6X6" => Some(Self::aruco_6x6()),
            _ => None,
        }
    }

    fn from_codes(n_bits: usize, tau_override: Option<u32>, codes_slice: &[u64]) -> Self {
        let mut code_list = Vec::with_capacity(codes_slice.len());
        let mut codes = HashMap::new();

        for (i, &val) in codes_slice.iter().enumerate() {
            code_list.push(val);
            codes.insert(val, i);
        }

        let tau = if let Some(t) = tau_override {
            t
        } else {
            // Calculate tau (min distance between any two codes)
            let mut t = u32::MAX;
            // Scan subset if too large? 5x5 has 1024 codes -> 1M comparisons. 
            // In JS they do full scan. 1M ops in Rust is fast.
            for i in 0..code_list.len() {
                for j in (i + 1)..code_list.len() {
                    let d = (code_list[i] ^ code_list[j]).count_ones();
                    if d < t { t = d; }
                }
            }
            t
        };

        Self { n_bits, tau, codes, code_list }
    }

    fn aruco_4x4_50() -> Self {
        let n_bits = 16;
        // manually converted to u64 for reuse with from_codes if needed, 
        // but keeping original byte definition for now to avoid huge diff, 
        // just converting it before call
        let raw_codes: Vec<&[u8]> = vec![
            &[181,50],&[15,154],&[51,45],&[153,70],&[84,158],&[121,205],&[158,46],&[196,242],&[254,218],&[207,86],
            &[249,145],&[17,167],&[14,183],&[42,15],&[36,177],&[38,62],&[70,101],&[102,0],&[108,94],&[118,175],
            &[134,139],&[176,43],&[204,213],&[221,130],&[254,71],&[148,113],&[172,228],&[165,84],&[33,35],&[52,111],
            &[68,21],&[87,178],&[158,207],&[240,203],&[8,174],&[9,41],&[24,117],&[4,255],&[13,246],&[28,90],
            &[23,24],&[42,40],&[50,140],&[56,178],&[36,232],&[46,235],&[45,63],&[75,100],&[80,46],&[80,19]
        ];

        let mut u64_codes = Vec::new();
        for bytes in raw_codes {
             u64_codes.push(Self::bytes_to_u64(bytes));
        }
        
        Self::from_codes(n_bits, None, &u64_codes)
    }

    fn aruco_5x5() -> Self {
        // ARUCO dictionary (5x5, 25 bits)
        // Tau is 3 explicitly in JS definition
        use crate::aruco::dict_data::ARUCO_5X5_CODES;
        Self::from_codes(25, Some(3), ARUCO_5X5_CODES)
    }

    fn aruco_6x6() -> Self {
        // ARUCO_MIP_36h12 (6x6, 36 bits)
        // Tau is 12 explicitly in JS definition
        use crate::aruco::dict_data::ARUCO_6X6_CODES;
        Self::from_codes(36, Some(12), ARUCO_6X6_CODES)
    }

    fn bytes_to_u64(bytes: &[u8]) -> u64 {
        let mut val = 0u64;
        for &b in bytes {
            val = (val << 8) | (b as u64);
        }
        val
    }

    pub fn find(&self, bits: u64, max_hamming: u32) -> Option<(usize, u32)> {
        // Direct lookup
        if let Some(&id) = self.codes.get(&bits) {
            return Some((id, 0));
        }

        // Find nearest within max_hamming
        let mut min_dist = u32::MAX;
        let mut min_id = 0;

        for (id, &code) in self.code_list.iter().enumerate() {
            let dist = (bits ^ code).count_ones();
            if dist < min_dist {
                min_dist = dist;
                min_id = id;
            }
        }

        if min_dist <= max_hamming {
            Some((min_id, min_dist))
        } else {
            None
        }
    }

    /// Recognize a marker from a warped+thresholded image.
    /// Port of AR.Detector.prototype.getMarker from JS.
    ///
    /// Key differences from previous Rust version:
    /// 1. Validates that border cells are BLACK (non-zero count < minZero)
    /// 2. Reads bits using countNonZero over the full cell area (not 3x3 sample)
    /// 3. Tries all 4 rotations to find best match
    pub fn recognize_marker(&self, img: &image::GrayImage, _mark_size: usize, max_hamming: u32, corners: &[crate::Point], mut debug_log: Option<&mut Vec<String>>) -> Option<crate::aruco::Marker> {
        let grid_size = (self.n_bits as f64).sqrt() as usize + 2; // e.g. 6 for 4x4
        let cell_w_f = (img.width() as f32) / (grid_size as f32);
        let cell_h_f = (img.height() as f32) / (grid_size as f32);
        let data = img.as_raw();
        let img_w = img.width() as usize;

        // --- Step 1: Validate border cells are black ---
        for i in 0..grid_size {
            let inc = if i == 0 || i == grid_size - 1 { 1 } else { grid_size - 1 };

            let mut j = 0;
            while j < grid_size {
                let x1 = (j as f32 * cell_w_f + cell_w_f * 0.2) as usize;
                let y1 = (i as f32 * cell_h_f + cell_h_f * 0.2) as usize;
                let x2 = ((j as f32 + 1.0) * cell_w_f - cell_w_f * 0.2) as usize;
                let y2 = ((i as f32 + 1.0) * cell_h_f - cell_h_f * 0.2) as usize;
                
                let w = (x2 - x1).max(1);
                let h = (y2 - y1).max(1);
                let min_zero = (w * h) >> 1; // half the sampled area

                let nz = crate::cv::geometry::count_non_zero(data, img_w, x1, y1, w, h);
                if nz > min_zero {
                    if let Some(ref mut logs) = debug_log {
                         logs.push(format!("DEBUG: Border check failed at {},{}. NZ={} > {}", j, i, nz, min_zero));
                    }
                    return None; // Border cell is not black
                }
                j += inc;
            }
        }

        // --- Step 2: Read inner bits using centered sampling ---
        let inner_size = grid_size - 2;
        let mut bits: Vec<Vec<u8>> = Vec::new();

        for i in 0..inner_size {
            let mut row = Vec::new();
            for j in 0..inner_size {
                // Use inner cells (skip first border)
                let x1 = ((j + 1) as f32 * cell_w_f + cell_w_f * 0.2) as usize;
                let y1 = ((i + 1) as f32 * cell_h_f + cell_h_f * 0.2) as usize;
                let x2 = ((j + 2) as f32 * cell_w_f - cell_w_f * 0.2) as usize;
                let y2 = ((i + 2) as f32 * cell_h_f - cell_h_f * 0.2) as usize;
                
                let w = (x2 - x1).max(1);
                let h = (y2 - y1).max(1);
                let min_zero = (w * h) >> 1;

                let nz = crate::cv::geometry::count_non_zero(data, img_w, x1, y1, w, h);
                row.push(if nz > min_zero { 1 } else { 0 });
            }
            bits.push(row);
        }

        // --- Step 3: Try all 4 rotations for best match ---
        let mut current_bits = bits;
        let mut best_found: Option<(usize, u32)> = None;
        let mut best_rot = 0;
        let mut debug_codes_tried = Vec::new();

        for rot in 0..4 {
            // Flatten to u64
            let mut val = 0u64;
            for row in &current_bits {
                for &b in row {
                    val = (val << 1) | (b as u64);
                }
            }
            
            if debug_log.is_some() { debug_codes_tried.push(format!("0x{:x}", val)); }

            if let Some((id, dist)) = self.find(val, max_hamming) {
                if best_found.is_none() || dist < best_found.unwrap().1 {
                    best_found = Some((id, dist));
                    best_rot = rot;
                    if dist == 0 { break; } // Perfect match
                }
            }

            // Rotate 90° CW for next iteration
            current_bits = rotate_grid_2d(&current_bits);
        }

        if let Some((id, dist)) = best_found {
            let shift = (4 - best_rot) % 4;
            let mut new_corners = vec![crate::Point { x: 0.0, y: 0.0 }; 4];
            for i in 0..4 {
                new_corners[i] = corners[(i + shift) % 4];
            }
            Some(crate::aruco::Marker::new(id as u32, new_corners, dist))
        } else {
            if let Some(ref mut logs) = debug_log {
                logs.push(format!("DEBUG: Recognition failed. Tried codes: {:?}. MaxHamming={}", debug_codes_tried, max_hamming));
            }
            None
        }
    }

    /// Genera una imagen (GrayImage) de un marcador ArUco dado su ID.
    /// La imagen tendrá un borde negro de 1 celda y el patrón NxN interno.
    pub fn generate_marker_image(&self, id: usize, cell_size: usize) -> Option<image::GrayImage> {
        if let Some(&code) = self.code_list.get(id) {
            let n = (self.n_bits as f64).sqrt() as usize;
            let grid_size = n + 2;
            let img_size = grid_size * cell_size;
            
            // GrayImage::new inicializa con 0 (negro), lo cual nos sirve para el borde.
            let mut img = image::GrayImage::new(img_size as u32, img_size as u32);
            
            // Dibujar bits internos blancos
            for row in 0..n {
                for col in 0..n {
                    let bit_idx = (n * n) - 1 - (row * n + col);
                    let bit = (code >> bit_idx) & 1;
                    if bit == 1 {
                        // Celda blanca
                        for py in 0..cell_size {
                            for px in 0..cell_size {
                                img.put_pixel(
                                    ((col + 1) * cell_size + px) as u32,
                                    ((row + 1) * cell_size + py) as u32,
                                    image::Luma([255])
                                );
                            }
                        }
                    }
                }
            }
            Some(img)
        } else {
            None
        }
    }

    /// Genera una cadena con el contenido de un archivo SVG para el marcador dado su ID.
    /// El SVG es vectorial y escala perfectamente sin pérdida de calidad.
    pub fn generate_marker_svg(&self, id: usize) -> Option<String> {
        if let Some(&code) = self.code_list.get(id) {
            let n = (self.n_bits as f64).sqrt() as usize;
            let grid_size = n + 2;
            
            let mut svg = format!(
                r#"<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<svg width="{0}" height="{0}" viewBox="0 0 {0} {0}" xmlns="http://www.w3.org/2000/svg" shape-rendering="crispEdges">
  <rect x="0" y="0" width="{0}" height="{0}" fill="black" />"#,
                grid_size
            );

            // Dibujar bits internos blancos
            for row in 0..n {
                for col in 0..n {
                    let bit_idx = (n * n) - 1 - (row * n + col);
                    let bit = (code >> bit_idx) & 1;
                    if bit == 1 {
                        svg.push_str(&format!(
                            r#"
  <rect x="{}" y="{}" width="1" height="1" fill="white" />"#,
                            col + 1,
                            row + 1
                        ));
                    }
                }
            }
            
            svg.push_str("\n</svg>");
            Some(svg)
        } else {
            None
        }
    }
}

/// Rotate a 2D grid 90° clockwise.
/// Port of AR.Detector.prototype.rotate from JS.
fn rotate_grid_2d(src: &[Vec<u8>]) -> Vec<Vec<u8>> {
    let len = src.len();
    let mut dst = vec![vec![0u8; len]; len];
    for i in 0..len {
        for j in 0..src[i].len() {
            dst[i][j] = src[src[i].len() - j - 1][i];
        }
    }
    dst
}