Skip to content

Instantly share code, notes, and snippets.

@Coro365

Coro365/waifu2x-wrapper.rb

Last active Mar 10, 2021
Embed
What would you like to do?
This script can process multiple files and directories with waifu2x-converter-cpp
# Dipendency waifu2x and imagemagic
require 'Open3'
require 'fileutils'
require 'ruby-progressbar'
require "#{__dir__}/get_files_from_ARGV.rb"
require "#{__dir__}/natural_sort.rb"
include Natural_sort
def waifu_run(image, mode='noise_scale')
image_2x = File.join(File.dirname(image), File.basename(image, ".*") + "-#{mode}.png")
system("docker run --rm -i sgsunder/waifu2x-converter-docker -m #{mode} < '#{image}' > '#{image_2x}'")
end
def status(file, total_num, i=1)
printf("\e[1A\e[2K") # line up and clear
print ("Convreting #{i} of #{total_num} item: #{file}")
printf("\r\e[1B") # begining of line and down
end
def image_size(image)
cmd = ['identify', '-format', '%w %h', image]
r, e, i = Open3.capture3(*cmd)
w, h = r.split("\s")
{width: w.to_i, hight: h.to_i}
end
def over_suw?(size)
5120 <= size[:width] && 1440 <= size[:hight]
end
def image_aspect_type(size)
w, h = size.values
if (w - h).abs < (w / 10)
return 'square'
elsif w > h
return 'width'
elsif h > w
return 'hight'
end
end
def anti_dupl(files)
files.select do |file|
file_name = File.basename(file)
scale_count = scale_count(file_name)
pure_name = file_name.gsub(/(-noise_scale)*/, '')
# Collect same images in different sizes
some_size_files = files.select { |e| File.basename(e).include?(File.basename(pure_name, ".*")) }
scale_count_max = some_size_files.map { |file| scale_count(File.basename(file)) }.max
scale_count == scale_count_max
end
end
def scale_count(file_name)
file_name.scan(/-noise_scale/).size
end
def only_width(images)
images.select { |image| 'width' == image_aspect_type(image_size(image)) }
end
def ignore_suw(images)
images.reject { |image| over_suw?(image_size(image)) }
end
def display_ignore(prev_files_size, files_size, ignore_type)
puts("Ignore #{prev_files_size - files_size} #{ignore_type} images (#{prev_files_size} => #{files_size})")
files_size
end
def scan_images(flags)
files = get_files_from_ARGV.select { |f| f.index(/jpg|jpeg|png/) }.natural_sort
prev_files_size = files.size
if flags[:mix_size_copy_flag]
copy_file(anti_dupl(files))
exit
end
if flags[:anti_duple_flag]
files = anti_dupl(files)
prev_files_size = display_ignore(prev_files_size, files.size, 'duplecate')
end
if flags[:only_width_flag]
files = only_width(files)
prev_files_size = display_ignore(prev_files_size, files.size, 'vertical')
end
if flags[:ignore_suw_flag]
files = ignore_suw(files)
prev_files_size = display_ignore(prev_files_size, files.size, 'over suw')
end
files = anti_override(files)
prev_files_size > files.size && display_ignore(prev_files_size, files.size, 'exist')
return files
end
def anti_override(files)
files.reject do |file|
upconverted_path = File.join(File.dirname(file), File.basename(file, ".*") + "-noise_scale.png")
File.exist?(upconverted_path)
end
end
def copy_file(files)
puts("Copy #{files.size}files")
files.each do |file|
dest_dir = File.dirname(file) + '-maximum'
FileUtils.mkdir_p(dest_dir)
FileUtils.copy(file, dest_dir)
end
end
def waifu2x(files)
puts('')
status(files.first, files.size)
pgbar = ProgressBar.create(:total => files.size, :format => '%e [%w>>%i] %a')
files.each_with_index do |file, i|
status(file, files.size, i+1) unless i.zero?
waifu_run(file)
pgbar.increment
end
puts('')
end
def app(flags)
flags[:loop_limit].times do |i|
flags[:loop_limit] > 1 && puts("loop #{i+1}")
files = scan_images(flags)
if files.size.zero?
puts ('x2 complete!')
exit 3
end
if flags[:dry_run_flag]
puts ('Dry run')
exit 3
else
waifu2x(files)
end
end
end
def print_help
puts('version 1.0')
puts('-ad: anti_duple_flag')
puts('-ow: only_width_flag')
puts('-is: ignore_suw_flag')
puts('-dr: dry_run_flag')
puts('-mc: mix_size_copy_flag')
puts('-lp: loop')
exit 3
end
# config
flags = {anti_duple_flag: true}
flags.store(:loop_limit, 1)
# option
ARGV.delete('-ad') && flags.store(:anti_duple_flag, true)
ARGV.delete('-ow') && flags.store(:only_width_flag, true)
ARGV.delete('-is') && flags.store(:ignore_suw_flag, true)
ARGV.delete('-mc') && flags.store(:mix_size_copy_flag, true)
ARGV.delete('-dr') && flags.store(:dry_run_flag, true)
ARGV.delete('-lp') && flags.store(:loop_limit, 3)
ARGV.delete('-h') && print_help
app(flags)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment