论文阅读：Adaptive Fusion for RGB-D Salient Object Detection
这篇代码的创新点在于使用了SW层，使用SW_logits * img_logits + (1 - SW_logits) * (1 - depth_logits) 来获得最终的预测结果
另外一个关键点是使用了3种loss损失值
第一种损失值，即经过归一化的标签g_t 与 输出的结果logits的sigmoid的损失值
第二种损失值, 即将im_logits进行sigmoid转换为0, 1之间，然后使用sigmoid_im * label + (1 - sigmoid_im) * (1 - label) # 获得标签值与图片值的交叉熵损失值
将计算好的交叉熵损失函数与SW_map 计算-log的交叉熵损失函数，个人认为这个loss存在问题
第三种损失值，即edge_loss即边界的损失值
将预测的结果进行sigmoid操作，转换为(0, 1)
使用tf.reshape(tf.constant([-1, 0, 1], tf.float32), [1, 3, 1, 1]) 构造x方向的边界卷积
使用tf.reshape(x_weight, [3, 1, 1, 1]) # 构造y方向的边界卷积
使用tf.nn.conv2d(g_t, x_weight, [1, 1, 1, 1], 'SAME') 进行标签的x轴方向和y轴方向上的边界卷积
使用tf.nn.conv2d(sigmoid_p, x_weight, [1, 1, 1, 1], 'SAME') 进行预测结果的x轴方向和y轴方向上的边界卷积
最后使用tf.losses.mean_squre_error(xgrad_gt, xgrad_sal) + tf.losses.mean_squre_error(ygrad_gt, ygrad_sal) 获得最终的mse损失函数
论文中的网络结构图
run_saliency.py 用于执行代码
from __future__ import print_function import tensorflow as tf import numpy as np import scipy.misc as misc import os import cv2 from net import * from loss import * FLAGS = tf.flags.FLAGS tf.flags.DEFINE_string('data_dir', './data/', 'path to dataset') tf.flags.DEFINE_string('ckpt_file', './model/AF-Net', 'checkpoint file') tf.flags.DEFINE_string('save_dir', './result', 'path to prediction direction') tf.flags.DEFINE_string('train_data', './train_data/', 'path to train_data') IMAGE_SIZE = 224 BATCH_SIZE = 1 train_num = 1500 num_epoch = 1000 def _transform(filename, _channels=True): image = misc.imread(filename) if _channels and len(image.shape) < 3: image = np.array([image for _ in range(3)]) resize_image = misc.imresize(image, [IMAGE_SIZE, IMAGE_SIZE], interp='nearest') return image def main(argv=None, is_training=True): image = tf.placeholder(tf.float32, [None, IMAGE_SIZE, IMAGE_SIZE, 3], name='input_image') depth_2 = tf.placeholder(tf.float32, [None, IMAGE_SIZE, IMAGE_SIZE], name='input_depth') depth = tf.expand_dims(depth_2, axis=3) processed_image = image - [123.64, 116.779, 103.939] # 减去最后一个维度的均值 gt_2 = tf.placeholder(tf.float32, [None, IMAGE_SIZE, IMAGE_SIZE], name='label') gt = tf.expand_dims(gt_2, axis=3) net_handler = NetHandler() logits, im_logits, SW_map = net_handler.RGBD_SW_net(processed_image, depth) pred_annotation = tf.sigmoid(logits) # 将其转换为0, 1 # 构造sal损失值 loss_sal = sigmoid_CEloss(logits, gt) # 构造SW损失值 loss_sw = SW_loss(im_logits, SW_map, gt) # 计算边缘损失值 loss_edge = edge_loss(logits, gt) # 计算总的损失值loss loss = loss_sal + loss_sw + loss_edge # 构造损失值的优化器 train_op = tf.train.AdamOptimizer(1e-6, beta1=0.5).minimize(loss) # 构造执行函数 sess = tf.Session() # 变量初始化 sess.run(tf.global_variables_initializer()) # 打印保存的参数的地址 print('Rreading params from {}'.format(FLAGS.ckpt_file)) # 如果已经保存了参数就加载 if tf.train.get_checkpoint_state('model'): saver = tf.train.Saver(None) saver.restore(sess, FLAGS.ckpt_file) # 进行图片结果的保存 if not os.path.exists(FLAGS.save_dir): os.makedirs(FLAGS.save_dir) if is_training == False: files = os.listdir(os.path.join(FLAGS.data_dir + '/RGB/')) test_num = len(files) test_RGB = np.array([_transform(os.path.join(FLAGS.data_dir + '/RGB/' + filename), _channels=True) for filename in files]) # 这里是不对的 test_depth = np.array([np.expand_dims(_transform(os.path.join(FLAGS.data_dir + '/depth/' + filename), _channels=True) for filename in files)]) # 进行测试操作 for k in range(test_num): # 进行结果的预测，这里结果的范围为0,1之间 test_prediction = sess.run(pred_annotation, feed_dict={image:test_RGB[k], depth:test_depth[k]}) test_origin_RGB = misc.imread(os.path.join(FLAGS.data_dir + '/RGB/' + files[k].split('.')[0] + '.jpg')) image_shape = test_origin_RGB.shape # 将图片转换为原来的图片的大小 test_pred = misc.imresize(test_prediction[0, :, :, 0], image_shape, interp='bilinear') misc.imsave('{}/{}'.format(FLAGS.save_dir, files[k].split('.')[0] + '.jpg'), test_pred.astype(np.uint8)) print('Save results in to %s' % (FLAGS.save_dir)) else: iter = 0 for epoch in range(num_epoch): # 载入数据 for i in range(train_num // BATCH_SIZE): deep_img, GT_image, Img_img = read_data_some('train_data.npy', BATCH_SIZE) _, _loss = sess.run([train_op, loss], feed_dict={image:Img_img, depth_2:deep_img, gt_2:GT_image}) if iter % 100 == 0 and iter != 0: print('iter', iter, 'loss', _loss) saver = tf.train.Saver() if epoch % 10 == 0: saver.save(sess, FLAGS.ckpt_file, write_meta_graph=FLAGS) test_deep, test_GT, test_RGB = read_data_some('test_data.npy', 1) test_prediction = sess.run(pred_annotation, feed_dict={image:test_RGB, depth_2:test_deep, gt_2:test_GT}) # 进行图片保存 cv2.imwrite('train_result/deep.png', test_deep) cv2.imwrite('train_result/GT.png', test_GT) cv2.imwrite('train_result/RGB.png', test_RGB) cv2.imwrite('train_result/pred.png', test_prediction) iter += 1 if __name__ == '__main__': tf.app.run()
net.py 网络结构
import tensorflow as tf import tensorflow.contrib.slim as slim class NetHandler(object): def __int__(self, weights_initializer = tf.contrib.layers.xavier_initializer(), weight_decay = 0.0001, padding='SAME'): self.padding = padding self.weight_initializer = weights_initializer self.weight_decay = weight_decay def vgg16_net(self, inputs, depth_suf = ''): layers = ( 'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2', 'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'pool5' ) kernel_size = 3 num_outputs = 64 net = {} current = inputs # 当前输入 for i, name in enumerate(layers): if depth_suf == '_d' and i == 0: current = slim.conv2d(current, 64, [3, 3], weights_initializer=self.weight_initializer, padding=self.padding, stride=1, activation_fn=None) net[name] = current continue kind = name[:4] if kind == 'conv': if name[:5] == 'conv1': num_outputs = 64 # 构造第一个卷积的输出fiter elif name[:5] == 'conv2': num_outputs = 128 elif name[:5] == 'conv3': num_outputs = 256 elif name[:5] == 'conv4': num_outputs = 512 elif name[:5] == 'conv5': num_outputs = 512 _, _, _, c = current.get_shape() kernels = tf.get_variable(name=name + '_w' + depth_suf, shape=[kernel_size, kernel_size, c, num_outputs], initializer=self.weight_initializer, regularizer=tf.contrib.layers.l2_regularizer(self.weight_decay), trainable=True) _, _, _, bias_size = kernels.get_shape() bias = tf.get_variable(name=name + '_b' + depth_suf, shape=[bias_size], initializer=tf.zeros_initializer(), trainable=True) conv = tf.nn.conv2d(current, kernels, strides=[1, 1, 1, 1], padding=self.padding) current = tf.nn.bias_add(conv, bias) elif kind == 'relu': current = tf.nn.relu(current, name=name) elif kind == 'pool': current = tf.nn.max_pool(current, kernel_size=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding=self.padding) net[name] = current return net def RGBD_SW_net(self, image, depth): image_net = self.vgg16_net(image) depth_net = self.vgg16_net(depth, depth_suf='_d') conv_5 = image_net['relu5_3'] # 获得第一层到最后一层卷积的结果 conv_4 = image_net['relu4_3'] conv_3 = image_net['relu3_3'] conv_2 = image_net['relu2_2'] conv_1 = image_net['relu1_2'] depth_5 = depth_net['relu5_3'] depth_4 = depth_net['relu4_3'] depth_3 = depth_net['relu3_3'] depth_2 = depth_net['relu2_2'] depth_1 = depth_net['relu1_2'] with slim.arg_scope([slim.conv2d], weights_initializer=self.weight_initializer, weight_regularizer=slim.l2_regularizer(self.weight_decay), padding=self.padding, stride=1, activation_fn=tf.nn.relu): conv5 = slim.repeat(conv_5, 2, slim.conv2d, 64, [3, 3], scope='conv5') # 2表示进行了两次的卷积操作 conv4 = slim.repeat(conv_4, 2, slim.conv2d, 64, [3, 3], scope='conv4') # 进行两次卷积 conv3 = slim.repeat(conv_3, 2, slim.conv2d, 64, [3, 3], scope='conv3') # conv2 = slim.repeat(conv_2, 2, slim.conv2d, 64, [3, 3], scope='conv2') conv1 = slim.repeat(conv_1, 2, slim.conv2d, 64, [3, 3], scope='conv1') depth5 = slim.repeat(depth_5, 2, slim.conv2d, 64, [3, 3], scope='depth5') depth4 = slim.repeat(depth_4, 2, slim.conv2d, 64, [3, 3], scope='depth4') depth3 = slim.repeat(depth_3, 2, slim.conv2d, 64, [3, 3], scope='depth3') depth2 = slim.repeat(depth_2, 2, slim.conv2d, 64, [3, 3], scope='depth2') depth1 = slim.repeat(depth_1, 2, slim.conv2d, 64, [3, 3], scope='depth1') conv5_up = tf.image.resize_images(conv5, [224, 224]) conv4_up = tf.image.resize_images(conv4, [224, 224]) conv3_up = tf.image.resize_images(conv3, [224, 224]) conv2_up = tf.image.resize_images(conv2, [224, 224]) depth5_up = tf.image.resize_images(depth5, [224, 224]) depth4_up = tf.image.resize_images(depth4, [224, 224]) depth3_up = tf.image.resize_images(depth3, [224, 224]) depth2_up = tf.image.resize_images(depth2, [224, 224]) # 将卷积层进行维度变化，卷积后的结果输入到下一层 concat4_im = tf.concat([conv5_up, conv4_up], 3) feat4_im = slim.conv2d(concat4_im, 64, [3, 3], scope='feat4_im') concat3_im = tf.concat([feat4_im, conv3_up], 3) feat3_im = slim.conv2d(concat3_im, 64, [3, 3], scope='feat3_im') concat2_im = tf.concat([feat3_im, conv2_up], 3) feat2_im = slim.conv2d(concat2_im, 64, [3, 3], scope='feat2_im') concat1_im = tf.concat([feat2_im, conv1], 3) feat1_im = slim.conv2d(concat1_im, 64, [3, 3], scope='feat1_im') # # 同理对深度图做相同的操作 concat4_d = tf.concat([depth4_up, depth5_up], 3) feat4_d = slim.conv2d(concat4_d, 64, [3, 3], scope='feat4_d') concat3_d = tf.concat([feat4_d, depth3]) feat3_d = slim.conv2d(concat3_d, 64, [3, 3], scope='feat3_d') concat2_d = tf.concat([feat3_d, depth2]) feat2_d = slim.conv2d(concat2_d, 64, [3, 3], scope='feat2_d') concat1_d = tf.concat([feat2_d, depth]) feat1_d = slim.conv2d(concat1_d, 64, [3, 3], scope='feat1_d') # 进行1*1的卷积, 时期维度变化为1 conv1_im_logits = slim.conv2d(feat1_im, 1, [1, 1], activation_fn=None, scope='conv1_im_logits') conv1_d_logits = slim.conv2d(feat1_d, 1, [1, 1], activation_fn=None, scope='conv1_d_logits') # 将图像卷积图与深度卷积图合并 feat1 = slim.conv2d(tf.concat([feat1_im, feat1_d], 3), 64, [3, 3], scope='feat1') SW_map = tf.nn.sigmoid(slim.conv2d(feat1, 1, [1, 1], activation_fn=None, scope='feat1_attn')) conv1_fused_logits = SW_map * conv1_im_logits + (1 - SW_map) * conv1_d_logits return conv1_fused_logits, conv1_im_logits, SW_map
loss.py 定义的损失值
import tensorflow as tf def sigmoid_CEloss(logits, gt): loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=tf.cast(gt, tf.float32)) ) def SW_loss(im_logits, SW_map, gt): label = tf.cast(gt, tf.float32) sigmoid_im = tf.nn.sigmoid(im_logits) SW_gt = label * sigmoid_im + (1 - label) * (1 - sigmoid_im) cost = -SW_gt * tf.log(tf.clip_by_value(SW_map, 1e-8, 1.0)) - (1 - SW_gt) * tf.log(tf.clip_by_value(1 - SW_map, 1e-8, 1.0)) return tf.reduce_mean(cost) # 边缘轮廓的损失值 def edge_loss(logits, gt): gt = tf.cast(gt, tf.float32) sigmoid_p = tf.nn.sigmoid(logits) x_weight = tf.reshape(tf.constant([-1, 0, +1], tf.float32), [1, 3, 1, 1]) # 构造了一个卷积核 y_weight = tf.reshape(x_weight, [3, 1, 1, 1]) # 构造了卷积核 # 获得其标签边缘的梯度值，获得x边缘的损失值 xgrad_gt = tf.nn.conv2d(gt, x_weight, [1, 1, 1, 1], 'SAME') ygrad_gt = tf.nn.conv2d(gt, y_weight, [1, 1, 1, 1], 'SAME') # 获得输出结果的边缘梯度值 xgrad_sal = tf.nn.conv2d(sigmoid_p, x_weight, [1, 1, 1, 1], 'SAME') ygrad_sal = tf.nn.conv2d(sigmoid_p, y_weight, [1, 1, 1, 1], 'SAME') # 计算平方根误差 loss = tf.losses.mean_squared_error(xgrad_gt, xgrad_sal) + tf.losses.mean_squared_error(ygrad_gt, ygrad_sal) return loss
read_data 读取一个batch_size的数据
import numpy as np import cv2 def read_data_some(path, bacth_size): data = np.array(np.load('npy/' + path)) num = len(data) indx = np.random.randint(0, num, bacth_size) deep_img, GT_img, Img_imgs = data[indx][:, 0], data[indx][:, 1], data[indx][:, 2] deep_imgs = [] GT_imgs = [] for i in range(bacth_size): deep_imgs.append(cv2.cvtColor(deep_img[i], cv2.COLOR_BGR2GRAY)) GT_imgs.append(cv2.cvtColor(GT_img[i], cv2.COLOR_BGR2GRAY)) return deep_imgs, GT_imgs, Img_imgs if __name__ == '__main__': read_train_data(64)
save_data 保存数据为.npy
import random import os import cv2 import numpy as np import glob def save_data(path): data = [] for root, dirs, files in os.walk(path): if len(dirs) != 0: file_names = glob.glob(path + dirs[0] + '/*.png') for deep_name in file_names: GT_name = deep_name.replace('deep', 'GT') Img_name = deep_name.replace('deep', 'Img').replace('png', 'jpg') # 图片的读取 deep_img = cv2.imread(deep_name) deep_img = cv2.resize(deep_img, (224, 224)) GT_img = cv2.imread(GT_name) GT_img = cv2.resize(GT_img, (224, 224)) Img_img = cv2.imread(Img_name) Img_img = cv2.resize(Img_img, (224, 224)) data.append((deep_img, GT_img, Img_img)) # 进行数据的清洗 random.shuffle(data) np.save('npy/' + path[:-1] + '.npy', data)