python实现左右、上下、对角线翻转的数据增扩

这是一个目标检测项目的数据增扩,所以不是仅仅翻转图片就完成了,还需要修改标签里的目标的位置信息。

#上下左右翻转
import json
import os
import cv2
import numpy as np
global x_min
x_min=0
global x_max
x_max=0
global y_min
y_min=0
global y_max
y_max=0
#读json文件
def read_json(json_file):
    json_content=[]
    with open(json_file,'r',encoding='utf8') as load_f:
        data=json.load(load_f)
        path = data['path']
        outputs = data['outputs']
        time_labeled = data['time_labeled']
        labeled = data['labeled']
        size = data['size']
        
        item_dict = {
            "path": path,
            "outputs": outputs,
            "time_labeled": time_labeled,
            "labeled": labeled ,
            "size": size
        }
        json_content.append(item_dict)
        json_list=json_content[0]
        return json_list
#修改标签里的坐标信息
def json_affine(json_son,M,img,mirror):
    print("before:",json_son["outputs"]["object"])
    list01 = json_son["outputs"]["object"]
    num=len(list01)
    h, w = img.shape[0], img.shape[1]
    global x_min
    global x_max
    global y_min
    global y_max
    for i in range(num):
        if mirror==1:  #水平翻转
            x_min=list01[i]["bndbox"]["xmin"]
            x_max=list01[i]["bndbox"]["xmax"]
            json_son["outputs"]["object"][i]["bndbox"]["xmin"]=w-x_max
            json_son["outputs"]["object"][i]["bndbox"]["xmax"]=w-x_min
        elif mirror==2:  #上下翻转
            y_min=list01[i]["bndbox"]["ymin"]
            y_max=list01[i]["bndbox"]["ymax"]
            json_son["outputs"]["object"][i]["bndbox"]["ymin"]=h-y_max
            json_son["outputs"]["object"][i]["bndbox"]["ymax"]=h-y_min
        elif mirror==3:  #对角线翻转
            x_min=list01[i]["bndbox"]["xmin"]
            x_max=list01[i]["bndbox"]["xmax"]
            y_min=list01[i]["bndbox"]["ymin"]
            y_max=list01[i]["bndbox"]["ymax"]
            json_son["outputs"]["object"][i]["bndbox"]["xmin"]=w-x_max
            json_son["outputs"]["object"][i]["bndbox"]["xmax"]=w-x_min
            json_son["outputs"]["object"][i]["bndbox"]["ymin"]=h-y_max
            json_son["outputs"]["object"][i]["bndbox"]["ymax"]=h-y_min
    json_son['size']['height']=h
    json_son['size']['width']=w
    print("after:",json_son["outputs"]["object"])
    return json_son

def mirror_matrix(mirror,w_son,h_son):
    if mirror == 0:
        return np.float32([[1,0,0],[0,1,0]])
    elif mirror == 1:  #水平翻转
        return np.float32([[-1,0,w_son],[0,1,0]])
    elif mirror == 2:  #上下翻转
        return np.float32([[1,0,0],[0,-1,h_son]])
    else:  #对角线翻转
        return np.float32([[-1,0,w_son],[0,-1,h_son]])
#翻转图像    
def adjust(M_flip,w_son,h_son,son_scale,json_son,mirror):
     #M_flip = np.float32([[-1,0,w_son],[0,1,0]])
        cos = np.abs(M_flip[0, 0])
        sin = np.abs(M_flip[0, 1])
        new_w = h_son * sin + w_son * cos
        new_h = h_son * cos + w_son * sin
        w = int(np.round(new_w))
        h = int(np.round(new_h))

        #son_flip = cv2.warpAffine(src=son_rotation, M=M_flip, dsize=(h_son, w_son), borderValue=(255, 255, 255))
        son_flip = cv2.warpAffine(src=son_scale, M=M_flip, dsize=(w, h), borderValue=(255, 255, 255))
        json_son = json_affine(json_son,M_flip,son_flip,mirror)
        h_son, w_son = son_flip.shape[0], son_flip.shape[1]
        return son_flip,json_son,w_son,h_son

def main():
    base_path="D:/data_augmentation/oringin/"
    path1="D:/data_augmentation/oringin/"
    path2="D:/data_augmentation/new/"
    filelist = os.listdir(base_path)
    filelist.sort()
    namelist=['种类1','种类2','种类3','种类4','种类5','种类6','种类7','种类8','种类9','种类10']
    

    
    for name in filelist:
        filename = os.path.splitext(name)[0]
        filename_suffix = os.path.splitext(name)[1]
        if filename_suffix == ".json":
            json_son=read_json(base_path + name)
            label_name=json_son["outputs"]["object"]
            son = cv2.imread(path1+filename+'.jpg')
            h_son, w_son = son.shape[0], son.shape[1]
            #print(label_name)
            
            for _ in label_name:
                if _["name"] in namelist:
                    M_flip = mirror_matrix(1,w_son,h_son)
                    son_flip,json_son,w_son,h_son=adjust(M_flip,w_son,h_son,son,json_son,1)
                    son_path1=path2+filename+'_level'+'.jpg'
                    json_path1=path2+filename+'_level'+'.json'
                    json_son["path"]=son_path1
                    cv2.imwrite(son_path1, son_flip)
                    with open(json_path1,'w') as dump_f:
                        json.dump(json_son, dump_f,indent=4)
                    
                    json_son=read_json(base_path + name)
                    M_flip = mirror_matrix(2,w_son,h_son)
                    son_flip,json_son,w_son,h_son=adjust(M_flip,w_son,h_son,son,json_son,2)
                    son_path1=path2+filename+'_turn'+'.jpg'
                    json_path1=path2+filename+'_turn'+'.json'
                    json_son["path"]=son_path1
                    cv2.imwrite(son_path1, son_flip)
                    with open(json_path1,'w') as dump_f:
                        json.dump(json_son, dump_f,indent=4)
                    
                    json_son=read_json(base_path + name)
                    M_flip = mirror_matrix(3,w_son,h_son)
                    son_flip,json_son,w_son,h_son=adjust(M_flip,w_son,h_son,son,json_son,3)
                    son_path1=path2+filename+'_diagonal'+'.jpg'
                    json_path1=path2+filename+'_diagonal'+'.json'
                    json_son["path"]=son_path1
                    cv2.imwrite(son_path1, son_flip)
                    with open(json_path1,'w') as dump_f:
                        json.dump(json_son, dump_f,indent=4)
                    
                    #break
    print("done")
     
if __name__ == '__main__':
    main()

 

posted @ 2022-12-06 14:58  实数集  阅读(229)  评论(0)    收藏  举报