圖像數據增強(縮放)

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)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章