CodeMirror/waifu2x
1
0
Fork 0
mirror of synced 2024-06-24 17:20:59 +12:00
Code Issues Releases Wiki Activity

Add support for noise level 3

Browse source
This commit is contained in:
nagadomi 2016-03-17 21:21:18 +09:00
parent 4f67f47d3b
commit 8a799e2d56
3 changed files with 18 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
lib/pairwise_transform.lua
View file

@ -168,7 +168,8 @@ function pairwise_transform.jpeg(src, style, level, size, offset, n, options)
if level == 1 then
return pairwise_transform.jpeg_(src, {torch.random(65, 85)},
size, offset, n, options)
elseif level == 2 then
elseif level == 2 or level == 3 then
-- level 2/3 adjusting by -nr_rate. for level3, -nr_rate=1
local r = torch.uniform()
if r > 0.6 then
return pairwise_transform.jpeg_(src, {torch.random(27, 70)},

2
lib/settings.lua
View file

@ -24,7 +24,7 @@ cmd:option("-backend", "cunn", '(cunn|cudnn)')
cmd:option("-test", "images/miku_small.png", 'path to test image')
cmd:option("-model_dir", "./models", 'model directory')
cmd:option("-method", "scale", 'method to training (noise|scale)')
cmd:option("-noise_level", 1, '(1|2)')
cmd:option("-noise_level", 1, '(1|2|3)')
cmd:option("-style", "art", '(art|photo)')
cmd:option("-color", 'rgb', '(y|rgb)')
cmd:option("-random_color_noise_rate", 0.0, 'data augmentation using color noise (0.0-1.0)')

51
waifu2x.lua
View file

@ -69,7 +69,8 @@ local function convert_image(opt)
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
end
local function convert_frames(opt)
local model_path, noise1_model, noise2_model, scale_model
local model_path, scale_model
local noise_model = {}
local scale_f, image_f
if opt.tta == 1 then
scale_f = reconstruct.scale_tta
@ -84,16 +85,10 @@ local function convert_frames(opt)
if not scale_model then
error("Load Error: " .. model_path)
end
elseif opt.m == "noise" and opt.noise_level == 1 then
model_path = path.join(opt.model_dir, "noise1_model.t7")
noise1_model = torch.load(model_path, "ascii")
if not noise1_model then
error("Load Error: " .. model_path)
end
elseif opt.m == "noise" and opt.noise_level == 2 then
model_path = path.join(opt.model_dir, "noise2_model.t7")
noise2_model = torch.load(model_path, "ascii")
if not noise2_model then
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] = torch.load(model_path, "ascii")
if not noise_model[opt.noise_level] then
error("Load Error: " .. model_path)
end
elseif opt.m == "noise_scale" then
@ -102,18 +97,10 @@ local function convert_frames(opt)
if not scale_model then
error("Load Error: " .. model_path)
end
if opt.noise_level == 1 then
model_path = path.join(opt.model_dir, "noise1_model.t7")
noise1_model = torch.load(model_path, "ascii")
if not noise1_model then
error("Load Error: " .. model_path)
end
elseif opt.noise_level == 2 then
model_path = path.join(opt.model_dir, "noise2_model.t7")
noise2_model = torch.load(model_path, "ascii")
if not noise2_model then
error("Load Error: " .. model_path)
end
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
noise_model[opt.noise_level] = torch.load(model_path, "ascii")
if not noise_model[opt.noise_level] then
error("Load Error: " .. model_path)
end
end
local fp = io.open(opt.l)
@ -130,24 +117,16 @@ local function convert_frames(opt)
if opt.resume == 0 or path.exists(string.format(opt.o, i)) == false then
local x, alpha = image_loader.load_float(lines[i])
local new_x = nil
if opt.m == "noise" and opt.noise_level == 1 then
new_x = image_f(noise1_model, x, opt.crop_size)
new_x = alpha_util.composite(new_x, alpha)
elseif opt.m == "noise" and opt.noise_level == 2 then
new_x = image_f(noise2_model, x, opt.crop_size)
if opt.m == "noise" then
new_x = image_f(noise_model[opt.noise_level], x, opt.crop_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)
new_x = alpha_util.composite(new_x, alpha, scale_model)
elseif opt.m == "noise_scale" and opt.noise_level == 1 then
elseif opt.m == "noise_scale" then
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
x = image_f(noise1_model, x, opt.crop_size)
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size)
new_x = alpha_util.composite(new_x, alpha, scale_model)
elseif opt.m == "noise_scale" and opt.noise_level == 2 then
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
x = image_f(noise2_model, x, opt.crop_size)
x = image_f(noise_model[opt.noise_level], x, opt.crop_size)
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size)
new_x = alpha_util.composite(new_x, alpha, scale_model)
else
@ -185,7 +164,7 @@ local function waifu2x()
cmd:option("-depth", 8, 'bit-depth of the output image (8|16)')
cmd:option("-model_dir", "./models/anime_style_art_rgb", 'path to model directory')
cmd:option("-m", "noise_scale", 'method (noise|scale|noise_scale)')
cmd:option("-noise_level", 1, '(1|2)')
cmd:option("-noise_level", 1, '(1|2|3)')
cmd:option("-crop_size", 128, 'patch size per process')
cmd:option("-resume", 0, "skip existing files (0|1)")
cmd:option("-thread", -1, "number of CPU threads")