diff --git a/assets/index.html b/assets/index.html index 6a18979..8ceb94a 100644 --- a/assets/index.html +++ b/assets/index.html @@ -2,51 +2,17 @@ - waifu2x - + - +

waifu2x

-
- Fork me on GitHub - +
+ Fork me on GitHub +
en/ja/ru
@@ -67,11 +33,18 @@
+ Style + + +
+
Noise Reduction (expect JPEG Artifact) -
When using 2x scaling, we never recommend to use high level of noise reduction, it almost always makes image worse, it makes sense for only some rare cases when image had really bad quality from the beginning.
+
+ When using 2x scaling, we never recommend to use high level of noise reduction, it almost always makes image worse, it makes sense for only some rare cases when image had really bad quality from the beginning. +
Upscaling @@ -82,7 +55,7 @@
-
diff --git a/assets/index.ja.html b/assets/index.ja.html index 357e82c..5200e51 100644 --- a/assets/index.ja.html +++ b/assets/index.ja.html @@ -2,51 +2,17 @@ - + waifu2x - - +

waifu2x

-
- Fork me on GitHub - +
+ Fork me on GitHub +
en/ja/ru
@@ -67,6 +33,11 @@
+ スタイル + + +
+
ノイズ除去 (JPEGノイズを想定) @@ -81,7 +52,7 @@
-
- Устранение шума (артефактов JPEG) + Стиль + + +
+
+ Устранение шума (артефактов JPEG) - +
Устранение шума нужно использовать, если на картинке действительно есть шум, иначе это даст противоположный эффект. Также не рекомендуется сильное устранение шума, оно даёт выгоду только в редких случаях, когда картинка изначально была сильно испорчена.
@@ -82,8 +54,9 @@
-
diff --git a/assets/style.css b/assets/style.css new file mode 100644 index 0000000..46342fe --- /dev/null +++ b/assets/style.css @@ -0,0 +1,46 @@ +body { + margin: 1em 2em 1em 2em; + background: LightGray; + width: 640px; +} +fieldset { + margin-top: 1em; + margin-bottom: 1em; +} +.about { + position: relative; + display: inline-block; + font-size: 0.9em; + padding: 1em 5px 0.2em 0; +} +.help { + font-size: 0.85em; + margin: 1em 0 0 0; +} +.github-banner { + position:absolute; + display:block; + top:0; + left:540px; + max-height:140px; +} +.github-banner-image { + position: absolute; + display: block; + left: 0; + top: 0; + width: 149px; + height: 149px; + border: 0; +} +.github-banner-link { + position: absolute; + display: block; + left:0; + top:0; + width:149px; + height:130px; +} +.padding-left { + padding-left: 15px; +} diff --git a/assets/ui.js b/assets/ui.js new file mode 100644 index 0000000..8fc27b9 --- /dev/null +++ b/assets/ui.js @@ -0,0 +1,24 @@ +$(function (){ + function clear_file() { + var new_file = $("#file").clone(); + new_file.change(clear_url); + $("#file").replaceWith(new_file); + } + function clear_url() { + $("#url").val("") + } + function on_change_style(e) { + var style = $("input[name=style]:checked").val() + if (style == "photo") { + $("input[name=noise]").prop("disabled", true); + $(".noise-field").hide() + } else { + $("input[name=noise]").prop("disabled", false); + $(".noise-field").show(); + } + } + + $("#url").change(clear_file); + $("#file").change(clear_url); + $("input[name=style]").change(on_change_style); +}) diff --git a/lib/cleanup_model.lua b/lib/cleanup_model.lua new file mode 100644 index 0000000..1784992 --- /dev/null +++ b/lib/cleanup_model.lua @@ -0,0 +1,48 @@ +-- ref: https://github.com/torch/nn/issues/112#issuecomment-64427049 + +local function zeroDataSize(data) + if type(data) == 'table' then + for i = 1, #data do + data[i] = zeroDataSize(data[i]) + end + elseif type(data) == 'userdata' then + data = torch.Tensor():typeAs(data) + end + return data +end +-- Resize the output, gradInput, etc temporary tensors to zero (so that the +-- on disk size is smaller) +local function cleanupModel(node) + if node.output ~= nil then + node.output = zeroDataSize(node.output) + end + if node.gradInput ~= nil then + node.gradInput = zeroDataSize(node.gradInput) + end + if node.finput ~= nil then + node.finput = zeroDataSize(node.finput) + end + if tostring(node) == "nn.LeakyReLU" or tostring(node) == "w2nn.LeakyReLU" then + if node.negative ~= nil then + node.negative = zeroDataSize(node.negative) + end + end + if tostring(node) == "nn.Dropout" then + if node.noise ~= nil then + node.noise = zeroDataSize(node.noise) + end + end + -- Recurse on nodes with 'modules' + if (node.modules ~= nil) then + if (type(node.modules) == 'table') then + for i = 1, #node.modules do + local child = node.modules[i] + cleanupModel(child) + end + end + end +end +function w2nn.cleanup_model(model) + cleanupModel(model) + return model +end diff --git a/lib/w2nn.lua b/lib/w2nn.lua index 6efb5f9..ce53823 100644 --- a/lib/w2nn.lua +++ b/lib/w2nn.lua @@ -21,5 +21,6 @@ else require 'DepthExpand2x' require 'WeightedMSECriterion' require 'WeightedHuberCriterion' + require 'cleanup_model' return w2nn end diff --git a/tools/cleanup_model.lua b/tools/cleanup_model.lua index 408ae5d..40d6b26 100644 --- a/tools/cleanup_model.lua +++ b/tools/cleanup_model.lua @@ -4,53 +4,6 @@ package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. packa require 'w2nn' torch.setdefaulttensortype("torch.FloatTensor") --- ref: https://github.com/torch/nn/issues/112#issuecomment-64427049 -local function zeroDataSize(data) - if type(data) == 'table' then - for i = 1, #data do - data[i] = zeroDataSize(data[i]) - end - elseif type(data) == 'userdata' then - data = torch.Tensor():typeAs(data) - end - return data -end - --- Resize the output, gradInput, etc temporary tensors to zero (so that the --- on disk size is smaller) -local function cleanupModel(node) - if node.output ~= nil then - node.output = zeroDataSize(node.output) - end - if node.gradInput ~= nil then - node.gradInput = zeroDataSize(node.gradInput) - end - if node.finput ~= nil then - node.finput = zeroDataSize(node.finput) - end - if tostring(node) == "nn.LeakyReLU" or tostring(node) == "w2nn.LeakyReLU" then - if node.negative ~= nil then - node.negative = zeroDataSize(node.negative) - end - end - if tostring(node) == "nn.Dropout" then - if node.noise ~= nil then - node.noise = zeroDataSize(node.noise) - end - end - -- Recurse on nodes with 'modules' - if (node.modules ~= nil) then - if (type(node.modules) == 'table') then - for i = 1, #node.modules do - local child = node.modules[i] - cleanupModel(child) - end - end - end - - collectgarbage() -end - local cmd = torch.CmdLine() cmd:text() cmd:text("cleanup model") @@ -62,7 +15,7 @@ cmd:option("-oformat", "binary", 'output format') local opt = cmd:parse(arg) local model = torch.load(opt.model, opt.iformat) if model then - cleanupModel(model) + w2nn.cleanup_model(model) torch.save(opt.model, model, opt.oformat) else error("model not found") diff --git a/web.lua b/web.lua index 2160da5..09696a8 100644 --- a/web.lua +++ b/web.lua @@ -1,10 +1,10 @@ local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)() -package.path = path.join(path.dirname(__FILE__), "lib", "?.lua;") .. package.path +require 'pl' +local ROOT = path.dirname(__FILE__) +package.path = path.join(ROOT, "lib", "?.lua;") .. package.path _G.TURBO_SSL = true -require 'pl' require 'w2nn' -local turbo = require 'turbo' local uuid = require 'uuid' local ffi = require 'ffi' local md5 = require 'md5' @@ -12,6 +12,11 @@ local iproc = require 'iproc' local reconstruct = require 'reconstruct' local image_loader = require 'image_loader' +-- Notes: turbo and xlua has different implementation of string:split(). +-- Therefore, string:split() has conflict issue. +-- In this script, use turbo's string:split(). +local turbo = require 'turbo' + local cmd = torch.CmdLine() cmd:text() cmd:text("waifu2x-api") @@ -29,14 +34,14 @@ if cudnn then cudnn.fastest = true cudnn.benchmark = false end +local ART_MODEL_DIR = path.join(ROOT, "models", "anime_style_art_rgb") +local PHOTO_MODEL_DIR = path.join(ROOT, "models", "ukbench") +local art_noise1_model = torch.load(path.join(ART_MODEL_DIR, "noise1_model.t7"), "ascii") +local art_noise2_model = torch.load(path.join(ART_MODEL_DIR, "noise2_model.t7"), "ascii") +local art_scale2_model = torch.load(path.join(ART_MODEL_DIR, "scale2.0x_model.t7"), "ascii") +local photo_scale2_model = torch.load(path.join(PHOTO_MODEL_DIR, "scale2.0x_model.t7"), "ascii") -local MODEL_DIR = "./models/anime_style_art_rgb" -local noise1_model = torch.load(path.join(MODEL_DIR, "noise1_model.t7"), "ascii") -local noise2_model = torch.load(path.join(MODEL_DIR, "noise2_model.t7"), "ascii") -local scale20_model = torch.load(path.join(MODEL_DIR, "scale2.0x_model.t7"), "ascii") - -local USE_CACHE = true -local CACHE_DIR = "./cache" +local CACHE_DIR = path.join(ROOT, "cache") local MAX_NOISE_IMAGE = 2560 * 2560 local MAX_SCALE_IMAGE = 1280 * 1280 local CURL_OPTIONS = { @@ -55,15 +60,6 @@ local function valid_size(x, scale) end end -local function apply_denoise1(x) - return reconstruct.image(noise1_model, x) -end -local function apply_denoise2(x) - return reconstruct.image(noise2_model, x) -end -local function apply_scale2x(x) - return reconstruct.scale(scale20_model, 2.0, x) -end local function cache_url(url) local hash = md5.sumhexa(url) local cache_file = path.join(CACHE_DIR, "url_" .. hash) @@ -91,15 +87,6 @@ local function cache_url(url) end return nil, nil, nil end -local function cache_do(cache, x, func) - if path.exists(cache) then - return image.load(cache) - else - x = func(x) - image.save(cache, x) - return x - end -end local function get_image(req) local file = req:get_argument("file", "") local url = req:get_argument("url", "") @@ -114,7 +101,30 @@ local function get_image(req) end return nil, nil, nil end - +local function convert(x, options) + local cache_file = path.join(CACHE_DIR, options.prefix .. ".png") + if path.exists(cache_file) then + return image.load(cache_file) + else + if options.style == "art" then + if options.method == "scale" then + x = reconstruct.scale(art_scale2_model, 2.0, x) + w2nn.cleanup_model(art_scale2_model) + elseif options.method == "noise1" then + x = reconstruct.image(art_noise1_model, x) + w2nn.cleanup_model(art_noise1_model) + else -- options.method == "noise2" + x = reconstruct.image(art_noise2_model, x) + w2nn.cleanup_model(art_noise2_model) + end + else -- photo + x = reconstruct.scale(photo_scale2_model, 2.0, x) + w2nn.cleanup_model(photo_scale2_model) + end + image.save(cache_file, x) + return x + end +end local function client_disconnected(handler) return not(handler.request and handler.request.connection and @@ -129,30 +139,28 @@ function APIHandler:post() self:write("client disconnected") return end - local x, alpha, src = get_image(self) + local x, alpha, blob = get_image(self) local scale = tonumber(self:get_argument("scale", "0")) local noise = tonumber(self:get_argument("noise", "0")) + local style = self:get_argument("style", "art") + if style ~= "art" then + style = "photo" -- style must be art or photo + end if x and valid_size(x, scale) then - if USE_CACHE and (noise ~= 0 or scale ~= 0) then - local hash = md5.sumhexa(src) - local cache_noise1 = path.join(CACHE_DIR, hash .. "_noise1.png") - local cache_noise2 = path.join(CACHE_DIR, hash .. "_noise2.png") - local cache_scale = path.join(CACHE_DIR, hash .. "_scale.png") - local cache_noise1_scale = path.join(CACHE_DIR, hash .. "_noise1_scale.png") - local cache_noise2_scale = path.join(CACHE_DIR, hash .. "_noise2_scale.png") - + if (noise ~= 0 or scale ~= 0) then + local hash = md5.sumhexa(blob) if noise == 1 then - x = cache_do(cache_noise1, x, apply_denoise1) + x = convert(x, {method = "noise1", style = style, prefix = style .. "_noise1_" .. hash}) elseif noise == 2 then - x = cache_do(cache_noise2, x, apply_denoise2) + x = convert(x, {method = "noise2", style = style, prefix = style .. "_noise2_" .. hash}) end if scale == 1 or scale == 2 then if noise == 1 then - x = cache_do(cache_noise1_scale, x, apply_scale2x) + x = convert(x, {method = "scale", style = style, prefix = style .. "_noise1_scale_" .. hash}) elseif noise == 2 then - x = cache_do(cache_noise2_scale, x, apply_scale2x) + x = convert(x, {method = "scale", style = style, prefix = style .. "_noise2_scale_" .. hash}) else - x = cache_do(cache_scale, x, apply_scale2x) + x = convert(x, {method = "scale", style = style, prefix = style .. "_scale_" .. hash}) end if scale == 1 then x = iproc.scale(x, @@ -161,23 +169,9 @@ function APIHandler:post() "Jinc") end end - elseif noise ~= 0 or scale ~= 0 then - if noise == 1 then - x = apply_denoise1(x) - elseif noise == 2 then - x = apply_denoise2(x) - end - if scale == 1 then - local x16 = {math.floor(x:size(3) * 1.6 + 0.5), math.floor(x:size(2) * 1.6 + 0.5)} - x = apply_scale2x(x) - x = iproc.scale(x, x16[1], x16[2], "Jinc") - elseif scale == 2 then - x = apply_scale2x(x) - end end local name = uuid() .. ".png" local blob, len = image_loader.encode_png(x, alpha) - self:set_header("Content-Disposition", string.format('filename="%s"', name)) self:set_header("Content-Type", "image/png") self:set_header("Content-Length", string.format("%d", len)) @@ -194,9 +188,9 @@ function APIHandler:post() collectgarbage() end local FormHandler = class("FormHandler", turbo.web.RequestHandler) -local index_ja = file.read("./assets/index.ja.html") -local index_ru = file.read("./assets/index.ru.html") -local index_en = file.read("./assets/index.html") +local index_ja = file.read(path.join(ROOT, "assets", "index.ja.html")) +local index_ru = file.read(path.join(ROOT, "assets", "index.ru.html")) +local index_en = file.read(path.join(ROOT, "assets", "index.html")) function FormHandler:get() local lang = self.request.headers:get("Accept-Language") if lang then @@ -226,9 +220,11 @@ turbo.log.categories = { local app = turbo.web.Application:new( { {"^/$", FormHandler}, - {"^/index.html", turbo.web.StaticFileHandler, path.join("./assets", "index.html")}, - {"^/index.ja.html", turbo.web.StaticFileHandler, path.join("./assets", "index.ja.html")}, - {"^/index.ru.html", turbo.web.StaticFileHandler, path.join("./assets", "index.ru.html")}, + {"^/style.css", turbo.web.StaticFileHandler, path.join(ROOT, "assets", "style.css")}, + {"^/ui.js", turbo.web.StaticFileHandler, path.join(ROOT, "assets", "ui.js")}, + {"^/index.html", turbo.web.StaticFileHandler, path.join(ROOT, "assets", "index.html")}, + {"^/index.ja.html", turbo.web.StaticFileHandler, path.join(ROOT, "assets", "index.ja.html")}, + {"^/index.ru.html", turbo.web.StaticFileHandler, path.join(ROOT, "assets", "index.ru.html")}, {"^/api$", APIHandler}, } )