279 lines
7.9 KiB
Lua
279 lines
7.9 KiB
Lua
require 'pl'
|
|
require 'cunn'
|
|
local iproc = require 'iproc'
|
|
local gm = {}
|
|
gm.Image = require 'graphicsmagick.Image'
|
|
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 random_rect_size(rect_min, rect_max)
|
|
local r = torch.Tensor(2):uniform():cmul(torch.Tensor({rect_max - rect_min, rect_max - rect_min})):int()
|
|
local rect_h = r[1] + rect_min
|
|
local rect_w = r[2] + rect_min
|
|
return rect_h, rect_w
|
|
end
|
|
function random_rect(height, width, rect_h, rect_w)
|
|
local r = torch.Tensor(2):uniform():cmul(torch.Tensor({height - 1 - rect_h, width-1 - rect_w})):int()
|
|
local rect_y1 = r[1] + 1
|
|
local rect_x1 = r[2] + 1
|
|
local rect_x2 = rect_x1 + rect_w
|
|
local rect_y2 = rect_y1 + rect_h
|
|
return {x1 = rect_x1, y1 = rect_y1, x2 = rect_x2, y2 = rect_y2}
|
|
end
|
|
function data_augmentation.erase(src, p, n, rect_min, rect_max)
|
|
if torch.uniform() < p then
|
|
local src, conversion = iproc.byte2float(src)
|
|
src = src:contiguous():clone()
|
|
local ch = src:size(1)
|
|
local height = src:size(2)
|
|
local width = src:size(3)
|
|
for i = 1, n do
|
|
local rect_h, rect_w = random_rect_size(rect_min, rect_max)
|
|
local rect1 = random_rect(height, width, rect_h, rect_w)
|
|
local rect2 = random_rect(height, width, rect_h, rect_w)
|
|
dest_rect = src:sub(1, ch, rect1.y1, rect1.y2, rect1.x1, rect1.x2)
|
|
src_rect = src:sub(1, ch, rect2.y1, rect2.y2, rect2.x1, rect2.x2)
|
|
dest_rect:copy(src_rect:clone())
|
|
end
|
|
if conversion then
|
|
src = iproc.float2byte(src)
|
|
end
|
|
return src
|
|
else
|
|
return src
|
|
end
|
|
end
|
|
function data_augmentation.color_noise(src, p, factor)
|
|
factor = factor or 0.1
|
|
if torch.uniform() < p then
|
|
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:clamp(0.0, 1.0)
|
|
|
|
if conversion then
|
|
dest = iproc.float2byte(dest)
|
|
end
|
|
return dest
|
|
else
|
|
return src
|
|
end
|
|
end
|
|
function data_augmentation.overlay(src, p)
|
|
if torch.uniform() < p then
|
|
local r = torch.uniform()
|
|
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
|
|
function data_augmentation.unsharp_mask(src, p)
|
|
if torch.uniform() < p then
|
|
local radius = 0 -- auto
|
|
local sigma = torch.uniform(0.5, 1.5)
|
|
local amount = torch.uniform(0.1, 0.9)
|
|
local threshold = torch.uniform(0.0, 0.05)
|
|
local unsharp = gm.Image(src, "RGB", "DHW"):
|
|
unsharpMask(radius, sigma, amount, threshold):
|
|
toTensor("float", "RGB", "DHW")
|
|
|
|
if src:type() == "torch.ByteTensor" then
|
|
return iproc.float2byte(unsharp)
|
|
else
|
|
return unsharp
|
|
end
|
|
else
|
|
return src
|
|
end
|
|
end
|
|
function data_augmentation.blur(src, p, size, sigma_min, sigma_max)
|
|
size = size or "3"
|
|
filters = utils.split(size, ",")
|
|
for i = 1, #filters do
|
|
local s = tonumber(filters[i])
|
|
filters[i] = s
|
|
end
|
|
if torch.uniform() < p then
|
|
local src, conversion = iproc.byte2float(src)
|
|
local kernel_size = filters[torch.random(1, #filters)]
|
|
local sigma
|
|
if sigma_min == sigma_max then
|
|
sigma = sigma_min
|
|
else
|
|
sigma = torch.uniform(sigma_min, sigma_max)
|
|
end
|
|
local kernel = iproc.gaussian2d(kernel_size, sigma)
|
|
local dest = image.convolve(src, kernel, 'same')
|
|
if conversion then
|
|
dest = iproc.float2byte(dest)
|
|
end
|
|
return dest
|
|
else
|
|
return src
|
|
end
|
|
end
|
|
function data_augmentation.pairwise_scale(x, y, p, scale_min, scale_max)
|
|
if torch.uniform() < p then
|
|
assert(x:size(2) == y:size(2) and x:size(3) == y:size(3))
|
|
local scale = torch.uniform(scale_min, scale_max)
|
|
local h = math.floor(x:size(2) * scale)
|
|
local w = math.floor(x:size(3) * scale)
|
|
local filters = {"Lanczos", "Catrom"}
|
|
local x_filter = filters[torch.random(1, 2)]
|
|
x = iproc.scale(x, w, h, x_filter)
|
|
y = iproc.scale(y, w, h, "Triangle")
|
|
return x, y
|
|
else
|
|
return x, y
|
|
end
|
|
end
|
|
function data_augmentation.pairwise_rotate(x, y, p, r_min, r_max)
|
|
if torch.uniform() < p then
|
|
assert(x:size(2) == y:size(2) and x:size(3) == y:size(3))
|
|
local r = torch.uniform(r_min, r_max) / 360.0 * math.pi
|
|
x = iproc.rotate(x, r)
|
|
y = iproc.rotate(y, r)
|
|
return x, y
|
|
else
|
|
return x, y
|
|
end
|
|
end
|
|
function data_augmentation.pairwise_negate(x, y, p)
|
|
if torch.uniform() < p then
|
|
assert(x:size(2) == y:size(2) and x:size(3) == y:size(3))
|
|
x = iproc.negate(x)
|
|
y = iproc.negate(y)
|
|
return x, y
|
|
else
|
|
return x, y
|
|
end
|
|
end
|
|
function data_augmentation.pairwise_negate_x(x, y, p)
|
|
if torch.uniform() < p then
|
|
assert(x:size(2) == y:size(2) and x:size(3) == y:size(3))
|
|
x = iproc.negate(x)
|
|
return x, y
|
|
else
|
|
return x, y
|
|
end
|
|
end
|
|
function data_augmentation.pairwise_flip(x, y)
|
|
local flip = torch.random(1, 4)
|
|
local tr = torch.random(1, 2)
|
|
local x, conversion = iproc.byte2float(x)
|
|
y = iproc.byte2float(y)
|
|
x = x:contiguous()
|
|
y = y:contiguous()
|
|
if tr == 1 then
|
|
-- pass
|
|
elseif tr == 2 then
|
|
x = x:transpose(2, 3):contiguous()
|
|
y = y:transpose(2, 3):contiguous()
|
|
end
|
|
if flip == 1 then
|
|
x = iproc.hflip(x)
|
|
y = iproc.hflip(y)
|
|
elseif flip == 2 then
|
|
x = iproc.vflip(x)
|
|
y = iproc.vflip(y)
|
|
elseif flip == 3 then
|
|
x = iproc.hflip(iproc.vflip(x))
|
|
y = iproc.hflip(iproc.vflip(y))
|
|
elseif flip == 4 then
|
|
end
|
|
if conversion then
|
|
x = iproc.float2byte(x)
|
|
y = iproc.float2byte(y)
|
|
end
|
|
return x, y
|
|
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 tr = torch.random(1, 2)
|
|
local src, conversion = iproc.byte2float(src)
|
|
local dest
|
|
src = src:contiguous()
|
|
if tr == 1 then
|
|
-- pass
|
|
elseif tr == 2 then
|
|
src = src:transpose(2, 3):contiguous()
|
|
end
|
|
if flip == 1 then
|
|
dest = iproc.hflip(src)
|
|
elseif flip == 2 then
|
|
dest = iproc.vflip(src)
|
|
elseif flip == 3 then
|
|
dest = iproc.hflip(iproc.vflip(src))
|
|
elseif flip == 4 then
|
|
dest = src
|
|
end
|
|
if conversion then
|
|
dest = iproc.float2byte(dest)
|
|
end
|
|
return dest
|
|
end
|
|
|
|
local function test_blur()
|
|
torch.setdefaulttensortype("torch.FloatTensor")
|
|
local image =require 'image'
|
|
local src = image.lena()
|
|
|
|
image.display({image = src, min=0, max=1})
|
|
local dest = data_augmentation.blur(src, 1.0, "3,5", 0.5, 0.6)
|
|
image.display({image = dest, min=0, max=1})
|
|
dest = data_augmentation.blur(src, 1.0, "3", 1.0, 1.0)
|
|
image.display({image = dest, min=0, max=1})
|
|
dest = data_augmentation.blur(src, 1.0, "5", 0.75, 0.75)
|
|
image.display({image = dest, min=0, max=1})
|
|
end
|
|
--test_blur()
|
|
|
|
return data_augmentation
|