support variable number of channels, png

Cargo.toml

@@ -6,3 +6,4 @@ edition = "2018"
 [dependencies]
 thiserror = "1.0.25"
 tiff = "0.7.0"
+png = "0.16"

src/codec.rs

@@ -137,25 +137,25 @@ mod tests {
 
     #[test]
     fn test_codec_12131() {
-        let frame = RGB48Frame::open("src/testdata/tears_of_steel_12130.tif").unwrap();
+        let frame = RGB48Frame::from_tiff("src/testdata/tears_of_steel_12130.tif").unwrap();
         assert_eq!(frame.data.len(), 4096 * 1714 * 3); // 42,123,264 bytes uncompressed
 
         let mut encoded = Vec::new();
         frame.encode::<Codec, _>(&mut encoded).unwrap();
-        assert_eq!(encoded.len(), 25526583);
+        assert_eq!(encoded.len(), 25526584);
 
         let decoded = RGB48Frame::decode::<Codec, _>(&*encoded, frame.width, frame.height).unwrap();
         assert_eq!(frame == decoded, true);
     }
 
     #[test]
     fn test_codec_12209() {
-        let frame = RGB48Frame::open("src/testdata/tears_of_steel_12209.tif").unwrap();
+        let frame = RGB48Frame::from_tiff("src/testdata/tears_of_steel_12209.tif").unwrap();
         assert_eq!(frame.data.len(), 4096 * 1714 * 3); // 42,123,264 bytes uncompressed
 
         let mut encoded = Vec::new();
         frame.encode::<Codec, _>(&mut encoded).unwrap();
-        assert_eq!(encoded.len(), 28270586);
+        assert_eq!(encoded.len(), 28270587);
 
         let decoded = RGB48Frame::decode::<Codec, _>(&*encoded, frame.width, frame.height).unwrap();
         assert_eq!(frame == decoded, true);

src/frame.rs

@@ -2,6 +2,8 @@ use std::{
     io::{self, Read, Write},
     path::Path,
 };
+
+use super::bitstream::{Bitstream, BitstreamWriter};
 use thiserror::Error;
 
 pub struct Plane<T> {
@@ -29,6 +31,8 @@ pub enum FrameOpenError {
     IO(#[from] io::Error),
     #[error(transparent)]
     TiffError(#[from] tiff::TiffError),
+    #[error(transparent)]
+    PngError(#[from] png::DecodingError),
     #[error("unsupported color type: {0:?}")]
     UnsupportedColorType(tiff::ColorType),
     #[error("unsupported sample type")]
@@ -43,7 +47,20 @@ pub struct RGB48Frame {
 }
 
 impl RGB48Frame {
-    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, FrameOpenError> {
+    pub fn from_png<P: AsRef<Path>>(path: P) -> Result<Self, FrameOpenError> {
+        let decoder = png::Decoder::new(std::fs::File::open(path)?);
+        let (info, mut reader) = decoder.read_info()?;
+        let mut data = vec![0; info.buffer_size()];
+        reader.next_frame(&mut data)?;
+        let data = data.iter().map(|&x| x as u16).collect();
+        Ok(RGB48Frame {
+            data,
+            width: info.width as _,
+            height: info.height as _,
+        })
+    }
+
+    pub fn from_tiff<P: AsRef<Path>>(path: P) -> Result<Self, FrameOpenError> {
         let f = std::fs::File::open(path)?;
         let mut dec =
             tiff::decoder::Decoder::new(f)?.with_limits(tiff::decoder::Limits::unlimited());
@@ -63,32 +80,24 @@ impl RGB48Frame {
     }
 
     pub fn planes(&self) -> Vec<Plane<&[u16]>> {
-        vec![
-            Plane {
-                data: &self.data,
-                width: self.width,
-                height: self.height,
-                row_stride: 3 * self.width,
-                sample_stride: 3,
-            },
-            Plane {
-                data: &self.data[1..],
-                width: self.width,
-                height: self.height,
-                row_stride: 3 * self.width,
-                sample_stride: 3,
-            },
-            Plane {
-                data: &self.data[2..],
+        let n_planes = self.data.len() / (self.width * self.height);
+        return (0..n_planes)
+            .map(|plane| Plane {
+                data: &self.data[plane..],
                 width: self.width,
                 height: self.height,
-                row_stride: 3 * self.width,
-                sample_stride: 3,
-            },
-        ]
+                row_stride: n_planes * self.width,
+                sample_stride: n_planes,
+            })
+            .collect();
     }
 
     pub fn encode<C: Codec, W: Write>(&self, mut dest: W) -> io::Result<()> {
+        {
+            let mut bitstream = BitstreamWriter::new(&mut dest);
+            bitstream.write_bits((self.planes().len() - 1) as _, 2)?;
+            bitstream.flush()?;
+        }
         for plane in self.planes() {
             C::encode(&plane, &mut dest)?;
         }
@@ -100,20 +109,24 @@ impl RGB48Frame {
         width: usize,
         height: usize,
     ) -> io::Result<Self> {
+        let n_planes = {
+            let mut bitstream = Bitstream::new(&mut source);
+            (bitstream.read_bits(2)? as usize) + 1
+        };
         let mut ret = Self {
-            data: vec![0; width * height * 3],
+            data: vec![0; width * height * n_planes],
             width,
             height,
         };
-        for plane in 0..3 {
+        for plane in 0..n_planes {
             C::decode(
                 &mut source,
                 &mut Plane {
                     data: &mut ret.data[plane..],
                     width: width,
                     height: height,
-                    row_stride: 3 * width,
-                    sample_stride: 3,
+                    row_stride: n_planes * width,
+                    sample_stride: n_planes,
                 },
             )?;
         }
@@ -127,6 +140,6 @@ mod tests {
 
     #[test]
     fn test_rgb48_frame_open() {
-        RGB48Frame::open("src/testdata/tears_of_steel_12130.tif").unwrap();
+        RGB48Frame::from_tiff("src/testdata/tears_of_steel_12130.tif").unwrap();
     }
 }