diff --git a/images/gen.sh b/images/gen.sh old mode 100644 new mode 100755 diff --git a/images/miku_CC_BY-NC_noisy_waifu2x.png b/images/miku_CC_BY-NC_noisy_waifu2x.png index f5a520e..5104e46 100644 Binary files a/images/miku_CC_BY-NC_noisy_waifu2x.png and b/images/miku_CC_BY-NC_noisy_waifu2x.png differ diff --git a/images/miku_small_noisy_waifu2x.png b/images/miku_small_noisy_waifu2x.png index 781553c..73a40d6 100644 Binary files a/images/miku_small_noisy_waifu2x.png and b/images/miku_small_noisy_waifu2x.png differ diff --git a/images/miku_small_waifu2x.png b/images/miku_small_waifu2x.png index 0d71947..dd8f36e 100644 Binary files a/images/miku_small_waifu2x.png and b/images/miku_small_waifu2x.png differ diff --git a/lib/reconstruct.lua b/lib/reconstruct.lua index f6ddb22..1b590a1 100644 --- a/lib/reconstruct.lua +++ b/lib/reconstruct.lua @@ -26,7 +26,8 @@ local function reconstruct_layer(model, x, block_size, offset) end return new_x end -local function reconstruct(model, x, offset, block_size) +local reconstruct = {} +function reconstruct.image(model, x, offset, block_size) block_size = block_size or 128 local output_size = block_size - offset * 2 local h_blocks = math.floor(x:size(2) / output_size) + @@ -54,5 +55,37 @@ local function reconstruct(model, x, offset, block_size) return output end +function reconstruct.scale(model, scale, x, offset, block_size) + block_size = block_size or 128 + local x_jinc = iproc.scale(x, x:size(3) * scale, x:size(2) * scale, "Jinc") + x = iproc.scale(x, x:size(3) * scale, x:size(2) * scale, "Box") + + local output_size = block_size - offset * 2 + local h_blocks = math.floor(x:size(2) / output_size) + + ((x:size(2) % output_size == 0 and 0) or 1) + local w_blocks = math.floor(x:size(3) / output_size) + + ((x:size(3) % output_size == 0 and 0) or 1) + + local h = offset + h_blocks * output_size + offset + local w = offset + w_blocks * output_size + offset + local pad_h1 = offset + local pad_w1 = offset + local pad_h2 = (h - offset) - x:size(2) + local pad_w2 = (w - offset) - x:size(3) + local yuv_nn = image.rgb2yuv(iproc.padding(x, pad_w1, pad_w2, pad_h1, pad_h2)) + local yuv_jinc = image.rgb2yuv(iproc.padding(x_jinc, pad_w1, pad_w2, pad_h1, pad_h2)) + local y = reconstruct_layer(model, yuv_nn[1], block_size, offset) + y[torch.lt(y, 0)] = 0 + y[torch.gt(y, 1)] = 1 + yuv_jinc[1]:copy(y) + local output = image.yuv2rgb(image.crop(yuv_jinc, + pad_w1, pad_h1, + yuv_jinc:size(3) - pad_w2, yuv_jinc:size(2) - pad_h2)) + output[torch.lt(output, 0)] = 0 + output[torch.gt(output, 1)] = 1 + collectgarbage() + + return output +end return reconstruct diff --git a/train.lua b/train.lua index 37a002f..4560418 100644 --- a/train.lua +++ b/train.lua @@ -13,15 +13,11 @@ local pairwise_transform = require './lib/pairwise_transform' local image_loader = require './lib/image_loader' local function save_test_scale(model, rgb, file) - local input = iproc.scale(rgb, - rgb:size(3) * settings.scale, - rgb:size(2) * settings.scale) - local up = reconstruct(model, input, settings.block_offset) - + local up = reconstruct.scale(model, settings.scale, rgb, settings.block_offset) image.save(file, up) end local function save_test_jpeg(model, rgb, file) - local im, count = reconstruct(model, rgb, settings.block_offset) + local im, count = reconstruct.image(model, rgb, settings.block_offset) image.save(file, im) end local function split_data(x, test_size) diff --git a/waifu2x.lua b/waifu2x.lua index 96f9948..d50854e 100644 --- a/waifu2x.lua +++ b/waifu2x.lua @@ -37,12 +37,11 @@ local function waifu2x() local model = torch.load(path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level)), "ascii") model:evaluate() - new_x = reconstruct(model, x, BLOCK_OFFSET) + new_x = reconstruct.image(model, x, BLOCK_OFFSET) elseif opt.m == "scale" then local model = torch.load(path.join(opt.model_dir, "scale2.0x_model.t7"), "ascii") model:evaluate() - x = iproc.scale(x, x:size(3) * 2.0, x:size(2) * 2.0) - new_x = reconstruct(model, x, BLOCK_OFFSET) + new_x = reconstruct.scale(model, 2.0, x, BLOCK_OFFSET) elseif opt.m == "noise_scale" then local noise_model = torch.load(path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level)), "ascii") @@ -50,9 +49,8 @@ local function waifu2x() noise_model:evaluate() scale_model:evaluate() - x = reconstruct(noise_model, x, BLOCK_OFFSET) - x = iproc.scale(x, x:size(3) * 2.0, x:size(2) * 2.0) - new_x = reconstruct(scale_model, x, BLOCK_OFFSET) + x = reconstruct.image(noise_model, x, BLOCK_OFFSET) + new_x = reconstruct.scale(scale_model, 2.0, x, BLOCK_OFFSET) else error("undefined method:" .. opt.method) end diff --git a/web.lua b/web.lua index 47b1a92..483e01a 100644 --- a/web.lua +++ b/web.lua @@ -74,15 +74,13 @@ local function get_image(req) end local function apply_denoise1(x) - return reconstruct(noise1_model, x, BLOCK_OFFSET) + return reconstruct.image(noise1_model, x, BLOCK_OFFSET) end local function apply_denoise2(x) - return reconstruct(noise2_model, x, BLOCK_OFFSET) + return reconstruct.image(noise2_model, x, BLOCK_OFFSET) end local function apply_scale2x(x) - return reconstruct(scale20_model, - iproc.scale(x, x:size(3) * 2.0, x:size(2) * 2.0), - BLOCK_OFFSET) + return reconstruct.scale(scale20_model, 2.0, x, BLOCK_OFFSET) end local function cache_do(cache, x, func) if path.exists(cache) then