Torch手札3:拆两个图像处理的轮子

Posted by Kriz on 2017-05-27

CV做的项目和图像着色有关,之前刷git的时候找到了几个项目里能直接拿来用的lua轮子,现在试着顺一遍源码稍微学习一下,也是为之后的文章梳理打一个基础。注释少到爆炸可以说是想打人了

系列文章传送门:

Torch手札1:安装及基础使用(以XOR问题为例)

Torch手札2:一小时Torch入门

DataLoader

是一个改良版DataLoader,用来读取YUV规范化的输入。

先提一句,YUV是一种颜色编码方法,Y通道保留亮度,UV通道保留色度。

1
2
3
4
5
6
require 'torch'
require 'hdf5'
require 'image'
local utils = require 'utils'
local DataLoader = torch.class('DataLoader')

导入torch库主要是为了方便引入Tensor,image库主要是为了将RGB图像转换成YUV。

至于这个hdf5,它是一种高效、简洁、全面的数据存储格式,广泛用于各大网络。这里导入的torch-hdf5是适配torch的hdf5库。一个hdf5文件由group和dataset组成,一个group如同一个文件夹,可以包含多个其他group和dataset;而dataset可以理解成数据集,和numpy中的数组非常相似。更详细的信息可以参考这里

utils是一个包含很多有用函数的第三方库,这里进行引入,但实例化好像也没什么用的样子。后面这个torch.class是torch用来实现面向对象编程的。

DataLoader类共有三个方法:__initresetgetBatch。我在原始文件的基础上进行了一系列删减和修改,主要是改掉了一些可选项,使代码结构更清晰。

__init

在lua中创建一个对象后,就会执行类的__init方法(从最顶层父类到最底层子类,依次调用定义的该函数)。具体的完整定义是:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
function DataLoader:__init(opt)
assert(opt.h5_file, 'Must provide h5_file')
assert(opt.batch_size, 'Must provide batch size')
self.h5_file = hdf5.open(opt.h5_file, 'r')
self.batch_size = opt.batch_size
self.split_idxs = {
train = 1,
val = 1,
}
self.image_paths = {
train = '/train',
val = '/validation',
}
local train_size = self.h5_file:read(self.image_paths.train):dataspaceSize()
self.split_sizes = {
train = train_size[1],
val = self.h5_file:read(self.image_paths.val):dataspaceSize()[1],
}
self.num_channels = train_size[2]
self.image_height = train_size[3]
self.image_width = train_size[4]
self.num_minibatches = {}
for k, v in pairs(self.split_sizes) do
self.num_minibatches[k] = math.floor(v / self.batch_size)
end
if opt.max_train and opt.max_train > 0 then
self.split_sizes.train = opt.max_train
end
end

不着急慢慢看。

首先是两个assert:

1
2
assert(opt.h5_file, 'Must provide h5_file')
assert(opt.batch_size, 'Must provide batch size')

这里提到的h5_file是提前建立好的存入图像数据的hdf5文件。

assert函数是检查错误的利器,它可以使lua在特定文件异常时抛出特定错误,比如非数字相加、调用非函数变量、访问表中不存在的值之类的。使用方法是assert(a,b),a是要检查是否有错误的一个参数,b是a错误时抛出的信息。参数b可选。

1
2
3
4
5
6
7
8
9
10
11
12
self.h5_file = hdf5.open(opt.h5_file, 'r')
self.batch_size = opt.batch_size
self.split_idxs = {
train = 1,
val = 1,
}
self.image_paths = {
train = '/train/images',
val = '/validation/images',
}

这一部分初始化了几个属性,包括存放灰度图数据的.h5文件,batch size,索引起始位置和图像路径(h5文件中的)。

1
2
3
4
5
6
7
8
local train_size = self.h5_file:read(self.image_paths.train):dataspaceSize()
self.split_sizes = {
train = train_size[1],
val = self.h5_file:read(self.image_paths.val):dataspaceSize()[1],
}
self.num_channels = train_size[2]
self.image_height = train_size[3]
self.image_width = train_size[4]

read函数读入hdf5文件中的对应数据,train_size获得dataset的size(点这里查看示例),train_size[1]是序列的数量,train_size[2]是序列的深度(这里即是通道数),train_size[3]和train_size[4]分别是图像的高度和宽度。validation这边也同理。

不过有个问题是不太清楚图库尺寸是否需要相同。传进了不同似乎也没什么问题,那么dataspaceSize的第三维和第四维难道是取最大?可是没什么意义啊。谜。有知道的小伙伴请告诉我。

1
2
3
4
5
6
7
8
self.num_minibatches = {}
for k, v in pairs(self.split_sizes) do
self.num_minibatches[k] = math.floor(v / self.batch_size)
end
if opt.max_train and opt.max_train > 0 then
self.split_sizes.train = opt.max_train
end

联系上下文感觉这里有种谜一般的没必要的优雅(?)……取split_sizes中的键和分成batch的值(math.floor用来劈掉小数部分)填入num_minibatches。最后一段调整实际训练图像数量,一般max_train都是-1,就不用太在意了。

reset

就一句code:

1
2
3
function DataLoader:reset(split)
self.split_idxs[split] = 1
end

在使用验证集check loss的时候会用到,仿佛就是为了从头开始计算所以置1。没什么可说的。

getBatch

先看一下整体效果:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
function DataLoader:getBatch(split)
local path = self.image_paths[split]
local start_idx = self.split_idxs[split]
local end_idx = math.min(start_idx + self.batch_size - 1,
self.split_sizes[split])
local images = self.h5_file:read(path):partial(
{start_idx, end_idx},
{1, self.num_channels},
{1, self.image_height},
{1, self.image_width}):float():div(255)
self.split_idxs[split] = end_idx + 1
if self.split_idxs[split] > self.split_sizes[split] then
self.split_idxs[split] = 1
end
local y = torch.Tensor(images:size(1),1,self.image_height,self.image_width)
local uv = torch.Tensor(images:size(1),2,self.image_height,self.image_width)
for t=1,images:size(1) do
local yuv = image.rgb2yuv(images[t])
y[t][1] = yuv[1]
uv[t][1] = torch.div(yuv[2],0.436)
uv[t][2] = torch.div(yuv[3],0.615)
end
y:add(-0.5)
return y, uv
end

好像还挺好懂的。

先看return。lua是可以返回多个值的,这里返回一个y一个uv,望文知义肯定是统一图像格式的函数了。

上半部分:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
local path = self.image_paths[split]
local start_idx = self.split_idxs[split]
local end_idx = math.min(start_idx + self.batch_size - 1,
self.split_sizes[split])
local images = self.h5_file:read(path):partial(
{start_idx, end_idx},
{1, self.num_channels},
{1, self.image_height},
{1, self.image_width}):float():div(255)
self.split_idxs[split] = end_idx + 1
if self.split_idxs[split] > self.split_sizes[split] then
self.split_idxs[split] = 1
end

取h5中的路径,确定好起止索引,读入范围内的数据内容,并规范化数值到0-1区间内。

下半部分就是进行灰度图着色部分的适配了。

1
2
local y = torch.Tensor(images:size(1),1,self.image_height,self.image_width)
local uv = torch.Tensor(images:size(1),2,self.image_height,self.image_width)

y和uv都是4维Tensor,分别为长度(图像数据数量)、深度、图像高度和图像宽度。

1
2
3
4
5
6
7
8
for t=1,images:size(1) do
local yuv = image.rgb2yuv(images[t])
y[t][1] = yuv[1]
uv[t][1] = torch.div(yuv[2],0.436)
uv[t][2] = torch.div(yuv[3],0.615)
end
y:add(-0.5)
return y, uv

这一部分要进行YUV的处理了。函数rgb2yuv的定义里,涉及转换的部分是这样一段:

1
2
3
outputY:zero():add(0.299, inputRed):add(0.587, inputGreen):add(0.114, inputBlue)
outputU:zero():add(-0.14713, inputRed):add(-0.28886, inputGreen):add(0.436, inputBlue)
outputV:zero():add(0.615, inputRed):add(-0.51499, inputGreen):add(-0.10001, inputBlue)

可以发现使用的公式是经典的

Y = 0.299R + 0.587G + 0.114B
U = -0.147R - 0.289G + 0.436B
V = 0.615R - 0.515G - 0.100B

因为之前对image进行过规范化,所以此处Y、U、V的取值范围分别是:Y [0,1],U [-0.436,0.436],V [-0.615,0.615]。使用一步div可以使uv值都规范到-1和1之间,方便之后的处理;Y值减0.5来进行中心对称。

ShaveImage

第一行:

1
local layer, parent = torch.class('nn.ShaveImage', 'nn.Module')

之前没见过这样的用法就去翻了一下官方文档。对应示例中的

1
local B = class('B', 'A')

即是实例化继承自A类的B类。ShaveImage包括构造函数在内共有3个函数,一个一个来看。

__init

首先仍然是__init。简单粗暴,既然是继承就直接调用父类的原构造函数,再设定一下size。

1
2
3
4
function layer:__init(size)
parent.__init(self)
self.size = size
end

父类的构造函数里都有些什么名堂呢?去官方文档学习观摩一下:

1
2
3
4
5
function Module:__init()
self.gradInput = torch.Tensor()
self.output = torch.Tensor()
self._type = self.output:type()
end

其实就是把gradInput(计算input对应的梯度)、output和type初始化一下。【不知道type()是干什么的不管了

updateOutput

1
2
3
4
5
6
7
8
function layer:updateOutput(input)
local N, C = input:size(1), input:size(2)
local H, W = input:size(3), input:size(4)
local s = self.size
self.output:resize(N, C, H - 2 * s, W - 2 * s)
self.output:copy(input[{{}, {}, {s + 1, H - s}, {s + 1, W - s}}])
return self.output
end

这个函数在前向传播的时候会被调用,这里进行了重写。核心步骤是:

1
2
self.output:resize(N, C, H - 2 * s, W - 2 * s)
self.output:copy(input[{{}, {}, {s + 1, H - s}, {s + 1, W - s}}])

也就是割掉图像的四条边。会在之后的系列文章里阐述它的用处。

updateGradInput

1
2
3
4
5
6
7
8
function layer:updateGradInput(input, gradOutput)
local N, C = input:size(1), input:size(2)
local H, W = input:size(3), input:size(4)
local s = self.size
self.gradInput:resizeAs(input):zero()
self.gradInput[{{}, {}, {s + 1, H - s}, {s + 1, W - s}}]:copy(gradOutput)
return self.gradInput
end

父类中的updateGradInput就只有一句return self.gradInput,这里配合上面的操作根据图像Tensor更新了一下梯度值Tensor的size。