pure_jpeg 0.1.0 → 0.3.0

data/lib/pure_jpeg/decoder.rb

@@ -13,7 +13,7 @@ module PureJPEG
     # @param path_or_data [String] a file path or raw JPEG bytes
     # @return [Image] decoded image with pixel access
     def self.decode(path_or_data)
-      data = if path_or_data.is_a?(String) && !path_or_data.include?("\x00") && File.exist?(path_or_data)
+      data = if path_or_data.is_a?(String) && !path_or_data.start_with?("\xFF\xD8".b) && File.exist?(path_or_data)
                File.binread(path_or_data)
              else
                path_or_data.b
@@ -27,6 +27,10 @@ module PureJPEG
 
     def decode
       jfif = JFIFReader.new(@data)
+      @icc_profile = jfif.icc_profile
+      validate_dimensions!(jfif.width, jfif.height)
+      return decode_progressive(jfif) if jfif.progressive
+
       width = jfif.width
       height = jfif.height
 
@@ -88,10 +92,8 @@ module PureJPEG
           end
 
           jfif.scan_components.each do |sc|
-            comp = comp_info[sc.id]
-            dc_tab = dc_tables[sc.dc_table_id]
-            ac_tab = ac_tables[sc.ac_table_id]
-            qt = jfif.quant_tables[comp.qt_id]
+            comp, dc_tab, ac_tab = resolve_scan_references!(sc, comp_info, dc_tables, ac_tables)
+            qt = fetch_quant_table!(jfif, comp)
             ch = channels[comp.id]
 
             comp.v_sampling.times do |bv|
@@ -120,13 +122,303 @@ module PureJPEG
       num_components = jfif.components.length
       if num_components == 1
         assemble_grayscale(width, height, channels, jfif.components[0])
-      else
+      elsif num_components == 3
         assemble_color(width, height, channels, jfif.components, max_h, max_v)
+      else
+        raise DecodeError, "Unsupported number of components: #{num_components}"
       end
     end
 
     private
 
+    def validate_dimensions!(width, height)
+      raise DecodeError, "Invalid image dimensions: #{width}x#{height}" if width <= 0 || height <= 0
+      raise DecodeError, "Image too large: #{width}x#{height} (max #{MAX_DIMENSION}x#{MAX_DIMENSION})" if width > MAX_DIMENSION || height > MAX_DIMENSION
+    end
+
+    # --- Progressive JPEG decoding ---
+
+    def decode_progressive(jfif)
+      width = jfif.width
+      height = jfif.height
+
+      comp_info = {}
+      jfif.components.each { |c| comp_info[c.id] = c }
+
+      max_h = jfif.components.map(&:h_sampling).max
+      max_v = jfif.components.map(&:v_sampling).max
+
+      mcu_px_w = max_h * 8
+      mcu_px_h = max_v * 8
+      mcus_x = (width + mcu_px_w - 1) / mcu_px_w
+      mcus_y = (height + mcu_px_h - 1) / mcu_px_h
+
+      # Coefficient buffers per component (zigzag order, pre-dequantization)
+      coeffs = {}
+      comp_blocks = {}
+      jfif.components.each do |c|
+        bx = mcus_x * c.h_sampling
+        by = mcus_y * c.v_sampling
+        coeffs[c.id] = Array.new(bx * by * 64, 0)
+        comp_blocks[c.id] = [bx, by]
+      end
+
+      restart_interval = jfif.restart_interval
+
+      jfif.scans.each do |scan|
+        # Build Huffman tables from this scan's snapshot (tables change between scans)
+        dc_tables = {}
+        ac_tables = {}
+        scan.huffman_tables.each do |(table_class, table_id), info|
+          table = Huffman::DecodeTable.new(info[:bits], info[:values])
+          if table_class == 0
+            dc_tables[table_id] = table
+          else
+            ac_tables[table_id] = table
+          end
+        end
+
+        reader = BitReader.new(scan.data)
+        ss = scan.spectral_start
+        se = scan.spectral_end
+        ah = scan.successive_high
+        al = scan.successive_low
+
+        if scan.components.length == 1
+          prog_scan_non_interleaved(reader, scan, comp_info, dc_tables, ac_tables,
+                                    coeffs, comp_blocks, restart_interval, ss, se, ah, al)
+        else
+          prog_scan_interleaved(reader, scan, comp_info, dc_tables, ac_tables,
+                                coeffs, comp_blocks, mcus_x, mcus_y, restart_interval, ss, se, ah, al)
+        end
+      end
+
+      # Reconstruct: unzigzag, dequantize, IDCT, write to channel buffers
+      padded_w = mcus_x * mcu_px_w
+      padded_h = mcus_y * mcu_px_h
+      channels = {}
+      jfif.components.each do |c|
+        ch_w = (padded_w * c.h_sampling) / max_h
+        ch_h = (padded_h * c.v_sampling) / max_v
+        channels[c.id] = { data: Array.new(ch_w * ch_h, 0), width: ch_w, height: ch_h }
+      end
+
+      zigzag = Array.new(64, 0)
+      raster = Array.new(64, 0.0)
+      dequant = Array.new(64, 0.0)
+      temp = Array.new(64, 0.0)
+      spatial = Array.new(64, 0.0)
+
+      jfif.components.each do |c|
+        qt = fetch_quant_table!(jfif, c)
+        ch = channels[c.id]
+        coeff_buf = coeffs[c.id]
+        bx_count, by_count = comp_blocks[c.id]
+
+        by_count.times do |block_y|
+          bx_count.times do |block_x|
+            offset = (block_y * bx_count + block_x) * 64
+            64.times { |i| zigzag[i] = coeff_buf[offset + i] }
+
+            Zigzag.unreorder!(zigzag, raster)
+            Quantization.dequantize!(raster, qt, dequant)
+            DCT.inverse!(dequant, temp, spatial)
+            write_block(spatial, ch[:data], ch[:width], block_x * 8, block_y * 8)
+          end
+        end
+      end
+
+      num_components = jfif.components.length
+      if num_components == 1
+        assemble_grayscale(width, height, channels, jfif.components[0])
+      elsif num_components == 3
+        assemble_color(width, height, channels, jfif.components, max_h, max_v)
+      else
+        raise DecodeError, "Unsupported number of components: #{num_components}"
+      end
+    end
+
+    def prog_scan_non_interleaved(reader, scan, comp_info, dc_tables, ac_tables,
+                                  coeffs, comp_blocks, restart_interval, ss, se, ah, al)
+      sc = scan.components[0]
+      comp, dc_tab, ac_tab = resolve_scan_references!(sc, comp_info, dc_tables, ac_tables, require_ac: ss > 0)
+      coeff_buf = coeffs[comp.id]
+      bx_count, by_count = comp_blocks[comp.id]
+
+      prev_dc = 0
+      eobrun = 0
+      mcu_count = 0
+
+      by_count.times do |block_y|
+        bx_count.times do |block_x|
+          if restart_interval > 0 && mcu_count > 0 && (mcu_count % restart_interval) == 0
+            reader.reset
+            prev_dc = 0
+            eobrun = 0
+          end
+
+          offset = (block_y * bx_count + block_x) * 64
+
+          if ss == 0
+            if ah == 0
+              prev_dc = prog_dc_first(reader, dc_tab, prev_dc, coeff_buf, offset, al)
+            else
+              prog_dc_refine(reader, coeff_buf, offset, al)
+            end
+          else
+            if ah == 0
+              eobrun = prog_ac_first(reader, ac_tab, coeff_buf, offset, ss, se, al, eobrun)
+            else
+              eobrun = prog_ac_refine(reader, ac_tab, coeff_buf, offset, ss, se, al, eobrun)
+            end
+          end
+
+          mcu_count += 1
+        end
+      end
+    end
+
+    def prog_scan_interleaved(reader, scan, comp_info, dc_tables, ac_tables,
+                              coeffs, comp_blocks, mcus_x, mcus_y, restart_interval, ss, se, ah, al)
+      prev_dc = Hash.new(0)
+      mcu_count = 0
+
+      mcus_y.times do |mcu_row|
+        mcus_x.times do |mcu_col|
+          if restart_interval > 0 && mcu_count > 0 && (mcu_count % restart_interval) == 0
+            reader.reset
+            prev_dc.clear
+          end
+
+          scan.components.each do |sc|
+            comp, dc_tab = resolve_scan_references!(sc, comp_info, dc_tables, ac_tables, require_ac: false)
+            coeff_buf = coeffs[comp.id]
+            bx_count = comp_blocks[comp.id][0]
+
+            comp.v_sampling.times do |bv|
+              comp.h_sampling.times do |bh|
+                block_x = mcu_col * comp.h_sampling + bh
+                block_y = mcu_row * comp.v_sampling + bv
+                offset = (block_y * bx_count + block_x) * 64
+
+                if ah == 0
+                  prev_dc[sc.id] = prog_dc_first(reader, dc_tab, prev_dc[sc.id], coeff_buf, offset, al)
+                else
+                  prog_dc_refine(reader, coeff_buf, offset, al)
+                end
+              end
+            end
+          end
+
+          mcu_count += 1
+        end
+      end
+    end
+
+    def prog_dc_first(reader, dc_tab, prev_dc, coeff_buf, offset, al)
+      cat = dc_tab.decode(reader)
+      diff = reader.receive_extend(cat)
+      dc_val = prev_dc + diff
+      coeff_buf[offset] = dc_val << al
+      dc_val
+    end
+
+    def prog_dc_refine(reader, coeff_buf, offset, al)
+      coeff_buf[offset] |= (reader.read_bit << al)
+    end
+
+    def prog_ac_first(reader, ac_tab, coeff_buf, offset, ss, se, al, eobrun)
+      return eobrun - 1 if eobrun > 0
+
+      k = ss
+      while k <= se
+        symbol = ac_tab.decode(reader)
+        run = (symbol >> 4) & 0x0F
+        size = symbol & 0x0F
+
+        if size == 0
+          if run == 15
+            k += 16
+          else
+            # EOBn
+            eobrun = (1 << run)
+            eobrun += reader.read_bits(run) if run > 0
+            return eobrun - 1
+          end
+        else
+          k += run
+          coeff_buf[offset + k] = reader.receive_extend(size) << al
+          k += 1
+        end
+      end
+
+      0
+    end
+
+    def prog_ac_refine(reader, ac_tab, coeff_buf, offset, ss, se, al, eobrun)
+      p1 = 1 << al
+      m1 = -(1 << al)
+
+      if eobrun > 0
+        ss.upto(se) do |k|
+          prog_refine_bit(reader, coeff_buf, offset + k, p1, m1) if coeff_buf[offset + k] != 0
+        end
+        return eobrun - 1
+      end
+
+      k = ss
+      while k <= se
+        symbol = ac_tab.decode(reader)
+        r = (symbol >> 4) & 0x0F
+        s = symbol & 0x0F
+
+        # Read the new coefficient value before processing the run
+        # (the value bits come before refinement bits in the bitstream)
+        new_value = nil
+        if s != 0
+          new_value = reader.receive_extend(s) << al
+        elsif r != 15
+          # EOBn: refine remaining nonzero coefficients in this block
+          eobrun = (1 << r)
+          eobrun += reader.read_bits(r) if r > 0
+          while k <= se
+            prog_refine_bit(reader, coeff_buf, offset + k, p1, m1) if coeff_buf[offset + k] != 0
+            k += 1
+          end
+          return eobrun - 1
+        end
+
+        # Advance through the band: refine nonzero coefficients, count zeros for run.
+        # Break when we've skipped `r` zeros and found the target zero position.
+        while k <= se
+          if coeff_buf[offset + k] != 0
+            prog_refine_bit(reader, coeff_buf, offset + k, p1, m1)
+          elsif r == 0
+            break
+          else
+            r -= 1
+          end
+          k += 1
+        end
+
+        # Place new coefficient at the target zero position
+        if new_value && k <= se
+          coeff_buf[offset + k] = new_value
+        end
+        k += 1
+      end
+
+      0
+    end
+
+    def prog_refine_bit(reader, coeff_buf, idx, p1, m1)
+      if reader.read_bit == 1
+        coeff_buf[idx] += coeff_buf[idx] > 0 ? p1 : m1
+      end
+    end
+
+    # --- Baseline decoding helpers ---
+
     def decode_block(reader, dc_tab, ac_tab, prev_dc, comp_id, out)
       # DC coefficient
       dc_cat = dc_tab.decode(reader)
@@ -177,6 +469,28 @@ module PureJPEG
       end
     end
 
+    def resolve_scan_references!(scan_component, comp_info, dc_tables, ac_tables, require_ac: true)
+      comp = comp_info[scan_component.id]
+      raise DecodeError, "Scan references unknown component id #{scan_component.id}" unless comp
+
+      dc_tab = dc_tables[scan_component.dc_table_id]
+      raise DecodeError, "Component #{scan_component.id} references missing DC Huffman table #{scan_component.dc_table_id}" unless dc_tab
+
+      if require_ac
+        ac_tab = ac_tables[scan_component.ac_table_id]
+        raise DecodeError, "Component #{scan_component.id} references missing AC Huffman table #{scan_component.ac_table_id}" unless ac_tab
+      end
+
+      [comp, dc_tab, ac_tab]
+    end
+
+    def fetch_quant_table!(jfif, comp)
+      qt = jfif.quant_tables[comp.qt_id]
+      raise DecodeError, "Component #{comp.id} references missing quantization table #{comp.qt_id}" unless qt
+
+      qt
+    end
+
     def assemble_grayscale(width, height, channels, comp)
       ch = channels[comp.id]
       pixels = Array.new(width * height)
@@ -185,20 +499,19 @@ module PureJPEG
         dst_row = y * width
         width.times do |x|
           v = ch[:data][src_row + x]
-          pixels[dst_row + x] = Source::Pixel.new(v, v, v)
+          pixels[dst_row + x] = (v << 16) | (v << 8) | v
         end
       end
-      Image.new(width, height, pixels)
+      Image.new(width, height, pixels, icc_profile: @icc_profile)
     end
 
     def assemble_color(width, height, channels, components, max_h, max_v)
       # Upsample chroma channels if needed and convert YCbCr to RGB
-      y_ch  = channels[components[0].id]
-      cb_ch = channels[components[1].id]
-      cr_ch = channels[components[2].id]
+      y_comp, cb_comp, cr_comp = resolve_color_components(components)
 
-      cb_comp = components[1]
-      cr_comp = components[2]
+      y_ch = channels[y_comp.id]
+      cb_ch = channels[cb_comp.id]
+      cr_ch = channels[cr_comp.id]
 
       pixels = Array.new(width * height)
 
@@ -228,11 +541,20 @@ module PureJPEG
           g = g < 0 ? 0 : (g > 255 ? 255 : g)
           b = b < 0 ? 0 : (b > 255 ? 255 : b)
 
-          pixels[dst_row + px] = Source::Pixel.new(r, g, b)
+          pixels[dst_row + px] = (r << 16) | (g << 8) | b
         end
       end
 
-      Image.new(width, height, pixels)
+      Image.new(width, height, pixels, icc_profile: @icc_profile)
+    end
+
+    def resolve_color_components(components)
+      by_id = components.each_with_object({}) { |comp, memo| memo[comp.id] = comp }
+      if by_id[1] && by_id[2] && by_id[3]
+        [by_id[1], by_id[2], by_id[3]]
+      else
+        components
+      end
     end
   end
 end