Mask-RCNN复现——实现copymove图像篡改检测

前言

1. 基于mask_rcnn网络图像篡改检测，对在网上阅读了许多相关内容发现都不利于小白复现mask_rcnn，因此记录一下自己实现的过程。
2. 神经网络的训练最好基于gpu, 楼主条件不允许所以直接拿自己电脑刚了。初次训练了10来个小时，效果还不错（心疼我的小风扇嘤嘤嘤）。
3. 第一次写博客，所以排版及表达可能不是特别严谨，万各位包容以及提出指正，谢谢～

准备工具

1. macos系统
2. anaconda——用于装Keras,tensorflow等机器学习必要的包
3. pycharm开发工具
4. labelme标注工具——用于制作train的数据集，打上自己想要的标签
   下载地址：https://github.com/wkentaro/labelme
   使用参考：https://blog.csdn.net/shwan_ma/article/details/77823281
5. mask-rcnn源码下载地址：https://github.com/matterport/Mask_RCNN 。在mask-rcnn根目录下建立logs文件夹，用于保存训练模型的历史信息
6. 所使用的图片集：coverage图片集和comofod图片集——这两个图片集都是经过copymove篡改后的图片，但是由于他们所提供的标注的信息我不会使用，所以我就自己使用labelme做了标注。所以其实在这种情况下，无论我们要使用mask-rcnn进行何种检测，都可以参考楼主的过程，因为此时只有标注分类不同。
7. 原模型下载：用的是coco的模型，一般来讲有一个初始模型算法会收敛的更快一点，所以使用了mask_rcnn_coco.h5。要楼主模型的可私信，只针对copymove图像篡改检测的。mask_rcnn_coco.h5下载地址： https://github.com/matterport/Mask_RCNN/releases/download/v2.0/mask_rcnn_coco.h5.
   将用于训练的原模型置于指定路径下。（由于我的代码找原模型是在根目录下找，所以将模型置于根目录下）

数据集制作

图片准备

1. 将coverage和comofod图片集中的图片合到一个文件夹下，并使所有图片都文件名为递增序列如1.png、2.png。
2. 最先开始跑代码的时候出现错误：

Exception: Image size must be dividable by 2 at least 6 times to avoid fractions when downscaling and upscaling.

马上查看了一下图片的大小，果然很多图片的个位数都是奇数。然后将代码中有关图片大小设置为512*512（这里可以由读者自行更改），并将准备的图片集进行resize一下，使用python代码进行resize。代码内容如下：

//resize.py
from PIL import Image
import os.path
import glob
def convertjpg(jpgfile,outdir,width=512,height=512):
    img=Image.open(jpgfile)
    try:
        new_img=img.resize((width,height),Image.BILINEAR)
        new_img.save(os.path.join(outdir,os.path.basename(jpgfile)))
    except Exception as e:
        print(e)
for jpgfile in glob.glob("/Users/empty/Desktop/dataset/*.png"): //对路径下适配‘*.png’的所有图片进行转换大小
    convertjpg(jpgfile,"/Users/empty/Desktop/dataset")

图片标注处理

1. 在终端下输入conda activate labelme，系统将会启动labelme程序，而后输入labelme打开图形化界面，而后如何使用labelme进行标注可以参考这篇文章：https://blog.csdn.net/shwan_ma/article/details/77823281
2. 参考完上篇文章使用labelme对resize后的图片进行标注，生成文件名.json数据对所有的json数据进行处理生成文件名_json文件。此时对json数据的处理过程是使用shell脚本进行批量化处理的，由于必须使用labelme_json_to_dataset，所以要进入到labelme应用程序（在终端下输入：conda activate labelme 进入labelme）。脚本代码如下：

// task.bash
#!/bin/bash
for((i=1;i<301;i++))//这里是文件名从小到大，如果训练集有500张图片就到500
do 
s=${i}
labelme_json_to_dataset ${s}.json
done

3. 此时对所有的文件名.json文件都会生成文件名_json文件夹，如下图所示：
   
   打开1_json文件夹，我们可看到如下文件：
   
   如上图所示，img.png为原图，label.png为生成掩膜的图片。

生成最后数据集

1. 为了为代码提供正确的数据，我们将数据集设置为4个文件：
   
   a)json文件夹下存储所有的标注后的json数据
   
   b)labelme_json文件夹下存储所有的文件名_json文件夹
   
   c)pic文件夹下存储所有的原图
   
   d)cv2_mask文件夹下存储所有的掩码标注图片，即上文中所说的文件名_json下的label.png图片。由于labelme生成的掩码标签 label.png为16位图，opencv默认处理8位图，所以需要将所有的掩码标签图由16位转8位，这里使用python实现，并将生成后的8位图直接保存到cv2_mask文件夹下。代码如下：

//img_16to8.py
def img_16to8():
    from PIL import Image
    import numpy as np
    import shutil
    import os

    src_dir = r'/Users/empty/Desktop/dataset/labelme_json'
    dest_dir = r'/Users/empty/Desktop/dataset/cv2_mask'
    for child_dir in os.listdir(src_dir):
        new_name = child_dir.split('_')[0] + '.png'
        old_mask = os.path.join(os.path.join(src_dir, child_dir), 'label.png')
        img = Image.open(old_mask)
        img = Image.fromarray(np.uint8(np.array(img)))
        new_mask = os.path.join(dest_dir, new_name)
        img.save(new_mask)

img_16to8()

代码运行成功后cv2_mask文件夹内容如下所示，此时所有的图片都是黑乎乎的，这是正常现象不用担心。

2. 完成所有文件夹创建并将内容归位后，数据集的制作就完成啦！

代码修改

代码准备

1. 在samples目录下建立文件夹copymove(或者你自己取的文件名)，并在copymove文件目录下建立新的文件train.py和test.py

train.py代码解析

1. train.py对训练数据集进行训练并生成模型，代码内容如下：

# -*- coding: utf-8 -*-

import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import tensorflow as tf
from mrcnn.config import Config
#import utils
from mrcnn import model as modellib,utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image


#os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Root directory of the project
#ROOT_DIR = os.getcwd()

ROOT_DIR = os.path.abspath("../../")
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

iter_num=0

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_shapes_0022.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)


class ShapesConfig(Config):
    """Configuration for training on the toy shapes dataset.
    Derives from the base Config class and overrides values specific
    to the toy shapes dataset.
    """
    # Give the configuration a recognizable name
    NAME = "shapes"

    # Train on 1 GPU and 2 images per GPU. We can put multiple images on each
    
    GPU_COUNT = 1
    IMAGES_PER_GPU = 2

    # Number of classes (including background)
    NUM_CLASSES = 1 + 1  # background + 1 shapes

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 512
    IMAGE_MAX_DIM = 512

    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels

    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE = 100

    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 100

    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50


config = ShapesConfig()
config.display()

class DrugDataset(utils.Dataset):
    # 得到该图中有多少个实例（物体）
    def get_obj_index(self, image):
        n = np.max(image)
        return n

    # 解析labelme中得到的yaml文件，从而得到mask每一层对应的实例标签
    def from_yaml_get_class(self, image_id):
        info = self.image_info[image_id]
        with open(info['yaml_path']) as f:
            temp = yaml.load(f.read())
            labels = temp['label_names']
            del labels[0]
        return labels

    # 重新写draw_mask
    def draw_mask(self, num_obj, mask, image,image_id):
        #print("draw_mask-->",image_id)
        #print("self.image_info",self.image_info)
        info = self.image_info[image_id]
        #print("info-->",info)
        #print("info[width]----->",info['width'],"-info[height]--->",info['height'])
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    #print("image_id-->",image_id,"-i--->",i,"-j--->",j)
                    #print("info[width]----->",info['width'],"-info[height]--->",info['height'])
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        return mask

    # 重新写load_shapes，里面包含自己的类别,可以任意添加
    # 并在self.image_info信息中添加了path、mask_path 、yaml_path
    # yaml_pathdataset_root_path = "/tongue_dateset/"
    # img_floder = dataset_root_path + "rgb"
    # mask_floder = dataset_root_path + "mask"
    # dataset_root_path = "/tongue_dateset/"
    def load_shapes(self, count, img_floder, mask_floder, imglist, dataset_root_path):
        """Generate the requested number of synthetic images.
        count: number of images to generate.
        height, width: the size of the generated images.
        """
        # Add classes,可通过这种方式扩展多个物体
        self.add_class("shapes", 1, "copymove") # 黑色素瘤

        for i in range(count):
            # 获取图片宽和高

            filestr = imglist[i].split(".")[0]
            if(filestr!=""):
            #print(imglist[i],"-->",cv_img.shape[1],"--->",cv_img.shape[0])
            #print("id-->", i, " imglist[", i, "]-->", imglist[i],"filestr-->",filestr)
            #filestr = filestr.split("_")[1]
                mask_path = mask_floder + "/" + filestr + ".png"
                yaml_path = dataset_root_path + "labelme_json/" + filestr + "_json/info.yaml"
                print(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
                cv_img = cv2.imread(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")

                self.add_image("shapes", image_id=i, path=img_floder + "/" + imglist[i],
                               width=cv_img.shape[1], height=cv_img.shape[0], mask_path=mask_path, yaml_path=yaml_path)

    # 重写load_mask
    def load_mask(self, image_id):
        """Generate instance masks for shapes of the given image ID.
        """
        global iter_num
        print("image_id",image_id)
        info = self.image_info[image_id]
        count = 1  # number of object
        img = Image.open(info['mask_path'])
        num_obj = self.get_obj_index(img)
        mask = np.zeros([info['height'], info['width'], num_obj], dtype=np.uint8)
        mask = self.draw_mask(num_obj, mask, img,image_id)
        occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)
        for i in range(count - 2, -1, -1):
            mask[:, :, i] = mask[:, :, i] * occlusion

            occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))
        labels = []
        labels = self.from_yaml_get_class(image_id)
        labels_form = []
        for i in range(len(labels)):
            if labels[i].find("copymove") != -1:
                # print "box"
                labels_form.append("copymove")

        class_ids = np.array([self.class_names.index(s) for s in labels_form])
        return mask, class_ids.astype(np.int32)

def get_ax(rows=1, cols=1, size=8):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.

    Change the default size attribute to control the size
    of rendered images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax

#基础设置
dataset_root_path="/Users/empty/Desktop/dataset/"
img_floder = dataset_root_path + "pic"
mask_floder = dataset_root_path + "cv2_mask"
#yaml_floder = dataset_root_path
imglist = os.listdir(img_floder)
count = len(imglist)

#train与val数据集准备
dataset_train = DrugDataset()
dataset_train.load_shapes(count, img_floder, mask_floder, imglist,dataset_root_path)
dataset_train.prepare()

#print("dataset_train-->",dataset_train._image_ids)

dataset_val = DrugDataset()
dataset_val.load_shapes(7, img_floder, mask_floder, imglist,dataset_root_path)
dataset_val.prepare()

#print("dataset_val-->",dataset_val._image_ids)

# Load and display random samples
#image_ids = np.random.choice(dataset_train.image_ids, 4)
#for image_id in image_ids:
#    image = dataset_train.load_image(image_id)
#    mask, class_ids = dataset_train.load_mask(image_id)
#    visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names)

# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config,
                          model_dir=MODEL_DIR)

# Which weights to start with?
init_with = "coco"  # imagenet, coco, or last

if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    # Load weights trained on MS COCO, but skip layers that
    # are different due to the different number of classes
    # See README for instructions to download the COCO weights
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last()[1], by_name=True)

# Train the head branches
# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=20,
            layers='heads')

# Fine tune all layers
# Passing layers="all" trains all layers. You can also
# pass a regular expression to select which layers to
# train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE / 10,
            epochs=40,
            layers="all")

2. 读者根据自己的训练集和环境有需要更改的内容如下：
   ROOT_DIR：整个mask_rcnn项目所在路径
   MODEL_DIR：用于存储训练出来的模型路径
   COCO_MODEL_PATH：原始模型路径及模型文件名
   GPU_COUNT：
   IMAGES_PER_GPU：
   NUM_CLASSES：
   IMAGE_MIN_DIM：
   IMAGE_MAX_DIM：
   dataset_root_path：数据集的根目录
   除上以外，由于我标注的种类只有copymove，因此读者还需将代码中含"copymove"的位置添加/修改为自己的标注种类
3. 运行代码，此时会迭代生成很多模型，将logs文件夹下最新生成的模型置于根目录下，为测试代码做准备。

test.py代码解析

1. test.py对测试数据进行标注，代码内容如下：

import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt

# Root directory of the project
ROOT_DIR = os.path.abspath("../../")

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize

# Import COCO config
sys.path.append(os.path.join(ROOT_DIR, "samples/copymove/"))  # To find local version
import copymove

# get_ipython().run_line_magic('matplotlib', 'inline')

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_shapes_0022.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)


class InferenceConfig(copymove.ShapesConfig):
    # Set batch size to 1 since we'll be running inference on
    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1


config = InferenceConfig()
config.display()

# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)

class_names = ['BG', 'copymove']

image = skimage.io.imread('2.png')

# Run detection
results = model.detect([image], verbose=1)

# Visualize results
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                            class_names, r['scores'])

2. 读者根据自己的测试数据和环境有需要更改的内容如下：
   ROOT_DIR：整个mask_rcnn项目所在路径
   MODEL_DIR：用于存储训练出来的模型路径
   COCO_MODEL_PATH：用于测试数据所使用的模型及模型文件名
   GPU_COUNT：
   IMAGES_PER_GPU：
   class_names：添加/修改为自己所标注的类
   image：测试的图片
   除上以外，由于我标注的种类只有copymove，因此读者还需将代码中含"copymove"的位置添加/修改为自己的标注种类
3. 此时可以根据测试数据来运行代码了

结果展示

最后测试结果如下图，基本实现最初试验目的。

总结

在此次试验中有一些心得体会：

1. shell脚本超好用！！！顶礼膜拜！！！
2. comofod可用图片不过才200张，但是整个rar包有3G多，并且由于是越狱下载，速度贼慢。所以针对copymove的数据集不推荐使用comofod。coverage好用，并且图片很清晰，但是量太少了。
3. 在初次训练的过程中，由于要下载mask_rcnn_coco.h5模型，所以运行时间特别长，但是把它放在那里就行了。
4. 每次进行训练的时候，记得要用最新的训练好的模型。。我有次忘记换模型了，结果做了好几个小时的无用功。
5. 每个框都有可信度，如果想要过滤掉可信值太低的，就去config.py中将DETECTION_MIN_CONFIDENCE改为理想中的最低值。
6. 用gpu跑吧，因为宿舍断电熬了一个通宵的人如是说。
7. 祝大家五一快乐～٩(๑o๑)۶