torchcodec 0.1.0

data/ext/torchcodec/ext.cpp

@@ -0,0 +1,12 @@
+#include <rice/rice.hpp>
+
+void init_core(Rice::Module m);
+
+extern "C"
+void Init_ext() {
+  auto rb_mTorchCodec = Rice::define_module("TorchCodec");
+
+  auto rb_mCore = Rice::define_module_under(rb_mTorchCodec, "Core");
+
+  init_core(rb_mCore);
+}

data/ext/torchcodec/extconf.rb

@@ -0,0 +1,73 @@
+require "mkmf-rice"
+
+$CXXFLAGS += " -std=c++17 $(optflags)"
+
+paths = []
+if RbConfig::CONFIG["host_os"] =~ /darwin/i
+  if RbConfig::CONFIG["host_cpu"] =~ /arm|aarch64/i
+    paths << "/opt/homebrew"
+  else
+    paths << "/usr/local"
+  end
+elsif !Gem.win_platform?
+  paths << "/home/linuxbrew/.linuxbrew"
+end
+
+inc, lib = dir_config("torch")
+inc ||= paths.map { |v| "#{v}/opt/pytorch/include" }.find { |v| Dir.exist?("#{v}/torch") }
+lib ||= paths.map { |v| "#{v}/opt/pytorch/lib" }.find { |v| Dir["#{v}/*torch_cpu*"].any? }
+
+unless inc && lib
+  abort "LibTorch not found"
+end
+
+cuda_inc, cuda_lib = dir_config("cuda")
+cuda_lib ||= "/usr/local/cuda/lib64"
+
+cudnn_inc, cudnn_lib = dir_config("cudnn")
+cudnn_lib ||= "/usr/local/cuda/lib"
+
+$LDFLAGS += " -L#{lib}" if Dir.exist?(lib)
+abort "LibTorch not found" unless have_library("torch")
+
+have_library("mkldnn")
+have_library("nnpack")
+
+with_cuda = false
+if Dir["#{lib}/*torch_cuda*"].any?
+  $LDFLAGS += " -L#{cuda_lib}" if Dir.exist?(cuda_lib)
+  $LDFLAGS += " -L#{cudnn_lib}" if Dir.exist?(cudnn_lib) && cudnn_lib != cuda_lib
+  with_cuda = have_library("cuda") && have_library("cudnn")
+end
+
+$INCFLAGS += " -I#{inc}"
+$INCFLAGS += " -I#{inc}/torch/csrc/api/include"
+
+$LDFLAGS += " -Wl,-rpath,#{lib}"
+$LDFLAGS += ":#{cuda_lib}/stubs:#{cuda_lib}" if with_cuda
+
+# https://github.com/pytorch/pytorch/blob/v2.9.0/torch/utils/cpp_extension.py#L1351-L1364
+$LDFLAGS += " -lc10 -ltorch_cpu -ltorch"
+if with_cuda
+  $LDFLAGS += " -lcuda -lnvrtc"
+  $LDFLAGS += " -lnvToolsExt" if File.exist?("#{cuda_lib}/libnvToolsExt.so")
+  $LDFLAGS += " -lcudart -lc10_cuda -ltorch_cuda -lcufft -lcurand -lcublas -lcudnn"
+  # TODO figure out why this is needed
+  $LDFLAGS += " -Wl,--no-as-needed,#{lib}/libtorch.so"
+end
+
+ffmpeg_inc, ffmpeg_lib = dir_config("ffmpeg")
+ffmpeg_inc ||= paths.map { |v| "#{v}/opt/ffmpeg/include" }.find { |v| File.exist?("#{v}/libavcodec/avcodec.h") }
+ffmpeg_lib ||= paths.map { |v| "#{v}/opt/ffmpeg/lib" }.find { |v| Dir["#{v}/*libavcodec*"].any? }
+
+ffmpeg_incs = ffmpeg_inc ? [ffmpeg_inc] : []
+ffmpeg_libs = ffmpeg_lib ? [ffmpeg_lib] : []
+
+abort "libavcodec not found" unless find_header("libavcodec/avcodec.h", *ffmpeg_incs)
+abort "libavcodec not found" unless find_library("avcodec", nil, *ffmpeg_libs)
+abort "libavdevice not found" unless find_library("avdevice", nil, *ffmpeg_libs)
+abort "libavfilter not found" unless find_library("avfilter", nil, *ffmpeg_libs)
+abort "libavutil not found" unless find_library("avutil", nil, *ffmpeg_libs)
+
+# create makefile
+create_makefile("torchcodec/ext")

data/lib/torchcodec/core/metadata.rb

@@ -0,0 +1,41 @@
+module TorchCodec
+  module Core
+    def self.get_container_metadata(decoder)
+      container_dict = JSON.parse(_get_container_json_metadata(decoder))
+      streams_metadata = []
+
+      container_dict["numStreams"].times do |stream_index|
+        stream_dict = JSON.parse(_get_stream_json_metadata(decoder, stream_index))
+        common_meta = {
+          duration_seconds_from_header: stream_dict["durationSecondsFromHeader"],
+          duration_seconds: stream_dict["durationSeconds"],
+          bit_rate: stream_dict["bitRate"],
+          begin_stream_seconds_from_header: stream_dict["beginStreamSecondsFromHeader"],
+          begin_stream_seconds: stream_dict["beginStreamSeconds"],
+          codec: stream_dict["codec"],
+          stream_index: stream_index
+        }
+        if stream_dict["mediaType"] == "video"
+          raise Todo
+        elsif stream_dict["mediaType"] == "audio"
+          streams_metadata << {
+            sample_rate: stream_dict["sampleRate"],
+            num_channels: stream_dict["numChannels"],
+            sample_format: stream_dict["sampleFormat"],
+            **common_meta
+          }
+        else
+          raise Todo
+        end
+      end
+
+      {
+        duration_seconds_from_header: container_dict["durationSecondsFromHeader"],
+        bit_rate_from_header: container_dict["bitRate"],
+        best_video_stream_index: container_dict["bestVideoStreamIndex"],
+        best_audio_stream_index: container_dict["bestAudioStreamIndex"],
+        streams: streams_metadata
+      }
+    end
+  end
+end

data/lib/torchcodec/decoders/audio_decoder.rb

@@ -0,0 +1,88 @@
+module TorchCodec
+  module Decoders
+    class AudioDecoder
+      attr_reader :metadata
+
+      def initialize(
+        source,
+        stream_index: nil,
+        sample_rate: nil,
+        num_channels: nil
+      )
+        @decoder = Decoders.create_decoder(source, "approximate")
+
+        container_metadata = Core.get_container_metadata(@decoder)
+        @stream_index =
+          stream_index.nil? ? container_metadata[:best_audio_stream_index] : stream_index
+        if @stream_index.nil?
+          raise ArgumentError, "The best audio stream is unknown and there is no specified stream."
+        end
+        if @stream_index >= container_metadata[:streams].length
+          raise ArgumentError, "The stream at index #{@stream_index} is not a valid stream."
+        end
+
+        @metadata = container_metadata[:streams][@stream_index]
+        if !@metadata.key?(:sample_rate)
+          raise ArgumentError, "The stream at index #{@stream_index} is not an audio stream."
+        end
+
+        @desired_sample_rate =
+          !sample_rate.nil? ? sample_rate : @metadata[:sample_rate]
+
+        Core.add_audio_stream(
+          @decoder,
+          stream_index,
+          sample_rate,
+          num_channels
+        )
+      end
+
+      def get_all_samples
+        get_samples_played_in_range
+      end
+
+      def get_samples_played_in_range(start_seconds: 0.0, stop_seconds: nil)
+        if !stop_seconds.nil? && !(start_seconds <= stop_seconds)
+          raise ArgumentError, "Invalid start seconds: #{start_seconds}. It must be less than or equal to stop seconds (#{stop_seconds})."
+        end
+        frames, first_pts = Core.get_frames_by_pts_in_range_audio(
+          @decoder,
+          start_seconds,
+          stop_seconds
+        )
+        first_pts = first_pts.item
+
+        sample_rate = @desired_sample_rate
+        # TODO: metadata's sample_rate should probably not be Optional
+        raise if sample_rate.nil?
+
+        if first_pts < start_seconds
+          offset_beginning = ((start_seconds - first_pts) * sample_rate).round
+          output_pts_seconds = start_seconds
+        else
+          # In normal cases we'll have first_pts <= start_pts, but in some
+          # edge cases it's possible to have first_pts > start_seconds,
+          # typically if the stream's first frame's pts isn't exactly 0.
+          offset_beginning = 0
+          output_pts_seconds = first_pts
+        end
+
+        num_samples = frames.shape[1]
+        last_pts = first_pts + num_samples / sample_rate
+        if !stop_seconds.nil? && stop_seconds < last_pts
+          offset_end = num_samples - ((last_pts - stop_seconds) * sample_rate).round
+        else
+          offset_end = num_samples
+        end
+
+        data = frames[0.., offset_beginning...offset_end]
+        {
+          data: data,
+          pts_seconds: output_pts_seconds,
+          duration_seconds: data.shape[1] / sample_rate.to_f,
+          sample_rate: sample_rate
+        }
+      end
+    end
+  end
+end

data/lib/torchcodec/decoders/decoder_utils.rb

@@ -0,0 +1,11 @@
+module TorchCodec
+  module Decoders
+    def self.create_decoder(source, seek_mode)
+      if source.is_a?(String)
+        Core.create_from_file(source, seek_mode)
+      else
+        raise TypeError, "Unknown source type: #{source.class.name}"
+      end
+    end
+  end
+end

data/lib/torchcodec/version.rb

@@ -0,0 +1,3 @@
+module TorchCodec
+  VERSION = "0.1.0"
+end

data/lib/torchcodec.rb

@@ -0,0 +1,28 @@
+# dependencies
+require "torch"
+
+# ext
+require "torchcodec/ext"
+
+# stdlib
+require "json"
+
+# modules
+require_relative "torchcodec/version"
+
+# core
+require_relative "torchcodec/core/metadata"
+
+# decoders
+require_relative "torchcodec/decoders/audio_decoder"
+require_relative "torchcodec/decoders/decoder_utils"
+
+module TorchCodec
+  class Error < StandardError; end
+
+  class Todo < Error
+    def message
+      "not implemented yet"
+    end
+  end
+end

metadata

@@ -0,0 +1,90 @@
+--- !ruby/object:Gem::Specification
+name: torchcodec
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Andrew Kane
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: torch-rb
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.23'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.23'
+email: andrew@ankane.org
+executables: []
+extensions:
+- ext/torchcodec/extconf.rb
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- ext/torchcodec/AVIOContextHolder.cpp
+- ext/torchcodec/AVIOContextHolder.h
+- ext/torchcodec/AVIOTensorContext.cpp
+- ext/torchcodec/AVIOTensorContext.h
+- ext/torchcodec/CpuDeviceInterface.cpp
+- ext/torchcodec/CpuDeviceInterface.h
+- ext/torchcodec/DeviceInterface.cpp
+- ext/torchcodec/DeviceInterface.h
+- ext/torchcodec/Encoder.cpp
+- ext/torchcodec/Encoder.h
+- ext/torchcodec/FFMPEGCommon.cpp
+- ext/torchcodec/FFMPEGCommon.h
+- ext/torchcodec/FilterGraph.cpp
+- ext/torchcodec/FilterGraph.h
+- ext/torchcodec/Frame.cpp
+- ext/torchcodec/Frame.h
+- ext/torchcodec/Metadata.cpp
+- ext/torchcodec/Metadata.h
+- ext/torchcodec/SingleStreamDecoder.cpp
+- ext/torchcodec/SingleStreamDecoder.h
+- ext/torchcodec/StableABICompat.h
+- ext/torchcodec/StreamOptions.h
+- ext/torchcodec/Transform.cpp
+- ext/torchcodec/Transform.h
+- ext/torchcodec/ValidationUtils.cpp
+- ext/torchcodec/ValidationUtils.h
+- ext/torchcodec/custom_ops.cpp
+- ext/torchcodec/ext.cpp
+- ext/torchcodec/extconf.rb
+- lib/torchcodec.rb
+- lib/torchcodec/core/metadata.rb
+- lib/torchcodec/decoders/audio_decoder.rb
+- lib/torchcodec/decoders/decoder_utils.rb
+- lib/torchcodec/version.rb
+homepage: https://github.com/ankane/torchcodec-ruby
+licenses:
+- BSD-3-Clause
+metadata: {}
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.2'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.3
+specification_version: 4
+summary: Media encoding and decoding for Torch.rb
+test_files: []