数据扩充--矩形框旋转

文件目录:
├── mark
│   └── dog.jpg
└── txt
└── dog.txt

dog.txt
392,226,804,222,776,556,508,548,354,514
1144,248,1704,112,1728,448,1236,536

对应的框:

import os
import cv2
import math
import random
import numpy as np

def rotation_point(img, angle=15, points=None):
    cols = img.shape[1]
    rows = img.shape[0]
    M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
    heightNew = int(cols * math.fabs(math.sin(math.radians(angle))) + rows * math.fabs(math.cos(math.radians(angle))))
    widthNew = int(rows * math.fabs(math.sin(math.radians(angle))) + cols * math.fabs(math.cos(math.radians(angle))))
    M[0, 2] += (widthNew - cols) / 2
    M[1, 2] += (heightNew - rows) / 2
    img = cv2.warpAffine(img, M, (widthNew, heightNew))
    a = M[:, :2]  ##a.shape (2,2)
    b = M[:, 2:]  ###b.shape(2,1)
    b = np.reshape(b, newshape=(1, 2))
    a = np.transpose(a)
    v_pt = []
    for point in points:
        len_1 = len(point)
        point = np.reshape(point, newshape=(len_1, 2))  # point = np.reshape(point, newshape=(len(point) * 4, 2))
        point = np.dot(point, a) + b
        point = np.reshape(point, newshape=(len_1, 2))
        point = point.astype(int)
        v_pt.append(point)
    return img, v_pt

def get_pt(txt_path):
    vv_pt_new = []
    with open(txt_path,'r') as fr:
        txt_line = fr.readlines()
        for line in txt_line:
            line = line.strip()
            pts = line.split(',')
            vv_pt = []
            len_pts = len(pts)
            for i in range(0,int(len_pts),2):
                v_pt = []
                x = pts[i]
                y = pts[i + 1]
                v_pt.append(int(x))
                v_pt.append(int(y))
                vv_pt.append(v_pt)
            vv_pt_new.append(vv_pt)
    return vv_pt_new

def write_pt(txt_path,vv_pt_rot):
    with open(txt_path,'a')as fw:
        for pts in vv_pt_rot:
            for i,pt in enumerate(pts):
                x = pt[0]
                y = pt[1]
                fw.write(str(x))
                fw.write(',')
                fw.write(str(y))
                if i != len(pts)-1:
                    fw.write(',')
                else:
                    fw.write('\n')

def show_debug(img_,vv_pt,m_rot,vv_pt_rot,path_show,b_show_save=False):
    img = img_.copy()
    for val in vv_pt:
        val = np.reshape(val,newshape=(len(val),2))
        val = val.astype(int)
        cv2.polylines(img, [val], True, (255, 0, 255), 2)
    for val in vv_pt_rot:
        cv2.polylines(m_rot, [val], True, (0, 255, 255), 2)

    if b_show_save:
        cv2.imwrite(path_show,m_rot)
        cv2.imwrite(path_show.replace('.jpg','_src_.jpg'),img)

    if False:
        cv2.namedWindow('src',0)
        cv2.imshow('src',img)
        cv2.namedWindow('m_rot',0)
        cv2.imshow('m_rot',m_rot)
        cv2.waitKey(0)


b_flg_show_rot_save = True
root_img = "/media/data_2/everyday/0801/5555/mark/"
root_pt_txt = "/media/data_2/everyday/0801/5555/txt/"
save_dir = os.path.dirname(os.path.dirname(root_img)) + '/rot/'
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

save_txt = save_dir + 'txt/'
save_img = save_dir + 'm_rot/'
if not os.path.exists(save_txt):
    os.makedirs(save_txt)
if not os.path.exists(save_img):
    os.makedirs(save_img)

save_show = ''
if b_flg_show_rot_save:
    save_show = save_dir + 'show/'
    if not os.path.exists(save_show):
        os.makedirs(save_show)


list_img = os.listdir(root_img)
for cnt_img,img_name in enumerate(list_img):
    print("cnt=%d,img_name=%s"%(cnt_img,img_name))
    img_path = root_img + img_name
    img = cv2.imread(img_path)
    txt_name = img_name.replace('.jpg','.txt')
    txt_path = root_pt_txt + txt_name
    vv_pt = get_pt(txt_path)
    for cnt in range(0,20):
        ang = random.randint(2,350)
        print("ang=%d"%(ang))
        m_rot, vv_pt_rot = rotation_point(img, ang, vv_pt)
        name_img_rot = img_name.replace('.jpg','_'+str(ang)+'.jpg')
        name_txt_rot = name_img_rot.replace('.jpg','.txt')
        save_img_path = save_img + name_img_rot
        save_txt_path = save_txt + name_txt_rot
        cv2.imwrite(save_img_path,m_rot)
        write_pt(save_txt_path, vv_pt_rot)

        if b_flg_show_rot_save:
            show_debug(img, vv_pt, m_rot, vv_pt_rot, save_show + name_img_rot, b_flg_show_rot_save)

运行脚本生成,目录如下:
├── mark
│   └── dog.jpg
├── rot
│   ├── m_rot
│   │   ├── dog_108.jpg
│   │   ├── dog_120.jpg
│   │   ├── dog_150.jpg
│   │   ├── dog_159.jpg
│   │   ├── dog_169.jpg
│   │   ├── dog_170.jpg
│   │   ├── dog_196.jpg
│   │   ├── dog_204.jpg
│   │   ├── dog_212.jpg
│   │   ├── dog_220.jpg
│   │   ├── dog_222.jpg
│   │   ├── dog_249.jpg
│   │   ├── dog_279.jpg
│   │   ├── dog_290.jpg
│   │   ├── dog_300.jpg
│   │   ├── dog_332.jpg
│   │   ├── dog_39.jpg
│   │   ├── dog_66.jpg
│   │   ├── dog_6.jpg
│   │   └── dog_72.jpg
│   ├── show
│   │   ├── dog_108.jpg
│   │   ├── dog_120.jpg
│   │   ├── dog_150.jpg
│   │   ├── dog_159.jpg
│   │   ├── dog_169.jpg
│   │   ├── dog_170.jpg
│   │   ├── dog_196.jpg
│   │   ├── dog_204.jpg
│   │   ├── dog_212.jpg
│   │   ├── dog_220.jpg
│   │   ├── dog_222.jpg
│   │   ├── dog_249.jpg
│   │   ├── dog_279.jpg
│   │   ├── dog_290.jpg
│   │   ├── dog_300.jpg
│   │   ├── dog_332.jpg
│   │   ├── dog_39.jpg
│   │   ├── dog_66.jpg
│   │   ├── dog_6.jpg
│   │   └── dog_72.jpg
│   └── txt
│   ├── dog_108.txt
│   ├── dog_120.txt
│   ├── dog_150.txt
│   ├── dog_159.txt
│   ├── dog_169.txt
│   ├── dog_170.txt
│   ├── dog_196.txt
│   ├── dog_204.txt
│   ├── dog_212.txt
│   ├── dog_220.txt
│   ├── dog_222.txt
│   ├── dog_249.txt
│   ├── dog_279.txt
│   ├── dog_290.txt
│   ├── dog_300.txt
│   ├── dog_332.txt
│   ├── dog_39.txt
│   ├── dog_66.txt
│   ├── dog_6.txt
│   └── dog_72.txt
└── txt
└── dog.txt

dog_65.txt 示例:
369,1479,540,1104,831,1271,710,1510,614,1635
707,807,821,242,1135,362,1007,845


猜你喜欢

转载自www.cnblogs.com/yanghailin/p/11285538.html