2015-11-08 22:31:46 +13:00
|
|
|
require 'pl'
|
2015-10-28 19:30:47 +13:00
|
|
|
local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)()
|
|
|
|
package.path = path.join(path.dirname(__FILE__), "lib", "?.lua;") .. package.path
|
2015-05-16 17:48:05 +12:00
|
|
|
require 'sys'
|
2015-10-28 19:30:47 +13:00
|
|
|
require 'w2nn'
|
|
|
|
local iproc = require 'iproc'
|
|
|
|
local reconstruct = require 'reconstruct'
|
|
|
|
local image_loader = require 'image_loader'
|
2015-12-01 21:26:45 +13:00
|
|
|
local alpha_util = require 'alpha_util'
|
2015-05-16 17:48:05 +12:00
|
|
|
|
|
|
|
torch.setdefaulttensortype('torch.FloatTensor')
|
|
|
|
|
2016-04-23 15:32:20 +12:00
|
|
|
local function format_output(opt, src, no)
|
|
|
|
no = no or 1
|
|
|
|
local name = path.basename(src)
|
|
|
|
local e = path.extension(name)
|
|
|
|
local basename = name:sub(0, name:len() - e:len())
|
|
|
|
|
|
|
|
if opt.o == "(auto)" then
|
|
|
|
return path.join(path.dirname(src), string.format("%s_%s.png", basename, opt.m))
|
|
|
|
else
|
|
|
|
local basename_pos = opt.o:find("%%s")
|
|
|
|
local no_pos = opt.o:find("%%%d*d")
|
|
|
|
if basename_pos ~= nil and no_pos ~= nil then
|
|
|
|
if basename_pos < no_pos then
|
|
|
|
return string.format(opt.o, basename, no)
|
|
|
|
else
|
|
|
|
return string.format(opt.o, no, basename)
|
|
|
|
end
|
|
|
|
elseif basename_pos ~= nil then
|
|
|
|
return string.format(opt.o, basename)
|
|
|
|
elseif no_pos ~= nil then
|
|
|
|
return string.format(opt.o, no)
|
|
|
|
else
|
|
|
|
return opt.o
|
|
|
|
end
|
|
|
|
end
|
|
|
|
end
|
|
|
|
|
2015-05-25 01:09:42 +12:00
|
|
|
local function convert_image(opt)
|
2016-04-15 12:13:37 +12:00
|
|
|
local x, meta = image_loader.load_float(opt.i)
|
|
|
|
local alpha = meta.alpha
|
2015-05-25 01:09:42 +12:00
|
|
|
local new_x = nil
|
2015-11-09 08:01:28 +13:00
|
|
|
local scale_f, image_f
|
2015-12-01 21:26:45 +13:00
|
|
|
|
2015-11-09 08:01:28 +13:00
|
|
|
if opt.tta == 1 then
|
2016-06-10 10:34:11 +12:00
|
|
|
scale_f = function(model, scale, x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
return reconstruct.scale_tta(model, opt.tta_level,
|
2016-06-10 10:34:11 +12:00
|
|
|
scale, x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
end
|
2016-06-09 17:03:18 +12:00
|
|
|
image_f = function(model, x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
return reconstruct.image_tta(model, opt.tta_level,
|
2016-06-09 17:03:18 +12:00
|
|
|
x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
end
|
2015-11-09 08:01:28 +13:00
|
|
|
else
|
|
|
|
scale_f = reconstruct.scale
|
|
|
|
image_f = reconstruct.image
|
|
|
|
end
|
2016-04-23 15:32:20 +12:00
|
|
|
opt.o = format_output(opt, opt.i)
|
2015-05-16 17:48:05 +12:00
|
|
|
if opt.m == "noise" then
|
2015-11-05 20:17:55 +13:00
|
|
|
local model_path = path.join(opt.model_dir, ("noise%d_model.t7"):format(opt.noise_level))
|
2016-06-12 19:33:50 +12:00
|
|
|
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
2015-11-05 20:17:55 +13:00
|
|
|
if not model then
|
|
|
|
error("Load Error: " .. model_path)
|
|
|
|
end
|
2016-06-02 13:13:07 +12:00
|
|
|
local t = sys.clock()
|
2016-06-09 17:03:18 +12:00
|
|
|
new_x = image_f(model, x, opt.crop_size, opt.batch_size)
|
2015-12-01 21:26:45 +13:00
|
|
|
new_x = alpha_util.composite(new_x, alpha)
|
2016-06-02 13:13:07 +12:00
|
|
|
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
2015-05-16 17:48:05 +12:00
|
|
|
elseif opt.m == "scale" then
|
2015-11-05 20:17:55 +13:00
|
|
|
local model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
2016-06-12 19:33:50 +12:00
|
|
|
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
2015-11-05 20:17:55 +13:00
|
|
|
if not model then
|
|
|
|
error("Load Error: " .. model_path)
|
|
|
|
end
|
2016-06-02 13:13:07 +12:00
|
|
|
local t = sys.clock()
|
2015-12-01 21:26:45 +13:00
|
|
|
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(model))
|
2016-06-10 10:34:11 +12:00
|
|
|
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size, opt.batch_size)
|
2015-12-01 21:26:45 +13:00
|
|
|
new_x = alpha_util.composite(new_x, alpha, model)
|
2016-06-02 13:13:07 +12:00
|
|
|
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
2015-05-16 17:48:05 +12:00
|
|
|
elseif opt.m == "noise_scale" then
|
2016-06-08 12:32:27 +12:00
|
|
|
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))
|
2016-06-12 19:33:50 +12:00
|
|
|
local t, scale_model = pcall(load_model, scale_model_path, opt.force_cudnn)
|
|
|
|
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
2016-06-12 18:56:44 +12:00
|
|
|
if not t then
|
|
|
|
scale_model = model
|
2016-06-08 12:32:27 +12:00
|
|
|
end
|
|
|
|
local t = sys.clock()
|
|
|
|
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
2016-06-10 10:34:11 +12:00
|
|
|
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size)
|
2016-06-08 12:32:27 +12:00
|
|
|
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))
|
2016-06-12 19:33:50 +12:00
|
|
|
local noise_model = w2nn.load_model(noise_model_path, opt.force_cudnn)
|
2016-06-08 12:32:27 +12:00
|
|
|
local scale_model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
2016-06-12 19:33:50 +12:00
|
|
|
local scale_model = w2nn.load_model(scale_model_path, opt.force_cudnn)
|
2016-06-08 12:32:27 +12:00
|
|
|
local t = sys.clock()
|
|
|
|
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
2016-06-09 17:03:18 +12:00
|
|
|
x = image_f(noise_model, x, opt.crop_size, opt.batch_size)
|
2016-06-10 10:34:11 +12:00
|
|
|
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
2016-06-08 12:32:27 +12:00
|
|
|
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
|
|
|
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
2015-11-05 20:17:55 +13:00
|
|
|
end
|
2015-05-16 17:48:05 +12:00
|
|
|
else
|
|
|
|
error("undefined method:" .. opt.method)
|
|
|
|
end
|
2016-04-15 16:29:50 +12:00
|
|
|
image_loader.save_png(opt.o, new_x, tablex.update({depth = opt.depth, inplace = true}, meta))
|
2015-05-16 17:48:05 +12:00
|
|
|
end
|
2015-05-25 01:09:42 +12:00
|
|
|
local function convert_frames(opt)
|
2016-06-12 18:56:44 +12:00
|
|
|
local model_path, scale_model, t
|
2016-06-08 12:32:27 +12:00
|
|
|
local noise_scale_model = {}
|
2016-03-18 01:21:18 +13:00
|
|
|
local noise_model = {}
|
2015-11-09 08:01:28 +13:00
|
|
|
local scale_f, image_f
|
|
|
|
if opt.tta == 1 then
|
2016-06-10 10:34:11 +12:00
|
|
|
scale_f = function(model, scale, x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
return reconstruct.scale_tta(model, opt.tta_level,
|
2016-06-10 10:34:11 +12:00
|
|
|
scale, x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
end
|
2016-06-09 17:03:18 +12:00
|
|
|
image_f = function(model, x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
return reconstruct.image_tta(model, opt.tta_level,
|
2016-06-09 17:03:18 +12:00
|
|
|
x, block_size, batch_size)
|
2016-06-09 16:08:11 +12:00
|
|
|
end
|
2015-11-09 08:01:28 +13:00
|
|
|
else
|
|
|
|
scale_f = reconstruct.scale
|
|
|
|
image_f = reconstruct.image
|
|
|
|
end
|
2015-11-05 20:17:55 +13:00
|
|
|
if opt.m == "scale" then
|
2015-11-09 08:01:28 +13:00
|
|
|
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
2016-06-12 19:33:50 +12:00
|
|
|
scale_model = w2nn.load_model(model_path, opt.force_cudnn)
|
2016-03-18 01:21:18 +13:00
|
|
|
elseif opt.m == "noise" then
|
|
|
|
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
|
2016-06-12 19:33:50 +12:00
|
|
|
noise_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
2015-11-09 08:01:28 +13:00
|
|
|
elseif opt.m == "noise_scale" then
|
2016-06-08 12:32:27 +12:00
|
|
|
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
|
2016-06-12 19:33:50 +12:00
|
|
|
noise_scale_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
2016-06-08 12:32:27 +12:00
|
|
|
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
2016-06-12 19:33:50 +12:00
|
|
|
t, scale_model = pcall(load_model, model_path, opt.force_cudnn)
|
2016-06-12 18:56:44 +12:00
|
|
|
if not t then
|
|
|
|
scale_model = noise_scale_model[opt.noise_level]
|
2016-06-08 12:32:27 +12:00
|
|
|
end
|
|
|
|
else
|
|
|
|
model_path = path.join(opt.model_dir, ("scale%.1fx_model.t7"):format(opt.scale))
|
2016-06-12 19:33:50 +12:00
|
|
|
scale_model = w2nn.load_model(model_path, opt.force_cudnn)
|
2016-06-08 12:32:27 +12:00
|
|
|
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
|
2016-06-12 19:33:50 +12:00
|
|
|
noise_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
2015-11-09 08:01:28 +13:00
|
|
|
end
|
2015-11-05 20:17:55 +13:00
|
|
|
end
|
2015-05-25 01:09:42 +12:00
|
|
|
local fp = io.open(opt.l)
|
2015-11-05 20:17:55 +13:00
|
|
|
if not fp then
|
|
|
|
error("Open Error: " .. opt.l)
|
|
|
|
end
|
2015-05-25 01:09:42 +12:00
|
|
|
local count = 0
|
|
|
|
local lines = {}
|
|
|
|
for line in fp:lines() do
|
|
|
|
table.insert(lines, line)
|
|
|
|
end
|
|
|
|
fp:close()
|
|
|
|
for i = 1, #lines do
|
2016-04-23 15:32:20 +12:00
|
|
|
local output = format_output(opt, lines[i], i)
|
|
|
|
if opt.resume == 0 or path.exists(output) == false then
|
2016-04-15 12:13:37 +12:00
|
|
|
local x, meta = image_loader.load_float(lines[i])
|
|
|
|
local alpha = meta.alpha
|
2015-06-16 23:41:48 +12:00
|
|
|
local new_x = nil
|
2016-03-18 01:21:18 +13:00
|
|
|
if opt.m == "noise" then
|
2016-06-09 17:03:18 +12:00
|
|
|
new_x = image_f(noise_model[opt.noise_level], x, opt.crop_size, opt.batch_size)
|
2015-12-01 21:26:45 +13:00
|
|
|
new_x = alpha_util.composite(new_x, alpha)
|
2015-06-16 23:41:48 +12:00
|
|
|
elseif opt.m == "scale" then
|
2015-12-01 21:26:45 +13:00
|
|
|
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
2016-06-10 10:34:11 +12:00
|
|
|
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
2015-12-01 21:26:45 +13:00
|
|
|
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
2016-03-18 01:21:18 +13:00
|
|
|
elseif opt.m == "noise_scale" then
|
2015-12-01 21:26:45 +13:00
|
|
|
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(scale_model))
|
2016-06-08 12:32:27 +12:00
|
|
|
if noise_scale_model[opt.noise_level] then
|
2016-06-10 10:34:11 +12:00
|
|
|
new_x = scale_f(noise_scale_model[opt.noise_level], opt.scale, x, opt.crop_size, opt.batch_size)
|
2016-06-08 12:32:27 +12:00
|
|
|
else
|
2016-06-09 17:03:18 +12:00
|
|
|
x = image_f(noise_model[opt.noise_level], x, opt.crop_size, opt.batch_size)
|
2016-06-10 10:34:11 +12:00
|
|
|
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
2016-06-08 12:32:27 +12:00
|
|
|
end
|
2015-12-01 21:26:45 +13:00
|
|
|
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
2015-06-16 23:41:48 +12:00
|
|
|
else
|
|
|
|
error("undefined method:" .. opt.method)
|
|
|
|
end
|
2016-04-15 12:13:37 +12:00
|
|
|
image_loader.save_png(output, new_x,
|
2016-04-15 16:29:50 +12:00
|
|
|
tablex.update({depth = opt.depth, inplace = true}, meta))
|
2015-06-16 23:41:48 +12:00
|
|
|
xlua.progress(i, #lines)
|
|
|
|
if i % 10 == 0 then
|
|
|
|
collectgarbage()
|
|
|
|
end
|
|
|
|
else
|
|
|
|
xlua.progress(i, #lines)
|
|
|
|
end
|
2015-05-26 02:22:29 +12:00
|
|
|
end
|
|
|
|
end
|
2015-05-25 01:09:42 +12:00
|
|
|
local function waifu2x()
|
|
|
|
local cmd = torch.CmdLine()
|
|
|
|
cmd:text()
|
|
|
|
cmd:text("waifu2x")
|
|
|
|
cmd:text("Options:")
|
2015-11-06 14:08:54 +13:00
|
|
|
cmd:option("-i", "images/miku_small.png", 'path to input image')
|
|
|
|
cmd:option("-l", "", 'path to image-list.txt')
|
2015-05-25 01:09:42 +12:00
|
|
|
cmd:option("-scale", 2, 'scale factor')
|
2015-11-06 14:08:54 +13:00
|
|
|
cmd:option("-o", "(auto)", 'path to output file')
|
2015-11-08 03:01:57 +13:00
|
|
|
cmd:option("-depth", 8, 'bit-depth of the output image (8|16)')
|
2016-05-15 06:04:08 +12:00
|
|
|
cmd:option("-model_dir", "./models/upconv_7/art", 'path to model directory')
|
2015-05-25 01:09:42 +12:00
|
|
|
cmd:option("-m", "noise_scale", 'method (noise|scale|noise_scale)')
|
2016-03-18 01:21:18 +13:00
|
|
|
cmd:option("-noise_level", 1, '(1|2|3)')
|
2015-05-25 01:09:42 +12:00
|
|
|
cmd:option("-crop_size", 128, 'patch size per process')
|
2016-06-09 17:03:18 +12:00
|
|
|
cmd:option("-batch_size", 1, 'batch_size')
|
2015-05-26 03:35:34 +12:00
|
|
|
cmd:option("-resume", 0, "skip existing files (0|1)")
|
2015-10-28 19:30:47 +13:00
|
|
|
cmd:option("-thread", -1, "number of CPU threads")
|
2015-11-09 08:01:28 +13:00
|
|
|
cmd:option("-tta", 0, '8x slower and slightly high quality (0|1)')
|
2016-06-09 16:08:11 +12:00
|
|
|
cmd:option("-tta_level", 8, 'TTA level (2|4|8)')
|
2016-06-12 19:33:50 +12:00
|
|
|
cmd:option("-force_cudnn", 0, 'use cuDNN backend (0|1)')
|
|
|
|
|
2015-05-25 01:09:42 +12:00
|
|
|
local opt = cmd:parse(arg)
|
2015-10-28 19:30:47 +13:00
|
|
|
if opt.thread > 0 then
|
|
|
|
torch.setnumthreads(opt.thread)
|
|
|
|
end
|
2015-11-01 02:09:21 +13:00
|
|
|
if cudnn then
|
|
|
|
cudnn.fastest = true
|
2016-06-12 19:33:50 +12:00
|
|
|
if opt.l:len() > 0 then
|
|
|
|
cudnn.benchmark = true -- find fastest algo
|
|
|
|
else
|
|
|
|
cudnn.benchmark = false
|
|
|
|
end
|
|
|
|
end
|
|
|
|
if opt.force_cudnn == 1 then
|
|
|
|
opt.force_cudnn = true
|
|
|
|
else
|
|
|
|
opt.force_cudnn = false
|
2015-11-01 02:09:21 +13:00
|
|
|
end
|
2015-05-25 01:09:42 +12:00
|
|
|
if string.len(opt.l) == 0 then
|
|
|
|
convert_image(opt)
|
|
|
|
else
|
|
|
|
convert_frames(opt)
|
|
|
|
end
|
|
|
|
end
|
2015-05-16 17:48:05 +12:00
|
|
|
waifu2x()
|