Ver Fonte

refactor

nagadomi há 9 anos atrás
pai
commit
4c691b4640
5 ficheiros alterados com 195 adições e 186 exclusões
  1. 2 21
      convert_data.lua
  2. 95 0
      lib/data_augmentation.lua
  3. 33 0
      lib/iproc.lua
  4. 64 159
      lib/pairwise_transform.lua
  5. 1 6
      tools/benchmark.lua

+ 2 - 21
convert_data.lua

@@ -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 - 0
lib/data_augmentation.lua

@@ -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

+ 33 - 0
lib/iproc.lua

@@ -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

+ 64 - 159
lib/pairwise_transform.lua

@@ -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.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))
+      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
-local function test_overlay()
-   torch.setdefaulttensortype('torch.FloatTensor')
-   local loader = require './image_loader'
-   local src = loader.load_byte("../images/miku_CC_BY-NC.jpg")
+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
-      image.display(overlay_augment(src, 1.0))
+      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})
    end
 end
-
---test_scale()
---test_jpeg()
---test_jpeg_scale()
---test_color_noise()
---test_overlay()
-
 return pairwise_transform

+ 1 - 6
tools/benchmark.lua

@@ -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