Skip to main content

openkspace_io/fastmri/
reader.rs

1//! FastMRI HDF5 reader.
2//!
3//! FastMRI files (from fastmri.med.nyu.edu) store a pre-assembled k-space
4//! tensor and a ground-truth RSS reconstruction directly in the HDF5 root:
5//!
6//! ```text
7//! /kspace              complex64  [slices, coils, ky, kx]
8//! /reconstruction_rss  float32    [slices, recon_y, recon_x]   (ground truth)
9//! /ismrmrd_header      str        ISMRMRD-compatible XML
10//!
11//! File-level attrs:
12//!   acquisition   str   e.g. "CORPDFS_FBK", "AXT2"
13//!   patient_id    str   SHA-256 anonymised subject ID
14//!   max           str   string-encoded f32 max value
15//!   norm          str   string-encoded f32 norm
16//! ```
17//!
18//! The `ismrmrd_header` XML is structurally identical to that in an ISMRMRD
19//! file, so [`IsmrmrdHeader`] is reused here for metadata.
20//!
21//! The k-space tensor axis order `[slices, coils, ky, kx]` is converted to
22//! the pipeline-standard `[coils, kz/slices, ky, kx]` on load, matching the
23//! shape returned by [`IsmrmrdFile::read_kspace`].
24
25use crate::error::{IoError, IoResult};
26use crate::ismrmrd::header::IsmrmrdHeader;
27
28use hdf5_metno::{File, H5Type};
29use ndarray::{Array2, Array3, Array4};
30use num_complex::Complex32;
31use std::path::Path;
32use tracing::info;
33
34/// HDF5 complex64 on-disk layout: two contiguous f32s (real, imag).
35///
36/// We derive `H5Type` so hdf5-metno maps the compound dtype by field name,
37/// then convert to `num_complex::Complex32` ourselves. This avoids the
38/// ndarray version mismatch that arises when asking hdf5-metno to return
39/// `ndarray::Array<Complex32, _>` directly (hdf5-metno 0.9 bundles its own
40/// ndarray 0.16 while the workspace uses 0.15).
41///
42/// FastMRI HDF5 files store complex64 with compound members named `r` and
43/// `i` (not `re`/`im`), so field names must match exactly for HDF5 name-
44/// based type conversion to succeed.
45#[derive(Debug, Clone, Copy, H5Type)]
46#[repr(C)]
47struct Cf32 {
48    r: f32,
49    i: f32,
50}
51
52impl From<Cf32> for Complex32 {
53    #[inline]
54    fn from(c: Cf32) -> Self {
55        Complex32::new(c.r, c.i)
56    }
57}
58
59// ── Helpers to read flat Vecs from HDF5 ──────────────────────────────────────
60
61fn read_complex_flat(file: &File, path: &str) -> IoResult<(Vec<Complex32>, Vec<usize>)> {
62    let ds = file.dataset(path)?;
63    let shape = ds.shape();
64    let raw: Vec<Cf32> = ds.read_raw()?;
65    let data: Vec<Complex32> = raw.into_iter().map(Into::into).collect();
66    Ok((data, shape))
67}
68
69fn read_f32_flat(file: &File, path: &str) -> IoResult<(Vec<f32>, Vec<usize>)> {
70    let ds = file.dataset(path)?;
71    let shape = ds.shape();
72    let data: Vec<f32> = ds.read_raw()?;
73    Ok((data, shape))
74}
75
76// ── Public types ──────────────────────────────────────────────────────────────
77
78/// Metadata parsed from a FastMRI file.
79#[non_exhaustive]
80#[derive(Debug, Clone)]
81pub struct FastmriMeta {
82    /// ISMRMRD-compatible XML header.
83    pub header: IsmrmrdHeader,
84    /// Acquisition type label, e.g. `"CORPDFS_FBK"`, `"AXT2"`.
85    pub acquisition: String,
86    /// Anonymised patient identifier (SHA-256 hex string).
87    pub patient_id: String,
88    /// Number of slices in this file.
89    pub n_slices: usize,
90    /// Number of coils.
91    pub n_coils: usize,
92    /// Encoded ky dimension.
93    pub n_ky: usize,
94    /// Encoded kx dimension.
95    pub n_kx: usize,
96    /// Recon y dimension (from `reconstruction_rss`).
97    pub recon_y: usize,
98    /// Recon x dimension (from `reconstruction_rss`).
99    pub recon_x: usize,
100}
101
102/// A handle to an opened FastMRI file.
103pub struct FastmriFile {
104    file: File,
105    pub meta: FastmriMeta,
106}
107
108impl FastmriFile {
109    /// Open the HDF5 file, parse the header, and validate the tensor shapes.
110    pub fn open<P: AsRef<Path>>(path: P) -> IoResult<Self> {
111        let path = path.as_ref();
112        let file = File::open(path)?;
113
114        // ── XML header ────────────────────────────────────────────────────
115        let xml_ds = file.dataset("ismrmrd_header")?;
116        let xml_str = xml_ds
117            .read_scalar::<hdf5_metno::types::VarLenUnicode>()
118            .map(|s| s.as_str().to_string())
119            .or_else(|_| {
120                xml_ds
121                    .read_scalar::<hdf5_metno::types::VarLenAscii>()
122                    .map(|s| s.as_str().to_string())
123            })
124            .map_err(|_| IoError::MissingField("ismrmrd_header"))?;
125
126        let header = IsmrmrdHeader::parse(&xml_str)?;
127
128        // ── k-space shape ─────────────────────────────────────────────────
129        let kshape = file.dataset("kspace")?.shape();
130        if kshape.len() != 4 {
131            return Err(IoError::Unsupported(format!(
132                "kspace has {} dims (expected 4: [slices, coils, ky, kx])",
133                kshape.len()
134            )));
135        }
136        let (n_slices, n_coils, n_ky, n_kx) = (kshape[0], kshape[1], kshape[2], kshape[3]);
137
138        // ── RSS reconstruction shape ──────────────────────────────────────
139        let rshape = file.dataset("reconstruction_rss")?.shape();
140        if rshape.len() != 3 {
141            return Err(IoError::Unsupported(format!(
142                "reconstruction_rss has {} dims (expected 3: [slices, y, x])",
143                rshape.len()
144            )));
145        }
146        if rshape[0] != n_slices {
147            return Err(IoError::Inconsistent(format!(
148                "kspace has {} slices but reconstruction_rss has {}",
149                n_slices, rshape[0]
150            )));
151        }
152        let (recon_y, recon_x) = (rshape[1], rshape[2]);
153
154        // ── File-level attributes ─────────────────────────────────────────
155        let attr_str = |key: &str| -> String {
156            file.attr(key)
157                .and_then(|a| a.read_scalar::<hdf5_metno::types::VarLenUnicode>())
158                .map(|s| s.as_str().to_string())
159                .unwrap_or_default()
160        };
161        let acquisition = attr_str("acquisition");
162        let patient_id = attr_str("patient_id");
163
164        info!(
165            "Opened {}  -- {} slices, {} coils, ky={}, kx={}, recon={}x{}, acq={:?}",
166            path.display(),
167            n_slices,
168            n_coils,
169            n_ky,
170            n_kx,
171            recon_y,
172            recon_x,
173            acquisition,
174        );
175
176        Ok(FastmriFile {
177            file,
178            meta: FastmriMeta {
179                header,
180                acquisition,
181                patient_id,
182                n_slices,
183                n_coils,
184                n_ky,
185                n_kx,
186                recon_y,
187                recon_x,
188            },
189        })
190    }
191
192    /// Read the full k-space tensor and return it as `[coils, slices, ky, kx]`.
193    ///
194    /// This matches the axis order produced by [`IsmrmrdFile::read_kspace`]
195    /// so both can be fed into the same reconstruction pipeline unchanged.
196    ///
197    /// The raw HDF5 layout is `[slices, coils, ky, kx]`; we permute axes on load.
198    pub fn read_kspace(&self) -> IoResult<Array4<Complex32>> {
199        let m = &self.meta;
200        let (data, _) = read_complex_flat(&self.file, "kspace")?;
201        // data is in [slices, coils, ky, kx] row-major order.
202        // Build output in [coils, slices, ky, kx] order.
203        let (ns, nc, nky, nkx) = (m.n_slices, m.n_coils, m.n_ky, m.n_kx);
204        let coil_stride = nky * nkx;
205        let slice_stride = nc * coil_stride;
206
207        let mut out = Array4::<Complex32>::zeros((nc, ns, nky, nkx));
208        for s in 0..ns {
209            for c in 0..nc {
210                let src_off = s * slice_stride + c * coil_stride;
211                out.slice_mut(ndarray::s![c, s, .., ..])
212                    .iter_mut()
213                    .zip(&data[src_off..src_off + coil_stride])
214                    .for_each(|(dst, src)| *dst = *src);
215            }
216        }
217        Ok(out)
218    }
219
220    /// Read a single slice's k-space as `[coils, ky, kx]`.
221    ///
222    /// More memory-efficient than [`read_kspace`](Self::read_kspace) when
223    /// only one slice is needed - still loads the full file but avoids
224    /// allocating the full permuted tensor.
225    pub fn read_kspace_slice(&self, slice_idx: usize) -> IoResult<Array3<Complex32>> {
226        let m = &self.meta;
227        if slice_idx >= m.n_slices {
228            return Err(IoError::Inconsistent(format!(
229                "slice {} out of range (file has {} slices)",
230                slice_idx, m.n_slices
231            )));
232        }
233        let (data, _) = read_complex_flat(&self.file, "kspace")?;
234        let coil_stride = m.n_ky * m.n_kx;
235        let slice_stride = m.n_coils * coil_stride;
236        let slice_base = slice_idx * slice_stride;
237
238        let mut out = Array3::<Complex32>::zeros((m.n_coils, m.n_ky, m.n_kx));
239        for c in 0..m.n_coils {
240            let src_off = slice_base + c * coil_stride;
241            out.slice_mut(ndarray::s![c, .., ..])
242                .iter_mut()
243                .zip(&data[src_off..src_off + coil_stride])
244                .for_each(|(dst, src)| *dst = *src);
245        }
246        Ok(out)
247    }
248
249    /// Read the ground-truth RSS reconstruction as `[slices, recon_y, recon_x]`.
250    pub fn read_reconstruction_rss(&self) -> IoResult<Array3<f32>> {
251        let m = &self.meta;
252        let (data, _) = read_f32_flat(&self.file, "reconstruction_rss")?;
253        Array3::from_shape_vec((m.n_slices, m.recon_y, m.recon_x), data)
254            .map_err(|e| IoError::Inconsistent(e.to_string()))
255    }
256
257    /// Read the ground-truth RSS for a single slice as `[recon_y, recon_x]`.
258    pub fn read_reconstruction_rss_slice(&self, slice_idx: usize) -> IoResult<Array2<f32>> {
259        let m = &self.meta;
260        if slice_idx >= m.n_slices {
261            return Err(IoError::Inconsistent(format!(
262                "slice {} out of range (file has {} slices)",
263                slice_idx, m.n_slices
264            )));
265        }
266        let (data, _) = read_f32_flat(&self.file, "reconstruction_rss")?;
267        let stride = m.recon_y * m.recon_x;
268        let off = slice_idx * stride;
269        Array2::from_shape_vec((m.recon_y, m.recon_x), data[off..off + stride].to_vec())
270            .map_err(|e| IoError::Inconsistent(e.to_string()))
271    }
272
273    /// Returns `true` if the acquisition label heuristically indicates a brain scan.
274    ///
275    /// The heuristic checks whether the label (e.g. `"AXT2"`, `"AXFLAIR"`) contains
276    /// the substring `"AX"` (axial) after upper-casing. Brain FastMRI acquisitions
277    /// are typically labelled with axial plane prefixes; knee acquisitions use
278    /// coronal (`"COR"`) or sagittal (`"SAG"`) prefixes. This is a best-effort
279    /// heuristic -- files with non-standard labels may be misclassified.
280    pub fn is_brain(&self) -> bool {
281        self.meta.acquisition.to_ascii_uppercase().contains("AX")
282    }
283}