Add the ability to scale down images during decode (IDCT scaling)

src/decoder.rs

@@ -50,6 +50,7 @@ pub struct ImageInfo {
 /// JPEG decoder
 pub struct Decoder<R> {
     reader: R,
+    requested_size: Option<Dimensions>,
 
     frame: Option<FrameInfo>,
     dc_huffman_tables: Vec<Option<HuffmanTable>>,
@@ -70,8 +71,26 @@ pub struct Decoder<R> {
 impl<R: Read> Decoder<R> {
     /// Creates a new `Decoder` using the reader `reader`.
     pub fn new(reader: R) -> Decoder<R> {
+        Decoder::init(reader, None)
+    }
+
+    /// Creates a new `Decoder` using the reader `reader` that returns a
+    /// scaled image that is equal to or larger than the requested size in at
+    /// least one axis, or the full size of the image if the requested size is
+    /// larger.
+    /// 
+    /// This efficiently scales down the image by one of a fixed set of
+    /// available ratios during decoding. To generate a thumbnail of an
+    /// exact size, pass the desired size or larger and then scale to the
+    /// final size using a traditional resampling algorithm.
+    pub fn scaled(reader: R, requested_width: u16, requested_height: u16) -> Decoder<R> {
+        Decoder::init(reader, Some(Dimensions{ width: requested_width, height: requested_height }))
+    }
+
+    fn init(reader: R, requested_size: Option<Dimensions>) -> Decoder<R> {
         Decoder {
             reader: reader,
+            requested_size,
             frame: None,
             dc_huffman_tables: vec![None, None, None, None],
             ac_huffman_tables: vec![None, None, None, None],
@@ -100,8 +119,8 @@ impl<R: Read> Decoder<R> {
                 };
 
                 Some(ImageInfo {
-                    width: frame.image_size.width,
-                    height: frame.image_size.height,
+                    width: frame.output_size.width,
+                    height: frame.output_size.height,
                     pixel_format: pixel_format,
                 })
             },
@@ -153,7 +172,7 @@ impl<R: Read> Decoder<R> {
                         return Err(Error::Unsupported(UnsupportedFeature::Hierarchical));
                     }
 
-                    let frame = parse_sof(&mut self.reader, marker)?;
+                    let frame = parse_sof(&mut self.reader, marker, self.requested_size)?;
                     let component_count = frame.components.len();
 
                     if frame.is_differential {
@@ -329,7 +348,7 @@ impl<R: Read> Decoder<R> {
         }
 
         let frame = self.frame.as_ref().unwrap();
-        compute_image(&frame.components, &planes, frame.image_size, self.is_jfif, self.color_transform)
+        compute_image(&frame.components, &planes, frame.output_size, self.is_jfif, self.color_transform)
     }
 
     fn read_marker(&mut self) -> Result<Marker> {
@@ -436,7 +455,7 @@ impl<R: Read> Decoder<R> {
                             let x = (block_num % blocks_per_row) as u16;
                             let y = (block_num / blocks_per_row) as u16;
 
-                            if x * 8 >= component.size.width || y * 8 >= component.size.height {
+                            if x * component.dct_scale as u16 >= component.size.width || y * component.dct_scale as u16 >= component.size.height {
                                 continue;
                             }
 
@@ -762,12 +781,15 @@ fn compute_image(components: &[Component],
             return Ok(data[0].clone())
         }
 
-        let mut buffer = vec![0u8; component.size.width as usize * component.size.height as usize];
-        let line_stride = component.block_size.width as usize * 8;
+        let width = component.size.width as usize;
+        let height = component.size.height as usize;
+
+        let mut buffer = vec![0u8; width * height];
+        let line_stride = width * component.dct_scale;
 
-        for y in 0 .. component.size.height as usize {
-            for x in 0 .. component.size.width as usize {
-                buffer[y * component.size.width as usize + x] = data[0][y * line_stride + x];
+        for y in 0 .. width {
+            for x in 0 .. height {
+                buffer[y * width + x] = data[0][y * line_stride + x];
             }
         }
 

src/idct.rs

@@ -1,9 +1,51 @@
 // Malicious JPEG files can cause operations in the idct to overflow.
 // One example is tests/crashtest/images/imagetestsuite/b0b8914cc5f7a6eff409f16d8cc236c5.jpg
 // That's why wrapping operators are needed.
+use crate::parser::Dimensions;
+
+pub fn choose_idct_size(full_size: Dimensions, requested_size: Dimensions) -> usize {
+    fn scaled(len: u16, scale: usize) -> u16 { ((len as u32 * scale as u32 - 1) / 8 + 1) as u16 }
+
+    for &scale in &[1, 2, 4] {
+        if scaled(full_size.width, scale) >= requested_size.width || scaled(full_size.height, scale) >= requested_size.height {
+            return scale;
+        }
+    }
+
+    return 8;
+}
+
+#[test]
+fn test_choose_idct_size() {
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 200, height: 200}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 500, height: 500}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 684, height: 456}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 999, height: 456}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 684, height: 999}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 500, height: 333}, Dimensions{width: 63, height: 42}), 1);
+
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 685, height: 999}), 2);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 1000, height: 1000}), 2);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 1400, height: 1400}), 4);
+
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 5472, height: 3648}), 8);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 16384, height: 16384}), 8);
+    assert_eq!(choose_idct_size(Dimensions{width: 1, height: 1}, Dimensions{width: 65535, height: 65535}), 8);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 16384, height: 16384}), 8);
+}
+
+pub fn dequantize_and_idct_block(scale: usize, coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+    match scale {
+        8 => dequantize_and_idct_block_8x8(coefficients, quantization_table, output_linestride, output),
+        4 => dequantize_and_idct_block_4x4(coefficients, quantization_table, output_linestride, output),
+        2 => dequantize_and_idct_block_2x2(coefficients, quantization_table, output_linestride, output),
+        1 => dequantize_and_idct_block_1x1(coefficients, quantization_table, output_linestride, output),
+        _ => panic!("Unsupported IDCT scale {}/8", scale),
+    }
+}
 
 // This is based on stb_image's 'stbi__idct_block'.
-pub fn dequantize_and_idct_block(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+fn dequantize_and_idct_block_8x8(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
     debug_assert_eq!(coefficients.len(), 64);
 
     let mut temp = [0i32; 64];
@@ -155,6 +197,103 @@ pub fn dequantize_and_idct_block(coefficients: &[i16], quantization_table: &[u16
     }
 }
 
+// 4x4 and 2x2 IDCT based on Rakesh Dugad and Narendra Ahuja: "A Fast Scheme for Image Size Change in the Compressed Domain" (2001).
+// http://sylvana.net/jpegcrop/jidctred/
+fn dequantize_and_idct_block_4x4(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+    debug_assert_eq!(coefficients.len(), 64);
+    let mut temp = [0i32; 4*4];
+
+    const CONST_BITS: u32 = 12;
+    const PASS1_BITS: u32 = 2;
+    const FINAL_BITS: u32 = CONST_BITS + PASS1_BITS + 3;
+
+    // columns
+    for i in 0 .. 4 {
+        let s0 = coefficients[i + 8*0] as i32 * quantization_table[i + 8*0] as i32;
+        let s1 = coefficients[i + 8*1] as i32 * quantization_table[i + 8*1] as i32;
+        let s2 = coefficients[i + 8*2] as i32 * quantization_table[i + 8*2] as i32;
+        let s3 = coefficients[i + 8*3] as i32 * quantization_table[i + 8*3] as i32;
+
+        let x0 = s0.wrapping_add(s2).wrapping_shl(PASS1_BITS);
+        let x2 = s0.wrapping_sub(s2).wrapping_shl(PASS1_BITS);
+
+        let p1 = s1.wrapping_add(s3).wrapping_mul(stbi_f2f(0.541196100));
+        let t0 = p1.wrapping_add(s3.wrapping_mul(stbi_f2f(-1.847759065))).wrapping_add(512).wrapping_shr(CONST_BITS - PASS1_BITS);
+        let t2 = p1.wrapping_add(s1.wrapping_mul(stbi_f2f( 0.765366865))).wrapping_add(512).wrapping_shr(CONST_BITS - PASS1_BITS);
+
+        temp[i + 4*0] = x0.wrapping_add(t2);
+        temp[i + 4*3] = x0.wrapping_sub(t2);
+        temp[i + 4*1] = x2.wrapping_add(t0);
+        temp[i + 4*2] = x2.wrapping_sub(t0);
+    }
+
+    for i in 0 .. 4 {
+        let s0 = temp[i * 4 + 0];
+        let s1 = temp[i * 4 + 1];
+        let s2 = temp[i * 4 + 2];
+        let s3 = temp[i * 4 + 3];
+
+        let x0 = s0.wrapping_add(s2).wrapping_shl(CONST_BITS);
+        let x2 = s0.wrapping_sub(s2).wrapping_shl(CONST_BITS);
+
+        let p1 = s1.wrapping_add(s3).wrapping_mul(stbi_f2f(0.541196100));
+        let t0 = p1.wrapping_add(s3.wrapping_mul(stbi_f2f(-1.847759065)));
+        let t2 = p1.wrapping_add(s1.wrapping_mul(stbi_f2f(0.765366865)));
+
+        // constants scaled things up by 1<<12, plus we had 1<<2 from first
+        // loop, plus horizontal and vertical each scale by sqrt(8) so together
+        // we've got an extra 1<<3, so 1<<17 total we need to remove.
+        // so we want to round that, which means adding 0.5 * 1<<17,
+        // aka 65536. Also, we'll end up with -128 to 127 that we want
+        // to encode as 0..255 by adding 128, so we'll add that before the shift
+        let x0 = x0.wrapping_add((1 << (FINAL_BITS - 1)) + (128 << FINAL_BITS));
+        let x2 = x2.wrapping_add((1 << (FINAL_BITS - 1)) + (128 << FINAL_BITS));
+
+        output[i * output_linestride + 0] = stbi_clamp(x0.wrapping_add(t2).wrapping_shr(FINAL_BITS));
+        output[i * output_linestride + 3] = stbi_clamp(x0.wrapping_sub(t2).wrapping_shr(FINAL_BITS));
+        output[i * output_linestride + 1] = stbi_clamp(x2.wrapping_add(t0).wrapping_shr(FINAL_BITS));
+        output[i * output_linestride + 2] = stbi_clamp(x2.wrapping_sub(t0).wrapping_shr(FINAL_BITS));
+    }
+}
+
+fn dequantize_and_idct_block_2x2(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+    debug_assert_eq!(coefficients.len(), 64);
+
+    const SCALE_BITS: u32 = 3;
+
+    // Column 0
+    let s00 = coefficients[8*0] as i32 * quantization_table[8*0] as i32;
+    let s10 = coefficients[8*1] as i32 * quantization_table[8*1] as i32;
+
+    let x0 = s00.wrapping_add(s10);
+    let x2 = s00.wrapping_sub(s10);
+
+    // Column 1
+    let s01 = coefficients[8*0+1] as i32 * quantization_table[8*0+1] as i32;
+    let s11 = coefficients[8*1+1] as i32 * quantization_table[8*1+1] as i32;
+
+    let x1 = s01.wrapping_add(s11);
+    let x3 = s01.wrapping_sub(s11);
+
+    let x0 = x0.wrapping_add((1 << (SCALE_BITS-1)) + (128 << SCALE_BITS));
+    let x2 = x2.wrapping_add((1 << (SCALE_BITS-1)) + (128 << SCALE_BITS));
+
+    // Row 0
+    output[0] = stbi_clamp(x0.wrapping_add(x1).wrapping_shr(SCALE_BITS));
+    output[1] = stbi_clamp(x0.wrapping_sub(x1).wrapping_shr(SCALE_BITS));
+
+    // Row 1
+    output[output_linestride + 0] = stbi_clamp(x2.wrapping_add(x3).wrapping_shr(SCALE_BITS));
+    output[output_linestride + 1] = stbi_clamp(x2.wrapping_sub(x3).wrapping_shr(SCALE_BITS));
+}
+
+fn dequantize_and_idct_block_1x1(coefficients: &[i16], quantization_table: &[u16; 64], _output_linestride: usize, output: &mut [u8]) {
+    debug_assert_eq!(coefficients.len(), 64);
+
+    let s0 = (coefficients[0] as i32 * quantization_table[0] as i32).wrapping_add(128 * 8) / 8;
+    output[0] = stbi_clamp(s0);
+}
+
 // take a -128..127 value and stbi__clamp it and convert to 0..255
 fn stbi_clamp(x: i32) -> u8
 {

src/parser.rs

@@ -34,6 +34,7 @@ pub struct FrameInfo {
     pub precision: u8,
 
     pub image_size: Dimensions,
+    pub output_size: Dimensions,
     pub mcu_size: Dimensions,
     pub components: Vec<Component>,
 }
@@ -58,6 +59,8 @@ pub struct Component {
 
     pub quantization_table_index: usize,
 
+    pub dct_scale: usize,
+
     pub size: Dimensions,
     pub block_size: Dimensions,
 }
@@ -104,7 +107,7 @@ fn skip_bytes<R: Read>(reader: &mut R, length: usize) -> Result<()> {
 }
 
 // Section B.2.2
-pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker) -> Result<FrameInfo> {
+pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker, requested_size: Option<Dimensions>) -> Result<FrameInfo> {
     let length = read_length(reader, marker)?;
 
     if length <= 6 {
@@ -156,6 +159,10 @@ pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker) -> Result<FrameInfo> {
         return Err(Error::Format("zero width in frame header".to_owned()));
     }
 
+    let scale = if let Some(req) = requested_size {
+        crate::idct::choose_idct_size(Dimensions { width, height }, req)
+    } else { 8 };
+
     let component_count = reader.read_u8()?;
 
     if component_count == 0 {
@@ -201,6 +208,7 @@ pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker) -> Result<FrameInfo> {
             horizontal_sampling_factor: horizontal_sampling_factor,
             vertical_sampling_factor: vertical_sampling_factor,
             quantization_table_index: quantization_table_index as usize,
+            dct_scale: scale,
             size: Dimensions {width: 0, height: 0},
             block_size: Dimensions {width: 0, height: 0},
         });
@@ -214,8 +222,8 @@ pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker) -> Result<FrameInfo> {
     };
 
     for component in &mut components {
-        component.size.width = (width as f32 * (component.horizontal_sampling_factor as f32 / h_max as f32)).ceil() as u16;
-        component.size.height = (height as f32 * (component.vertical_sampling_factor as f32 / v_max as f32)).ceil() as u16;
+        component.size.width = (width as f32 * component.horizontal_sampling_factor as f32 * component.dct_scale as f32 / (h_max as f32 * 8.0)).ceil() as u16;
+        component.size.height = (height as f32 * component.vertical_sampling_factor as f32 * component.dct_scale as f32 / (v_max as f32 * 8.0)).ceil() as u16;
 
         component.block_size.width = mcu_size.width * component.horizontal_sampling_factor as u16;
         component.block_size.height = mcu_size.height * component.vertical_sampling_factor as u16;
@@ -228,6 +236,7 @@ pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker) -> Result<FrameInfo> {
         entropy_coding: entropy_coding,
         precision: precision,
         image_size: Dimensions {width: width, height: height},
+        output_size: Dimensions {width: (width as f32 * scale as f32 / 8.0).ceil() as u16, height: (height as f32 * scale as f32 / 8.0).ceil() as u16},
         mcu_size: mcu_size,
         components: components,
     })

src/upsampler.rs

@@ -29,7 +29,7 @@ impl Upsampler {
                 upsampler: upsampler,
                 width: component.size.width as usize,
                 height: component.size.height as usize,
-                row_stride: component.block_size.width as usize * 8,
+                row_stride: component.block_size.width as usize * component.dct_scale,
             });
         }
 

src/worker/immediate.rs

@@ -26,7 +26,7 @@ impl ImmediateWorker {
         assert!(self.results[data.index].is_empty());
 
         self.offsets[data.index] = 0;
-        self.results[data.index].resize(data.component.block_size.width as usize * data.component.block_size.height as usize * 64, 0u8);
+        self.results[data.index].resize(data.component.block_size.width as usize * data.component.block_size.height as usize * data.component.dct_scale * data.component.dct_scale, 0u8);
         self.components[data.index] = Some(data.component);
         self.quantization_tables[data.index] = Some(data.quantization_table);
     }
@@ -36,20 +36,21 @@ impl ImmediateWorker {
         let component = self.components[index].as_ref().unwrap();
         let quantization_table = self.quantization_tables[index].as_ref().unwrap();
         let block_count = component.block_size.width as usize * component.vertical_sampling_factor as usize;
-        let line_stride = component.block_size.width as usize * 8;
+        let line_stride = component.block_size.width as usize * component.dct_scale;
 
         assert_eq!(data.len(), block_count * 64);
 
         for i in 0..block_count {
-            let x = (i % component.block_size.width as usize) * 8;
-            let y = (i / component.block_size.width as usize) * 8;
-            dequantize_and_idct_block(&data[i * 64..(i + 1) * 64],
-                                    quantization_table,
-                                    line_stride,
-                                    &mut self.results[index][self.offsets[index] + y * line_stride + x..]);
+            let x = (i % component.block_size.width as usize) * component.dct_scale;
+            let y = (i / component.block_size.width as usize) * component.dct_scale;
+
+            let coefficients = &data[i * 64..(i + 1) * 64];
+            let output = &mut self.results[index][self.offsets[index] + y * line_stride + x..];
+
+            dequantize_and_idct_block(component.dct_scale, coefficients, quantization_table, line_stride, output);
         }
 
-        self.offsets[index] += data.len();
+        self.offsets[index] += block_count * component.dct_scale * component.dct_scale;
     }
     pub fn get_result_immediate(&mut self, index: usize) -> Vec<u8> {
         mem::replace(&mut self.results[index], Vec::new())