RGB-D显著性突出物体(学习)

论文阅读：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)