module TorchVision::Utils
Public Class Methods
image_from_array(array)
click to toggle source
private Ruby-specific method TODO use Numo when bridge available
# File lib/torchvision/utils.rb, line 89 def image_from_array(array) case array when Torch::Tensor # TODO support more dtypes raise "Type not supported yet: #{array.dtype}" unless array.dtype == :uint8 array = array.contiguous unless array.contiguous? width, height = array.shape bands = array.shape[2] || 1 data = FFI::Pointer.new(:uint8, array._data_ptr) data = data.slice(0, array.numel * array.element_size) Vips::Image.new_from_memory(data, width, height, bands, :uchar) when Numo::NArray # TODO support more types raise "Type not supported yet: #{array.class.name}" unless array.is_a?(Numo::UInt8) width, height = array.shape bands = array.shape[2] || 1 data = array.to_binary Vips::Image.new_from_memory(data, width, height, bands, :uchar) else raise "Expected Torch::Tensor or Numo::NArray, not #{array.class.name}" end end
make_grid(tensor, nrow: 8, padding: 2, normalize: false, range: nil, scale_each: false, pad_value: 0)
click to toggle source
# File lib/torchvision/utils.rb, line 4 def make_grid(tensor, nrow: 8, padding: 2, normalize: false, range: nil, scale_each: false, pad_value: 0) unless Torch.tensor?(tensor) || (tensor.is_a?(Array) && tensor.all? { |t| Torch.tensor?(t) }) raise ArgumentError, "tensor or list of tensors expected, got #{tensor.class.name}" end # if list of tensors, convert to a 4D mini-batch Tensor if tensor.is_a?(Array) tensor = Torch.stack(tensor, dim: 0) end if tensor.dim == 2 # single image H x W tensor = tensor.unsqueeze(0) end if tensor.dim == 3 # single image if tensor.size(0) == 1 # if single-channel, convert to 3-channel tensor = Torch.cat([tensor, tensor, tensor], 0) end tensor = tensor.unsqueeze(0) end if tensor.dim == 4 && tensor.size(1) == 1 # single-channel images tensor = Torch.cat([tensor, tensor, tensor], 1) end if normalize tensor = tensor.clone # avoid modifying tensor in-place if !range.nil? && !range.is_a?(Array) raise "range has to be an array (min, max) if specified. min and max are numbers" end norm_ip = lambda do |img, min, max| img.clamp!(min, max) img.add!(-min).div!(max - min + 1e-5) end norm_range = lambda do |t, range| if !range.nil? norm_ip.call(t, range[0], range[1]) else norm_ip.call(t, t.min.to_f, t.max.to_f) end end if scale_each tensor.each do |t| # loop over mini-batch dimension norm_range.call(t, range) end else norm_range.call(tensor, range) end end if tensor.size(0) == 1 return tensor.squeeze(0) end # make the mini-batch of images into a grid nmaps = tensor.size(0) xmaps = [nrow, nmaps].min ymaps = (nmaps.to_f / xmaps).ceil height, width = (tensor.size(2) + padding), (tensor.size(3) + padding) num_channels = tensor.size(1) grid = tensor.new_full([num_channels, height * ymaps + padding, width * xmaps + padding], pad_value) k = 0 ymaps.times do |y| xmaps.times do |x| break if k >= nmaps grid.narrow(1, y * height + padding, height - padding).narrow(2, x * width + padding, width - padding).copy!(tensor[k]) k += 1 end end grid end
save_image(tensor, fp, nrow: 8, padding: 2, normalize: false, range: nil, scale_each: false, pad_value: 0)
click to toggle source
# File lib/torchvision/utils.rb, line 78 def save_image(tensor, fp, nrow: 8, padding: 2, normalize: false, range: nil, scale_each: false, pad_value: 0) grid = make_grid(tensor, nrow: nrow, padding: padding, pad_value: pad_value, normalize: normalize, range: range, scale_each: scale_each) # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer ndarr = grid.mul(255).add!(0.5).clamp!(0, 255).permute(1, 2, 0).to("cpu", dtype: :uint8) im = image_from_array(ndarr) im.write_to_file(fp) end