Deep Leraningで線画抽出をする方法として有名なHEDという論文がある。
ネットワークの中間層で抽出したエッジを、さらに多階層分合わせることで精度良くエッジを抽出する手法である。
線画抽出として使えそうなため試してみた。
オリジナルのCaffe版は、Caffe自体を改造する必要があるため動かすのに苦労したが、
最近有志が実装したKerasコードは簡単に動作させることができた。
動かしてみた際の備忘録をここに残しておく。


002



○論文
https://arxiv.org/abs/1504.06375

○リポジトリ
https://github.com/lc82111/Keras_HED
○導入
1.HED_KERASのリポジトリをダウンロード
https://github.com/lc82111/Keras_HED

2.VGG16のモデルをダウンロード
https://github.com/MinerKasch/applied_deep_learning/blob/master/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5
または
https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5

3. BSDSのデータセットのダウンロード
http://vcl.ucsd.edu/hed/HED-BSDS

4. ソースコード修正
i) src/networks/hed.py 67行目
VGG16の重みファイルパスを修正（例）
filepath = 'D:/dataset/VGG16/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'

ii) src/utils/HED_data_parser.py 27行目～
BSDSデータセットへのパスを修正（例）
self.train_file = os.path.join('D:/dataset/HED-BSDS/', 'train_pair.lst')
self.train_data_dir = 'D:/dataset/HED-BSDS/'

iii) main_segmentation.py 56行目
エポックが大きすぎて終わらないので30程度にしておく。
epochs = 30

5. 学習
main_segmentation.py を実行

6. 結果（自前画像）

002

001

003

004


7. 結果の確認用コード
テストデータで確認する場合



%matplotlib inline
from __future__ import print_function
import os
from src.utils.HED_data_parser import DataParser
from src.networks.hed import hed
from keras.utils import plot_model
from keras import backend as K
from keras import callbacks
import numpy as np
import pdb
from matplotlib import pyplot as plt
import cv2

def generate_minibatches(dataParser, train=True):
    #pdb.set_trace()
    while True:
        if train:
            batch_ids = np.random.choice(dataParser.training_ids, dataParser.batch_size_train)
        else:
            batch_ids = np.random.choice(dataParser.validation_ids, dataParser.batch_size_train*2)
        ims, ems, _ = dataParser.get_batch(batch_ids)
        yield(ims, [ems, ems, ems, ems, ems, ems])

# params
model_name = 'HEDSeg'
model_dir     = os.path.join('checkpoints', model_name)
batch_size = 10

# environment
K.set_image_data_format('channels_last')
K.image_data_format()
os.environ["CUDA_VISIBLE_DEVICES"]= '0'
if not os.path.isdir(model_dir): os.makedirs(model_dir)

# prepare data
dataParser = DataParser(batch_size)

# model
model = hed()
model.load_weights('checkpoints/HEDSeg/checkpoint.24-0.12.hdf5')

    
def show_edge_image(ed):
    plt.gray()
    plt.imshow(ed)
    plt.show()

def show_test_image(im):
    plt.imshow((im + np.array([103.939, 116.779,123.68]))[..., ::-1].astype(np.uint8))
    plt.show()
    
def show_edge_image(ed):
    plt.gray()
    plt.imshow(ed)
    plt.show()

batch_ids = np.random.choice(dataParser.validation_ids, 1)
ims, ems, _ = dataParser.get_batch(batch_ids)

est = labels = model.predict(ims)

show_test_image(ims[0])
show_edge_image(ems[0].reshape(480,480))
show_edge_image(est[0][0].reshape(480,480))
show_edge_image(est[1][0].reshape(480,480))
show_edge_image(est[2][0].reshape(480,480))
show_edge_image(est[3][0].reshape(480,480))
show_edge_image(est[4][0].reshape(480,480))
show_edge_image(est[5][0].reshape(480,480))

自前のデータで試す場合
def generate_batch(im):
    ims = cv2.resize(im, (480, 480)).reshape(1, 480, 480, 3)
    ims = np.array(ims, dtype=np.float32)
    ims[..., 0] -= 103.939
    ims[..., 1] -= 116.779
    ims[..., 2] -= 123.68
    shape = im.shape[:2][::-1]
    return ims, shape

def show_edge_image_resize(im, shape):
    plt.imshow(cv2.resize(im.reshape(480,480), shape[:2]))
    plt.show()

im = cv2.imread('001.jpg')
ims, shape = generate_batch(im)
est = labels = model.predict(ims)    

plt.imshow(im[..., ::-1])
plt.show()
plt.gray()
show_edge_image_resize(est[0][0].reshape(480,480), shape)
show_edge_image_resize(est[1][0].reshape(480,480), shape)
show_edge_image_resize(est[2][0].reshape(480,480), shape)
show_edge_image_resize(est[3][0].reshape(480,480), shape)
show_edge_image_resize(est[4][0].reshape(480,480), shape)
show_edge_image_resize(est[5][0].reshape(480,480), shape)






タグ ：
#Keras
#HED
#DeepLearning