image_augmentation.py

# -*- coding: utf-8 -*-
"""image_augmentation.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1iHvkV1IOUOCNQH2MeZLeq6bX4--PZkXA
"""

# Commented out IPython magic to ensure Python compatibility.
# importing all the required libraries
import warnings
warnings.filterwarnings('ignore')
import numpy as np
import skimage.io as io
from skimage.transform import rotate, AffineTransform, warp
import os

# %matplotlib inline

dir_path_0 = r'/content/drive/MyDrive/Hackerthon_data/0_img' #path to file for scenario 0
dir_path_1 = r'/content/drive/MyDrive/Hackerthon_data/1_img' #path to file for scenario 1
dir_path_2 = r'/content/drive/MyDrive/Hackerthon_data/2_img' #path to file for scenario 2

def percentage_matcher(path0, path1, path2, n_iter):
  '''change the percentages of files in each scenario to be similar to 2 dp.
  input        path0: path to scenario 0
  input        path1: path to scenario 1
  input        path2: path to scenario 2
  input       n_iter: numbner of iterations for the for loop
  return     lengths: the length of each directory
  return percentages: the percentage of evenets in each directory'''

  len_0 = (len([entry for entry in os.listdir(dir_path_0) if os.path.isfile(os.path.join(dir_path_0, entry))])) #calculate the number of files in the directory
  len_1 = (len([entry for entry in os.listdir(dir_path_1) if os.path.isfile(os.path.join(dir_path_1, entry))])) #calculate the number of files in the directory
  len_2 = (len([entry for entry in os.listdir(dir_path_2) if os.path.isfile(os.path.join(dir_path_2, entry))])) #calculate the number of files in the directory
  lengths = np.array([len_0,len_1,len_2]) #put the starting number of flies for each scenario into an array
  percentages = (lengths)/np.sum(lengths) #calculate the starting percentages of each image in the three scenarios
  for i in range(n_iter):
    min_length = min(lengths)
    boolarr = lengths == min_length
    lengths[boolarr] = lengths[boolarr] + 5
    percentages = (lengths)/np.sum(lengths)
  return lengths, percentages

lengths, percentages = percentage_matcher(dir_path_0,dir_path_1,dir_path_2,200)