Merge 837e9830cc into 60e22a05bd
This commit is contained in:
Junhyuk Park
GitHub
commit
dc5d76b730
225
convert_data.lua
225
convert_data.lua
|
|
@ -11,126 +11,129 @@ local image_loader = require 'image_loader'
|
|||
local iproc = require 'iproc'
|
||||
local alpha_util = require 'alpha_util'
|
||||
|
||||
local function crop_if_large(src, max_size)
|
||||
if max_size < 0 then
|
||||
return src
|
||||
end
|
||||
local tries = 4
|
||||
if src:size(2) >= max_size and src:size(3) >= max_size then
|
||||
local rect
|
||||
for i = 1, tries do
|
||||
local yi = torch.random(0, src:size(2) - max_size)
|
||||
local xi = torch.random(0, src:size(3) - max_size)
|
||||
rect = iproc.crop(src, xi, yi, xi + max_size, yi + max_size)
|
||||
-- ignore simple background
|
||||
if rect:float():std() >= 0 then
|
||||
break
|
||||
end
|
||||
end
|
||||
return rect
|
||||
else
|
||||
return src
|
||||
end
|
||||
end
|
||||
local function crop_if_large_pair(x, y, max_size)
|
||||
if max_size < 0 then
|
||||
return x, y
|
||||
end
|
||||
local scale_y = y:size(2) / x:size(2)
|
||||
local mod = 4
|
||||
assert(x:size(3) == (y:size(3) / scale_y))
|
||||
|
||||
local tries = 4
|
||||
if y:size(2) > max_size and y:size(3) > max_size then
|
||||
assert(max_size % 4 == 0)
|
||||
local rect_x, rect_y
|
||||
for i = 1, tries do
|
||||
local yi = torch.random(0, y:size(2) - max_size)
|
||||
local xi = torch.random(0, y:size(3) - max_size)
|
||||
if mod then
|
||||
yi = yi - (yi % mod)
|
||||
xi = xi - (xi % mod)
|
||||
end
|
||||
rect_y = iproc.crop(y, xi, yi, xi + max_size, yi + max_size)
|
||||
rect_x = iproc.crop(y, xi / scale_y, yi / scale_y, xi / scale_y + max_size / scale_y, yi / scale_y + max_size / scale_y)
|
||||
-- ignore simple background
|
||||
if rect_y:float():std() >= 0 then
|
||||
break
|
||||
end
|
||||
end
|
||||
return rect_x, rect_y
|
||||
else
|
||||
return x, y
|
||||
end
|
||||
local function crop_if_large(src, max_size)
|
||||
if max_size < 0 then
|
||||
return src
|
||||
end
|
||||
local tries = 4
|
||||
if src:size(2) >= max_size and src:size(3) >= max_size then
|
||||
local rect
|
||||
for i = 1, tries do
|
||||
local yi = torch.random(0, src:size(2) - max_size)
|
||||
local xi = torch.random(0, src:size(3) - max_size)
|
||||
rect = iproc.crop(src, xi, yi, xi + max_size, yi + max_size)
|
||||
-- ignore simple background
|
||||
if rect:float():std() >= 0 then
|
||||
break
|
||||
end
|
||||
end
|
||||
return rect
|
||||
else
|
||||
return src
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
local function crop_if_large_pair(x, y, max_size)
|
||||
if max_size < 0 then
|
||||
return x, y
|
||||
end
|
||||
local scale_y = y:size(2) / x:size(2)
|
||||
local mod = 4
|
||||
assert(x:size(3) == (y:size(3) / scale_y))
|
||||
local tries = 4
|
||||
if y:size(2) > max_size and y:size(3) > max_size then
|
||||
assert(max_size % 4 == 0)
|
||||
local rect_x, rect_y
|
||||
for i = 1, tries do
|
||||
local yi = torch.random(0, y:size(2) - max_size)
|
||||
local xi = torch.random(0, y:size(3) - max_size)
|
||||
if mod then
|
||||
yi = yi - (yi % mod)
|
||||
xi = xi - (xi % mod)
|
||||
end
|
||||
rect_y = iproc.crop(y, xi, yi, xi + max_size, yi + max_size)
|
||||
rect_x = iproc.crop(y, xi / scale_y, yi / scale_y, xi / scale_y + max_size / scale_y, yi / scale_y + max_size / scale_y)
|
||||
-- ignore simple background
|
||||
if rect_y:float():std() >= 0 then
|
||||
break
|
||||
end
|
||||
end
|
||||
return rect_x, rect_y
|
||||
else
|
||||
return x, y
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
local function load_images(list)
|
||||
local MARGIN = 32
|
||||
local csv = csvigo.load({path = list, verbose = false, mode = "raw"})
|
||||
local x = {}
|
||||
local skip_notice = false
|
||||
for i = 1, #csv do
|
||||
local filename = csv[i][1]
|
||||
local csv_meta = csv[i][2]
|
||||
if csv_meta and csv_meta:len() > 0 then
|
||||
csv_meta = cjson.decode(csv_meta)
|
||||
end
|
||||
if csv_meta and csv_meta.filters then
|
||||
filters = csv_meta.filters
|
||||
end
|
||||
local im, meta = image_loader.load_byte(filename)
|
||||
local skip = false
|
||||
local alpha_color = torch.random(0, 1)
|
||||
|
||||
if meta and meta.alpha then
|
||||
if settings.use_transparent_png then
|
||||
im = alpha_util.fill(im, meta.alpha, alpha_color)
|
||||
else
|
||||
skip = true
|
||||
end
|
||||
end
|
||||
if skip then
|
||||
if not skip_notice then
|
||||
io.stderr:write("skip transparent png (settings.use_transparent_png=0)\n")
|
||||
skip_notice = true
|
||||
end
|
||||
local MARGIN = 32
|
||||
local csv = csvigo.load({path = list, verbose = false, mode = "raw"})
|
||||
local x = {}
|
||||
local skip_notice = false
|
||||
for i = 1, #csv do
|
||||
local filename = csv[i][1]
|
||||
local csv_meta = csv[i][2]
|
||||
if csv_meta and csv_meta:len() > 0 then
|
||||
csv_meta = cjson.decode(csv_meta)
|
||||
end
|
||||
if csv_meta and csv_meta.filters then
|
||||
filters = csv_meta.filters
|
||||
end
|
||||
local im, meta = image_loader.load_byte(filename)
|
||||
local skip = false
|
||||
local alpha_color = torch.random(0, 1)
|
||||
if meta and meta.alpha then
|
||||
if settings.use_transparent_png then
|
||||
im = alpha_util.fill(im, meta.alpha, alpha_color)
|
||||
else
|
||||
if csv_meta and csv_meta.x then
|
||||
-- method == user
|
||||
local yy = im
|
||||
local xx, meta2 = image_loader.load_byte(csv_meta.x)
|
||||
if meta2 and meta2.alpha then
|
||||
xx = alpha_util.fill(xx, meta2.alpha, alpha_color)
|
||||
end
|
||||
xx, yy = crop_if_large_pair(xx, yy, settings.max_training_image_size)
|
||||
table.insert(x, {{y = compression.compress(yy), x = compression.compress(xx)},
|
||||
{data = {filters = filters, has_x = true}}})
|
||||
else
|
||||
im = crop_if_large(im, settings.max_training_image_size)
|
||||
im = iproc.crop_mod4(im)
|
||||
local scale = 1.0
|
||||
if settings.random_half_rate > 0.0 then
|
||||
scale = 2.0
|
||||
end
|
||||
if im then
|
||||
if im:size(2) > (settings.crop_size * scale + MARGIN) and im:size(3) > (settings.crop_size * scale + MARGIN) then
|
||||
table.insert(x, {compression.compress(im), {data = {filters = filters}}})
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: image is too small (%d > size).\n", filename, settings.crop_size * scale + MARGIN))
|
||||
end
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: load error.\n", filename))
|
||||
end
|
||||
end
|
||||
skip = true
|
||||
end
|
||||
xlua.progress(i, #csv)
|
||||
if i % 10 == 0 then
|
||||
collectgarbage()
|
||||
end
|
||||
if skip then
|
||||
if not skip_notice then
|
||||
io.stderr:write("skip transparent png (settings.use_transparent_png=0)\n")
|
||||
skip_notice = true
|
||||
end
|
||||
end
|
||||
return x
|
||||
else
|
||||
if csv_meta and csv_meta.x then
|
||||
-- method == user
|
||||
local yy = im
|
||||
local xx, meta2 = image_loader.load_byte(csv_meta.x)
|
||||
if meta2 and meta2.alpha then
|
||||
xx = alpha_util.fill(xx, meta2.alpha, alpha_color)
|
||||
end
|
||||
xx, yy = crop_if_large_pair(xx, yy, settings.max_training_image_size)
|
||||
table.insert(x, {{y = compression.compress(yy), x = compression.compress(xx)},
|
||||
{data = {filters = filters, has_x = true}}})
|
||||
else
|
||||
im = crop_if_large(im, settings.max_training_image_size)
|
||||
im = iproc.crop_mod4(im)
|
||||
local scale = 1.0
|
||||
if settings.random_half_rate > 0.0 then
|
||||
scale = 2.0
|
||||
end
|
||||
if im then
|
||||
if im:size(2) > (settings.crop_size * scale + MARGIN) and im:size(3) > (settings.crop_size * scale + MARGIN) then
|
||||
table.insert(x, {compression.compress(im), {data = {filters = filters}}})
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: image is too small (%d > size).\n", filename, settings.crop_size * scale + MARGIN))
|
||||
end
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: load error.\n", filename))
|
||||
end
|
||||
end
|
||||
end
|
||||
xlua.progress(i, #csv)
|
||||
if i % 10 == 0 then
|
||||
collectgarbage()
|
||||
end
|
||||
end
|
||||
return x
|
||||
end
|
||||
|
||||
|
||||
torch.manualSeed(settings.seed)
|
||||
print(settings)
|
||||
local x = load_images(settings.image_list)
|
||||
|
|
|
|||
518
waifu2x.lua
518
waifu2x.lua
|
|
@ -10,287 +10,285 @@ local alpha_util = require 'alpha_util'
|
|||
|
||||
torch.setdefaulttensortype('torch.FloatTensor')
|
||||
|
||||
|
||||
local function format_output(opt, src, no)
|
||||
no = no or 1
|
||||
local name = path.basename(src)
|
||||
local e = path.extension(name)
|
||||
local basename = name:sub(0, name:len() - e:len())
|
||||
|
||||
if opt.o == "(auto)" then
|
||||
return path.join(path.dirname(src), string.format("%s_%s.png", basename, opt.m))
|
||||
else
|
||||
local basename_pos = opt.o:find("%%s")
|
||||
local no_pos = opt.o:find("%%%d*d")
|
||||
if basename_pos ~= nil and no_pos ~= nil then
|
||||
if basename_pos < no_pos then
|
||||
return string.format(opt.o, basename, no)
|
||||
else
|
||||
return string.format(opt.o, no, basename)
|
||||
end
|
||||
elseif basename_pos ~= nil then
|
||||
return string.format(opt.o, basename)
|
||||
elseif no_pos ~= nil then
|
||||
return string.format(opt.o, no)
|
||||
no = no or 1
|
||||
local name = path.basename(src)
|
||||
local e = path.extension(name)
|
||||
local basename = name:sub(0, name:len() - e:len())
|
||||
if opt.o == "(auto)" then
|
||||
return path.join(path.dirname(src), string.format("%s_%s.png", basename, opt.m))
|
||||
else
|
||||
local basename_pos = opt.o:find("%%s")
|
||||
local no_pos = opt.o:find("%%%d*d")
|
||||
if basename_pos ~= nil and no_pos ~= nil then
|
||||
if basename_pos < no_pos then
|
||||
return string.format(opt.o, basename, no)
|
||||
else
|
||||
return opt.o
|
||||
return string.format(opt.o, no, basename)
|
||||
end
|
||||
end
|
||||
elseif basename_pos ~= nil then
|
||||
return string.format(opt.o, basename)
|
||||
elseif no_pos ~= nil then
|
||||
return string.format(opt.o, no)
|
||||
else
|
||||
return opt.o
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
local function convert_image(opt)
|
||||
local x, meta = image_loader.load_float(opt.i)
|
||||
if not x then
|
||||
error(string.format("failed to load image: %s", opt.i))
|
||||
end
|
||||
local alpha = meta.alpha
|
||||
local new_x = nil
|
||||
local scale_f, image_f
|
||||
|
||||
if opt.tta == 1 then
|
||||
scale_f = function(model, scale, x, block_size, batch_size)
|
||||
return reconstruct.scale_tta(model, opt.tta_level,
|
||||
scale, x, block_size, batch_size)
|
||||
end
|
||||
image_f = function(model, x, block_size, batch_size)
|
||||
return reconstruct.image_tta(model, opt.tta_level,
|
||||
x, block_size, batch_size)
|
||||
end
|
||||
else
|
||||
scale_f = reconstruct.scale
|
||||
image_f = reconstruct.image
|
||||
end
|
||||
opt.o = format_output(opt, opt.i)
|
||||
if opt.m == "noise" then
|
||||
local model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level))
|
||||
local x, meta = image_loader.load_float(opt.i)
|
||||
if not x then
|
||||
error(string.format("failed to load image: %s", opt.i))
|
||||
end
|
||||
local alpha = meta.alpha
|
||||
local new_x = nil
|
||||
local scale_f, image_f
|
||||
if opt.tta == 1 then
|
||||
scale_f = function(model, scale, x, block_size, batch_size)
|
||||
return reconstruct.scale_tta(model, opt.tta_level, scale, x, block_size, batch_size)
|
||||
end
|
||||
image_f = function(model, x, block_size, batch_size)
|
||||
return reconstruct.image_tta(model, opt.tta_level, x, block_size, batch_size)
|
||||
end
|
||||
else
|
||||
scale_f = reconstruct.scale
|
||||
image_f = reconstruct.image
|
||||
end
|
||||
opt.o = format_output(opt, opt.i)
|
||||
if opt.m == "noise" then
|
||||
local model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level))
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
end
|
||||
local t = sys.clock()
|
||||
new_x = image_f(model, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
elseif opt.m == "scale" then
|
||||
local model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
end
|
||||
local t = sys.clock()
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(model))
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, model)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
elseif opt.m == "noise_scale" then
|
||||
local model_path = path.join(opt.model_dir, ("noise%d_scale%.1fx_model.t7"):format(opt.noise_level, opt.scale))
|
||||
if path.exists(model_path) then
|
||||
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
local t, scale_model = pcall(w2nn.load_model, scale_model_path, opt.force_cudnn)
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
if not t then
|
||||
scale_model = model
|
||||
end
|
||||
local t = sys.clock()
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
else
|
||||
local noise_model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level))
|
||||
local noise_model = w2nn.load_model(noise_model_path, opt.force_cudnn)
|
||||
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
local scale_model = w2nn.load_model(scale_model_path, opt.force_cudnn)
|
||||
local t = sys.clock()
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
x = image_f(noise_model, x, opt.crop_size, opt.batch_size)
|
||||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
end
|
||||
elseif opt.m == "user" then
|
||||
local model_path = opt.model_path
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
end
|
||||
local t = sys.clock()
|
||||
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(model))
|
||||
if opt.scale == 1 then
|
||||
new_x = image_f(model, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
elseif opt.m == "scale" then
|
||||
local model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
end
|
||||
local t = sys.clock()
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(model))
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, model)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
elseif opt.m == "noise_scale" then
|
||||
local model_path = path.join(opt.model_dir, ("noise%d_scale%.1fx_model.t7"):format(opt.noise_level, opt.scale))
|
||||
if path.exists(model_path) then
|
||||
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
local t, scale_model = pcall(w2nn.load_model, scale_model_path, opt.force_cudnn)
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not t then
|
||||
scale_model = model
|
||||
end
|
||||
local t = sys.clock()
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
else
|
||||
local noise_model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level))
|
||||
local noise_model = w2nn.load_model(noise_model_path, opt.force_cudnn)
|
||||
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
local scale_model = w2nn.load_model(scale_model_path, opt.force_cudnn)
|
||||
local t = sys.clock()
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
x = image_f(noise_model, x, opt.crop_size, opt.batch_size)
|
||||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
end
|
||||
elseif opt.m == "user" then
|
||||
local model_path = opt.model_path
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
end
|
||||
local t = sys.clock()
|
||||
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(model))
|
||||
if opt.scale == 1 then
|
||||
new_x = image_f(model, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha) -- TODO: should it use model?
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
else
|
||||
error("undefined method:" .. opt.method)
|
||||
end
|
||||
image_loader.save_png(opt.o, new_x, tablex.update({depth = opt.depth, inplace = true}, meta))
|
||||
else
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha) -- TODO: should it use model?
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
else
|
||||
error("undefined method:" .. opt.method)
|
||||
end
|
||||
image_loader.save_png(opt.o, new_x, tablex.update({depth = opt.depth, inplace = true}, meta))
|
||||
end
|
||||
|
||||
|
||||
local function convert_frames(opt)
|
||||
local model_path, scale_model, t
|
||||
local noise_scale_model = {}
|
||||
local noise_model = {}
|
||||
local user_model = nil
|
||||
local scale_f, image_f
|
||||
if opt.tta == 1 then
|
||||
scale_f = function(model, scale, x, block_size, batch_size)
|
||||
return reconstruct.scale_tta(model, opt.tta_level,
|
||||
scale, x, block_size, batch_size)
|
||||
local model_path, scale_model, t
|
||||
local noise_scale_model = {}
|
||||
local noise_model = {}
|
||||
local user_model = nil
|
||||
local scale_f, image_f
|
||||
if opt.tta == 1 then
|
||||
scale_f = function(model, scale, x, block_size, batch_size)
|
||||
return reconstruct.scale_tta(model, opt.tta_level, scale, x, block_size, batch_size)
|
||||
end
|
||||
image_f = function(model, x, block_size, batch_size)
|
||||
return reconstruct.image_tta(model, opt.tta_level, x, block_size, batch_size)
|
||||
end
|
||||
else
|
||||
scale_f = reconstruct.scale
|
||||
image_f = reconstruct.image
|
||||
end
|
||||
if opt.m == "scale" then
|
||||
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
scale_model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
elseif opt.m == "noise" then
|
||||
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
|
||||
noise_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
elseif opt.m == "noise_scale" then
|
||||
local model_path = path.join(opt.model_dir, ("noise%d_scale%.1fx_model.t7"):format(opt.noise_level, opt.scale))
|
||||
if path.exists(model_path) then
|
||||
noise_scale_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
t, scale_model = pcall(w2nn.load_model, model_path, opt.force_cudnn)
|
||||
if not t then
|
||||
scale_model = noise_scale_model[opt.noise_level]
|
||||
end
|
||||
image_f = function(model, x, block_size, batch_size)
|
||||
return reconstruct.image_tta(model, opt.tta_level,
|
||||
x, block_size, batch_size)
|
||||
end
|
||||
else
|
||||
scale_f = reconstruct.scale
|
||||
image_f = reconstruct.image
|
||||
end
|
||||
if opt.m == "scale" then
|
||||
else
|
||||
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
scale_model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
elseif opt.m == "noise" then
|
||||
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
|
||||
noise_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
elseif opt.m == "noise_scale" then
|
||||
local model_path = path.join(opt.model_dir, ("noise%d_scale%.1fx_model.t7"):format(opt.noise_level, opt.scale))
|
||||
if path.exists(model_path) then
|
||||
noise_scale_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
t, scale_model = pcall(w2nn.load_model, model_path, opt.force_cudnn)
|
||||
if not t then
|
||||
scale_model = noise_scale_model[opt.noise_level]
|
||||
end
|
||||
end
|
||||
elseif opt.m == "user" then
|
||||
user_model = w2nn.load_model(opt.model_path, opt.force_cudnn)
|
||||
end
|
||||
local fp = io.open(opt.l)
|
||||
if not fp then
|
||||
error("Open Error: " .. opt.l)
|
||||
end
|
||||
local count = 0
|
||||
local lines = {}
|
||||
for line in fp:lines() do
|
||||
table.insert(lines, line)
|
||||
end
|
||||
fp:close()
|
||||
for i = 1, #lines do
|
||||
local output = format_output(opt, lines[i], i)
|
||||
if opt.resume == 0 or path.exists(output) == false then
|
||||
local x, meta = image_loader.load_float(lines[i])
|
||||
if not x then
|
||||
io.stderr:write(string.format("failed to load image: %s\n", lines[i]))
|
||||
else
|
||||
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
||||
scale_model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
|
||||
noise_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
local alpha = meta.alpha
|
||||
local new_x = nil
|
||||
if opt.m == "noise" then
|
||||
new_x = image_f(noise_model[opt.noise_level], x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
elseif opt.m == "scale" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
elseif opt.m == "noise_scale" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
if noise_scale_model[opt.noise_level] then
|
||||
new_x = scale_f(noise_scale_model[opt.noise_level], opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
x = image_f(noise_model[opt.noise_level], x, opt.crop_size, opt.batch_size)
|
||||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
elseif opt.m == "user" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(user_model))
|
||||
if opt.scale == 1 then
|
||||
new_x = image_f(user_model, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
new_x = scale_f(user_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
else
|
||||
error("undefined method:" .. opt.method)
|
||||
end
|
||||
image_loader.save_png(output, new_x, tablex.update({depth = opt.depth, inplace = true}, meta))
|
||||
end
|
||||
elseif opt.m == "user" then
|
||||
user_model = w2nn.load_model(opt.model_path, opt.force_cudnn)
|
||||
end
|
||||
local fp = io.open(opt.l)
|
||||
if not fp then
|
||||
error("Open Error: " .. opt.l)
|
||||
end
|
||||
local count = 0
|
||||
local lines = {}
|
||||
for line in fp:lines() do
|
||||
table.insert(lines, line)
|
||||
end
|
||||
fp:close()
|
||||
|
||||
for i = 1, #lines do
|
||||
local output = format_output(opt, lines[i], i)
|
||||
if opt.resume == 0 or path.exists(output) == false then
|
||||
local x, meta = image_loader.load_float(lines[i])
|
||||
if not x then
|
||||
io.stderr:write(string.format("failed to load image: %s\n", lines[i]))
|
||||
else
|
||||
local alpha = meta.alpha
|
||||
local new_x = nil
|
||||
if opt.m == "noise" then
|
||||
new_x = image_f(noise_model[opt.noise_level], x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
elseif opt.m == "scale" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
elseif opt.m == "noise_scale" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
||||
if noise_scale_model[opt.noise_level] then
|
||||
new_x = scale_f(noise_scale_model[opt.noise_level], opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
x = image_f(noise_model[opt.noise_level], x, opt.crop_size, opt.batch_size)
|
||||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
elseif opt.m == "user" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(user_model))
|
||||
if opt.scale == 1 then
|
||||
new_x = image_f(user_model, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
new_x = scale_f(user_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
else
|
||||
error("undefined method:" .. opt.method)
|
||||
end
|
||||
image_loader.save_png(output, new_x,
|
||||
tablex.update({depth = opt.depth, inplace = true}, meta))
|
||||
end
|
||||
if not opt.q then
|
||||
xlua.progress(i, #lines)
|
||||
end
|
||||
if i % 10 == 0 then
|
||||
collectgarbage()
|
||||
end
|
||||
else
|
||||
if not opt.q then
|
||||
xlua.progress(i, #lines)
|
||||
end
|
||||
if not opt.q then
|
||||
xlua.progress(i, #lines)
|
||||
end
|
||||
end
|
||||
if i % 10 == 0 then
|
||||
collectgarbage()
|
||||
end
|
||||
else
|
||||
if not opt.q then
|
||||
xlua.progress(i, #lines)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
local function waifu2x()
|
||||
local cmd = torch.CmdLine()
|
||||
cmd:text()
|
||||
cmd:text("waifu2x")
|
||||
cmd:text("Options:")
|
||||
cmd:option("-i", "images/miku_small.png", 'path to input image')
|
||||
cmd:option("-l", "", 'path to image-list.txt')
|
||||
cmd:option("-scale", 2, 'scale factor')
|
||||
cmd:option("-o", "(auto)", 'path to output file')
|
||||
cmd:option("-depth", 8, 'bit-depth of the output image (8|16)')
|
||||
cmd:option("-model_dir", "./models/upconv_7/art", 'path to model directory')
|
||||
cmd:option("-name", "user", 'model name for user method')
|
||||
cmd:option("-m", "noise_scale", 'method (noise|scale|noise_scale|user)')
|
||||
cmd:option("-method", "", 'same as -m')
|
||||
cmd:option("-noise_level", 1, '(1|2|3)')
|
||||
cmd:option("-crop_size", 128, 'patch size per process')
|
||||
cmd:option("-batch_size", 1, 'batch_size')
|
||||
cmd:option("-resume", 0, "skip existing files (0|1)")
|
||||
cmd:option("-thread", -1, "number of CPU threads")
|
||||
cmd:option("-tta", 0, 'use TTA mode. It is slow but slightly high quality (0|1)')
|
||||
cmd:option("-tta_level", 8, 'TTA level (2|4|8). A higher value makes better quality output but slow')
|
||||
cmd:option("-force_cudnn", 0, 'use cuDNN backend (0|1)')
|
||||
cmd:option("-q", 0, 'quiet (0|1)')
|
||||
|
||||
local opt = cmd:parse(arg)
|
||||
if opt.method:len() > 0 then
|
||||
opt.m = opt.method
|
||||
end
|
||||
if opt.thread > 0 then
|
||||
torch.setnumthreads(opt.thread)
|
||||
end
|
||||
if cudnn then
|
||||
cudnn.fastest = true
|
||||
if opt.l:len() > 0 then
|
||||
cudnn.benchmark = true -- find fastest algo
|
||||
else
|
||||
cudnn.benchmark = false
|
||||
end
|
||||
end
|
||||
opt.force_cudnn = opt.force_cudnn == 1
|
||||
opt.q = opt.q == 1
|
||||
opt.model_path = path.join(opt.model_dir, string.format("%s_model.t7", opt.name))
|
||||
|
||||
if string.len(opt.l) == 0 then
|
||||
convert_image(opt)
|
||||
else
|
||||
convert_frames(opt)
|
||||
end
|
||||
local cmd = torch.CmdLine()
|
||||
cmd:text()
|
||||
cmd:text("waifu2x")
|
||||
cmd:text("Options:")
|
||||
cmd:option("-i", "images/miku_small.png", 'path to input image')
|
||||
cmd:option("-l", "", 'path to image-list.txt')
|
||||
cmd:option("-scale", 2, 'scale factor')
|
||||
cmd:option("-o", "(auto)", 'path to output file')
|
||||
cmd:option("-depth", 8, 'bit-depth of the output image (8|16)')
|
||||
cmd:option("-model_dir", "./models/upconv_7/art", 'path to model directory')
|
||||
cmd:option("-name", "user", 'model name for user method')
|
||||
cmd:option("-m", "noise_scale", 'method (noise|scale|noise_scale|user)')
|
||||
cmd:option("-method", "", 'same as -m')
|
||||
cmd:option("-noise_level", 1, '(1|2|3)')
|
||||
cmd:option("-crop_size", 128, 'patch size per process')
|
||||
cmd:option("-batch_size", 1, 'batch_size')
|
||||
cmd:option("-resume", 0, "skip existing files (0|1)")
|
||||
cmd:option("-thread", -1, "number of CPU threads")
|
||||
cmd:option("-tta", 0, 'use TTA mode. It is slow but slightly high quality (0|1)')
|
||||
cmd:option("-tta_level", 8, 'TTA level (2|4|8). A higher value makes better quality output but slow')
|
||||
cmd:option("-force_cudnn", 0, 'use cuDNN backend (0|1)')
|
||||
cmd:option("-q", 0, 'quiet (0|1)')
|
||||
local opt = cmd:parse(arg)
|
||||
if opt.method:len() > 0 then
|
||||
opt.m = opt.method
|
||||
end
|
||||
if opt.thread > 0 then
|
||||
torch.setnumthreads(opt.thread)
|
||||
end
|
||||
if cudnn then
|
||||
cudnn.fastest = true
|
||||
if opt.l:len() > 0 then
|
||||
cudnn.benchmark = true -- find fastest algo
|
||||
else
|
||||
cudnn.benchmark = false
|
||||
end
|
||||
end
|
||||
opt.force_cudnn = opt.force_cudnn == 1
|
||||
opt.q = opt.q == 1
|
||||
opt.model_path = path.join(opt.model_dir, string.format("%s_model.t7", opt.name))
|
||||
if string.len(opt.l) == 0 then
|
||||
convert_image(opt)
|
||||
else
|
||||
convert_frames(opt)
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
waifu2x()
|
||||
|
|
|
|||
Loading…
Reference in a new issue