[代码记录]制作篡改图像数据集

之前制作篡改图代码，翻出来记录下，以防忘记。

import cv2
import os
import numpy as np
import random
import matplotlib.pyplot as plt


def get_Y_and_mask():
    # 获取路径和图片名称
    pixelmap_Path = 'Image/pixelmap/'
    train_Path = 'Image/train/'
    pixelmap_names = os.listdir(pixelmap_Path)
    train_names = os.listdir(train_Path)
    # 截取黄色部分
    color = [([0, 183, 183], [0, 255, 255])]  # 黄色范围
    for i in range(len(pixelmap_names)):
        img_path = pixelmap_Path + pixelmap_names[i]
        train_path = train_Path + train_names[i]
        pixelmap = cv2.imread(img_path)
        train_img = cv2.imread(train_path)
        for (lower, upper) in color:
            lower = np.array(lower, dtype="uint8")  # 颜色下限
            upper = np.array(upper, dtype="uint8")  # 颜色上限
            mask = cv2.inRange(pixelmap, lower, upper)
            output = cv2.bitwise_and(train_img, train_img, mask=mask)
            bin_img_save = np.copy(mask)
            (_, contoures, hierarchy) = cv2.findContours(bin_img_save, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
            for idx in range(len(contoures)):
                poly_img = np.zeros(output.shape, dtype=np.uint8)
                cv2.drawContours(poly_img, contoures, idx, [255, 255, 255], -1)
                poly_img = poly_img & output
                # poly_img = cv.add(poly_img,output)
                gray = cv2.cvtColor(poly_img, cv2.COLOR_BGR2GRAY)
                ret, im_fixed = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
                mask_part = ~im_fixed
                # 腐蚀膨胀
                erode = cv2.erode(mask_part, None, iterations=5)
                dilate = cv2.dilate(erode, None, iterations=4)
                # 保存图片
                cv2.imwrite("Image/Y/" + train_names[i].split(".")[0] + "_Y_" + str((idx + 1)) + ".jpg", poly_img)
                cv2.imwrite("Image/Ymask/" + train_names[i].split(".")[0] + "_Y_mask_" + str((idx + 1)) + ".jpg",
                            dilate)


def Y_train_splice():
    # 图片路径和名称获取
    Y_path = "Image/Y/"
    train_path = "Image/train/"
    Y_mask_path = "Image/Ymask/"
    Y_names = os.listdir(Y_path)
    Y_mask_name = os.listdir(Y_mask_path)
    train_names = os.listdir(train_path)

    # 批量处理Y，即将Y中的图片分别粘贴到train中
    for i in range(len(Y_names)):
        # 导入图片

        Y_pathimg = Y_path + Y_names[i]
        Ymask_pathimg = Y_mask_path + Y_mask_name[i]
        Y_img = cv2.imread(Y_pathimg)
        Ymask = cv2.imread(Ymask_pathimg, cv2.THRESH_BINARY)
        # mask 做取反 腐蚀膨胀操作，用于获取轮廓，意图是获取Y的最小矩阵
        mash_inv = cv2.bitwise_not(Ymask)
        closed = cv2.erode(mash_inv, None, iterations=1)
        closed = cv2.dilate(closed, None, iterations=1)
        _, contours, hierarchy = cv2.findContours(closed, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        x, y, w, h = cv2.boundingRect(contours[0])
        new_Y = Y_img[y:y + h, x:x + w]
        new_Y_mask = Ymask[y:y + h, x:x + w]

        ret, new_Y_mask = cv2.threshold(new_Y_mask, 100, 255, cv2.THRESH_BINARY)
        # print(new_Y_mask)
        # cv2.imshow('img', new_Y_mask)
        # cv2.waitKey(0)
        new_mash_inv = cv2.bitwise_not(new_Y_mask)
        # 每个Y都循环粘贴到不一样的train中，形成新的篡改图像
        for j in range(len(train_names)):
            # 导入train图像
            train_pathimg = train_path + train_names[j]
            train_img = cv2.imread(train_pathimg)
            # 判断：如果Y比train大或者是同一张图，则continue
            if (train_img.shape[0] < new_Y.shape[0] or
                    train_img.shape[1] < new_Y.shape[1] or
                    train_names[j].split(".")[0] == Y_names[i].split("_")[0]):
                continue
            else:
                Y_areas = new_Y.shape[0] * new_Y.shape[0]
                train_areas = train_img.shape[0] * train_img.shape[0]
                if Y_areas < train_areas * 0.05:
                    size = random.randint(int(train_img.shape[0] * 0.2), int(train_img.shape[0] * 0.3))
                    r = (size) / train_img.shape[1]
                    dim = (size, int(train_img.shape[0] * r))
                    new_Y = cv2.resize(new_Y, dim, interpolation=cv2.INTER_AREA)
                    new_Y_mask = cv2.resize(new_Y_mask, dim, interpolation=cv2.INTER_AREA)
                    new_mash_inv = cv2.bitwise_not(new_Y_mask)
                elif Y_areas > train_areas * 0.3:
                    size = random.randint(int(train_img.shape[0] * 0.3), int(train_img.shape[0] * 0.4))
                    r = (size) / train_img.shape[1]
                    dim = (size, int(train_img.shape[0] * r))
                    new_Y = cv2.resize(new_Y, dim, interpolation=cv2.INTER_AREA)
                    new_Y_mask = cv2.resize(new_Y_mask, dim, interpolation=cv2.INTER_AREA)
                    new_mash_inv = cv2.bitwise_not(new_Y_mask)

                shift_x = random.randint(0, train_img.shape[0] - new_Y.shape[0])
                shift_y = random.randint(0, train_img.shape[1] - new_Y.shape[1])
                # print(shift_x,shift_y)
                # 计算中间位置
                # row, col = (train_img.shape[0] - new_Y.shape[0]) // 2, \
                #            (train_img.shape[1] - new_Y.shape[1]) // 2
                row, col = shift_x, shift_y
                roi = train_img[row:row + new_Y.shape[0], col:col + new_Y.shape[1]]
                img1_bg = cv2.bitwise_and(roi, roi, mask=new_Y_mask)
                img2_fg = cv2.bitwise_and(new_Y, new_Y, mask=new_mash_inv)
                dst = img1_bg + img2_fg
                new_img1_bg = img1_bg
                new_img2_fg = img2_fg
                new_img1_bg[:, :] = 255
                new_img1_bg = cv2.cvtColor(new_img1_bg, cv2.COLOR_BGR2GRAY)
                new_img2_fg = cv2.cvtColor(new_img2_fg, cv2.COLOR_BGR2GRAY)
                # new_img2_fg[:,:]=255

                dst2 = cv2.bitwise_and(new_img1_bg, new_img2_fg)
                ret, dst2 = cv2.threshold(dst2, 0, 255, cv2.THRESH_BINARY)
                erode = cv2.erode(dst2, None, iterations=2)
                dst2 = cv2.dilate(erode, None, iterations=2)

                Y_masksplice = np.zeros((train_img.shape[0], train_img.shape[1]), dtype=np.int)
                Y_masksplice[row:row + new_Y.shape[0], col:col + new_Y.shape[1]] = dst2
                train_img[row:row + new_Y.shape[0], col:col + new_Y.shape[1]] = dst
                Y_masksplice = 255 - Y_masksplice

                new_img_path = "Image/Ysplicetrain/" + train_names[j].split(".")[0] + Y_names[i].split(".")[
                    0] + "_splice.jpg"
                cv2.imwrite(new_img_path, train_img)
                new_mask_path = "Image/Ysplicemask/" + train_names[j].split(".")[0] + Y_names[i].split(".")[
                    0] + "_splice_mask.jpg"
                cv2.imwrite(new_mask_path, Y_masksplice)


if __name__ == "__main__":
    # get_Y_and_mask()
    Y_train_splice()