Initial commit

crates/specialized/rtx-neuro-anatomy/src/annotation.rs

@@ -0,0 +1,430 @@
+//! FreeSurfer annotation file reading
+//!
+//! Reads .annot files containing parcellation labels and color tables.
+//! Common atlases include:
+//! - aparc (Desikan-Killiany atlas)
+//! - aparc.a2009s (Destrieux atlas)
+//! - aparc.DKTatlas (DKT atlas)
+
+use std::collections::HashMap;
+use std::fs::File;
+use std::io::{BufReader, Read};
+use std::path::Path;
+
+use byteorder::{LittleEndian, ReadBytesExt};
+use serde::{Deserialize, Serialize};
+
+use crate::error::{AnatomyError, Result};
+use crate::surface::Hemisphere;
+
+/// A color entry in the annotation color table
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AnnotationEntry {
+    /// Region name
+    pub name: String,
+    /// Red component (0-255)
+    pub r: u8,
+    /// Green component (0-255)
+    pub g: u8,
+    /// Blue component (0-255)
+    pub b: u8,
+    /// Alpha component (0-255, usually 0 = opaque)
+    pub a: u8,
+}
+
+impl AnnotationEntry {
+    /// Get RGB color as [0-1] floats
+    pub fn color_rgb(&self) -> [f32; 3] {
+        [
+            self.r as f32 / 255.0,
+            self.g as f32 / 255.0,
+            self.b as f32 / 255.0,
+        ]
+    }
+
+    /// Get RGBA color as [0-1] floats
+    pub fn color_rgba(&self) -> [f32; 4] {
+        [
+            self.r as f32 / 255.0,
+            self.g as f32 / 255.0,
+            self.b as f32 / 255.0,
+            1.0 - (self.a as f32 / 255.0), // FreeSurfer uses 0=opaque
+        ]
+    }
+}
+
+/// Color table mapping label IDs to entries
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ColorTable {
+    /// Entries indexed by label ID
+    pub entries: HashMap<i32, AnnotationEntry>,
+}
+
+impl ColorTable {
+    /// Get entry by label ID
+    pub fn get(&self, label: i32) -> Option<&AnnotationEntry> {
+        self.entries.get(&label)
+    }
+
+    /// Get all region names
+    pub fn region_names(&self) -> Vec<&str> {
+        self.entries.values().map(|e| e.name.as_str()).collect()
+    }
+
+    /// Number of regions
+    pub fn n_regions(&self) -> usize {
+        self.entries.len()
+    }
+}
+
+/// Parcellation annotation for a brain surface
+#[derive(Debug, Clone)]
+pub struct Annotation {
+    /// Per-vertex label IDs (using the packed RGBA format)
+    pub labels: Vec<i32>,
+    /// Color table with region info
+    pub color_table: ColorTable,
+    /// Atlas name (e.g., "aparc", "aparc.a2009s")
+    pub atlas: String,
+    /// Hemisphere
+    pub hemisphere: Hemisphere,
+}
+
+impl Annotation {
+    /// Number of vertices
+    pub fn n_vertices(&self) -> usize {
+        self.labels.len()
+    }
+
+    /// Get the region name for a vertex
+    pub fn vertex_region(&self, vertex: usize) -> Option<&str> {
+        let label = self.labels.get(vertex)?;
+        self.color_table.get(*label).map(|e| e.name.as_str())
+    }
+
+    /// Get the color for a vertex
+    pub fn vertex_color(&self, vertex: usize) -> Option<[f32; 3]> {
+        let label = self.labels.get(vertex)?;
+        self.color_table.get(*label).map(AnnotationEntry::color_rgb)
+    }
+
+    /// Get all labels as region indices (for lookup in color table)
+    pub fn labels_as_indices(&self) -> Vec<i32> {
+        self.labels.clone()
+    }
+
+    /// Get vertex colors as flat array for visualization
+    pub fn vertex_colors(&self) -> Vec<[f32; 3]> {
+        self.labels
+            .iter()
+            .map(|label| {
+                self.color_table
+                    .get(*label)
+                    .map_or([0.5, 0.5, 0.5], AnnotationEntry::color_rgb) // Gray for unknown
+            })
+            .collect()
+    }
+
+    /// Get unique labels present in this annotation
+    pub fn unique_labels(&self) -> Vec<i32> {
+        let mut unique: Vec<i32> = self.labels.clone();
+        unique.sort_unstable();
+        unique.dedup();
+        unique
+    }
+}
+
+/// Read a FreeSurfer annotation file
+///
+/// # Arguments
+///
+/// * `path` - Path to the .annot file
+/// * `atlas` - Atlas name (e.g., "aparc")
+/// * `hemisphere` - Hemisphere of the annotation
+///
+/// # Format
+///
+/// The .annot file format (all little-endian):
+/// 1. n_vertices (int32)
+/// 2. For each vertex: vertex_index (int32), label (int32)
+/// 3. has_colortable (int32, 0 or 1)
+/// 4. If has_colortable: color table data
+///
+/// # Example
+///
+/// ```ignore
+/// use rtx_neuro_anatomy::annotation::{read_annotation, Hemisphere};
+///
+/// let annot = read_annotation(
+///     "/path/to/subject/label/lh.aparc.annot",
+///     "aparc",
+///     Hemisphere::Left,
+/// )?;
+/// println!("Loaded {} vertices", annot.n_vertices());
+/// ```
+pub fn read_annotation(
+    path: impl AsRef<Path>,
+    atlas: &str,
+    hemisphere: Hemisphere,
+) -> Result<Annotation> {
+    let path = path.as_ref();
+    let file = File::open(path).map_err(|_| AnatomyError::AnnotationNotFound {
+        path: path.to_path_buf(),
+    })?;
+    let mut reader = BufReader::new(file);
+
+    // Read number of vertices
+    let n_vertices = reader.read_i32::<LittleEndian>()? as usize;
+
+    // Read vertex labels
+    let mut labels = vec![0i32; n_vertices];
+    let mut max_vertex = 0usize;
+
+    for _ in 0..n_vertices {
+        let vertex_idx = reader.read_i32::<LittleEndian>()? as usize;
+        let label = reader.read_i32::<LittleEndian>()?;
+
+        if vertex_idx < n_vertices {
+            labels[vertex_idx] = label;
+            max_vertex = max_vertex.max(vertex_idx);
+        }
+    }
+
+    // Check if there's a color table
+    let has_colortable = reader.read_i32::<LittleEndian>()? != 0;
+
+    let color_table = if has_colortable {
+        read_colortable(&mut reader)?
+    } else {
+        ColorTable::default()
+    };
+
+    Ok(Annotation {
+        labels,
+        color_table,
+        atlas: atlas.to_string(),
+        hemisphere,
+    })
+}
+
+/// Read the color table from an annotation file
+fn read_colortable(reader: &mut BufReader<File>) -> Result<ColorTable> {
+    // Read number of entries
+    let num_entries = reader.read_i32::<LittleEndian>()? as usize;
+
+    if num_entries == 0 {
+        return Ok(ColorTable::default());
+    }
+
+    // Check for new format (version > 0)
+    let version_or_len = reader.read_i32::<LittleEndian>()?;
+
+    let entries = if version_or_len > 0 {
+        // New format: version_or_len is the length of the original filename
+        read_colortable_new(reader, num_entries, version_or_len)?
+    } else {
+        // Old format (not commonly used)
+        read_colortable_old(reader, num_entries)?
+    };
+
+    Ok(ColorTable { entries })
+}
+
+/// Read new format color table
+fn read_colortable_new(
+    reader: &mut BufReader<File>,
+    num_entries: usize,
+    filename_len: i32,
+) -> Result<HashMap<i32, AnnotationEntry>> {
+    // Skip the original filename
+    let mut filename_buf = vec![0u8; filename_len as usize];
+    reader.read_exact(&mut filename_buf)?;
+
+    // Read number of entries again (redundant in new format)
+    let _num_entries_check = reader.read_i32::<LittleEndian>()?;
+
+    let mut entries = HashMap::new();
+
+    for _ in 0..num_entries {
+        // Read structure number (label ID)
+        let struct_id = reader.read_i32::<LittleEndian>()?;
+
+        // Read name length and name
+        let name_len = reader.read_i32::<LittleEndian>()? as usize;
+        let mut name_buf = vec![0u8; name_len];
+        reader.read_exact(&mut name_buf)?;
+
+        // Remove null terminator if present
+        let name = String::from_utf8_lossy(&name_buf)
+            .trim_end_matches('\0')
+            .to_string();
+
+        // Read RGBA values
+        let r = reader.read_i32::<LittleEndian>()? as u8;
+        let g = reader.read_i32::<LittleEndian>()? as u8;
+        let b = reader.read_i32::<LittleEndian>()? as u8;
+        let a = reader.read_i32::<LittleEndian>()? as u8;
+
+        // Compute the packed label value (same as what's stored per-vertex)
+        let label = pack_rgba(r, g, b, a);
+
+        entries.insert(label, AnnotationEntry { name, r, g, b, a });
+
+        // Also insert by struct_id for direct lookup
+        if struct_id != label {
+            entries.insert(
+                struct_id,
+                AnnotationEntry {
+                    name: entries[&label].name.clone(),
+                    r,
+                    g,
+                    b,
+                    a,
+                },
+            );
+        }
+    }
+
+    Ok(entries)
+}
+
+/// Read old format color table (rarely used)
+fn read_colortable_old(
+    reader: &mut BufReader<File>,
+    num_entries: usize,
+) -> Result<HashMap<i32, AnnotationEntry>> {
+    let mut entries = HashMap::new();
+
+    // Skip filename length that was already read as version
+    let filename_len = reader.read_i32::<LittleEndian>()? as usize;
+    let mut filename_buf = vec![0u8; filename_len];
+    reader.read_exact(&mut filename_buf)?;
+
+    for _ in 0..num_entries {
+        // Read name length and name
+        let name_len = reader.read_i32::<LittleEndian>()? as usize;
+        let mut name_buf = vec![0u8; name_len];
+        reader.read_exact(&mut name_buf)?;
+
+        let name = String::from_utf8_lossy(&name_buf)
+            .trim_end_matches('\0')
+            .to_string();
+
+        // Read RGBA values
+        let r = reader.read_i32::<LittleEndian>()? as u8;
+        let g = reader.read_i32::<LittleEndian>()? as u8;
+        let b = reader.read_i32::<LittleEndian>()? as u8;
+        let a = reader.read_i32::<LittleEndian>()? as u8;
+
+        let label = pack_rgba(r, g, b, a);
+
+        entries.insert(label, AnnotationEntry { name, r, g, b, a });
+    }
+
+    Ok(entries)
+}
+
+/// Pack RGBA values into a single i32 label value
+///
+/// FreeSurfer packs colors as: R + G*256 + B*65536 + A*16777216
+fn pack_rgba(r: u8, g: u8, b: u8, a: u8) -> i32 {
+    (r as i32) + (g as i32) * 256 + (b as i32) * 65536 + (a as i32) * 16777216
+}
+
+/// Unpack i32 label value into RGBA components
+pub fn unpack_rgba(label: i32) -> (u8, u8, u8, u8) {
+    let r = (label & 0xFF) as u8;
+    let g = ((label >> 8) & 0xFF) as u8;
+    let b = ((label >> 16) & 0xFF) as u8;
+    let a = ((label >> 24) & 0xFF) as u8;
+    (r, g, b, a)
+}
+
+/// Get the annotation filename for an atlas
+pub fn annotation_filename(hemisphere: &Hemisphere, atlas: &str) -> String {
+    format!("{}.{}.annot", hemisphere.prefix(), atlas)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pack_unpack_rgba() {
+        let r = 128u8;
+        let g = 64u8;
+        let b = 32u8;
+        let a = 0u8;
+
+        let packed = pack_rgba(r, g, b, a);
+        let (ur, ug, ub, ua) = unpack_rgba(packed);
+
+        assert_eq!(r, ur);
+        assert_eq!(g, ug);
+        assert_eq!(b, ub);
+        assert_eq!(a, ua);
+    }
+
+    #[test]
+    fn test_annotation_entry_colors() {
+        let entry = AnnotationEntry {
+            name: "test".to_string(),
+            r: 255,
+            g: 128,
+            b: 0,
+            a: 0,
+        };
+
+        let rgb = entry.color_rgb();
+        assert!((rgb[0] - 1.0).abs() < 1e-6);
+        assert!((rgb[1] - 0.502).abs() < 0.01);
+        assert!((rgb[2] - 0.0).abs() < 1e-6);
+    }
+
+    #[test]
+    fn test_annotation_filename() {
+        assert_eq!(
+            annotation_filename(&Hemisphere::Left, "aparc"),
+            "lh.aparc.annot"
+        );
+        assert_eq!(
+            annotation_filename(&Hemisphere::Right, "aparc.a2009s"),
+            "rh.aparc.a2009s.annot"
+        );
+    }
+
+    #[test]
+    fn test_color_table_operations() {
+        let mut entries = HashMap::new();
+        entries.insert(
+            1,
+            AnnotationEntry {
+                name: "region1".to_string(),
+                r: 255,
+                g: 0,
+                b: 0,
+                a: 0,
+            },
+        );
+        entries.insert(
+            2,
+            AnnotationEntry {
+                name: "region2".to_string(),
+                r: 0,
+                g: 255,
+                b: 0,
+                a: 0,
+            },
+        );
+
+        let table = ColorTable { entries };
+
+        assert_eq!(table.n_regions(), 2);
+        assert!(table.get(1).is_some());
+        assert!(table.get(3).is_none());
+
+        let names = table.region_names();
+        assert!(names.contains(&"region1"));
+        assert!(names.contains(&"region2"));
+    }
+}