gtcrn 0.0.3 → 0.0.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e838b8e452d988facb9cf6cfaa99c8a1b4e9c48073d7a035a07d3b3e42461ec8
-  data.tar.gz: 433c73cbe706d29786d62499ad17eb8b629ab5f5fc640aa3c8e999a56aa46397
+  metadata.gz: 996bfc8a312128c55b4eee2b91f908a612ccbcc619711168bd0468d92018a79b
+  data.tar.gz: 63831dc111ee4d3016c99915aa4dc9cf81b440bb04374c0f64f36402ce8ea1c8
 SHA512:
-  metadata.gz: 21fd283b28ca5b35b7f9f02dab22cef8d4e342af47e7556594b97e70643351e18ff77a9da542e6b5d3963fbb812922740f00f6067b780e95409eb081111b8806
-  data.tar.gz: bc8d8aa599c3b53ef4801bebcb1ac5ef49e11349c71753292d4eeeecfc65bfb495d40a65df5cbe4ba8d14e6db17fe7e16ff25ff01417484a8db0f981ab8dac36
+  metadata.gz: 55836e22a48b395dc403b0340d14d3a7e8ac107902756a6b685bc09514290b8571c2b50e37ac6dcb67cfcc6c716ad3476be847e1fbff9648f324a0071da701ff
+  data.tar.gz: 38558c34bc1ef02473b84c6d74a0397030153aa6cd09393c408f704eebb22e91898e920e6a72617b7904e840489d8d94b27216418e13ffe04b54b332f2a1ee8c

data/README.md

@@ -23,7 +23,8 @@ 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
+    # See https://pytorch.org/get-started/locally/ for download URI for your environment
+    % wget https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.10.0.zip
     % 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

data/bin/gtcrn

@@ -1,4 +1,5 @@
 require "optparse"
+require "optparse/pathname"
 require "gtcrn"
 
 def main(argv)
@@ -18,7 +19,7 @@ def parse_options(argv)
 
       Usage: gtcrn [options] INPUT
     EOB
-    opt.on "-o", "--output=PATH", "Specify output file or directory" do |path|
+    opt.on "-o", "--output=PATH", "Specify output file or directory", Pathname do |path|
       options[:output] = path
     end
   }.parse!(argv)

data/gtcrn.gemspec

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

data/lib/gtcrn.rb

@@ -2,9 +2,11 @@ require "pathname"
 require "torchaudio"
 require "torch"
 require "onnxruntime"
-require "numo/narray/alt"
+require "onnxruntime/torch/tensor"
 
 class GTCRN
+  include NDAV::Converter
+
   MODEL_PATH = File.join(__dir__, "../vendor/gtcrn/stream/onnx_models/gtcrn_simple.onnx").freeze
   ISTFT_OPTS = {
     n_fft: 512,
@@ -19,6 +21,7 @@ class GTCRN
 
   def initialize
     @session = OnnxRuntime::InferenceSession.new(MODEL_PATH)
+    @cache_shapes = @session.inputs[1..].collect {|input| input[:shape]}
     @output_names = @session.outputs.collect {|output| output[:name]}
   end
 
@@ -44,7 +47,7 @@ class GTCRN
   def enhance_speech_waveform(waveform)
     ndim = waveform.ndim
     unless ndim == 1 or ndim == 2
-      raise ArgumentError, "wrong dimension of argment (given #{ndim}, expected 1D or 2D"
+      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)}
@@ -52,33 +55,25 @@ class GTCRN
   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])
-      )
+    conv_cache, tra_cache, inter_cache = @cache_shapes.collect {|shape|
+      OrtValue(Torch.zeros(*shape, dtype: :float32))
     }
     inputs = Torch.view_as_real(
       Torch.stft(channel, **STFT_OPTS)[nil]
-    ).numo
+    )
     outputs = []
     inputs.shape[-2].times do |i|
+      input = inputs[0.., 0.., i..i, 0..]
       enh, conv_cache, tra_cache, inter_cache = @session.run(
         @output_names,
-        {
-          mix: OnnxRuntime::OrtValue.from_numo(inputs[0.., 0.., i..i, 0..]),
-          conv_cache:, tra_cache:, inter_cache:,
-        },
+        {mix: OrtValue(input), conv_cache:, tra_cache:, inter_cache:},
         output_type: :ort_value
       )
-      outputs << enh.numo
+      outputs << TorchTensor(enh)
     end
-    concated = Numo::NArray.concatenate(outputs, axis: 2)
-    real = concated[0.., 0.., 0.., 0]
-    imag = concated[0.., 0.., 0.., 1]
-    enhanced = Torch.istft(
-      Torch.complex(Torch.from_numo(real), Torch.from_numo(imag)),
-      **ISTFT_OPTS
-    )
-    enhanced.squeeze(0)
+    concated = Torch.cat(outputs, dim: 2)
+    Torch
+      .istft(Torch.view_as_complex(concated), **ISTFT_OPTS)
+      .squeeze(0)
   end
 end

data/test/test_gtcrn.rb

@@ -26,7 +26,7 @@ class TestGTCRN < Test::Unit::TestCase
     assert_equal waveform.ndim, enhanced.ndim
   end
 
-  def test_enhance_speech_waveform_two_dim
+  def test_enhance_speech_waveform_multi_dim
     channels = 5
     gtcrn = GTCRN.new
     waveform = Torch.rand(channels, 16000)
@@ -36,7 +36,7 @@ class TestGTCRN < Test::Unit::TestCase
 
     0.upto(channels - 1) do |i|
       enh = gtcrn.enhance_speech_waveform(waveform[i])
-      assert enh.equal(enhanced[i])
+      assert enh.equal(enhanced[i]), "channel #{i}"
     end
   end
 

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gtcrn
 version: !ruby/object:Gem::Version
-  version: 0.0.3
+  version: 0.0.4
 platform: ruby
 authors:
 - Kitaiti Makoto
@@ -66,7 +66,21 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '0'
 - !ruby/object:Gem::Dependency
-  name: numo-narray-alt
+  name: optparse-pathname
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: onnxruntime-torch-tensor
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="