gtcrn 0.0.2 → 0.0.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e62dd7f2ba7da50ab841769e0ccc9d0aeef8f0a31499a2a59b80ffc749b1ca9b
-  data.tar.gz: 555a5c91d412822e1e066bb1ea4bccdfd1cdf9e334adad025b886aea3a2b62f3
+  metadata.gz: e838b8e452d988facb9cf6cfaa99c8a1b4e9c48073d7a035a07d3b3e42461ec8
+  data.tar.gz: 433c73cbe706d29786d62499ad17eb8b629ab5f5fc640aa3c8e999a56aa46397
 SHA512:
-  metadata.gz: c82eb21211da591054d65de5aa131219432f20ce161b53b750d961baed6bd5e3f840076ff17dcdf9568a059257544de648a33c1f7d51911e0513c2e855bb3e2e
-  data.tar.gz: 50814826a17fb42e874f7ea7bd8fae2366f4b08d2961c2e354f4b8b9e4346199d37e61495c4a236cafc46f9886c2a1a9ba6f4b4e9b1b5f5e6f837a0a16265f03
+  metadata.gz: 21fd283b28ca5b35b7f9f02dab22cef8d4e342af47e7556594b97e70643351e18ff77a9da542e6b5d3963fbb812922740f00f6067b780e95409eb081111b8806
+  data.tar.gz: bc8d8aa599c3b53ef4801bebcb1ac5ef49e11349c71753292d4eeeecfc65bfb495d40a65df5cbe4ba8d14e6db17fe7e16ff25ff01417484a8db0f981ab8dac36

data/README.md

@@ -16,7 +16,32 @@ output = GTCRN.new.enhance_speech("path/to/audio.wav")
 # => <Pathname:path/to/audio.enhanced.wav>
 ````
 
-Audio file must be mono WAV with 16kHz sampling rate and 16-bit per sample.
+Audio file must be with 16kHz sampling rate and 16-bit per sample. Currently, file formats supported by [TorchAudio Ruby][] ([TorchCodec Ruby][]) are available.
+
+INSTALLATION
+------------
+
+This gem depends on [Torch.rb][], [TorchAudio Ruby][] and [TorchCodec Ruby][] which require precompiled libtorch and being built with it.
+
+    % wget  https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.10.0.zip # See https://pytorch.org/get-started/locally/ for download URI for your environment
+    % unzip -d path/to/libtorch libtorch-macos-arm64-2.10.0.zip
+    % gem install torch-rb -- --with-torch-dir=path/to/libtorch
+    % gem install torchaudio -- --with-torch-dir=path/to/libtorch
+    % gem install torchcodec -- --with-torch-dir=path/to/libtorch
+    % gem install gtcrn
+
+Or,
+
+    % bundle config set --local build.torch-rb --with-torch-dir=path/to/libtorch
+    % bundle config set --local build.torchaudio --with-torch-dir=path/to/libtorch
+    % bundle config set --local build.torchcodec --with-torch-dir=path/to/libtorch
+    % bundle install
+
+These instructions might be outdated. Refer to each library's instruction if you have trouble.
+
+[Torch.rb]: https://github.com/ankane/torch.rb
+[TorchAudio Ruby]: https://github.com/ankane/torchaudio-ruby
+[TorchCodec Ruby]: https://github.com/ankane/torchcodec-ruby
 
 CLI
 ---
@@ -44,6 +69,8 @@ enhanced = GTCRN.new.enhance_speech_waveform(waveform)
 TorchAudio.save("path/to/output.wav", enhanced.squeeze, sample_rate)
 ```
 
+`GTCRN#enhance_speech_waveform` enhances each channel separately if you pass multi-channel audio.
+
 LICENSE
 -------
 

data/gtcrn.gemspec

@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name = "gtcrn"
-  s.version = "0.0.2"
+  s.version = "0.0.3"
   s.authors = ["Kitaiti Makoto"]
   s.email = ["KitaitiMakoto@gmail.com"]
   s.summary = "Denoises audio"
@@ -16,7 +16,8 @@ Gem::Specification.new do |s|
 
   s.add_runtime_dependency "onnxruntime"
   s.add_runtime_dependency "torch-rb"
-  s.add_runtime_dependency "torchaudio"
+  s.add_runtime_dependency "torchaudio", ">= 0.5.0"
+  s.add_runtime_dependency "torchcodec"
   s.add_runtime_dependency "numo-narray-alt"
 
   s.add_development_dependency "rake"

data/lib/gtcrn.rb

@@ -19,6 +19,7 @@ class GTCRN
 
   def initialize
     @session = OnnxRuntime::InferenceSession.new(MODEL_PATH)
+    @output_names = @session.outputs.collect {|output| output[:name]}
   end
 
   def enhance_speech(path, dest=nil)
@@ -41,31 +42,43 @@ class GTCRN
   end
 
   def enhance_speech_waveform(waveform)
-    conv_cache, tra_cache, inter_cache = 1.upto(3).collect {|i| Numo::SFloat.zeros(*@session.inputs[i][:shape]) }
+    ndim = waveform.ndim
+    unless ndim == 1 or ndim == 2
+      raise ArgumentError, "wrong dimension of argment (given #{ndim}, expected 1D or 2D"
+    end
+    waveform = [waveform] if ndim == 1
+    channels = waveform.collect {|channel| enhance_speech_waveform_channel(channel)}
+    ndim == 1 ? channels[0] : Torch.stack(channels)
+  end
+
+  def enhance_speech_waveform_channel(channel)
+    conv_cache, tra_cache, inter_cache = 1.upto(3).collect {|i|
+      OnnxRuntime::OrtValue.from_numo(
+        Numo::SFloat.zeros(*@session.inputs[i][:shape])
+      )
+    }
     inputs = Torch.view_as_real(
-      Torch.stft(waveform[0], **STFT_OPTS)[nil]
+      Torch.stft(channel, **STFT_OPTS)[nil]
     ).numo
     outputs = []
     inputs.shape[-2].times do |i|
       enh, conv_cache, tra_cache, inter_cache = @session.run(
-        @session.outputs.collect {|output| output[:name]},
+        @output_names,
         {
           mix: OnnxRuntime::OrtValue.from_numo(inputs[0.., 0.., i..i, 0..]),
-          conv_cache: OnnxRuntime::OrtValue.from_numo(conv_cache),
-          tra_cache: OnnxRuntime::OrtValue.from_numo(tra_cache),
-          inter_cache: OnnxRuntime::OrtValue.from_numo(inter_cache)
+          conv_cache:, tra_cache:, inter_cache:,
         },
-        output_type: :numo
+        output_type: :ort_value
       )
-      outputs << enh
+      outputs << enh.numo
     end
     concated = Numo::NArray.concatenate(outputs, axis: 2)
     real = concated[0.., 0.., 0.., 0]
     imag = concated[0.., 0.., 0.., 1]
     enhanced = Torch.istft(
-      Torch.from_numo(real) + 1i * Torch.from_numo(imag),
+      Torch.complex(Torch.from_numo(real), Torch.from_numo(imag)),
       **ISTFT_OPTS
     )
-    enhanced.squeeze
+    enhanced.squeeze(0)
   end
 end

data/test/test_gtcrn.rb

@@ -19,4 +19,31 @@ class TestGTCRN < Test::Unit::TestCase
       assert source.total_sample_frames - enhanced.total_sample_frames < 512
     end
   end
+
+  def test_enhance_speech_waveform_one_dim
+    waveform = Torch.rand(16000)
+    enhanced = GTCRN.new.enhance_speech_waveform(waveform)
+    assert_equal waveform.ndim, enhanced.ndim
+  end
+
+  def test_enhance_speech_waveform_two_dim
+    channels = 5
+    gtcrn = GTCRN.new
+    waveform = Torch.rand(channels, 16000)
+    enhanced = gtcrn.enhance_speech_waveform(waveform)
+
+    assert_equal waveform.shape[0..-2], enhanced.shape[0..-2]
+
+    0.upto(channels - 1) do |i|
+      enh = gtcrn.enhance_speech_waveform(waveform[i])
+      assert enh.equal(enhanced[i])
+    end
+  end
+
+  def test_enhance_speech_waveform_channel
+    gtcrn = GTCRN.new
+    channel = Torch.rand(16000)
+    enhanced = gtcrn.enhance_speech_waveform_channel(channel)
+    assert_equal channel.ndim, enhanced.ndim
+  end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gtcrn
 version: !ruby/object:Gem::Version
-  version: 0.0.2
+  version: 0.0.3
 platform: ruby
 authors:
 - Kitaiti Makoto
@@ -39,6 +39,20 @@ dependencies:
         version: '0'
 - !ruby/object:Gem::Dependency
   name: torchaudio
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.5.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.5.0
+- !ruby/object:Gem::Dependency
+  name: torchcodec
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
@@ -188,7 +202,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 4.0.3
+rubygems_version: 4.0.6
 specification_version: 4
 summary: Denoises audio
 test_files: []