深度学习 数据增强之图片及标注框旋转

参考：python 简单图像处理（4） 旋转

图片旋转使用pytorch的torchvision 库

#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import random
import glob
import numpy as np
import os
import math
import cv2

import torch
import torchvision
from torchvision import transforms as tfs
from torchvision.transforms import RandomRotation

from PIL import Image

img_Lists = [] 

#dir for rotating image
img_dir = "/data2/myfile/cwb/data"
#dir for rotated image
img_target_dir = "/data2/myfile/cwb/rotate_images"
#dir for rotating label
txt_dir = "/data2/myfile/cwb/txt"
#dir for rotated label 
txt_target_dir = "/data2/myfile/cwb/rotate_txts"

img_box_dir = "/data2/myfile/cwb/box_images"

img_Lists = glob.glob(img_dir + '/*.jpg')

img_basenames = []  # e.g. 100.jpg
for item in img_Lists:
    img_basenames.append(os.path.basename(item))

img_names = []  # e.g. 100
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)
    img_names.append(temp1)

#rotate per image
for img in img_names:
    im = Image.open((img_dir + '/' + img + '.jpg'))
    #rotate image
    rot_im = tfs.RandomRotation((-30,-30),expand=True)(im)
    rot_im=rot_im.convert('RGB')
    rot_im.save((img_target_dir + '/' + img + '.jpg'))

    # open the crospronding txt file
    gt = open((txt_dir + '/' + img + '.txt')).read().splitlines()

    txt_file = open((txt_target_dir + '/' + img + '.txt') ,'w')

    image=np.array(rot_im.convert('RGB'))

    for img_each_label in gt:
        gt_ind = img_each_label.split(',')
        point_coordinates = []
        bbox = []
        line = ''
        for i in xrange(8):
            point_coordinates.append(float(gt_ind[i]))   
        #coordinate rotate
        for i in range(0, 8, 2):
            x = point_coordinates[i]
            y = point_coordinates[i + 1]
            angle = 30 * math.pi / 180.0
            rotated_x = x * math.cos(angle) - y * math.sin(angle) - 0.5 * im.width * math.cos(angle) + 0.5 * im.height * math.sin(
            angle) + 0.5 * rot_im.width
            rotated_y = x * math.sin(angle) + y * math.cos(angle) - 0.5 * im.width * math.sin(angle) - 0.5 * im.height * math.cos(
            angle) + 0.5 * rot_im.height
            line = line  + str(rotated_x) + ',' + str(rotated_y) + ','
            bbox.append(int(rotated_x))
            bbox.append(int(rotated_y))

        line = line + str(gt_ind[8])
        print(line)
        pts = np.array([[bbox[0], bbox[1]],  [bbox[2], bbox[3]], [bbox[4], bbox[5]], [bbox[6], bbox[7]]], np.int32)
        pts = pts.reshape((-1, 1, 2))
        
        cv2.polylines(image, [pts], True, (0, 255, 255))
        txt_file.write(line+'\n')

    cv2.imwrite((img_box_dir + '/' + img + '.jpg'),image)

方法二（使用opencv,旋转后的图片更加清晰）

#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import random
import glob
import numpy as np
import os
import math
from math import *
import cv2

# import torch
# import torchvision
# from torchvision import transforms as tfs
# from torchvision.transforms import RandomRotation

from PIL import Image

img_Lists = [] 

#dir for rotating image
img_dir = "/data2/myfile/cwb/data"
#dir for rotated image
img_target_dir = "/data2/myfile/cwb/rotate_images"
#dir for rotating label
txt_dir = "/data2/myfile/cwb/txt"
#dir for rotated label 
txt_target_dir = "/data2/myfile/cwb/rotate_txts"

img_box_dir = "/data2/myfile/cwb/box_images"

img_Lists = glob.glob(img_dir + '/*.jpg')

img_basenames = []  # e.g. 100.jpg
for item in img_Lists:
    img_basenames.append(os.path.basename(item))

img_names = []  # e.g. 100
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)
    img_names.append(temp1)

#rotate per image
for img in img_names:
    # im = Image.open((img_dir + '/' + img + '.jpg'))
    #rotate image
    # rot_im = tfs.RandomRotation((-30,-30),expand=True)(im)
    # rot_im=rot_im.convert('RGB')
    # rot_im.save((img_target_dir + '/' + img + '.jpg'))


    im = cv2.imread((img_dir + '/' + img + '.jpg'))

    height,width=im.shape[:2]
     
    degree= -30
    #旋转后的尺寸
    heightNew=int(width*fabs(sin(radians(degree)))+height*fabs(cos(radians(degree)))) # 这个公式参考之前内容
    widthNew=int(height*fabs(sin(radians(degree)))+width*fabs(cos(radians(degree))))
     
    matRotation=cv2.getRotationMatrix2D((width/2,height/2),degree,1)
     
    matRotation[0,2] +=(widthNew-width)/2  # 因为旋转之后,坐标系原点是新图像的左上角,所以需要根据原图做转化
    matRotation[1,2] +=(heightNew-height)/2  
     
    rot_im=cv2.warpAffine(im,matRotation,(widthNew,heightNew),borderValue=(255,255,255))
    cv2.imwrite((img_target_dir + '/' + img + '.jpg'),rot_im)

    # open the crospronding txt file
    gt = open((txt_dir + '/' + img + '.txt')).read().splitlines()

    txt_file = open((txt_target_dir + '/' + img + '.txt') ,'w')

    image=rot_im

    for img_each_label in gt:
        gt_ind = img_each_label.split(',')
        point_coordinates = []
        bbox = []
        line = ''
        for i in xrange(8):
            point_coordinates.append(float(gt_ind[i]))   
        #coordinate rotate
        for i in range(0, 8, 2):
            x = point_coordinates[i]
            y = point_coordinates[i + 1]
            angle = 30 * math.pi / 180.0
            rotated_x = x * math.cos(angle) - y * math.sin(angle) - 0.5 * width * math.cos(angle) + 0.5 * height * math.sin(
            angle) + 0.5 * widthNew
            rotated_y = x * math.sin(angle) + y * math.cos(angle) - 0.5 * width * math.sin(angle) - 0.5 * height * math.cos(
            angle) + 0.5 * heightNew
            line = line  + str(int(rotated_x)) + ',' + str(int(rotated_y)) + ','
            bbox.append(int(rotated_x))
            bbox.append(int(rotated_y))

        line = line + str(gt_ind[8])
        print(line)
        pts = np.array([[bbox[0], bbox[1]],  [bbox[2], bbox[3]], [bbox[4], bbox[5]], [bbox[6], bbox[7]]], np.int32)
        pts = pts.reshape((-1, 1, 2))
        #可视化结果
        cv2.polylines(image, [pts], True, (0, 255, 255))
        #写入txt文件
        txt_file.write(line+'\n')

    cv2.imwrite((img_box_dir + '/' + img + '.jpg'),image)