CopyForEqualDistribution

# -*- coding: utf-8 -*-
"""
Created on Sun Oct  4 13:35:21 2020

@author: Kevin

Equal distribution by copying. 
"""

import os
import glob
import shutil
import argparse
from matplotlib import pyplot as plt
import numpy as np
import time
np.random.seed(3121991)

def parse_args():
    
    parser = argparse.ArgumentParser(description="Will copy images to create an equal dsitribtuion")
    parser.add_argument("--path",type=str)
    args = parser.parse_args()
    
    return args

if __name__ == "__main__":
    
    
    args = parse_args()
    path_to_data  = args.path

    classes = os.listdir(path_to_data)
    lengths = [len(os.listdir(path_to_data + "/" + c)) for c in classes]
    
    plt.figure(1)
    plt.bar(np.arange(0,len(lengths)),lengths)    
    
    max_class = max(lengths)     
    print("The maximum is",max_class)
    for length,c  in zip(lengths,classes):
        
        base_path = path_to_data + "/" + c
        
        need = max_class - length 
        if need == 0:
            continue
        
        # get random files to copy
        idx = np.random.randint(0,length,need)
        
        possible = np.array(os.listdir(base_path))        
        
        to_copy = possible[idx]
        
        for element_to_copy in to_copy:
             
            print("copying:",base_path+ "/" + element_to_copy)            

            shutil.copy(base_path + "/" + element_to_copy, base_path + "/" + str(np.random.randint(0,1000000)) + str(time.time()) + ".jpg")

    lengths = [len(os.listdir(path_to_data + "/" + c)) for c in classes]
    
    plt.figure(2)
    plt.bar(np.arange(0,len(lengths)),lengths)