Ruby PNG Reader

r-png.rb

require 'zlib'

module RImage
  class InvalidSignature < Exception
  end

  class InvalidChecksum < Exception
  end

  class RImage
    FORMAT_PNG = 1

    FILE_SIGNATURE_PNG = [137, 80, 78, 71, 13, 10, 26, 10]

    KNOWN_FILE_EXTENSIONS = {
      'PNG' => FORMAT_PNG
    }

    # Creates an RImage object from a filesystem node
    def self.from_file(filename)
      file_contents = open(filename, "rb") {|io| io.read } 
      format = self.format_from_filename(filename)
      self.factory(file_contents, format)
    end

    # Gueses the file type from a file's name, using its extension
    def self.format_from_filename(filename)
      if (filename =~ /\.([^\.]+)$/)
        self.format_from_extension($1)
      end
    end

    # Guesses the file type from a file's extension
    def self.format_from_extension(extension)
      KNOWN_FILE_EXTENSIONS[extension.upcase]
    end

    # Guesses the file type by examining the first few bytes of a file
    def self.format_from_file_contents(contents)
      if contents[0,8].unpack("C*") == self::FILE_SIGNATURE_PNG
        self::FORMAT_PNG
      end
    end

    # Factory method for creating RImageXyz objects
    def self.factory(data, format)
      if format == self::FORMAT_PNG
        RImagePng.new(data, format)
      end
    end

    def copy
      self.dup
    end

    def to_file(filename)
      file_handle = open(filename, "w")
      file_handle.write @contents
      file_handle.close
    end

    def contents
      @contents
    end

    # Simply compares the contents of two RImage objects
    def ==(other_rimage)
      @contents == other_rimage.contents
    end

    def initialize(data, format)
      @format    = format
      @contents  = data
      @play_head = 0
      verify_signature!
    end

    def read_bytes(count)
      bytes = @contents[@play_head, count]
      @play_head += count
      bytes
    end
  end

  class RImageUncompressed < RImage
    CHANNEL_GREY  = 0
    CHANNEL_RED   = 1
    CHANNEL_GREEN = 2
    CHANNEL_BLUE  = 3
    CHANNEL_ALPHA = 4

    def initialize(bit_depth, color_channels, width, height)
      @bit_depth = bit_depth
      @width     = width
      @height    = height
      @contents  = {}
      color_channels.each do |channel|
        @contents[channel] = []
      end
      @play_head = 0
    end

    def contents
      @contents
    end

    # Get an arbitrary line by its index (first line is zero, last is -1...)
    def line(index)
      ret = {}
      @contents.keys.each do |channel|
        ret[channel] = @contents[channel][index * @width, @width]
      end
      ret
    end
  end

  class RImagePng < RImage

    FILTER_TYPES = [
      "NONE",
      "SUB",
      "UP",
      "AVERAGE",
      "PAETH",
    ]

    def verify_signature!
      signature = read_bytes(8).unpack('C*')
      if signature != FILE_SIGNATURE_PNG
        raise InvalidSignature.new
      end
    end

    def decode
      chunks do |chunk|
        if chunk[:type] == 'IDAT'
          d = read_chunk_IDAT(chunk[:data])
          # buffer += d
          # while (bufferlines > 1 || last_byte_read)
          #   process buffer
          # end
        else
          send "read_chunk_%s" % chunk[:type], chunk[:data]
        end
      end
      uncompressed
    end

    def uncompressed
      @uncompressed
    end

    def read_chunk
      length   = read_bytes(4).unpack('N')[0]
      type     = read_bytes(4).unpack('M')[0]
      data     = read_bytes(length)
      checksum = read_bytes(4).unpack('N')[0]

      if checksum != Zlib.crc32(data, Zlib.crc32(type))
        raise InvalidChecksum.new
      end
      {:type => type, :data => data}
    end

    def self.color_channels_from_color_type(color_type)
      case color_type
      when 0
        [
          RImageUncompressed::CHANNEL_GREYSCALE
        ]
      when 2..3
        [
          RImageUncompressed::CHANNEL_RED,
          RImageUncompressed::CHANNEL_GREEN,
          RImageUncompressed::CHANNEL_BLUE,
        ]
      when 4
        [
          RImageUncompressed::CHANNEL_GREYSCALE,
          RImageUncompressed::CHANNEL_ALPHA,
        ]
      when 6
        [
          RImageUncompressed::CHANNEL_RED,
          RImageUncompressed::CHANNEL_GREEN,
          RImageUncompressed::CHANNEL_BLUE,
          RImageUncompressed::CHANNEL_ALPHA,
        ]
      end
    end

    def self.read_chunk_IHDR(data)
      @meta = { 
        :width              => data[0,4].unpack('N')[0],
        :height             => data[4,4].unpack('N')[0],
        :bit_depth          => data[8],
        :color_type         => data[9],
        :compression_method => data[10],
        :filter_method      => data[11],
        :interlace_method   => data[12],
      }

      @meta
    end

    def read_chunk_IHDR(data)
      @meta         = self.class.read_chunk_IHDR(data)
      @uncompressed = RImageUncompressed.new(
        @meta[:bit_depth],
        self.class.color_channels_from_color_type(@meta[:color_type]),
        @meta[:width],
        @meta[:height]
      )
    end

    def self.read_chunk_IDAT(data, bit_depth, channels, width, previous_bytes)
      data_chunk      = Zlib::Inflate.inflate(data)
      line_bytes      = 1 + 8.0 / bit_depth * channels.length * width
      lines_left      = (data_chunk.length / line_bytes).ceil
      scanlines       = channels.collect { |channel| { channel => [] } }
      
      while (lines_left > 0) do
        line_position   = data_chunk.length - lines_left * line_bytes
        line            = data_chunk[line_position, line_bytes]
			  unfilter_method = "unfilter_%s" % FILTER_TYPES[line[0]]
        line            = line[1, line.length - 1]
       
        # call the unfiltering method and return a raw binary scanline
				scanline = send(
          unfilter_method, 
          line, 
          width, 
          channels, 
          previous_bytes)

        # index every channel/point and divide the values into arrays according
        # to the image's bit depth 
        points = []
				if    bit_depth < 8
					bitmask = bit_depth ^ 2 - 1
					line.each_byte do |byte|
						n = 0
						while(n < 8) do
							points.push(byte & (bitmask << n) >> n)
							n += bit_depth
						end
					end
				elsif bit_depth == 8
					points = line.unpack("C*")
				elsif bit_depth == 16
					points = line.unpack("n*")
				end

        # Add all the scanline pixels to the scanlines array
        scanline.each do |channel|
          scanlines[channel] += previous_scanline[channel]
        end

        # If we have a full scanline, save it in case the next line's filter
        # requires it
        if scanline.first.length == width
          previous_scanline = scanline
        end

        lines_left -= 1
      end
      scanlines
    end

    def read_chunk_IDAT(data)
      previous_bytes = @uncompressed.line(-1)
      if previous_bytes.first.length < @meta[:width]
        previous_bytes = @uncompressed.line(-2) + previous_bytes
      end

      self.class.read_chunk_IDAT(
        data, 
        @meta[:bit_depth], 
        @uncompressed.contents.keys, 
        @meta[:width],
        previous_bytes)
    end

    def unfilter_NONE(points)
      n = 0
      @uncompressed.contents.keys.cycle do |channel|
        break if n >= points.length
        @uncompressed.contents[channel].push(points[n])
        n += 1
      end
    end

    def self.unfilter_SUB(points, width, channels, scanline_position)
      n = 0
      unfiltered_pixels = channels.collect { |c| { c => [] } }
      channels.cycle do |channel|
        break if n >= points.length
        a = (unfiltered_pixels[channel].length + scanline_position) % width == 0 ? 0 : unfiltered_pixels[channel][-1]
        unfiltered_pixels[channel].push((points[n] + (a || 0)) % 256)
        n += 1
      end
      unfiltered_pixels
    end

    def unfilter_SUB(points)
      scanline_position = @uncompressed.contents[channel].length
      self.class.unfilter_SUB(points, @meta[:width], @uncompressed.contents.keys, scanline_position)
    end

    def unfilter_UP(points)
      n = 0
      @uncompressed.contents.keys.cycle do |channel|
        break if n >= points.length
        b = @uncompressed.contents[channel].length < @meta[:width] ? 0 : @uncompressed.contents[channel][-@meta[:width]]
        @uncompressed.contents[channel].push((points[n] + (b || 0)) % 256)
        n += 1
      end
    end

    def chunks
      while @play_head < @contents.length
        if block_given?
          yield read_chunk
        else
          chunk_array = []
          chunks do |chunk|
            chunk_array.push chunk
          end
        end
      end
      chunk_array
    end

    def method_missing(symbol, *args)
      if symbol.to_s =~ /read_chunk_(.+)/
        #puts "Unimplemented handler for chunk of type '#{$1}'"
      else
        raise NoMethodError.new(symbol.to_s)
      end
    end
  end
end