使用coco資料集建立賦值黏貼篡改資料集

haohulala發表於2021-01-01
賦值

為了能夠訓練賦值黏貼篡改檢測網路，我們需要建立一個複製黏貼篡改資料集，為了簡單起見，我們使用coco資料集來進行構建。

第一步，建立json檔案檔案

這裡我為了能夠按照自己的思路寫程式，將coco的json檔案格式做了點修改，程式碼如下

import json
import time
"""
file = open("instances_val2017.json", "r")
img_info = json.load(file)
print(img_info.keys())
tmp_dict = {}
tmp_dict["categories"] = img_info["categories"][:2]
tmp_dict = json.dumps(tmp_dict, indent=4)
"""
#tmp_dict = json.loads(tmp_dict)
#print(tmp_dict)
#img_info = json.dumps(img_info, indent=4)
#print(img_info)

# 獲取類別資訊
# 第一個返回值是列表
# 第二個返回值是字典
def get_cate():
    file = open("instances_val2017.json", "r")
    img_info = json.load(file)
    tmp_dict = {}
    tmp_dict["categories"] = img_info["categories"]
    #tmp_dict = json.dumps(tmp_dict, indent=4)
    list = tmp_dict["categories"]
    cate = []
    # 第一個位置是背景
    cate.append("background")
    cate_dict = {}
    for item in list:
        cate.append(item["name"])
        cate_dict[item["name"]] = item["id"]

    print(cate)
    print(cate_dict)
    return cate, cate_dict

""""""""""""""""""""""""""""
我自定義的資料格式
{
    "images":{
        "(file_name)":{
            "height": int,
            "width":  int,
            "id":     int,
            "annotation": {
                "segmentation":[
                    [
                    int,
                    int,
                    ...
                    ]
                ],
                "iscrowd": int,
                "bbox":[
                    int,
                    int,
                    int,
                    int
                ],
                "category_id": int,
            }
        }
    }
}

演算法步驟：
1.找出所有的filename和對應的id
2.對於每張圖片，找尋他的所有標註資訊
"""""""""""""""""""""""""""""""""""
def get_info(read_file, dump_file):
    # 儲存資訊的字典
    img_anno = {}
    img_anno["images"] = {}
    file = open(read_file, "r")
    img_info = json.load(file)
    tmp_dict = {}
    tmp_dict["images"] = img_info["images"]
    # 讀取所有的圖片資訊，建立鍵值對
    images_list = tmp_dict["images"]
    for item in images_list:
        file_name = item["file_name"]
        height = item["height"]
        width = item["width"]
        img_id = item["id"]
        # 建立鍵值對
        img_anno["images"][file_name] = {}
        img_anno["images"][file_name]["height"] = height
        img_anno["images"][file_name]["width"] = width
        img_anno["images"][file_name]["img_id"] = img_id
        # 標註資訊的值是列表，裡面每個元素都是一個字典
        img_anno["images"][file_name]["annotations"] = []
    print("讀取完所有圖片資訊✿✿ヽ(°▽°)ノ✿")
    # 看看結果
    #print("一共載入了%d張圖片資訊✿✿ヽ(°▽°)ノ✿"%len(img_anno["images"].keys()))
    #tmp = json.dumps(img_anno, indent=4)
    #print(tmp)
    st = time.time()
    # 接下來根據圖片id查詢他的標註資訊
    tmp_dict["annotations"] = img_info["annotations"]
    list_anno = tmp_dict["annotations"]
    num = 0
    for key in img_anno["images"]:
        # print(key)
        tmp_dict = img_anno["images"][key]
        image_id = tmp_dict["img_id"]
        for item in list_anno:
            img_id = item["image_id"]
            if(image_id == img_id):
                tmp_dict["annotations"].append(item)
                # 將使用過的標註資訊刪除
                list_anno.remove(item)
        num = num + 1
        if(num%1000 == 0):
            print("讀取完%d張圖片資訊，時間：%ds"%(num, time.time()-st))

    #tmp = json.dumps(img_anno, indent=4)
    #print(tmp)

    # 儲存json檔案
    filename = dump_file
    with open(filename, "w") as f:
        json.dump(img_anno, f, indent=4)

    return img_anno





if __name__ == "__main__":
    #get_cate()
    read_file = "instances_train2017.json"
    dump_file = "instances_train"
    get_info(read_file, dump_file)

第二步，根據圖片和json文件資訊選取複製黏貼的物體

目錄結構如下

img ：存放原始圖片

img_t ：存放篡改後的影像

info：存放被複制物體的資訊

mask_s：被賦值物體的mask

mask_t：黏貼物體的mask

tmp_img：從原圖中將目標物體摳出來的圖

這一步我們的主要目的是從圖片中選取複製的物體

import json
import cv2
import matplotlib.pyplot as plt
import numpy as np
import os
import random
import torchvision.transforms as tfs

"""
獲取coco影像中的物體，繪製mask，並且摳圖
"""

# 原始圖片地址
PREFIX = "coco/img/"
# 源物體mask
PREFIX_MASK_S = "coco/mask_s/"
# 圖片資訊地址
PREFIX_INFO = "coco/info/"
# 摳圖結果
PREFIX_TMP = "coco/tmp_img/"
# json檔案
JSON_FILE = "instances_val.json"

def change_geshi(seg_list):
    output = []
    tmp = seg_list[0]
    for i in range(len(tmp)//2):
        point = []
        point.append(tmp[2*i])
        point.append(tmp[2*i+1])
        output.append(point)

    #output = np.array(output, np.int32)
    return output


def get_random(max):
    return random.randint(0,max)


def getMask_s(img_name, f):
    img = cv2.imread(PREFIX + img_name)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    file = open(JSON_FILE, "r")
    img_info = json.load(file)
    tmp_dict = img_info["images"]
    info = tmp_dict[img_name]
    img_height = info["height"]
    img_width = info["width"]
    anno_list = info["annotations"]
    if(len(anno_list)==0):
        return 0
    idx = get_random(len(anno_list)-1)

    item = anno_list[idx]
    # 每個item都是一個字典
    bbox = item["bbox"]
    center_x = int(float(bbox[0]))
    center_y = int(float(bbox[1]))
    width = int(float(bbox[2]))
    height = int(float(bbox[3]))
    seg_points = item["segmentation"]
    if(isinstance(seg_points, dict)):
        return 0
    seg_points = change_geshi(seg_points)
    seg_points = np.array(seg_points, np.int32)


    mask = np.zeros_like(img)
    cv2.fillPoly(mask, [seg_points], (255, 255, 255))
    if(width<(int(img_width)//2) or height<(int(height)//2)):
        #cv2.imwrite(PREFIX_MASK_S + img_name.split(".")[0]+".png", mask)

        # 從原圖中把物體摳出來
        img = tfs.ToTensor()(img)
        mask = tfs.ToTensor()(mask)
        # mask不為0的地方為1
        mask[~(mask==0)] = 1


        img = img*mask
        img = tfs.ToPILImage()(img)
        img.save(PREFIX_TMP + img_name.split(".")[0]+".png")
        print(PREFIX_TMP + img_name.split(".")[0]+".png")
        mask = mask[0]
        mask = tfs.ToPILImage()(mask)
        mask.save(PREFIX_MASK_S + img_name.split(".")[0] + ".png")
        print(PREFIX_MASK_S + img_name.split(".")[0] + ".png")

        # 源物體資訊
        info_mask_s = PREFIX_MASK_S+img_name.split(".")[0]+".png"+" "+\
                        str(int(img_width))+" "+str(int(img_height))+" "\
                      +" "+str(width)+" "+str(height)+\
                      " "+str(center_x)+" "+str(center_y) + "\n"

        f.write(info_mask_s)
    else:
        print("所選目標太大o(￣ヘ￣o＃)")
        # 刪除源圖片
        # os.remove(PREFIX + img_name)

    #img_add = cv2.addWeighted(mask, 0.3, img, 0.7, 0)
    #plt.imshow(img_add)
    #plt.show()



def boot():
    # 獲取所有img的mask_s
    with open(PREFIX_INFO+"mask_s.txt", "w") as f:
        list_img = os.listdir(PREFIX)
        for img in list_img:
            getMask_s(img, f)

if __name__ == "__main__":
    boot()

第三步，對被賦值物體進行操作後黏貼

對物體進行操作的方法如下：

（1）對物體水平翻轉

（2）以一定概率對物體垂直翻轉

（3）以一定縮放率縮小物體

（4）將物體圖片填充為原來的形狀，並且黏貼到原圖上

import torch
import numpy as np
import os
import matplotlib.pyplot as plt
import cv2
import torchvision.transforms as tfs
import torch.nn.functional as f
from PIL import Image
"""
根據摳圖和源圖進行復制黏貼篡改
"""

# 原始圖片地址
PREFIX = "coco/img/"
# 篡改後圖片存放地址
PREFIX_T = "coco/img_t/"
# 源物體mask
PREFIX_MASK_S = "coco/mask_s/"
# 目標物體mask
PREFIX_MASK_T = "coco/mask_t/"
# 圖片資訊地址
PREFIX_INFO = "coco/info/"
# 摳圖結果
PREFIX_TMP = "coco/tmp_img/"

def tensor(img):
    return tfs.ToTensor()(img)

def PIL(img):
    return tfs.ToPILImage()(img)

def resize(input, factor=0.5):
    return f.interpolate(input, scale_factor=factor)

# 讀取圖片資訊，並且進行左右和上下映象
# 然後隨機縮小並填充圖片
# 返回三個引數，第一個引數是原圖
# 第二個引數是處理後的摳圖資訊
# 第三個引數是處理後的mask資訊
# 都是tensor格式的
def readImg(filename):
    # 源圖片
    img = Image.open(PREFIX + filename.split(".")[0]+".jpg")
    # 目標物體的mask
    mask_s = Image.open(PREFIX_MASK_S + filename)
    # 摳圖資訊
    img_tmp = Image.open(PREFIX_TMP + filename)

    # 記錄下原圖資訊引數
    img = tensor(img)
    height = img.shape[1]
    width = img.shape[2]

    # 對摳圖資訊和目標物體mask操作

    # 左右翻轉
    img_tmp = img_tmp.transpose(Image.FLIP_LEFT_RIGHT)
    mask_s = mask_s.transpose(Image.FLIP_LEFT_RIGHT)

    # 以一般的概率上下翻轉
    b = torch.rand(1).item()
    if(b>0.5):
        img_tmp = img_tmp.transpose(Image.FLIP_TOP_BOTTOM)
        mask_s = mask_s.transpose(Image.FLIP_TOP_BOTTOM)

    img_tmp = tensor(img_tmp)
    mask_s = tensor(mask_s)

    # 進行縮小
    factor = torch.rand(1).item()
    factor = np.clip(factor, a_max=0.95, a_min=0.6)
    print(factor)
    img_tmp = resize(img_tmp, factor.__float__())
    mask_s = resize(mask_s, factor.__float__())

    # 進行擴大，擴大成和源圖一樣大小
    # 填充的list分別是 左右上下 的填充數量
    img_tmp = f.pad(img_tmp, [0, width-img_tmp.shape[2], 0, height-img_tmp.shape[1]], value=0.)
    mask_s = f.pad(mask_s, [0, width-mask_s.shape[2], 0, height-mask_s.shape[1]], value=0.)

    """
    plt.subplot(1,3,1), plt.imshow(PIL(img))
    plt.subplot(1,3,2), plt.imshow(PIL(img_tmp))
    plt.subplot(1,3,3), plt.imshow(PIL(mask_s))
    plt.show()

    print(img.shape)
    print(img_tmp.shape)
    print(mask_s.shape)
    """
    return PIL(img), PIL(img_tmp), PIL(mask_s)

# 對圖片進行復制黏貼篡改，並且儲存圖片和篡改的mask資訊
def tamper(img, img_tmp, mask_t):
    img = tensor(img)
    img_tmp = tensor(img_tmp)
    mask_t = tensor(mask_t)
    img = img*~(mask_t==1)
    img = img+img_tmp
    return PIL(img), PIL(mask_t)


def boot():
    list_img = os.listdir(PREFIX_TMP)
    for item in list_img:
        img, img_tmp, mask_t = readImg(item)
        img_t, mask_t = tamper(img, img_tmp, mask_t)
        # 儲存圖片
        img_t.save(PREFIX_T + item)
        mask_t.save(PREFIX_MASK_T + item)
        print(PREFIX_T + item)
        print(PREFIX_MASK_T + item)


if __name__ == "__main__":
    boot()

經過這三步，我們的資料就做好了，下面來看看效果吧

原圖

篡改圖片

相關文章