54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
import os, sys
|
|
from pathlib import Path
|
|
|
|
url = "http://cs231n.stanford.edu/tiny-imagenet-200.zip"
|
|
|
|
def load_val():
|
|
path = 'tiny-imagenet-200/val/val_annotations.txt'
|
|
cfile = open(path, 'r')
|
|
content = cfile.readlines()
|
|
content = [x.strip().split('\t') for x in content]
|
|
cfile.close()
|
|
images = [x[0] for x in content]
|
|
labels = [x[1] for x in content]
|
|
return images, labels
|
|
|
|
def main():
|
|
os.system("wget {:}".format(url))
|
|
os.system("rm -rf tiny-imagenet-200")
|
|
os.system("unzip -o tiny-imagenet-200.zip")
|
|
images, labels = load_val()
|
|
savedir = 'tiny-imagenet-200/new_val'
|
|
if not os.path.exists(savedir): os.makedirs(savedir)
|
|
for image, label in zip(images, labels):
|
|
cdir = savedir + '/' + label
|
|
if not os.path.exists(cdir): os.makedirs(cdir)
|
|
ori_path = 'tiny-imagenet-200/val/images/' + image
|
|
os.system("cp {:} {:}".format(ori_path, cdir))
|
|
os.system("rm -rf tiny-imagenet-200/val")
|
|
os.system("mv {:} tiny-imagenet-200/val".format(savedir))
|
|
|
|
def generate_salt_pepper():
|
|
targetdir = Path('tiny-imagenet-200/val')
|
|
noisedir = Path('tiny-imagenet-200/val-noise')
|
|
assert targetdir.exists(), '{:} does not exist'.format(targetdir)
|
|
from imgaug import augmenters as iaa
|
|
import cv2
|
|
aug = iaa.SaltAndPepper(p=0.2)
|
|
|
|
for sub in targetdir.iterdir():
|
|
if not sub.is_dir(): continue
|
|
subdir = noisedir / sub.name
|
|
if not subdir.exists(): os.makedirs('{:}'.format(subdir))
|
|
images = sub.glob('*.JPEG')
|
|
for image in images:
|
|
I = cv2.imread(str(image))
|
|
Inoise = aug.augment_image(I)
|
|
savepath = subdir / image.name
|
|
cv2.imwrite(str(savepath), Inoise)
|
|
print ('{:} done'.format(sub))
|
|
|
|
if __name__ == "__main__":
|
|
#main()
|
|
generate_salt_pepper()
|