allow access to raw bytes on image extension for cursor readers

cds-astro · Feb 25, 2025 · 7fbc73d · 7fbc73d
1 parent fa5843f
commit 7fbc73d
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 88 deletions.
diff --git a/README.md b/README.md
@@ -72,7 +72,7 @@ Example
 ```rust
 use std::fs::File;
 use std::io::Cursor;
-use fitsrs::{Fits, ImageData, HDU, hdu::header::Xtension};
+use fitsrs::{Fits, ImageData, Pixels, HDU, hdu::header::Xtension};
 use fitsrs::wcs::{ImgXY, LonLat};
 
 use std::io::{BufReader, Read};
@@ -110,29 +110,29 @@ while let Some(Ok(hdu)) = hdu_list.next() {
                 assert!((xy.x() - xy_2.x()).abs() <= 1e-9);
                 assert!((xy.y() - xy_2.y()).abs() <= 1e-9);
             }
-
-            match hdu_list.get_data(&hdu) {
-                ImageData::U8(it) => {
+            let image = hdu_list.get_data(&hdu);
+            match image.pixels() {
+                Pixels::U8(it) => {
                     let data = it.collect::<Vec<_>>();
                     assert_eq!(num_pixels, data.len())
                 },
-                ImageData::I16(it) => {
+                Pixels::I16(it) => {
                     let data = it.collect::<Vec<_>>();
                     assert_eq!(num_pixels, data.len())
                 },
-                ImageData::I32(it) => {
+                Pixels::I32(it) => {
                     let data = it.collect::<Vec<_>>();
                     assert_eq!(num_pixels, data.len())
                 },
-                ImageData::I64(it) => {
+                Pixels::I64(it) => {
                     let data = it.collect::<Vec<_>>();
                     assert_eq!(num_pixels, data.len())
                 },
-                ImageData::F32(it) => {
+                Pixels::F32(it) => {
                     let data = it.collect::<Vec<_>>();
                     assert_eq!(num_pixels, data.len())
                 },
-                ImageData::F64(it) => {
+                Pixels::F64(it) => {
                     let data = it.collect::<Vec<_>>();
                     assert_eq!(num_pixels, data.len())
                 },

diff --git a/src/hdu/data/bintable/data.rs b/src/hdu/data/bintable/data.rs
@@ -144,7 +144,7 @@ pub struct TableData<R> {
     heap: bool,
 }
 
-impl<R> TableData<Cursor<R>>
+impl<'a, R> TableData<&'a mut Cursor<R>>
 where
     R: AsRef<[u8]>,
 {
@@ -161,7 +161,7 @@ where
 
 use std::io::Cursor;
 /// Get a reference to the inner reader for in-memory readers
-impl<R> TableData<Cursor<R>> {
+impl<'a, R> TableData<&'a mut Cursor<R>> {
     pub const fn get_ref(&self) -> &R {
         self.reader.get_ref()
     }

diff --git a/src/hdu/data/image.rs b/src/hdu/data/image.rs
@@ -20,34 +20,17 @@ where
 {
     type Data = ImageData<&'a mut Self>;
 
-    fn read_data_unit(&'a mut self, header: &Header<Image>, _start_pos: u64) -> Self::Data {
-        ImageData::new(header.get_xtension(), self)
+    fn read_data_unit(&'a mut self, header: &Header<Image>, start_pos: u64) -> Self::Data {
+        ImageData::new(header.get_xtension(), self, start_pos)
     }
 }
 
-use super::stream;
-#[async_trait(?Send)]
-impl<'a, R> AsyncDataBufRead<'a, Image> for futures::io::BufReader<R>
-where
-    R: AsyncReadExt + Debug + 'a + Unpin,
-{
-    type Data = DataStream<'a, Self>;
 
-    fn prepare_data_reading(
-        ctx: &Image,
-        num_remaining_bytes_in_cur_hdu: &'a mut usize,
-        reader: &'a mut Self,
-    ) -> Self::Data {
-        let bitpix = ctx.get_bitpix();
-        match bitpix {
-            Bitpix::U8 => DataStream::U8(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
-            Bitpix::I16 => DataStream::I16(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
-            Bitpix::I32 => DataStream::I32(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
-            Bitpix::I64 => DataStream::I64(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
-            Bitpix::F32 => DataStream::F32(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
-            Bitpix::F64 => DataStream::F64(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
-        }
-    }
+#[derive(Serialize, Debug)]
+pub struct ImageData<R> {
+    start_pos: u64,
+    num_bytes_data_block: u64,
+    pixels: Pixels<R>,
 }
 
 /// An iterator on the data array
@@ -58,7 +41,7 @@ where
 /// for non in-memory readers (typically BufReader) that ensures
 /// a file may not fit in memory
 #[derive(Serialize, Debug)]
-pub enum ImageData<R> {
+pub enum Pixels<R> {
     U8(It<R, u8>),
     I16(It<R, i16>),
     I32(It<R, i32>),
@@ -71,16 +54,72 @@ impl<R> ImageData<R>
 where
     R: Read,
 {
-    pub(crate) fn new(ctx: &Image, reader: R) -> Self {
+    pub(crate) fn new(ctx: &Image, reader: R, start_pos: u64) -> Self {
         let limit = ctx.get_num_bytes_data_block();
 
-        match ctx.get_bitpix() {
-            Bitpix::U8 => ImageData::U8(It::new(reader, limit)),
-            Bitpix::I16 => ImageData::I16(It::new(reader, limit)),
-            Bitpix::I32 => ImageData::I32(It::new(reader, limit)),
-            Bitpix::I64 => ImageData::I64(It::new(reader, limit)),
-            Bitpix::F32 => ImageData::F32(It::new(reader, limit)),
-            Bitpix::F64 => ImageData::F64(It::new(reader, limit)),
-        }
+        let pixels = match ctx.get_bitpix() {
+            Bitpix::U8 => Pixels::U8(It::new(reader, limit)),
+            Bitpix::I16 => Pixels::I16(It::new(reader, limit)),
+            Bitpix::I32 => Pixels::I32(It::new(reader, limit)),
+            Bitpix::I64 => Pixels::I64(It::new(reader, limit)),
+            Bitpix::F32 => Pixels::F32(It::new(reader, limit)),
+            Bitpix::F64 => Pixels::F64(It::new(reader, limit)),
+        };
+
+        Self { start_pos, num_bytes_data_block: limit, pixels }
+    }
+
+    /// Get the pixels iterator of the image
+    pub fn pixels(self) -> Pixels<R> {
+        self.pixels
     }
 }
+
+use std::io::Cursor;
+impl<'a, R> ImageData<&'a mut Cursor<R>>
+where
+    R: AsRef<[u8]>,
+{
+    /// For in memory buffers, access the raw bytes of the image.
+    /// You might need to convert the data from big to little endian at some point
+    pub fn raw_bytes(&self) -> &[u8] {
+        let inner = match &self.pixels {
+            Pixels::U8(It { reader, ..}) => reader.get_ref(),
+            Pixels::I16(It { reader, ..}) => reader.get_ref(),
+            Pixels::I32(It { reader, ..}) => reader.get_ref(),
+            Pixels::I64(It { reader, ..}) => reader.get_ref(),
+            Pixels::F32(It { reader, ..}) => reader.get_ref(),
+            Pixels::F64(It { reader, ..}) => reader.get_ref(),
+        };
+        let raw_bytes = inner.as_ref();
+
+        let s = self.start_pos as usize;
+        let e = s + (self.num_bytes_data_block as usize);
+        &raw_bytes[s..e]
+    }
+}
+
+use super::stream;
+#[async_trait(?Send)]
+impl<'a, R> AsyncDataBufRead<'a, Image> for futures::io::BufReader<R>
+where
+    R: AsyncReadExt + Debug + 'a + Unpin,
+{
+    type Data = DataStream<'a, Self>;
+
+    fn prepare_data_reading(
+        ctx: &Image,
+        num_remaining_bytes_in_cur_hdu: &'a mut usize,
+        reader: &'a mut Self,
+    ) -> Self::Data {
+        let bitpix = ctx.get_bitpix();
+        match bitpix {
+            Bitpix::U8 => DataStream::U8(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
+            Bitpix::I16 => DataStream::I16(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
+            Bitpix::I32 => DataStream::I32(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
+            Bitpix::I64 => DataStream::I64(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
+            Bitpix::F32 => DataStream::F32(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
+            Bitpix::F64 => DataStream::F64(stream::St::new(reader, num_remaining_bytes_in_cur_hdu)),
+        }
+    }
+}
diff --git a/src/hdu/data/iter.rs b/src/hdu/data/iter.rs
@@ -48,7 +48,7 @@ impl Value for f64 {
 #[derive(Debug, Serialize)]
 pub struct It<R, T> {
     /// The reader
-    reader: R,
+    pub(crate) reader: R,
     /// The number of items the reader must read
     num_items: usize,
     /// Number of item read

diff --git a/src/hdu/data/mod.rs b/src/hdu/data/mod.rs
@@ -6,7 +6,7 @@ pub mod stream;
 
 pub use bintable::BinaryTableData;
 pub use bintable::TableData;
-pub use image::ImageData;
+pub use image::{ImageData, Pixels};
 
 pub use iter::It;