From 37bc7a5eea8095dfa102debe5672f28fbcb0b422 Mon Sep 17 00:00:00 2001 From: nagadomi Date: Wed, 8 Jun 2016 09:32:27 +0900 Subject: [PATCH] Add support for new noise_scale method --- waifu2x.lua | 90 +++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 63 insertions(+), 27 deletions(-) diff --git a/waifu2x.lua b/waifu2x.lua index 0bc5c7e..2a60f08 100644 --- a/waifu2x.lua +++ b/waifu2x.lua @@ -73,23 +73,41 @@ local function convert_image(opt) new_x = alpha_util.composite(new_x, alpha, model) print(opt.o .. ": " .. (sys.clock() - t) .. " sec") elseif opt.m == "noise_scale" then - local noise_model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level)) - 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 = torch.load(scale_model_path, "ascii") - - if not noise_model then - error("Load Error: " .. noise_model_path) + 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 scale_model = torch.load(scale_model_path, "ascii") + local model = torch.load(model_path, "ascii") + if not model then + error("Load Error: " .. model_path) + end + if not scale_model then + error("Load Error: " .. model_path) + 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.upsampling_filter) + new_x = alpha_util.composite(new_x, alpha, scale_model) + print(opt.o .. ": " .. (sys.clock() - t) .. " sec") + else + local noise_model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level)) + 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 = 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() + x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model)) + x = image_f(noise_model, x, opt.crop_size) + new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.upsampling_filter) + new_x = alpha_util.composite(new_x, alpha, scale_model) + print(opt.o .. ": " .. (sys.clock() - t) .. " sec") end - if not scale_model then - error("Load Error: " .. scale_model_path) - end - 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) - new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.upsampling_filter) - new_x = alpha_util.composite(new_x, alpha, scale_model) - print(opt.o .. ": " .. (sys.clock() - t) .. " sec") else error("undefined method:" .. opt.method) end @@ -97,6 +115,7 @@ local function convert_image(opt) end local function convert_frames(opt) local model_path, scale_model + local noise_scale_model = {} local noise_model = {} local scale_f, image_f if opt.tta == 1 then @@ -119,15 +138,28 @@ local function convert_frames(opt) error("Load Error: " .. model_path) end elseif opt.m == "noise_scale" then - model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) - 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)) - noise_model[opt.noise_level] = torch.load(model_path, "ascii") - if not noise_model[opt.noise_level] then - error("Load Error: " .. model_path) + 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] = 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)) + scale_model = torch.load(model_path, "ascii") + if not scale_model then + error("Load Error: " .. model_path) + end + else + model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale)) + 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)) + 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 local fp = io.open(opt.l) @@ -155,8 +187,12 @@ local function convert_frames(opt) 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)) - x = image_f(noise_model[opt.noise_level], x, opt.crop_size) - new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, upsampling_filter) + if noise_scale_model[opt.noise_level] then + new_x = scale_f(noise_scale_model[opt.noise_level], opt.scale, x, opt.crop_size, upsampling_filter) + else + x = image_f(noise_model[opt.noise_level], x, opt.crop_size) + new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, upsampling_filter) + end new_x = alpha_util.composite(new_x, alpha, scale_model) else error("undefined method:" .. opt.method)