lib/dnn/lib/image_io.rb in ruby-dnn-0.3.2 vs lib/dnn/lib/image_io.rb in ruby-dnn-0.4.0

@@ -1,27 +1,31 @@
 require "numo/narray"
-require "dnn/ext/image_io/image_io_ext"
+require "dnn/ext/rb_stb_image/rb_stb_image"
 
 module DNN
   module ImageIO
-    private_class_method :_read
-    private_class_method :_write_bmp
-    private_class_method :_write_png
-    private_class_method :_write_jpg
-
     def self.read(file_name)
       raise ImageIO::ReadError.new("#{file_name} is not found.") unless File.exist?(file_name)
-      _read(file_name)
+      img, = Stb.stbi_load(file_name, 3)
+      img
     end
 
-    def self.write(file_name, nary, quality: 100)
+    def self.write(file_name, img, quality: 100)
+      img = img.clone
+      if img.shape.length == 2
+        img = Numo::UInt8[img, img, img].transpose(1, 2, 0).clone
+      elsif img.shape[2] == 1
+        img = img.shape(img.shape[0], img.shape[1])
+        img = Numo::UInt8[img, img, img].transpose(1, 2, 0).clone
+      end
       case file_name
       when /\.png$/
-        _write_png(file_name, nary)
+        stride_in_bytes = img.shape[0] * img.shape[2]
+        Stb.stbi_write_png(file_name, *img.shape, img, stride_in_bytes)
       when /\.bmp$/
-        _write_bmp(file_name, nary)
+        Stb.stbi_write_bmp(file_name, *img.shape, img)
       when /\.jpg$/
-        _write_jpg(file_name, nary, quality)
+        Stb.stbi_write_jpg(file_name, *img.shape, img, quality)
       end
     rescue => ex
       raise ImageIO::WriteError.new(ex.message)
     end
   end