CodeMirror/waifu2x
1
0
Fork 0
mirror of synced 2024-06-28 19:20:32 +12:00
Code Issues Releases Wiki Activity

refactor

Browse source
This commit is contained in:
nagadomi 2016-06-12 15:56:44 +09:00
parent 6be1479710
commit 599da6a665
1 changed files with 7 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
waifu2x.lua
View file

@ -82,13 +82,10 @@ local function convert_image(opt)
local model_path = path.join(opt.model_dir, ("noise%d_scale%.1fx_model.t7"):format(opt.noise_level, opt.scale)) 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 if path.exists(model_path) then
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
local scale_model = torch.load(scale_model_path, "ascii") local t, scale_model = pcall(torch.load, scale_model_path, "ascii")
local model = torch.load(model_path, "ascii") local model = torch.load(model_path, "ascii")
if not model then if not t then
error("Load Error: " .. model_path) scale_model = model
end
if not scale_model then
error("Load Error: " .. model_path)
end end
local t = sys.clock() local t = sys.clock()
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model)) x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
@ -100,13 +97,6 @@ local function convert_image(opt)
local noise_model = torch.load(noise_model_path, "ascii") local noise_model = torch.load(noise_model_path, "ascii")
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
local scale_model = torch.load(scale_model_path, "ascii") local scale_model = torch.load(scale_model_path, "ascii")
if not noise_model then
error("Load Error: " .. noise_model_path)
end
if not scale_model then
error("Load Error: " .. scale_model_path)
end
local t = sys.clock() local t = sys.clock()
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model)) x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
x = image_f(noise_model, x, opt.crop_size, opt.batch_size) x = image_f(noise_model, x, opt.crop_size, opt.batch_size)
@ -120,7 +110,7 @@ local function convert_image(opt)
image_loader.save_png(opt.o, new_x, tablex.update({depth = opt.depth, inplace = true}, meta)) image_loader.save_png(opt.o, new_x, tablex.update({depth = opt.depth, inplace = true}, meta))
end end
local function convert_frames(opt) local function convert_frames(opt)
local model_path, scale_model local model_path, scale_model, t
local noise_scale_model = {} local noise_scale_model = {}
local noise_model = {} local noise_model = {}
local scale_f, image_f local scale_f, image_f
@ -140,38 +130,23 @@ local function convert_frames(opt)
if opt.m == "scale" then if opt.m == "scale" then
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
scale_model = torch.load(model_path, "ascii") scale_model = torch.load(model_path, "ascii")
if not scale_model then
error("Load Error: " .. model_path)
end
elseif opt.m == "noise" then elseif opt.m == "noise" then
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level)) 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") 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 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)) 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 if path.exists(model_path) then
noise_scale_model[opt.noise_level] = torch.load(model_path, "ascii") noise_scale_model[opt.noise_level] = torch.load(model_path, "ascii")
if not noise_scale_model[opt.noise_level] then
error("Load Error: " .. model_path)
end
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
scale_model = torch.load(model_path, "ascii") t, scale_model = pcall(torch.load, model_path, "ascii")
if not scale_model then if not t then
error("Load Error: " .. model_path) scale_model = noise_scale_model[opt.noise_level]
end end
else else
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
scale_model = torch.load(model_path, "ascii") scale_model = torch.load(model_path, "ascii")
if not scale_model then
error("Load Error: " .. model_path)
end
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level)) 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") 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 end
end end
local fp = io.open(opt.l) local fp = io.open(opt.l)