图片分类-数据增强

import osimport PIL.Image as Imageimport shutilfrom torchvision import transformsfrom numpy.random import randint,randomdef zengqiang(root_path):for img in os.listdir(root_path):i = 0.if img[-3:] in [

只会搬砖的李富贵

459人浏览 · 2021-04-24 21:24:24

只会搬砖的李富贵 · 2021-04-24 21:24:24 发布

import os
import PIL.Image as Image
import shutil
from torchvision import transforms
from numpy.random import randint,random

def zengqiang(root_path):
    for img in os.listdir(root_path):
        i = 0.
        if img[-3:] in ["jpg", "png", "JPG","bmp"]:
            img_path = root_path + img[:-4]
            if not os.path.exists(img_path):
                os.mkdir(img_path)
            shutil.copy(root_path + img, img_path + "/" + img[:-4] + str(i) + ".jpg")
            print(img_path)
            im = Image.open(root_path + img)
            for p in range(5):
                # 随机裁剪
                new_im = transforms.RandomCrop(randint(1024,2048),padding = sequence,pad_if_needed=True, fill=0)(new_im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                # 随机水平翻转
                new_im = transforms.RandomHorizontalFlip(p=0.9)(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                # 随机竖直翻转
                new_im = transforms.RandomVerticalFlip(p=0.9)(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                # 随机旋转
                new_im = transforms.RandomRotation(30)(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                new_im = transforms.RandomRotation(60)(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1

            
            # 上下左右中心裁
            new_im = transforms.FiveCrop()(im)
            new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
            i += 1


            # 调节亮度对比度啥的
            for p in range(3):
                new_im = transforms.ColorJitter(brightness=0)(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                new_im = transforms.RandomCrop(randint(1024,2048),padding = sequence,pad_if_needed=True, fill=0)(new_im)
                new_im = transforms.ColorJitter(brightness=0)(new_im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                new_im = transforms.ColorJitter(hue=0.5)(new_im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                new_im = transforms.RandomRotation(45)(im)
                new_im = transforms.ColorJitter(contrast=0)(new_im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                new_im = transforms.ColorJitter(contrast=2)(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1
                # 
                new_im = transforms.CenteCrop(random(512,1500))(im)
                new_im.save(img_path + "/" + img[:-4] + str(i) + ".jpg")
                i += 1   
        else:
            continue