jxl 1.0.0

data/lib/jxl/vardct/lf_group.rb

@@ -0,0 +1,227 @@
+# frozen_string_literal: true
+
+module JXL
+  module VarDCT
+    LFGroup = Data.define(:dc, :quant_dc, :strategies, :strategy_first, :quant_field, :sharpness,
+                          :ytox_map, :ytob_map, :used_strategies, :x, :y) do
+      def self.read(reader, frame, metadata, global, group_index: 0, group_count: 1, dc_frame: nil)
+        group_dim = 128 << frame.group_size_shift
+        groups_x = Num.ceil_div(frame.blocks_width, group_dim)
+        x = (group_index % groups_x) * group_dim
+        y = (group_index / groups_x) * group_dim
+        width = [group_dim, frame.blocks_width - x].min
+        height = [group_dim, frame.blocks_height - y].min
+        dc, quant_dc = if frame.flags.anybits?(32)
+                         dc_from_frame(dc_frame, frame, width, height, x, y)
+                       else
+                         read_dc(reader, frame, metadata, global, group_index, width, height)
+                       end
+        if global.modular_group
+          Modular::Decoder.decode_grouped_dc!(reader, global.modular_group,
+                                              group_index:, groups_x:, stream_id: 1 + group_count + group_index)
+        end
+        stream_id = 1 + (2 * group_count) + group_index
+        metadata_result = read_metadata(reader, width, height, metadata, global.modular_state, stream_id)
+        build(dc, quant_dc, metadata_result.channels, width, height, x, y)
+      end
+
+      def self.read_dc(reader, frame, metadata, global, group_index, width, height)
+        extra_precision = reader.read(2)
+        shifts = [frame.chroma_subsampling[1], frame.chroma_subsampling[0], frame.chroma_subsampling[2]]
+        shapes = shifts.map { |hshift, vshift| [width >> hshift, height >> vshift] }
+        result = Modular::Stream.decode(reader, shapes,
+                                        state: global.modular_state, stream_id: 1 + group_index, metadata:)
+        dequantize_dc(result.channels, global, extra_precision, frame)
+      end
+      private_class_method :read_dc
+
+      def self.dc_from_frame(image, frame, width, height, x, y)
+        raise CorruptError, "missing progressive DC frame" unless image
+
+        channels = image.channels.first(3)
+        raise CorruptError, "invalid progressive DC frame" unless channels.length == 3
+
+        dc = frame.chroma_subsampling.map.with_index do |(hshift, vshift), channel|
+          source = channels.fetch(channel)
+          plane = Plane.new(width >> hshift, height >> vshift)
+          plane.height.times do |row|
+            source_x = x >> hshift
+            source_y = (y >> vshift) + row
+            plane.data[row * plane.width, plane.width] = source.data[(source_y * source.width) + source_x, plane.width]
+          end
+          plane
+        end
+        [dc, Plane.new(width, height)]
+      end
+      private_class_method :dc_from_frame
+
+      def self.dequantize_dc(channels, global, extra_precision, frame)
+        steps = global.quantizer.dc_steps(global.dequant).map { _1 / (1 << extra_precision) }
+        factors = global.chroma_from_luma.dc_factors
+        output = [channels[1], channels[0], channels[2]].map { Plane.new(_1.width, _1.height) }
+        unless frame.max_hshift.zero? && frame.max_vshift.zero?
+          channels[0].data.each_index { |index| output[1].data[index] = channels[0].data[index] * steps[1] }
+          channels[1].data.each_index { |index| output[0].data[index] = channels[1].data[index] * steps[0] }
+          channels[2].data.each_index { |index| output[2].data[index] = channels[2].data[index] * steps[2] }
+          return [output, dc_contexts(channels, global.block_context_map, frame)]
+        end
+
+        channels[0].data.length.times do |index|
+          y = channels[0].data[index] * steps[1]
+          output[1].data[index] = y
+          output[0].data[index] = (channels[1].data[index] * steps[0]) + (y * factors[0])
+          output[2].data[index] = (channels[2].data[index] * steps[2]) + (y * factors[2])
+        end
+        [output, dc_contexts(channels, global.block_context_map, frame)]
+      end
+      private_class_method :dequantize_dc
+
+      def self.dc_contexts(channels, map, frame)
+        plane = Plane.new(channels[0].width, channels[0].height)
+        return plane if map.num_dc_contexts == 1
+
+        plane.data.length.times do |index|
+          x = index % plane.width
+          y = index / plane.width
+          buckets = [channels[1], channels[0], channels[2]].each_with_index.map do |channel, c|
+            hshift, vshift = frame.chroma_subsampling[c]
+            value = channel[x >> hshift, y >> vshift]
+            map.dc_thresholds[c].count { value > _1 }
+          end
+          plane.data[index] = (((buckets[0] * (map.dc_thresholds[2].length + 1)) + buckets[2]) *
+                               (map.dc_thresholds[1].length + 1)) + buckets[1]
+        end
+        plane
+      end
+      private_class_method :dc_contexts
+
+      def self.read_metadata(reader, width, height, metadata, state, stream_id)
+        count = reader.read(((width * height) - 1).bit_length) + 1
+        shapes = [[Num.ceil_div(width, 8), Num.ceil_div(height, 8)],
+                  [Num.ceil_div(width, 8), Num.ceil_div(height, 8)], [count, 2], [width, height]]
+        Modular::Stream.decode(reader, shapes, state:, stream_id:, metadata:)
+      end
+      private_class_method :read_metadata
+
+      def self.build(dc, quant_dc, channels, width, height, x = 0, y = 0)
+        strategies = Array.new(width * height)
+        strategy_first = Array.new(width * height, false)
+        quant_field = Plane.new(width, height)
+        index = 0
+        height.times do |y|
+          width.times do |x|
+            next if strategies[(y * width) + x]
+
+            strategy = AcStrategy[channels[2][index, 0]]
+            place_strategy!(strategies, width, height, x, y, strategy)
+            strategy_first[(y * width) + x] = true
+            quant_field[x, y] = 1 + Num.clamp(channels[2][index, 1], 0, 255)
+            index += 1
+          end
+        end
+        raise CorruptError, "invalid AC metadata count" unless index == channels[2].width
+        raise CorruptError, "invalid EPF sharpness" unless channels[3].data.all? { _1.between?(0, 7) }
+
+        new(dc:, quant_dc:, strategies:, strategy_first:, quant_field:, sharpness: channels[3],
+            ytox_map: channels[0], ytob_map: channels[1],
+            used_strategies: strategies.compact.map(&:id).uniq.sort.freeze, x:, y:)
+      end
+      private_class_method :build
+
+      def self.merge(groups, frame)
+        return groups.first if groups.one?
+
+        width = frame.blocks_width
+        height = frame.blocks_height
+        dc = frame.chroma_subsampling.map do |hshift, vshift|
+          Plane.new(width >> hshift, height >> vshift)
+        end
+        quant_dc = Plane.new(width, height)
+        strategies = Array.new(width * height)
+        strategy_first = Array.new(width * height, false)
+        quant_field = Plane.new(width, height)
+        sharpness = Plane.new(width, height)
+        ytox_map = Plane.new(Num.ceil_div(width, 8), Num.ceil_div(height, 8))
+        ytob_map = Plane.new(ytox_map.width, ytox_map.height)
+        groups.each do |group|
+          group.dc.each_with_index do |plane, channel|
+            hshift, vshift = frame.chroma_subsampling[channel]
+            blit_plane!(plane, dc[channel], group.x >> hshift, group.y >> vshift)
+          end
+          blit_plane!(group.quant_dc, quant_dc, group.x, group.y)
+          blit_plane!(group.quant_field, quant_field, group.x, group.y)
+          blit_plane!(group.sharpness, sharpness, group.x, group.y)
+          blit_plane!(group.ytox_map, ytox_map, group.x / 8, group.y / 8)
+          blit_plane!(group.ytob_map, ytob_map, group.x / 8, group.y / 8)
+          group.strategies.each_with_index do |strategy, index|
+            local_x = index % group.quant_field.width
+            local_y = index / group.quant_field.width
+            target = ((group.y + local_y) * width) + group.x + local_x
+            strategies[target] = strategy
+            strategy_first[target] = group.strategy_first[index]
+          end
+        end
+        new(dc:, quant_dc:, strategies:, strategy_first:, quant_field:, sharpness:, ytox_map:, ytob_map:,
+            used_strategies: groups.flat_map(&:used_strategies).uniq.sort.freeze, x: 0, y: 0)
+      end
+
+      def self.smooth!(group, dc_factors)
+        width = group.dc.first.width
+        height = group.dc.first.height
+        return group if width <= 2 || height <= 2
+
+        w1 = 0.20345139757231578
+        w2 = 0.0334829185968739
+        w0 = 1.0 - (4.0 * (w1 + w2))
+        output = group.dc.map do |source|
+          Plane.new(width, height).tap { |target| target.data.replace(source.data) }
+        end
+        1.upto(height - 2) do |y|
+          1.upto(width - 2) do |x|
+            centres = group.dc.map { _1[x, y] }
+            smoothed = group.dc.map do |plane|
+              corners = plane[x - 1, y - 1] + plane[x + 1, y - 1] +
+                        plane[x - 1, y + 1] + plane[x + 1, y + 1]
+              sides = plane[x - 1, y] + plane[x + 1, y] + plane[x, y - 1] + plane[x, y + 1]
+              (w2 * corners) + (w1 * sides) + (w0 * plane[x, y])
+            end
+            gap = 3.times.reduce(0.5) do |value, channel|
+              [value, ((centres[channel] - smoothed[channel]) /
+                                                             dc_factors[channel]).abs].max
+            end
+            factor = [0.0, 3.0 - (4.0 * gap)].max
+            3.times do |channel|
+              output[channel][x, y] = centres[channel] +
+                                      ((smoothed[channel] - centres[channel]) * factor)
+            end
+          end
+        end
+        group.dc.each_index { |channel| group.dc[channel].data.replace(output[channel].data) }
+        group
+      end
+
+      def self.blit_plane!(source, target, x, y)
+        source.height.times do |row|
+          target.data[((y + row) * target.width) + x, source.width] = source.data[row * source.width, source.width]
+        end
+      end
+      private_class_method :blit_plane!
+
+      def self.place_strategy!(grid, width, height, x, y, strategy)
+        if x + strategy.blocks_x > width || y + strategy.blocks_y > height
+          raise CorruptError, "AC strategy exceeds group bounds"
+        end
+
+        strategy.blocks_y.times do |dy|
+          strategy.blocks_x.times do |dx|
+            index = ((y + dy) * width) + x + dx
+            raise CorruptError, "overlapping AC strategies" if grid[index]
+
+            grid[index] = strategy
+          end
+        end
+      end
+      private_class_method :place_strategy!
+    end
+  end
+end

data/lib/jxl/vardct/pass_group.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module JXL
+  module VarDCT
+    COEFFICIENT_FREQUENCY_CONTEXT = [
+      0xBAD, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
+      15, 15, 16, 16, 17, 17, 18, 18, 19, 19, 20, 20, 21, 21, 22, 22,
+      23, 23, 23, 23, 24, 24, 24, 24, 25, 25, 25, 25, 26, 26, 26, 26,
+      27, 27, 27, 27, 28, 28, 28, 28, 29, 29, 29, 29, 30, 30, 30, 30
+    ].freeze
+    COEFFICIENT_NONZERO_CONTEXT = [
+      0xBAD, 0, 31, 62, 62, 93, 93, 93, 93, 123, 123, 123, 123,
+      152, 152, 152, 152, 152, 152, 152, 152, 180, 180, 180, 180, 180,
+      180, 180, 180, 180, 180, 180, 180, 206, 206, 206, 206, 206, 206,
+      206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206,
+      206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206
+    ].freeze
+    QuantizedBlock = Data.define(:x, :y, :strategy, :coefficients)
+    PassGroup = Data.define(:blocks) do
+      def self.merge(groups)
+        return groups.first if groups.one?
+
+        blocks = groups.first.blocks.each_index.map do |index|
+          pass_blocks = groups.map { _1.blocks.fetch(index) }
+          first = pass_blocks.first
+          unless pass_blocks.all? { _1.x == first.x && _1.y == first.y && _1.strategy == first.strategy }
+            raise CorruptError, "progressive AC blocks differ between passes"
+          end
+
+          coefficients = 3.times.map do |channel|
+            first.coefficients[channel].each_index.map do |coefficient|
+              pass_blocks.sum { _1.coefficients[channel][coefficient] }
+            end
+          end
+          QuantizedBlock.new(x: first.x, y: first.y, strategy: first.strategy, coefficients:)
+        end
+        new(blocks:)
+      end
+
+      def self.read(reader, frame, lf_global, lf_group, hf_global, pass_index: 0, group_index: 0)
+        pass = hf_global.passes.fetch(pass_index)
+        selector_bits = (hf_global.histogram_count - 1).bit_length
+        selector = selector_bits.zero? ? 0 : reader.read(selector_bits)
+        raise CorruptError, "invalid AC histogram selector" if selector >= hf_global.histogram_count
+
+        context_offset = selector * lf_global.block_context_map.ac_context_count
+        decoder = pass.entropy.fork(reader, distance_multiplier: 0)
+        blocks = decode_blocks(decoder, frame, lf_global.block_context_map, lf_group,
+                               pass.orders, context_offset, frame.passes.shifts.fetch(pass_index), group_index)
+        decoder.final_state!
+        if lf_global.modular_group
+          group_dim = 128 << frame.group_size_shift
+          groups_x = Num.ceil_div(frame.encoded_width, group_dim)
+          dc_groups = Num.ceil_div(frame.blocks_width, group_dim) * Num.ceil_div(frame.blocks_height, group_dim)
+          group_count = groups_x * Num.ceil_div(frame.encoded_height, group_dim)
+          stream_id = 1 + (3 * dc_groups) + 17 + (group_count * pass_index) + group_index
+          Modular::Decoder.decode_grouped_pass!(reader, lf_global.modular_group, frame.passes,
+                                                pass_index:, group_index:, groups_x:, stream_id:)
+        end
+        new(blocks:)
+      end
+
+      def self.decode_blocks(decoder, frame, context_map, lf_group, orders, context_offset, shift, group_index)
+        frame_width = frame.blocks_width
+        frame_height = frame.blocks_height
+        group_dim = 128 << frame.group_size_shift
+        groups_x = Num.ceil_div(frame.encoded_width, group_dim)
+        block_group_dim = group_dim / 8
+        x0 = (group_index % groups_x) * block_group_dim
+        y0 = (group_index / groups_x) * block_group_dim
+        width = [block_group_dim, frame_width - x0].min
+        height = [block_group_dim, frame_height - y0].min
+        nonzeros = frame.chroma_subsampling.map do |hshift, vshift|
+          Plane.new(width >> hshift, height >> vshift)
+        end
+        blocks = []
+        height.times do |y|
+          width.times do |x|
+            absolute_x = x0 + x
+            absolute_y = y0 + y
+            next unless lf_group.strategy_first[(absolute_y * frame_width) + absolute_x]
+
+            strategy = lf_group.strategies[(absolute_y * frame_width) + absolute_x]
+            coefficients = Array.new(3) { Array.new(strategy.coefficient_count, 0) }
+            [1, 0, 2].each do |channel|
+              hshift, vshift = frame.chroma_subsampling[channel]
+              next unless absolute_x.nobits?((1 << hshift) - 1) && absolute_y.nobits?((1 << vshift) - 1)
+
+              decode_channel!(decoder, context_map, lf_group, nonzeros[channel], coefficients[channel],
+                              orders.fetch([strategy.order, channel]), strategy, channel, x >> hshift, y >> vshift,
+                              absolute_x, absolute_y, context_offset, shift)
+            end
+            blocks << QuantizedBlock.new(x: absolute_x, y: absolute_y, strategy:, coefficients:)
+          end
+        end
+        blocks
+      end
+      private_class_method :decode_blocks
+
+      def self.decode_channel!(decoder, map, lf_group, nonzeros, coefficients, order,
+                               strategy, channel, x, y, absolute_x, absolute_y, context_offset, shift)
+        predicted = predict_nonzeros(nonzeros, x, y)
+        block_context = map.context(lf_group.quant_dc[absolute_x, absolute_y],
+                                    lf_group.quant_field[absolute_x, absolute_y], strategy.order, channel)
+        count = decoder.read_uint(context_offset + map.nonzero_context(predicted, block_context))
+        covered = strategy.blocks_x * strategy.blocks_y
+        raise CorruptError, "too many nonzero AC coefficients" if count > strategy.coefficient_count - covered
+
+        average = (count + covered - 1) >> strategy.log2_blocks
+        strategy.blocks_y.times do |dy|
+          strategy.blocks_x.times { |dx| nonzeros[x + dx, y + dy] = average }
+        end
+        decode_coefficients!(decoder, map, coefficients, order, covered, strategy.log2_blocks,
+                             block_context, context_offset, count, shift)
+      end
+      private_class_method :decode_channel!
+
+      def self.decode_coefficients!(decoder, map, coefficients, order, covered, log2_blocks,
+                                    block_context, context_offset, count, shift)
+        previous = count > coefficients.length / 16 ? 0 : 1
+        k = covered
+        while k < coefficients.length && count.positive?
+          context = coefficient_context(count, k, covered, log2_blocks, previous)
+          packed = decoder.read_uint(context_offset + map.coefficient_context_offset(block_context) + context)
+          coefficients[order[k]] += Num.unpack_signed(packed) << shift
+          previous = packed.zero? ? 0 : 1
+          count -= previous
+          k += 1
+        end
+        raise CorruptError, "nonzero AC coefficients exceed scan" unless count.zero?
+      end
+      private_class_method :decode_coefficients!
+
+      def self.predict_nonzeros(plane, x, y)
+        return 32 if x.zero? && y.zero?
+        return plane[x, y - 1] if x.zero?
+        return plane[x - 1, y] if y.zero?
+
+        (plane[x, y - 1] + plane[x - 1, y] + 1) / 2
+      end
+      private_class_method :predict_nonzeros
+
+      def self.coefficient_context(nonzeros, k, covered, log2_blocks, previous)
+        nonzeros = (nonzeros + covered - 1) >> log2_blocks
+        k >>= log2_blocks
+        raise CorruptError, "invalid AC coefficient context" unless nonzeros.between?(1, 63) && k.between?(1, 63)
+
+        ((COEFFICIENT_NONZERO_CONTEXT[nonzeros] + COEFFICIENT_FREQUENCY_CONTEXT[k]) * 2) + previous
+      end
+      private_class_method :coefficient_context
+    end
+  end
+end

data/lib/jxl/vardct/quantizer.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module JXL
+  module VarDCT
+    QUANTIZER_SCALE_DENOMINATOR = 1 << 16
+    QUANTIZER_SCALE_DISTRIBUTION = [Bit::Field.bits_offset(11, 1), Bit::Field.bits_offset(11, 2049),
+                                    Bit::Field.bits_offset(12, 4097), Bit::Field.bits_offset(16, 8193)].freeze
+    QUANTIZER_DC_DISTRIBUTION = [Bit::Field.val(16), Bit::Field.bits_offset(5, 1),
+                                 Bit::Field.bits_offset(8, 1), Bit::Field.bits_offset(16, 1)].freeze
+
+    Quantizer = Data.define(:global_scale, :quant_dc) do
+      def self.read(reader)
+        new(global_scale: Bit::Field.u32(reader, *QUANTIZER_SCALE_DISTRIBUTION),
+            quant_dc: Bit::Field.u32(reader, *QUANTIZER_DC_DISTRIBUTION))
+      end
+
+      def scale = global_scale.fdiv(QUANTIZER_SCALE_DENOMINATOR)
+      def inverse_global_scale = QUANTIZER_SCALE_DENOMINATOR.fdiv(global_scale)
+      def inverse_dc = inverse_global_scale / quant_dc
+      def inverse_ac(quant) = inverse_global_scale / quant
+
+      def dc_steps(dequant)
+        dequant.dc.map { |value| inverse_dc * value }
+      end
+    end
+  end
+end