data_augmentation.py 文件源码

def load_and_augmentate(self, root):
        angles = [45, 90, 135, 180, 225, 270, 315]
        scale = 1.0
        for img_dir in os.listdir(root):
            img_dir_path = os.path.join(root, img_dir)
            for img in os.listdir(img_dir_path):
                img_path = os.path.join(img_dir_path, img)
                image = caffe.io.load_image(img_path,color=True)
                w = image.shape[1]
                h = image.shape[0]
                img_name = img.split(".")[0]
                img_type = img.split(".")[-1]
                img_vmirror = cv2.flip(image,1)
                img_vmirror_path = os.path.join(img_dir_path,img_name+"_v."+img_type)
                skimage.io.imsave(img_vmirror_path, img_vmirror )
                for angle in angles:
                    rot_mat = cv2.getRotationMatrix2D((w*0.5, h*0.5), angle, scale)
                    new_img = cv2.warpAffine(image, rot_mat, (int(math.ceil(w)), int(math.ceil(h))), flags=cv2.INTER_LANCZOS4)
                    new_img_path = os.path.join(img_dir_path,img_name+"_"+str(angle)+"."+img_type)
                    skimage.io.imsave(new_img_path, new_img)
                    new_img_vmirror = cv2.flip(new_img, 1)
                    new_img_vmirror_path = os.path.join(img_dir_path, img_name+"_"+str(angle)+"_v."+img_type)
                    skimage.io.imsave(new_img_vmirror_path, new_img_vmirror)