图像数据增强（缩放）

import cv2
from PIL import Image
import math
import numpy as np
import os
import pdb
import xml.etree.ElementTree as ET
from skimage import transform
from skimage.io import imread, imsave


class ImgAugemention():
    def __init__(self):
        self.formats = ['.png', '.jpg', '.jpeg']

    def zoom(self, img, p1x, p1y, p2x, p2y):
        w = img.shape[1]
        h = img.shape[0]    # print(w, h)    [1280, 720]
        crop_p1x = max(p1x, 0)
        crop_p1y = max(p1y, 0)
        crop_p2x = min(p2x, w)
        crop_p2y = min(p2y, h)

        cropped_img = img[crop_p1y:crop_p2y, crop_p1x:crop_p2x]

        x_pad_before = -min(0, p1x)
        x_pad_after = max(0, p2x - w)
        y_pad_before = -min(0, p1y)
        y_pad_after = max(0, p2y - h)

        padding = [(y_pad_before, y_pad_after), (x_pad_before, x_pad_after)]
        is_colour = len(img.shape) == 3
        if is_colour:
            padding.append((0, 0))  # colour images have an extra dimension

        padded_img = np.pad(cropped_img, padding, 'constant')
        return transform.resize(padded_img, (h, w))

    def vertical_xml(self, src, loc, zoom_loc):
        w = src.shape[1]
        h = src.shape[0]
        w1 = abs(zoom_loc[2] - zoom_loc[0])
        h1 = abs(zoom_loc[3] - zoom_loc[1])
        n_xmin = int(abs(loc[0] - zoom_loc[0]) * w / w1)
        n_ymin = int(abs(loc[1] - zoom_loc[1]) * h / h1)
        n_xmax = n_xmin + int(abs(loc[2] - loc[0]) * w / w1)
        n_ymax = n_ymin + int(abs(loc[3] - loc[1]) * h / h1)

        return n_xmin, n_ymin, n_xmax, n_ymax

    def process_img(self, imgs_path, xmls_path, img_save_path, xml_save_path, zoom_loc):
        for img_name in os.listdir(imgs_path):
            # split filename and suffix
            n, s = os.path.splitext(img_name)
            img_path = os.path.join(imgs_path, img_name)
            img = imread(img_path)
            zoom = self.zoom(img, zoom_loc[0], zoom_loc[1], zoom_loc[2], zoom_loc[3])
            # 写入图像
            imsave(img_save_path + n + "_" + "v_stretch.jpg", zoom)
            xml_url = img_name.split('.')[0] + ".xml"
            xml_path = os.path.join(xmls_path, xml_url)
            tree = ET.parse(xml_path)
            root = tree.getroot()
            for obj in root.iter('object'):
                loc = []
                for bnd in obj.iter('bndbox'):
                    for xmin in bnd.iter('xmin'):
                        loc.append(int(xmin.text))
                    for ymin in bnd.iter('ymin'):
                        loc.append(int(ymin.text))
                    for xmax in bnd.iter('xmax'):
                        loc.append(int(xmax.text))
                    for ymax in bnd.iter('ymax'):
                        loc.append(int(ymax.text))

                    n_xmin, n_ymin, n_xmax, n_ymax = self.vertical_xml(img, loc, zoom_loc)   # change locs in xml file
                    # update locs in xml file
                    for xmin in bnd.iter('xmin'):
                        xmin.text = str(n_xmin)
                    for ymin in bnd.iter('ymin'):
                        ymin.text = str(n_ymin)
                    for xmax in bnd.iter('xmax'):
                        xmax.text = str(n_xmax)
                    for ymax in bnd.iter('ymax'):
                        ymax.text = str(n_ymax)
                    bnd.set('updated', 'yes')
            # write into new xml file
            tree.write(xml_save_path + n + "_" + "v_stretch.xml")


if __name__ == "__main__":
    img_aug = ImgAugemention()
    root_dir_path1 = "E:\\singel_zoom\\tmp\\"
    root_dir_path2 = "E:\\singel_zoom1\\tmp\\"
    zoom_loc = [0, 40, 1280, 660]         # [0],[1]为左上角坐标， [2][3]为右下角坐标， 以此两点作伸缩变换(原图尺寸为(1280， 720))
                                          # zoom: [276, 6, 994, 633]  h_stretch: [250, 0, 900, 720]  v_stretch: [0, 40, 1280, 660]

    path_list = [""]                      # path_list = ["", "1\\", "2\\", "3\\", "4\\", "5\\"]
    for n in path_list:
        path_image = root_dir_path1 + n
        path_xml = root_dir_path1 + n
        path_image_dst = root_dir_path2 + n
        path_xml_dst = root_dir_path2 + n
        img_aug.process_img(path_image, path_xml, path_image_dst, path_xml_dst, zoom_loc)