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()