Add support for user specified pairwise data for universal filter
This commit is contained in:
parent
9ec1f5159b
commit
edac608f18
|
@ -91,6 +91,7 @@ luarocks install graphicsmagick # upgrade
|
|||
luarocks install lua-csnappy
|
||||
luarocks install md5
|
||||
luarocks install uuid
|
||||
luarocks install csvigo
|
||||
PREFIX=$HOME/torch/install luarocks install turbo # if you need to use web application
|
||||
```
|
||||
|
||||
|
|
|
@ -3,6 +3,8 @@ local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^
|
|||
package.path = path.join(path.dirname(__FILE__), "lib", "?.lua;") .. package.path
|
||||
|
||||
require 'image'
|
||||
local cjson = require 'cjson'
|
||||
local csvigo = require 'csvigo'
|
||||
local compression = require 'compression'
|
||||
local settings = require 'settings'
|
||||
local image_loader = require 'image_loader'
|
||||
|
@ -10,6 +12,9 @@ local iproc = require 'iproc'
|
|||
local alpha_util = require 'alpha_util'
|
||||
|
||||
local function crop_if_large(src, max_size)
|
||||
if max_size < 0 then
|
||||
return src
|
||||
end
|
||||
local tries = 4
|
||||
if src:size(2) >= max_size and src:size(3) >= max_size then
|
||||
local rect
|
||||
|
@ -27,25 +32,59 @@ local function crop_if_large(src, max_size)
|
|||
return src
|
||||
end
|
||||
end
|
||||
local function crop_if_large_pair(x, y, max_size)
|
||||
if max_size < 0 then
|
||||
return x, y
|
||||
end
|
||||
local scale_y = y:size(2) / x:size(2)
|
||||
local mod = 4
|
||||
assert(x:size(3) == (y:size(3) / scale_y))
|
||||
|
||||
local tries = 4
|
||||
if y:size(2) > max_size and y:size(3) > max_size then
|
||||
assert(max_size % 4 == 0)
|
||||
local rect_x, rect_y
|
||||
for i = 1, tries do
|
||||
local yi = torch.random(0, y:size(2) - max_size)
|
||||
local xi = torch.random(0, y:size(3) - max_size)
|
||||
if mod then
|
||||
yi = yi - (yi % mod)
|
||||
xi = xi - (xi % mod)
|
||||
end
|
||||
rect_y = iproc.crop(y, xi, yi, xi + max_size, yi + max_size)
|
||||
rect_x = iproc.crop(y, xi / scale_y, yi / scale_y, xi / scale_y + max_size / scale_y, yi / scale_y + max_size / scale_y)
|
||||
-- ignore simple background
|
||||
if rect_y:float():std() >= 0 then
|
||||
break
|
||||
end
|
||||
end
|
||||
return rect_x, rect_y
|
||||
else
|
||||
return x, y
|
||||
end
|
||||
end
|
||||
|
||||
local function load_images(list)
|
||||
local MARGIN = 32
|
||||
local lines = utils.split(file.read(list), "\n")
|
||||
local csv = csvigo.load({path = list, verbose = false, mode = "raw"})
|
||||
local x = {}
|
||||
local skip_notice = false
|
||||
for i = 1, #lines do
|
||||
local line = lines[i]
|
||||
local v = utils.split(line, ",")
|
||||
local filename = v[1]
|
||||
local filters = v[2]
|
||||
if filters then
|
||||
filters = utils.split(filters, ":")
|
||||
for i = 1, #csv do
|
||||
local filename = csv[i][1]
|
||||
local csv_meta = csv[i][2]
|
||||
if csv_meta and csv_meta:len() > 0 then
|
||||
csv_meta = cjson.decode(csv_meta)
|
||||
end
|
||||
if csv_meta.filters then
|
||||
filters = csv_meta.filters
|
||||
end
|
||||
local im, meta = image_loader.load_byte(filename)
|
||||
local skip = false
|
||||
local alpha_color = torch.random(0, 1)
|
||||
|
||||
if meta and meta.alpha then
|
||||
if settings.use_transparent_png then
|
||||
im = alpha_util.fill(im, meta.alpha, torch.random(0, 1))
|
||||
im = alpha_util.fill(im, meta.alpha, alpha_color)
|
||||
else
|
||||
skip = true
|
||||
end
|
||||
|
@ -56,25 +95,35 @@ local function load_images(list)
|
|||
skip_notice = true
|
||||
end
|
||||
else
|
||||
if settings.max_training_image_size > 0 then
|
||||
im = crop_if_large(im, settings.max_training_image_size)
|
||||
end
|
||||
im = iproc.crop_mod4(im)
|
||||
local scale = 1.0
|
||||
if settings.random_half_rate > 0.0 then
|
||||
scale = 2.0
|
||||
end
|
||||
if im then
|
||||
if im:size(2) > (settings.crop_size * scale + MARGIN) and im:size(3) > (settings.crop_size * scale + MARGIN) then
|
||||
table.insert(x, {compression.compress(im), {data = {filters = filters}}})
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: image is too small (%d > size).\n", filename, settings.crop_size * scale + MARGIN))
|
||||
if csv_meta.x then
|
||||
-- method == user
|
||||
local yy = im
|
||||
local xx, meta2 = image_loader.load_byte(csv_meta.x)
|
||||
if meta2 and meta2.alpha then
|
||||
xx = alpha_util.fill(xx, meta2.alpha, alpha_color)
|
||||
end
|
||||
xx, yy = crop_if_large_pair(xx, yy, settings.max_training_image_size)
|
||||
table.insert(x, {{y = compression.compress(yy), x = compression.compress(xx)},
|
||||
{data = {filters = filters, has_x = true}}})
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: load error.\n", filename))
|
||||
im = crop_if_large(im, settings.max_training_image_size)
|
||||
im = iproc.crop_mod4(im)
|
||||
local scale = 1.0
|
||||
if settings.random_half_rate > 0.0 then
|
||||
scale = 2.0
|
||||
end
|
||||
if im then
|
||||
if im:size(2) > (settings.crop_size * scale + MARGIN) and im:size(3) > (settings.crop_size * scale + MARGIN) then
|
||||
table.insert(x, {compression.compress(im), {data = {filters = filters}}})
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: image is too small (%d > size).\n", filename, settings.crop_size * scale + MARGIN))
|
||||
end
|
||||
else
|
||||
io.stderr:write(string.format("\n%s: skip: load error.\n", filename))
|
||||
end
|
||||
end
|
||||
end
|
||||
xlua.progress(i, #lines)
|
||||
xlua.progress(i, #csv)
|
||||
if i % 10 == 0 then
|
||||
collectgarbage()
|
||||
end
|
||||
|
|
|
@ -4,5 +4,6 @@ local pairwise_transform = {}
|
|||
pairwise_transform = tablex.update(pairwise_transform, require('pairwise_transform_scale'))
|
||||
pairwise_transform = tablex.update(pairwise_transform, require('pairwise_transform_jpeg'))
|
||||
pairwise_transform = tablex.update(pairwise_transform, require('pairwise_transform_jpeg_scale'))
|
||||
pairwise_transform = tablex.update(pairwise_transform, require('pairwise_transform_user'))
|
||||
|
||||
return pairwise_transform
|
||||
|
|
60
lib/pairwise_transform_user.lua
Normal file
60
lib/pairwise_transform_user.lua
Normal file
|
@ -0,0 +1,60 @@
|
|||
local pairwise_utils = require 'pairwise_transform_utils'
|
||||
local iproc = require 'iproc'
|
||||
local gm = require 'graphicsmagick'
|
||||
local pairwise_transform = {}
|
||||
|
||||
local function crop_if_large(x, y, scale_y, max_size, mod)
|
||||
local tries = 4
|
||||
if y:size(2) > max_size and y:size(3) > max_size then
|
||||
assert(max_size % 4 == 0)
|
||||
local rect_x, rect_y
|
||||
for i = 1, tries do
|
||||
local yi = torch.random(0, y:size(2) - max_size)
|
||||
local xi = torch.random(0, y:size(3) - max_size)
|
||||
if mod then
|
||||
yi = yi - (yi % mod)
|
||||
xi = xi - (xi % mod)
|
||||
end
|
||||
rect_y = iproc.crop(y, xi, yi, xi + max_size, yi + max_size)
|
||||
rect_x = iproc.crop(x, xi / scale_y, yi / scale_y, xi / scale_y + max_size / scale_y, yi / scale_y + max_size / scale_y)
|
||||
-- ignore simple background
|
||||
if rect_y:float():std() >= 0 then
|
||||
break
|
||||
end
|
||||
end
|
||||
return rect_x, rect_y
|
||||
else
|
||||
return x, y
|
||||
end
|
||||
end
|
||||
function pairwise_transform.user(x, y, size, offset, n, options)
|
||||
assert(x:size(1) == y:size(1))
|
||||
|
||||
local scale_y = y:size(2) / x:size(2)
|
||||
assert(x:size(3) == y:size(3) / scale_y)
|
||||
|
||||
x, y = crop_if_large(x, y, scale_y, options.max_size, options.scale, 2)
|
||||
assert(x:size(3) == y:size(3) / scale_y and x:size(2) == y:size(2) / scale_y)
|
||||
local batch = {}
|
||||
local lowres_y = gm.Image(y, "RGB", "DHW"):
|
||||
size(y:size(3) * 0.5, y:size(2) * 0.5, "Box"):
|
||||
size(y:size(3), y:size(2), "Box"):
|
||||
toTensor(t, "RGB", "DHW")
|
||||
local xs, ys, ls = pairwise_utils.flip_augmentation(x, y, lowres_y)
|
||||
for i = 1, n do
|
||||
local t = (i % #xs) + 1
|
||||
local xc, yc = pairwise_utils.active_cropping(xs[t], ys[t], ls[t], size, scale_y,
|
||||
options.active_cropping_rate,
|
||||
options.active_cropping_tries)
|
||||
xc = iproc.byte2float(xc)
|
||||
yc = iproc.byte2float(yc)
|
||||
if options.rgb then
|
||||
else
|
||||
yc = image.rgb2yuv(yc)[1]:reshape(1, yc:size(2), yc:size(3))
|
||||
xc = image.rgb2yuv(xc)[1]:reshape(1, xc:size(2), xc:size(3))
|
||||
end
|
||||
table.insert(batch, {xc, iproc.crop(yc, offset, offset, size - offset, size - offset)})
|
||||
end
|
||||
return batch
|
||||
end
|
||||
return pairwise_transform
|
41
train.lua
41
train.lua
|
@ -22,6 +22,13 @@ local function save_test_jpeg(model, rgb, file)
|
|||
local im, count = reconstruct.image(model, rgb)
|
||||
image.save(file, im)
|
||||
end
|
||||
local function save_test_user(model, rgb, file)
|
||||
if settings.scale == 1 then
|
||||
save_test_jpeg(model, rgb, file)
|
||||
else
|
||||
save_test_scale(model, rgb, file)
|
||||
end
|
||||
end
|
||||
local function split_data(x, test_size)
|
||||
local index = torch.randperm(#x)
|
||||
local train_size = #x - test_size
|
||||
|
@ -117,9 +124,15 @@ local function create_criterion(model, loss)
|
|||
end
|
||||
local function transformer(model, x, is_validation, n, offset)
|
||||
local meta = {data = {}}
|
||||
local y = nil
|
||||
if type(x) == "table" and type(x[2]) == "table" then
|
||||
meta = x[2]
|
||||
x = compression.decompress(x[1])
|
||||
if x[1].x and x[1].y then
|
||||
y = compression.decompress(x[1].y)
|
||||
x = compression.decompress(x[1].x)
|
||||
else
|
||||
x = compression.decompress(x[1])
|
||||
end
|
||||
else
|
||||
x = compression.decompress(x)
|
||||
end
|
||||
|
@ -197,6 +210,15 @@ local function transformer(model, x, is_validation, n, offset)
|
|||
settings.noise_level,
|
||||
settings.crop_size, offset,
|
||||
n, conf)
|
||||
elseif settings.method == "user" then
|
||||
local conf = tablex.update({
|
||||
max_size = settings.max_size,
|
||||
active_cropping_rate = active_cropping_rate,
|
||||
active_cropping_tries = active_cropping_tries,
|
||||
rgb = (settings.color == "rgb")}, meta)
|
||||
return pairwise_transform.user(x, y,
|
||||
settings.crop_size, offset,
|
||||
n, conf)
|
||||
end
|
||||
end
|
||||
|
||||
|
@ -248,8 +270,12 @@ local function remove_small_image(x)
|
|||
for i = 1, #x do
|
||||
local xe, meta, x_s
|
||||
xe = x[i]
|
||||
if type(xe) == "table" and type(xe[2]) == "table" then
|
||||
x_s = compression.size(xe[1])
|
||||
if type(x) == "table" and type(x[2]) == "table" then
|
||||
if xe[1].x and xe[1].y then
|
||||
x_s = compression.size(xe[1].y) -- y size
|
||||
else
|
||||
x_s = compression.size(xe[1])
|
||||
end
|
||||
else
|
||||
x_s = compression.size(xe)
|
||||
end
|
||||
|
@ -394,6 +420,11 @@ local function train()
|
|||
settings.scale,
|
||||
epoch, i))
|
||||
save_test_scale(model, test_image, log)
|
||||
elseif settings.method == "user" then
|
||||
local log = path.join(settings.model_dir,
|
||||
("%s_best.%d-%d.png"):format(settings.name,
|
||||
epoch, i))
|
||||
save_test_user(model, test_image, log)
|
||||
end
|
||||
else
|
||||
torch.save(settings.model_file, model:clearState(), "ascii")
|
||||
|
@ -410,6 +441,10 @@ local function train()
|
|||
("noise%d_scale%.1f_best.png"):format(settings.noise_level,
|
||||
settings.scale))
|
||||
save_test_scale(model, test_image, log)
|
||||
elseif settings.method == "user" then
|
||||
local log = path.join(settings.model_dir,
|
||||
("%s_best.png"):format(settings.name))
|
||||
save_test_user(model, test_image, log)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
33
waifu2x.lua
33
waifu2x.lua
|
@ -112,6 +112,24 @@ local function convert_image(opt)
|
|||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
end
|
||||
elseif opt.m == "user" then
|
||||
local model_path = opt.model_path
|
||||
local model = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
if not model then
|
||||
error("Load Error: " .. model_path)
|
||||
end
|
||||
local t = sys.clock()
|
||||
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(model))
|
||||
if opt.scale == 1 then
|
||||
new_x = image_f(model, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
new_x = scale_f(model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha) -- TODO: should it use model?
|
||||
if not opt.q then
|
||||
print(opt.o .. ": " .. (sys.clock() - t) .. " sec")
|
||||
end
|
||||
else
|
||||
error("undefined method:" .. opt.method)
|
||||
end
|
||||
|
@ -121,6 +139,7 @@ local function convert_frames(opt)
|
|||
local model_path, scale_model, t
|
||||
local noise_scale_model = {}
|
||||
local noise_model = {}
|
||||
local user_model = nil
|
||||
local scale_f, image_f
|
||||
if opt.tta == 1 then
|
||||
scale_f = function(model, scale, x, block_size, batch_size)
|
||||
|
@ -156,6 +175,8 @@ local function convert_frames(opt)
|
|||
model_path = path.join(opt.model_dir, string.format("noise%d_model.t7", opt.noise_level))
|
||||
noise_model[opt.noise_level] = w2nn.load_model(model_path, opt.force_cudnn)
|
||||
end
|
||||
elseif opt.m == "user" then
|
||||
user_model = w2nn.load_model(opt.model_path, opt.force_cudnn)
|
||||
end
|
||||
local fp = io.open(opt.l)
|
||||
if not fp then
|
||||
|
@ -189,6 +210,14 @@ local function convert_frames(opt)
|
|||
new_x = scale_f(scale_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha, scale_model)
|
||||
elseif opt.m == "user" then
|
||||
x = alpha_util.make_border(x, alpha, reconstruct.offset_size(user_model))
|
||||
if opt.scale == 1 then
|
||||
new_x = image_f(user_model, x, opt.crop_size, opt.batch_size)
|
||||
else
|
||||
new_x = scale_f(user_model, opt.scale, x, opt.crop_size, opt.batch_size)
|
||||
end
|
||||
new_x = alpha_util.composite(new_x, alpha)
|
||||
else
|
||||
error("undefined method:" .. opt.method)
|
||||
end
|
||||
|
@ -218,7 +247,8 @@ local function waifu2x()
|
|||
cmd:option("-o", "(auto)", 'path to output file')
|
||||
cmd:option("-depth", 8, 'bit-depth of the output image (8|16)')
|
||||
cmd:option("-model_dir", "./models/upconv_7/art", 'path to model directory')
|
||||
cmd:option("-m", "noise_scale", 'method (noise|scale|noise_scale)')
|
||||
cmd:option("-name", "user", 'model name for user method')
|
||||
cmd:option("-m", "noise_scale", 'method (noise|scale|noise_scale|user)')
|
||||
cmd:option("-method", "", 'same as -m')
|
||||
cmd:option("-noise_level", 1, '(1|2|3)')
|
||||
cmd:option("-crop_size", 128, 'patch size per process')
|
||||
|
@ -247,6 +277,7 @@ local function waifu2x()
|
|||
end
|
||||
opt.force_cudnn = opt.force_cudnn == 1
|
||||
opt.q = opt.q == 1
|
||||
opt.model_path = path.join(opt.model_dir, string.format("%s_model.t7", opt.name))
|
||||
|
||||
if string.len(opt.l) == 0 then
|
||||
convert_image(opt)
|
||||
|
|
Loading…
Reference in a new issue