libav-ruby 0.1.0

data/lib/libav/video_encoder.rb

@@ -0,0 +1,257 @@
+# frozen_string_literal: true
+
+module LibAV
+  class VideoEncoder
+    CODECS = {
+      h264: %w[libx264 h264],
+      hevc: %w[libx265 hevc],
+      vp9: %w[libvpx-vp9 vp9],
+      prores: %w[prores_ks prores_aw prores],
+      png: %w[png]
+    }.freeze
+    CODEC_IDS = {
+      h264: 27,
+      png: 61,
+      prores: 147,
+      vp9: 167,
+      hevc: 173
+    }.freeze
+    CRF_CODECS = %i[h264 hevc vp9].freeze
+    PRESET_CODECS = %i[h264 hevc].freeze
+    DEFAULT_PIXEL_FORMATS = {
+      png: :rgba,
+      prores: :yuv422p10le
+    }.freeze
+    PIXEL_FORMAT_ALIASES = {
+      prores: {
+        "rgba" => "yuva444p10le",
+        "yuv444p" => "yuv444p10le"
+      }
+    }.freeze
+    MAX_DIMENSION = 32_768
+    MAX_PIXELS = 134_217_728
+
+    attr_reader :stream, :time_base
+
+    def initialize(output, width:, height:, fps:, codec: :h264, bitrate: nil, crf: 18,
+                   preset: "medium", pixel_format: nil, gop: nil, **codec_options)
+      @output = output
+      @width = dimension(width, :width)
+      @height = dimension(height, :height)
+      validate_pixel_count!
+      @fps = positive_rational(fps, :fps)
+      @codec_name = codec.to_sym
+      @time_base = Rational(@fps.denominator, @fps.numerator)
+      @pixel_format = normalize_pixel_format(pixel_format)
+      validate_subsampled_dimensions!
+
+      @codec, @named_encoder = find_encoder
+      @stream = output.new_stream(@codec)
+      @context = Native.codec.avcodec_alloc_context3(@codec)
+      raise EncodingError, "avcodec_alloc_context3 returned NULL" if @context.null?
+
+      configure_context(bitrate:, gop:)
+      configure_private_options(crf:, preset:, codec_options:)
+      open_context
+      allocate_conversion
+      @frame_index = 0
+    rescue StandardError
+      close
+      raise
+    end
+
+    def encode(rgba)
+      raise ClosedError, "video encoder is closed" if @closed
+
+      Native.check(Native.util.av_frame_make_writable(@frame), "av_frame_make_writable")
+      source = FFI::MemoryPointer.from_string(rgba)
+      @source_data.put_pointer(0, source)
+      scaled = Native.scale.sws_scale(
+        @scale, @source_data, @source_stride, 0, @height,
+        Native::Layout.field_pointer(@frame, :frame, :data),
+        Native::Layout.field_pointer(@frame, :frame, :linesize)
+      )
+      Native.check(scaled, "sws_scale(video encode)") if scaled.negative?
+      raise EncodingError, "sws_scale converted #{scaled} rows, expected #{@height}" unless scaled == @height
+
+      Native::Layout.write_int64(@frame, :frame, :pts, @frame_index)
+      Native::Layout.write_int64(@frame, :frame, :duration, 1)
+      @frame_index += 1
+      send_frame(@frame)
+    end
+
+    def flush
+      return if @flushed || @closed
+
+      send_frame(nil)
+      @flushed = true
+    end
+
+    def close
+      return if @closed
+
+      Native.scale.sws_freeContext(@scale) if @scale && !@scale.null?
+      free_pointer(@frame, :av_frame_free)
+      free_pointer(@packet, :av_packet_free, Native.codec)
+      free_pointer(@context, :avcodec_free_context, Native.codec)
+      @scale = @frame = @packet = @context = nil
+      @closed = true
+    end
+
+    private
+
+    def configure_context(bitrate:, gop:)
+      pixel_format = Native.util.av_get_pix_fmt(@pixel_format)
+      raise ArgumentError, "unknown pixel format: #{@pixel_format}" if pixel_format.negative?
+
+      Native::Layout.write_int(@context, :codec_context, :width, @width)
+      Native::Layout.write_int(@context, :codec_context, :height, @height)
+      Native::Layout.write_int(@context, :codec_context, :pix_fmt, pixel_format)
+      Native::Layout.write_int(@context, :codec_context, :gop_size, gop || default_gop)
+      Native::Layout.write_rational(@context, :codec_context, :time_base, time_base)
+      Native::Layout.write_rational(@context, :codec_context, :framerate, @fps)
+      Native::Layout.write_int64(@context, :codec_context, :bit_rate, bitrate) if bitrate
+      Native::Layout.write_rational(stream, :stream, :time_base, time_base)
+
+      return unless @output.global_header?
+
+      flags = Native::Layout.read_int(@context, :codec_context, :flags)
+      Native::Layout.write_int(
+        @context, :codec_context, :flags, flags | Native::AV_CODEC_FLAG_GLOBAL_HEADER
+      )
+    end
+
+    def configure_private_options(crf:, preset:, codec_options:)
+      private_data = Native::Layout.read_pointer(@context, :codec_context, :priv_data)
+      set_option(private_data, "crf", crf) if
+        @named_encoder && CRF_CODECS.include?(@codec_name) && !crf.nil?
+      set_option(private_data, "preset", preset) if
+        @named_encoder && PRESET_CODECS.include?(@codec_name) && !preset.nil?
+      codec_options.each { |name, value| set_option(private_data, name, value) }
+    end
+
+    def open_context
+      Native.check(Native.codec.avcodec_open2(@context, @codec, nil), "avcodec_open2")
+      codec_parameters = Native::Layout.read_pointer(stream, :stream, :codecpar)
+      Native.check(
+        Native.codec.avcodec_parameters_from_context(codec_parameters, @context),
+        "avcodec_parameters_from_context"
+      )
+    end
+
+    def allocate_conversion
+      @frame = Native.util.av_frame_alloc
+      @packet = Native.codec.av_packet_alloc
+      raise EncodingError, "failed to allocate an AVFrame or AVPacket" if @frame.null? || @packet.null?
+
+      Native::Layout.write_int(@frame, :frame, :format,
+                               Native::Layout.read_int(@context, :codec_context, :pix_fmt))
+      Native::Layout.write_int(@frame, :frame, :width, @width)
+      Native::Layout.write_int(@frame, :frame, :height, @height)
+      Native.check(Native.util.av_frame_get_buffer(@frame, 0), "av_frame_get_buffer")
+
+      @scale = Native.scale.sws_getContext(
+        @width, @height, Native.util.av_get_pix_fmt("rgba"),
+        @width, @height, Native::Layout.read_int(@context, :codec_context, :pix_fmt),
+        Native::SWS_BICUBIC, nil, nil, nil
+      )
+      raise EncodingError, "sws_getContext returned NULL" if @scale.null?
+
+      @source_data = FFI::MemoryPointer.new(:pointer)
+      @source_stride = FFI::MemoryPointer.new(:int)
+      @source_stride.write_int(@width * 4)
+    end
+
+    def send_frame(frame)
+      result = Native.codec.avcodec_send_frame(@context, frame)
+      drain_packets if result == Native::AVERROR_EAGAIN
+      result = Native.codec.avcodec_send_frame(@context, frame) if result == Native::AVERROR_EAGAIN
+      Native.check(result, "avcodec_send_frame")
+      drain_packets
+    end
+
+    def drain_packets
+      loop do
+        result = Native.codec.avcodec_receive_packet(@context, @packet)
+        break if Native.control_flow?(result)
+
+        Native.check(result, "avcodec_receive_packet")
+        begin
+          packet_duration = Native::Layout.read_int64(@packet, :packet, :duration)
+          Native::Layout.write_int64(@packet, :packet, :duration, 1) unless packet_duration.positive?
+          @output.write_packet(@packet, time_base, stream)
+        ensure
+          Native.codec.av_packet_unref(@packet)
+        end
+      end
+    end
+
+    def find_encoder
+      names = CODECS.fetch(@codec_name) do
+        raise ArgumentError, "unsupported video codec: #{@codec_name.inspect}"
+      end
+      names.each do |name|
+        encoder = Native.codec.avcodec_find_encoder_by_name(name)
+        return [encoder, true] unless encoder.null?
+      end
+      encoder = Native.codec.avcodec_find_encoder(CODEC_IDS.fetch(@codec_name))
+      return [encoder, false] unless encoder.null?
+
+      raise EncodingError, "FFmpeg does not provide a #{names.join(" or ")} encoder"
+    end
+
+    def set_option(target, name, value)
+      raise ArgumentError, "#{@codec_name} encoder does not accept private options" if target.null?
+
+      option_name = name.to_s.tr("_", "-")
+      Native.check(
+        Native.util.av_opt_set(target, option_name, value.to_s, 0),
+        "av_opt_set(#{option_name})"
+      )
+    end
+
+    def dimension(value, name)
+      value = Integer(value)
+      return value if value.positive? && value <= MAX_DIMENSION
+
+      raise ArgumentError, "#{name} must be between 1 and #{MAX_DIMENSION}"
+    end
+
+    def positive_rational(value, name)
+      value = value.to_r
+      return value if value.positive?
+
+      raise ArgumentError, "#{name} must be positive"
+    end
+
+    def validate_subsampled_dimensions!
+      return unless @pixel_format == "yuv420p"
+      return if @width.even? && @height.even?
+
+      raise ArgumentError, "yuv420p requires even width and height"
+    end
+
+    def normalize_pixel_format(pixel_format)
+      requested = (pixel_format || DEFAULT_PIXEL_FORMATS.fetch(@codec_name, :yuv420p)).to_s
+      PIXEL_FORMAT_ALIASES.fetch(@codec_name, {}).fetch(requested, requested)
+    end
+
+    def validate_pixel_count!
+      return if @width * @height <= MAX_PIXELS
+
+      raise ArgumentError, "frame dimensions exceed the #{MAX_PIXELS}-pixel limit"
+    end
+
+    def default_gop
+      [(@fps * 2).round, 1].max
+    end
+
+    def free_pointer(pointer, function, library = Native.util)
+      return unless pointer && !pointer.null?
+
+      holder = FFI::MemoryPointer.new(:pointer)
+      holder.write_pointer(pointer)
+      library.public_send(function, holder)
+    end
+  end
+end

data/lib/libav/video_frame.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module LibAV
+  class VideoFrame
+    attr_reader :width, :height, :data, :pts
+
+    def initialize(width:, height:, data:, pts: nil)
+      @width = Integer(width)
+      @height = Integer(height)
+      @data = data.to_str
+      @data = @data.dup unless @data.encoding == Encoding::BINARY
+      @data.force_encoding(Encoding::BINARY)
+      @pts = pts
+    end
+
+    def to_str
+      data
+    end
+
+    def bytesize
+      data.bytesize
+    end
+  end
+end

data/lib/libav/video_reader.rb

@@ -0,0 +1,333 @@
+# frozen_string_literal: true
+
+module LibAV
+  class VideoReader
+    FRAME_TYPES = %i[video_frame texel].freeze
+
+    attr_reader :width, :height, :fps, :duration, :frame_count, :frame_type
+
+    def self.open(path, **options)
+      reader = new(path, **options)
+      return reader unless block_given?
+
+      begin
+        yield reader
+      ensure
+        reader.close
+      end
+    end
+
+    def initialize(path, frame_type: :video_frame)
+      @path = File.path(path)
+      @frame_type = normalize_frame_type(frame_type)
+      TexelSupport.load! if @frame_type == :texel
+      @format_holder = FFI::MemoryPointer.new(:pointer)
+      Native.check(
+        Native.format.avformat_open_input(@format_holder, @path, nil, nil),
+        "avformat_open_input"
+      )
+      @format_context = @format_holder.read_pointer
+      Native.check(
+        Native.format.avformat_find_stream_info(@format_context, nil),
+        "avformat_find_stream_info"
+      )
+      open_video_stream
+      allocate_decode_state
+      load_metadata
+      @reuse_buffer = false
+      reset_decode_state
+    rescue StandardError
+      close
+      raise
+    end
+
+    def read_frame
+      raise ClosedError, "video reader is closed" if closed?
+
+      loop do
+        return nil unless receive_next_frame
+
+        frame = convert_frame
+        next if before_seek_target?(@decoded_pts)
+
+        @seek_target = nil
+        @last_read_pts = @decoded_pts
+        return frame
+      end
+    end
+
+    def each_frame
+      return enum_for(__method__) unless block_given?
+
+      while (frame = read_frame)
+        yield frame, @last_read_pts
+      end
+      self
+    end
+
+    def seek(seconds)
+      raise ClosedError, "video reader is closed" if closed?
+
+      seconds = Float(seconds)
+      raise ArgumentError, "seek time must be nonnegative" if seconds.negative? || !seconds.finite?
+
+      timestamp = (seconds / @time_base).round
+      Native.check(
+        Native.format.av_seek_frame(
+          @format_context, @stream_index, timestamp, Native::AVSEEK_FLAG_BACKWARD
+        ),
+        "av_seek_frame"
+      )
+      Native.codec.avcodec_flush_buffers(@codec_context)
+      Native.util.av_frame_unref(@frame)
+      Native.codec.av_packet_unref(@packet)
+      reset_decode_state
+      @seek_target = seconds
+      self
+    end
+
+    def reuse_buffer
+      @reuse_buffer == true
+    end
+
+    def reuse_buffer=(enabled)
+      enabled = !!enabled
+      raise ArgumentError, "reuse_buffer is incompatible with immutable Texel::Image output" if
+        enabled && frame_type == :texel
+
+      @reuse_buffer = enabled
+    end
+
+    def close
+      return if closed?
+
+      Native.scale.sws_freeContext(@scale) if @scale && !@scale.null?
+      free_pointer(@frame, :av_frame_free, Native.util)
+      free_pointer(@packet, :av_packet_free, Native.codec)
+      free_pointer(@codec_context, :avcodec_free_context, Native.codec)
+      if @format_context && !@format_context.null?
+        @format_holder.write_pointer(@format_context)
+        Native.format.avformat_close_input(@format_holder)
+      end
+      @format_context = @codec_context = @frame = @packet = @scale = nil
+      @closed = true
+      nil
+    end
+
+    def closed?
+      @closed == true
+    end
+
+    private
+
+    def open_video_stream
+      decoder_holder = FFI::MemoryPointer.new(:pointer)
+      @stream_index = Native.format.av_find_best_stream(
+        @format_context, Native::AVMEDIA_TYPE_VIDEO, -1, -1, decoder_holder, 0
+      )
+      Native.check(@stream_index, "av_find_best_stream(video)")
+      decoder = decoder_holder.read_pointer
+      raise Error, "FFmpeg did not return a video decoder" if decoder.null?
+
+      @stream = stream_at(@stream_index)
+      @codec_context = Native.codec.avcodec_alloc_context3(decoder)
+      raise Error, "avcodec_alloc_context3 returned NULL" if @codec_context.null?
+
+      parameters = Native::Layout.read_pointer(@stream, :stream, :codecpar)
+      Native.check(
+        Native.codec.avcodec_parameters_to_context(@codec_context, parameters),
+        "avcodec_parameters_to_context"
+      )
+      Native.check(
+        Native.codec.avcodec_open2(@codec_context, decoder, nil),
+        "avcodec_open2"
+      )
+    end
+
+    def allocate_decode_state
+      @frame = Native.util.av_frame_alloc
+      @packet = Native.codec.av_packet_alloc
+      raise Error, "failed to allocate an AVFrame or AVPacket" if @frame.null? || @packet.null?
+
+      @destination_data = FFI::MemoryPointer.new(:pointer)
+      @destination_stride = FFI::MemoryPointer.new(:int)
+    end
+
+    def load_metadata
+      @width = Native::Layout.read_int(@codec_context, :codec_context, :width)
+      @height = Native::Layout.read_int(@codec_context, :codec_context, :height)
+      raise Error, "invalid decoded dimensions #{@width}x#{@height}" unless
+        @width.positive? && @height.positive?
+
+      @time_base = Native::Layout.read_rational(@stream, :stream, :time_base)
+      real_rate = Native::Layout.read_rational(@stream, :stream, :r_frame_rate)
+      rate = real_rate.positive? ? real_rate : Native::Layout.read_rational(
+        @stream, :stream, :avg_frame_rate
+      )
+      @fps = rate.positive? ? rate : nil
+
+      stream_duration = Native::Layout.read_int64(@stream, :stream, :duration)
+      @duration = if valid_timestamp?(stream_duration)
+                    stream_duration * @time_base
+                  else
+                    format_duration
+                  end
+      frames = Native::Layout.read_int64(@stream, :stream, :nb_frames)
+      @frame_count = if frames.positive?
+                       frames
+                     elsif @duration && @fps
+                       (@duration * @fps).round
+                     end
+    end
+
+    def receive_next_frame
+      loop do
+        result = Native.codec.avcodec_receive_frame(@codec_context, @frame)
+        return true if result.zero?
+        return false if result == Native::AVERROR_EOF
+
+        Native.check(result, "avcodec_receive_frame") unless result == Native::AVERROR_EAGAIN
+        return false unless feed_decoder
+      end
+    end
+
+    def feed_decoder
+      if @input_eof
+        return true if send_flush_packet
+
+        return false
+      end
+
+      loop do
+        result = Native.format.av_read_frame(@format_context, @packet)
+        if result == Native::AVERROR_EOF
+          @input_eof = true
+          return send_flush_packet
+        end
+        Native.check(result, "av_read_frame")
+
+        packet_stream = Native::Layout.read_int(@packet, :packet, :stream_index)
+        if packet_stream == @stream_index
+          begin
+            Native.check(
+              Native.codec.avcodec_send_packet(@codec_context, @packet),
+              "avcodec_send_packet"
+            )
+          ensure
+            Native.codec.av_packet_unref(@packet)
+          end
+          return true
+        end
+        Native.codec.av_packet_unref(@packet)
+      end
+    end
+
+    def send_flush_packet
+      return false if @flush_sent
+
+      result = Native.codec.avcodec_send_packet(@codec_context, nil)
+      Native.check(result, "avcodec_send_packet(flush)") unless result == Native::AVERROR_EOF
+      @flush_sent = true
+      true
+    end
+
+    def convert_frame
+      frame_width = Native::Layout.read_int(@frame, :frame, :width)
+      frame_height = Native::Layout.read_int(@frame, :frame, :height)
+      frame_format = Native::Layout.read_int(@frame, :frame, :format)
+      prepare_conversion(frame_width, frame_height, frame_format)
+
+      rows = Native.scale.sws_scale(
+        @scale,
+        Native::Layout.field_pointer(@frame, :frame, :data),
+        Native::Layout.field_pointer(@frame, :frame, :linesize),
+        0, frame_height, @destination_data, @destination_stride
+      )
+      Native.check(rows, "sws_scale(video decode)") if rows.negative?
+      raise Error, "sws_scale converted #{rows} rows, expected #{frame_height}" unless rows == frame_height
+
+      data = frame_data(@rgba_size)
+      timestamp = Native::Layout.read_int64(@frame, :frame, :best_effort_timestamp)
+      timestamp = Native::Layout.read_int64(@frame, :frame, :pts) unless valid_timestamp?(timestamp)
+      @decoded_pts = valid_timestamp?(timestamp) ? timestamp * @time_base : nil
+      build_frame(width: frame_width, height: frame_height, data:, pts: @decoded_pts)
+    end
+
+    def prepare_conversion(frame_width, frame_height, frame_format)
+      signature = [frame_width, frame_height, frame_format]
+      return if signature == @conversion_signature
+
+      Native.scale.sws_freeContext(@scale) if @scale && !@scale.null?
+      @scale = Native.scale.sws_getContext(
+        frame_width, frame_height, frame_format,
+        frame_width, frame_height, Native.util.av_get_pix_fmt("rgba"),
+        Native::SWS_BICUBIC, nil, nil, nil
+      )
+      raise Error, "sws_getContext returned NULL" if @scale.null?
+
+      @rgba_size = frame_width * frame_height * 4
+      @rgba_native = FFI::MemoryPointer.new(:uint8, @rgba_size)
+      @destination_data.put_pointer(0, @rgba_native)
+      @destination_stride.write_int(frame_width * 4)
+      @conversion_signature = signature
+      @width = frame_width
+      @height = frame_height
+    end
+
+    def frame_data(size)
+      decoded = @rgba_native.read_string_length(size)
+      return decoded unless reuse_buffer
+
+      @reused_data ||= String.new(encoding: Encoding::BINARY)
+      @reused_data.replace(decoded)
+    end
+
+    def build_frame(width:, height:, data:, pts:)
+      return TexelSupport.image(width:, height:, data:) if frame_type == :texel
+
+      VideoFrame.new(width:, height:, data:, pts:)
+    end
+
+    def normalize_frame_type(value)
+      value = value.to_sym
+      return value if FRAME_TYPES.include?(value)
+
+      raise ArgumentError, "frame_type must be :video_frame or :texel"
+    rescue NoMethodError
+      raise ArgumentError, "frame_type must be :video_frame or :texel"
+    end
+
+    def before_seek_target?(pts)
+      @seek_target && pts && pts < @seek_target
+    end
+
+    def reset_decode_state
+      @input_eof = false
+      @flush_sent = false
+    end
+
+    def stream_at(index)
+      streams = Native::Layout.read_pointer(@format_context, :format_context, :streams)
+      streams.get_pointer(index * FFI.type_size(:pointer))
+    end
+
+    def format_duration
+      value = Native::Layout.read_int64(@format_context, :format_context, :duration)
+      return unless valid_timestamp?(value)
+
+      Rational(value, Native::AV_TIME_BASE)
+    end
+
+    def valid_timestamp?(value)
+      value && value != Native::AV_NOPTS_VALUE && value >= 0
+    end
+
+    def free_pointer(pointer, function, library)
+      return unless pointer && !pointer.null?
+
+      holder = FFI::MemoryPointer.new(:pointer)
+      holder.write_pointer(pointer)
+      library.public_send(function, holder)
+    end
+  end
+end