refactor
This commit is contained in:
nagadomi
parent
21ea5dd858
commit
4c691b4640
|
|
@ -6,25 +6,7 @@ require 'image'
|
|||
local compression = require 'compression'
|
||||
local settings = require 'settings'
|
||||
local image_loader = require 'image_loader'
|
||||
|
||||
local MAX_SIZE = 1440
|
||||
|
||||
local function crop_if_large(src, max_size)
|
||||
if max_size > 0 and (src:size(2) > max_size or src:size(3) > max_size) then
|
||||
local sx = torch.random(0, src:size(3) - math.min(max_size, src:size(3)))
|
||||
local sy = torch.random(0, src:size(2) - math.min(max_size, src:size(2)))
|
||||
return image.crop(src, sx, sy,
|
||||
math.min(sx + max_size, src:size(3)),
|
||||
math.min(sy + max_size, src:size(2)))
|
||||
else
|
||||
return src
|
||||
end
|
||||
end
|
||||
local function crop_4x(x)
|
||||
local w = x:size(3) % 4
|
||||
local h = x:size(2) % 4
|
||||
return image.crop(x, 0, 0, x:size(3) - w, x:size(2) - h)
|
||||
end
|
||||
local iproc = require 'iproc'
|
||||
|
||||
local function load_images(list)
|
||||
local MARGIN = 32
|
||||
|
|
@ -36,8 +18,7 @@ local function load_images(list)
|
|||
if alpha then
|
||||
io.stderr:write(string.format("\n%s: skip: image has alpha channel.\n", line))
|
||||
else
|
||||
im = crop_if_large(im, settings.max_size)
|
||||
im = crop_4x(im)
|
||||
im = iproc.crop_mod4(im)
|
||||
local scale = 1.0
|
||||
if settings.random_half then
|
||||
scale = 2.0
|
||||
|
|
|
|||
95
lib/data_augmentation.lua
Normal file
95
lib/data_augmentation.lua
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
require 'image'
|
||||
local iproc = require 'iproc'
|
||||
|
||||
local data_augmentation = {}
|
||||
|
||||
local function pcacov(x)
|
||||
local mean = torch.mean(x, 1)
|
||||
local xm = x - torch.ger(torch.ones(x:size(1)), mean:squeeze())
|
||||
local c = torch.mm(xm:t(), xm)
|
||||
c:div(x:size(1) - 1)
|
||||
local ce, cv = torch.symeig(c, 'V')
|
||||
return ce, cv
|
||||
end
|
||||
function data_augmentation.color_noise(src, factor)
|
||||
factor = factor or 0.1
|
||||
local src, conversion = iproc.byte2float(src)
|
||||
local src_t = src:reshape(src:size(1), src:nElement() / src:size(1)):t():contiguous()
|
||||
local ce, cv = pcacov(src_t)
|
||||
local color_scale = torch.Tensor(3):uniform(1 / (1 + factor), 1 + factor)
|
||||
|
||||
pca_space = torch.mm(src_t, cv):t():contiguous()
|
||||
for i = 1, 3 do
|
||||
pca_space[i]:mul(color_scale[i])
|
||||
end
|
||||
local dest = torch.mm(pca_space:t(), cv:t()):t():contiguous():resizeAs(src)
|
||||
dest[torch.lt(dest, 0.0)] = 0.0
|
||||
dest[torch.gt(dest, 1.0)] = 1.0
|
||||
|
||||
if conversion then
|
||||
dest = iproc.float2byte(dest)
|
||||
end
|
||||
return dest
|
||||
end
|
||||
function data_augmentation.shift_1px(src)
|
||||
-- reducing the even/odd issue in nearest neighbor scaler.
|
||||
local direction = torch.random(1, 4)
|
||||
local x_shift = 0
|
||||
local y_shift = 0
|
||||
if direction == 1 then
|
||||
x_shift = 1
|
||||
y_shift = 0
|
||||
elseif direction == 2 then
|
||||
x_shift = 0
|
||||
y_shift = 1
|
||||
elseif direction == 3 then
|
||||
x_shift = 1
|
||||
y_shift = 1
|
||||
elseif flip == 4 then
|
||||
x_shift = 0
|
||||
y_shift = 0
|
||||
end
|
||||
local w = src:size(3) - x_shift
|
||||
local h = src:size(2) - y_shift
|
||||
w = w - (w % 4)
|
||||
h = h - (h % 4)
|
||||
local dest = iproc.crop(src, x_shift, y_shift, x_shift + w, y_shift + h)
|
||||
return dest
|
||||
end
|
||||
function data_augmentation.flip(src)
|
||||
local flip = torch.random(1, 4)
|
||||
local src, conversion = iproc.byte2float(src)
|
||||
local dest
|
||||
|
||||
src = src:contiguous()
|
||||
if flip == 1 then
|
||||
dest = image.hflip(src)
|
||||
elseif flip == 2 then
|
||||
dest = image.vflip(src)
|
||||
elseif flip == 3 then
|
||||
dest = image.hflip(image.vflip(src))
|
||||
elseif flip == 4 then
|
||||
dest = src
|
||||
end
|
||||
if conversion then
|
||||
dest = iproc.float2byte(dest)
|
||||
end
|
||||
return dest
|
||||
end
|
||||
function data_augmentation.overlay(src, p)
|
||||
p = p or 0.25
|
||||
if torch.uniform() < p then
|
||||
local r = torch.uniform(0.2, 0.8)
|
||||
local src, conversion = iproc.byte2float(src)
|
||||
src = src:contiguous()
|
||||
local flip = data_augmentation.flip(src)
|
||||
flip:mul(r):add(src * (1.0 - r))
|
||||
if conversion then
|
||||
flip = iproc.float2byte(flip)
|
||||
end
|
||||
return flip
|
||||
else
|
||||
return src
|
||||
end
|
||||
end
|
||||
return data_augmentation
|
||||
|
|
@ -2,6 +2,38 @@ local gm = require 'graphicsmagick'
|
|||
local image = require 'image'
|
||||
local iproc = {}
|
||||
|
||||
function iproc.crop_mod4(src)
|
||||
local w = src:size(3) % 4
|
||||
local h = src:size(2) % 4
|
||||
return image.crop(src, 0, 0, src:size(3) - w, src:size(2) - h)
|
||||
end
|
||||
function iproc.crop(src, w1, h1, w2, h2)
|
||||
local dest
|
||||
if src:dim() == 3 then
|
||||
dest = src[{{}, { h1 + 1, h2 }, { w1 + 1, w2 }}]:clone()
|
||||
else -- dim == 2
|
||||
dest = src[{{ h1 + 1, h2 }, { w1 + 1, w2 }}]:clone()
|
||||
end
|
||||
return dest
|
||||
end
|
||||
function iproc.byte2float(src)
|
||||
local conversion = false
|
||||
local dest = src
|
||||
if src:type() == "torch.ByteTensor" then
|
||||
conversion = true
|
||||
dest = src:float():div(255.0)
|
||||
end
|
||||
return dest, conversion
|
||||
end
|
||||
function iproc.float2byte(src)
|
||||
local conversion = false
|
||||
local dest = src
|
||||
if src:type() == "torch.FloatTensor" then
|
||||
conversion = true
|
||||
dest = (src * 255.0):byte()
|
||||
end
|
||||
return dest, conversion
|
||||
end
|
||||
function iproc.scale(src, width, height, filter)
|
||||
local t = "float"
|
||||
if src:type() == "torch.ByteTensor" then
|
||||
|
|
@ -22,4 +54,5 @@ function iproc.padding(img, w1, w2, h1, h2)
|
|||
flow[2]:add(-w1)
|
||||
return image.warp(img, flow, "simple", false, "clamp")
|
||||
end
|
||||
|
||||
return iproc
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
require 'image'
|
||||
local gm = require 'graphicsmagick'
|
||||
local iproc = require 'iproc'
|
||||
local reconstruct = require 'reconstruct'
|
||||
local data_augmentation = require 'data_augmentation'
|
||||
|
||||
local pairwise_transform = {}
|
||||
|
||||
local function random_half(src, p)
|
||||
|
|
@ -14,43 +15,52 @@ local function random_half(src, p)
|
|||
return src
|
||||
end
|
||||
end
|
||||
local function pcacov(x)
|
||||
local mean = torch.mean(x, 1)
|
||||
local xm = x - torch.ger(torch.ones(x:size(1)), mean:squeeze())
|
||||
local c = torch.mm(xm:t(), xm)
|
||||
c:div(x:size(1) - 1)
|
||||
local ce, cv = torch.symeig(c, 'V')
|
||||
return ce, cv
|
||||
end
|
||||
local function crop_if_large(src, max_size)
|
||||
if src:size(2) > max_size and src:size(3) > max_size then
|
||||
local yi = torch.random(0, src:size(2) - max_size)
|
||||
local xi = torch.random(0, src:size(3) - max_size)
|
||||
return image.crop(src, xi, yi, xi + max_size, yi + max_size)
|
||||
return iproc.crop(src, xi, yi, xi + max_size, yi + max_size)
|
||||
else
|
||||
return src
|
||||
end
|
||||
end
|
||||
local function active_cropping(x, y, size, offset, p, tries)
|
||||
local function preprocess(src, crop_size, options)
|
||||
local dest = src
|
||||
if options.random_half then
|
||||
dest = random_half(dest)
|
||||
end
|
||||
dest = crop_if_large(dest, math.max(crop_size * 4, 512))
|
||||
dest = data_augmentation.flip(dest)
|
||||
if options.color_noise then
|
||||
dest = data_augmentation.color_noise(dest)
|
||||
end
|
||||
if options.overlay then
|
||||
dest = data_augmentation.overlay(dest)
|
||||
end
|
||||
dest = data_augmentation.shift_1px(dest)
|
||||
|
||||
return dest
|
||||
end
|
||||
local function active_cropping(x, y, size, p, tries)
|
||||
assert("x:size == y:size", x:size(2) == y:size(2) and x:size(3) == y:size(3))
|
||||
local r = torch.uniform()
|
||||
if p < r then
|
||||
local xi = torch.random(offset, y:size(3) - (size + offset + 1))
|
||||
local yi = torch.random(offset, y:size(2) - (size + offset + 1))
|
||||
local xc = image.crop(x, xi, yi, xi + size, yi + size)
|
||||
local yc = image.crop(y, xi, yi, xi + size, yi + size)
|
||||
yc = yc:float():div(255)
|
||||
xc = xc:float():div(255)
|
||||
local xi = torch.random(0, y:size(3) - (size + 1))
|
||||
local yi = torch.random(0, y:size(2) - (size + 1))
|
||||
local xc = iproc.crop(x, xi, yi, xi + size, yi + size)
|
||||
local yc = iproc.crop(y, xi, yi, xi + size, yi + size)
|
||||
return xc, yc
|
||||
else
|
||||
local samples = {}
|
||||
local sum_mse = 0
|
||||
for i = 1, tries do
|
||||
local xi = torch.random(offset, y:size(3) - (size + offset + 1))
|
||||
local yi = torch.random(offset, y:size(2) - (size + offset + 1))
|
||||
local xc = image.crop(x, xi, yi, xi + size, yi + size):float():div(255)
|
||||
local yc = image.crop(y, xi, yi, xi + size, yi + size):float():div(255)
|
||||
local mse = (xc - yc):pow(2):mean()
|
||||
local xi = torch.random(0, y:size(3) - (size + 1))
|
||||
local yi = torch.random(0, y:size(2) - (size + 1))
|
||||
local xc = iproc.crop(x, xi, yi, xi + size, yi + size)
|
||||
local yc = iproc.crop(y, xi, yi, xi + size, yi + size)
|
||||
local xcf = iproc.byte2float(xc)
|
||||
local ycf = iproc.byte2float(yc)
|
||||
local mse = (xcf - ycf):pow(2):mean()
|
||||
sum_mse = sum_mse + mse
|
||||
table.insert(samples, {xc = xc, yc = yc, mse = mse})
|
||||
end
|
||||
|
|
@ -63,87 +73,6 @@ local function active_cropping(x, y, size, offset, p, tries)
|
|||
return samples[1].xc, samples[1].yc
|
||||
end
|
||||
end
|
||||
|
||||
local function color_noise(src)
|
||||
local p = 0.1
|
||||
src = src:float():div(255)
|
||||
local src_t = src:reshape(src:size(1), src:nElement() / src:size(1)):t():contiguous()
|
||||
local ce, cv = pcacov(src_t)
|
||||
local color_scale = torch.Tensor(3):uniform(1 / (1 + p), 1 + p)
|
||||
|
||||
pca_space = torch.mm(src_t, cv):t():contiguous()
|
||||
for i = 1, 3 do
|
||||
pca_space[i]:mul(color_scale[i])
|
||||
end
|
||||
x = torch.mm(pca_space:t(), cv:t()):t():contiguous():resizeAs(src)
|
||||
x[torch.lt(x, 0.0)] = 0.0
|
||||
x[torch.gt(x, 1.0)] = 1.0
|
||||
|
||||
return x:mul(255):byte()
|
||||
end
|
||||
local function shift_1px(src)
|
||||
-- reducing the even/odd issue in nearest neighbor.
|
||||
local r = torch.random(1, 4)
|
||||
|
||||
end
|
||||
local function flip_augment(x, y)
|
||||
local flip = torch.random(1, 4)
|
||||
if y then
|
||||
if flip == 1 then
|
||||
x = image.hflip(x)
|
||||
y = image.hflip(y)
|
||||
elseif flip == 2 then
|
||||
x = image.vflip(x)
|
||||
y = image.vflip(y)
|
||||
elseif flip == 3 then
|
||||
x = image.hflip(image.vflip(x))
|
||||
y = image.hflip(image.vflip(y))
|
||||
elseif flip == 4 then
|
||||
end
|
||||
return x, y
|
||||
else
|
||||
if flip == 1 then
|
||||
x = image.hflip(x)
|
||||
elseif flip == 2 then
|
||||
x = image.vflip(x)
|
||||
elseif flip == 3 then
|
||||
x = image.hflip(image.vflip(x))
|
||||
elseif flip == 4 then
|
||||
end
|
||||
return x
|
||||
end
|
||||
end
|
||||
local function overlay_augment(src, p)
|
||||
p = p or 0.25
|
||||
if torch.uniform() > (1.0 - p) then
|
||||
local r = torch.uniform(0.2, 0.8)
|
||||
local t = "float"
|
||||
if src:type() == "torch.ByteTensor" then
|
||||
src = src:float():div(255)
|
||||
t = "byte"
|
||||
end
|
||||
local flip = flip_augment(src)
|
||||
flip:mul(r):add(src * (1.0 - r))
|
||||
if t == "byte" then
|
||||
flip = flip:mul(255):byte()
|
||||
end
|
||||
return flip
|
||||
else
|
||||
return src
|
||||
end
|
||||
end
|
||||
local function data_augment(y, options)
|
||||
y = flip_augment(y)
|
||||
if options.color_noise then
|
||||
y = color_noise(y)
|
||||
end
|
||||
if options.overlay then
|
||||
y = overlay_augment(y)
|
||||
end
|
||||
return y
|
||||
end
|
||||
|
||||
local INTERPOLATION_PADDING = 16
|
||||
function pairwise_transform.scale(src, scale, size, offset, n, options)
|
||||
local filters = {
|
||||
"Box","Box", -- 0.012756949974688
|
||||
|
|
@ -152,13 +81,11 @@ function pairwise_transform.scale(src, scale, size, offset, n, options)
|
|||
--"Hanning", -- 0.013761314529647
|
||||
--"Hermite", -- 0.013850225205266
|
||||
"SincFast", -- 0.014095824314306
|
||||
--"Jinc", -- 0.014244299255442
|
||||
"Jinc", -- 0.014244299255442
|
||||
}
|
||||
if options.random_half then
|
||||
src = random_half(src)
|
||||
end
|
||||
local downscale_filter = filters[torch.random(1, #filters)]
|
||||
local y = data_augment(crop_if_large(src, math.max(size * 4, 512)), options)
|
||||
local y = preprocess(src, size, options)
|
||||
assert(y:size(2) % 4 == 0 and y:size(3) % 4 == 0)
|
||||
local down_scale = 1.0 / scale
|
||||
local x = iproc.scale(iproc.scale(y, y:size(3) * down_scale,
|
||||
y:size(2) * down_scale, downscale_filter),
|
||||
|
|
@ -167,20 +94,21 @@ function pairwise_transform.scale(src, scale, size, offset, n, options)
|
|||
for i = 1, n do
|
||||
local xc, yc = active_cropping(x, y,
|
||||
size,
|
||||
INTERPOLATION_PADDING,
|
||||
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, image.crop(yc, offset, offset, size - offset, size - offset)})
|
||||
table.insert(batch, {xc, iproc.crop(yc, offset, offset, size - offset, size - offset)})
|
||||
end
|
||||
return batch
|
||||
end
|
||||
function pairwise_transform.jpeg_(src, quality, size, offset, n, options)
|
||||
local y = data_augment(crop_if_large(src, math.max(size * 4, 512)), options)
|
||||
local y = preprocess(src, size, options)
|
||||
local x = y
|
||||
for i = 1, #quality do
|
||||
x = gm.Image(x, "RGB", "DHW")
|
||||
|
|
@ -194,20 +122,21 @@ function pairwise_transform.jpeg_(src, quality, size, offset, n, options)
|
|||
x:fromBlob(blob, len)
|
||||
x = x:toTensor("byte", "RGB", "DHW")
|
||||
end
|
||||
-- TODO: use shift_1px after compression?
|
||||
|
||||
local batch = {}
|
||||
for i = 1, n do
|
||||
local xc, yc = active_cropping(x, y, size, 0,
|
||||
local xc, yc = active_cropping(x, y, size,
|
||||
options.active_cropping_rate,
|
||||
options.active_cropping_tries)
|
||||
xc, yc = flip_augment(xc, yc)
|
||||
|
||||
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, image.crop(yc, offset, offset, size - offset, size - offset)})
|
||||
table.insert(batch, {xc, iproc.crop(yc, offset, offset, size - offset, size - offset)})
|
||||
end
|
||||
return batch
|
||||
end
|
||||
|
|
@ -276,60 +205,36 @@ function pairwise_transform.jpeg(src, category, level, size, offset, n, options)
|
|||
error("unknown category: " .. category)
|
||||
end
|
||||
end
|
||||
local function test_jpeg()
|
||||
local loader = require './image_loader'
|
||||
local src = loader.load_byte("../images/miku_CC_BY-NC.jpg")
|
||||
for i = 2, 9 do
|
||||
local xy = pairwise_transform.jpeg_(random_half(src),
|
||||
{i * 10}, 128, 0, 2, {color_noise = false, random_half = true, overlay = true, rgb = true})
|
||||
for i = 1, #xy do
|
||||
image.display({image = xy[i][1], legend = "y:" .. (i * 10), max=1,min=0})
|
||||
image.display({image = xy[i][2], legend = "x:" .. (i * 10),max=1,min=0})
|
||||
end
|
||||
--print(x:mean(), y:mean())
|
||||
end
|
||||
end
|
||||
|
||||
local function test_scale()
|
||||
torch.setdefaulttensortype('torch.FloatTensor')
|
||||
local loader = require './image_loader'
|
||||
local src = loader.load_byte("../images/miku_CC_BY-NC.jpg")
|
||||
function pairwise_transform.test_jpeg(src)
|
||||
local options = {color_noise = true,
|
||||
random_half = true,
|
||||
overlay = false,
|
||||
active_cropping_rate = 1.5,
|
||||
overlay = true,
|
||||
active_cropping_rate = 0.5,
|
||||
active_cropping_tries = 10,
|
||||
rgb = true
|
||||
}
|
||||
for i = 1, 9 do
|
||||
local xy = pairwise_transform.jpeg(src,
|
||||
"anime_style_art",
|
||||
torch.random(1, 2),
|
||||
128, 7, 1, options)
|
||||
image.display({image = xy[1][1], legend = "y:" .. (i * 10), min=0, max=1})
|
||||
image.display({image = xy[1][2], legend = "x:" .. (i * 10), min=0, max=1})
|
||||
end
|
||||
end
|
||||
function pairwise_transform.test_scale(src)
|
||||
local options = {color_noise = true,
|
||||
random_half = true,
|
||||
overlay = true,
|
||||
active_cropping_rate = 0.5,
|
||||
active_cropping_tries = 10,
|
||||
rgb = true
|
||||
}
|
||||
for i = 1, 10 do
|
||||
local xy = pairwise_transform.scale(src, 2.0, 128, 7, 1, options)
|
||||
image.display({image = xy[1][1], legend = "y:" .. (i * 10), min = 0, max = 1})
|
||||
image.display({image = xy[1][2], legend = "x:" .. (i * 10), min = 0, max = 1})
|
||||
print(xy[1][1]:size(), xy[1][2]:size())
|
||||
--print(x:mean(), y:mean())
|
||||
end
|
||||
end
|
||||
local function test_color_noise()
|
||||
torch.setdefaulttensortype('torch.FloatTensor')
|
||||
local loader = require './image_loader'
|
||||
local src = loader.load_byte("../images/miku_CC_BY-NC.jpg")
|
||||
for i = 1, 10 do
|
||||
image.display(color_noise(src))
|
||||
end
|
||||
end
|
||||
local function test_overlay()
|
||||
torch.setdefaulttensortype('torch.FloatTensor')
|
||||
local loader = require './image_loader'
|
||||
local src = loader.load_byte("../images/miku_CC_BY-NC.jpg")
|
||||
for i = 1, 10 do
|
||||
image.display(overlay_augment(src, 1.0))
|
||||
end
|
||||
end
|
||||
|
||||
--test_scale()
|
||||
--test_jpeg()
|
||||
--test_jpeg_scale()
|
||||
--test_color_noise()
|
||||
--test_overlay()
|
||||
|
||||
return pairwise_transform
|
||||
|
|
|
|||
|
|
@ -117,16 +117,11 @@ local function benchmark(color_weight, x, input_func, v1_noise, v2_noise)
|
|||
end
|
||||
io.stdout:write("\n")
|
||||
end
|
||||
local function crop_4x(x)
|
||||
local w = x:size(3) % 4
|
||||
local h = x:size(2) % 4
|
||||
return image.crop(x, 0, 0, x:size(3) - w, x:size(2) - h)
|
||||
end
|
||||
local function load_data(test_dir)
|
||||
local test_x = {}
|
||||
local files = dir.getfiles(test_dir, "*.*")
|
||||
for i = 1, #files do
|
||||
table.insert(test_x, crop_4x(image_loader.load_byte(files[i])))
|
||||
table.insert(test_x, iproc.crop_mod4(image_loader.load_byte(files[i])))
|
||||
xlua.progress(i, #files)
|
||||
end
|
||||
return test_x
|
||||
|
|
|
|||
Loading…
Reference in a new issue