Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
# -*- 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment