使用OpenCV合成训练图片，同时生成labelme兼容格式的标注文件

Posted 2022-12-14 henreash

 思路很简单，加载背景图片，加载目标图片，随机抽取N个目标图像，将目标图像的像素copy到背景图片上，拷贝的位置可根据要求指定；同时合成标注json文件。代码如下：

# coding=utf-8

from cv2 import cv2
import os
import numpy as np
import random

from gen_ann_json import AnnotationsJson

XXX_WIDTH = 48
XXX_HEIGHT = 72

#加载源图像，多个，随机抽取后复制到背景图像上
def load_src_images():
    path = './src_img'
    files = os.listdir(path)
    xxxs = []
    for f in files:
        if '000' in f:
            continue
        if os.path.isdir(f):
            continue
        img_file = os.path.join(path, f)
        xxx = cv2.imread(img_file)
        xxx = cv2.resize(xxx,(XXX_WIDTH, XXX_HEIGHT))
        xxxs.append(xxx)
    return xxxs

#图像混合（将src拷贝到dst上，位置：x、y）
def img_mix(dst, src, x, y):
    output = dst.copy()
    output[y:y + src.shape[0], x:x + src.shape[1]] = src
    return output

#功能同上，多了一次90°旋转操作
def img_mix_horv(dst, src, x, y):
    #transposedImage = cv2.transpose(src)
    transposedImage=np.rot90(src)
    output = img_mix(dst, transposedImage, x, y)
    return output

#合成一个大的子区域，上面copy多个src图像，最后在将子区域拷贝到背景上
def mix_xxx_region(offset_x, offset_y, show_region_img, xxxs):
    fixed_xxx_regions = [[85, 12], [138, 12], [208, 12], [261, 12]]
    random.shuffle(xxxs)#将样本打乱，抽取前4--6个
    outputImg = show_region_img.copy()
    pointsArr = []
    for i in range(4):
        xxx = xxxs[i]
        x, y = fixed_xxx_regions[i][0], fixed_xxx_regions[i][1]
        outputImg = img_mix(outputImg, xxx, x, y)
        pointsArr.append([[offset_x + x, offset_y + y], [offset_x + x + XXX_WIDTH, offset_y + y + XXX_HEIGHT]])
    if random.random() > 0.5:
        outputImg = img_mix_horv(outputImg, xxxs[4], 9, 35)
        pointsArr.append([[offset_x+9, offset_y+35], [offset_x+9 + XXX_HEIGHT, offset_y+35 + XXX_WIDTH]])
    if random.random() > 0.5:
        outputImg= img_mix_horv(outputImg, xxxs[5], 314, 35)
        pointsArr.append([[offset_x+314, offset_y+35], [offset_x+314 + XXX_HEIGHT, offset_y+35 + XXX_WIDTH]])   
    return outputImg, pointsArr

#保存json格式的标注文件
def saveJsonFile(jpgFileName, jsonFileName, pointsArr):
    annObj = AnnotationsJson(jpgFileName, 811, 833)
    for points in pointsArr:
        annObj.AppendLabel("xxx", points)
    annObj.SaveToJsonFile(jsonFileName)

#合成图片
def mix_image(dst, show_region_img, xxxs):
    #6个子区域所在位置
    show_region_pos_list = [[2, 157],[2, 435],[2, 713],[410, 157],[410, 435],[410, 713]]
    output = dst.copy()
    pointsArr = []
    subRegion = [0,1,2,3,4,5]
    random.shuffle(subRegion)#将位置打散，选取前4个
    for i in range(4):
        index = subRegion[i]
        x = show_region_pos_list[index][0]
        y = show_region_pos_list[index][1]
        #先将随机的牌复制到显示区域
        xxx_region, pointsArr2 = mix_xxx_region(x, y, show_region_img, xxxs)
        pointsArr.extend(pointsArr2)
        #再将显示区域复制到桌位面板上
        output = img_mix(output, xxx_region, x, y)
    return output, pointsArr

xxxs = load_xxx_images()
dst = cv2.imread('./back1.jpg')
show_region_img = cv2.imread('./show_region.jpg')
for i in range(1, 1001):
    output, pointsArr = mix_image(dst, show_region_img, xxxs)
    cv2.imwrite('./images_ann/.jpg'.format(i), output)
    saveJsonFile('.jpg'.format(i), './images_ann/.json'.format(i), pointsArr)

生成json标注文件代码：

import json

class AnnotationsJson(object):
    def __init__(self, imagePath, imgWidht, imgHeight):
        self.version = "3.16.7"
        self.flags = ,
        self.shapes = []
        self.fillColor = [255,0,0,128]
        self.lineColor = [0, 255, 0, 128]
        self.imagePath = imagePath
        self.imageData = None
        self.imageHeight = imgHeight
        self.imageWidth = imgWidht
    
    def AppendLabel(self, label, points):
        self.shapes.append("label":label, "line_color":None, "fill_color":None, "points":points, "shape_type":"rectangle", "flag":)

    def ToJson(self):
        outObj = "version": self.version, "flags": , "shapes":[],   "fillColor": self.fillColor,"lineColor":self.lineColor, 
                  "imagePath":self.imagePath, "imageData":None, "imageHeight":self.imageHeight, "imageWidth":self.imageWidth
        for shape in self.shapes:
            outObj["shapes"].append("label":shape["label"], "line_color":None, "fill_color":None, 
                                     "points":shape["points"], "shape_type":"rectangle", "flags":)
        res = json.dumps(outObj)
        return res

    def SaveToJsonFile(self, filename):
        jsonString = self.ToJson()
        with open(filename, "w") as f:
            f.write(jsonString)

if __name__ == "__main__":
    annObj = AnnotationsJson("1.jpg", 811, 833)
    annObj.AppendLabel("xxx", [[11,22], [33,44]])
    print(annObj.ToJson())
    annObj.SaveToJsonFile("1.json")

合成几千张图片后，即可使用unet等网络进行训练、验证。