import os, glob
import random
import shutil
from tqdm.auto import tqdm
def print_info(string, writepath):
print(string)
print(string, file=open(writepath, "a"))
## Use when all you care is just random split
def split_image_folder(source_root, target_root, portion=2, ratio= [0.8, 0.2]):
random.seed(73)
assert len(ratio) == portion
assert sum(ratio) == 1
os.makedirs(target_root, exist_ok = True)
class_folder = [n for n in os.listdir(source_root)]
for cf in class_folder:
ipath = os.path.join(source_root, cf, "*.png")
imnames = [os.path.basename(x) for x in glob.glob(ipath)]
random.shuffle(imnames)
lsize = [round(len(imnames)*r) for r in ratio ] ## floor will be cleaner
begin = 0; outnames = []
# print(lsize)
for i in range(portion):
outnames.append(imnames[begin:begin+lsize[i]])
begin += lsize[i]
## handle leftouts due to rounding
for j in range(begin, len(imnames)):
outnames[j%portion].append(imnames[j])
print_info(f"ClassFolder: {cf}, Count: {len(imnames)}", target_root+"/split-info.txt")
ipath_dir = os.path.dirname(ipath)
for o, grp in tqdm(enumerate(outnames)):
print_info(f'Group-{o+1}:',len(grp), target_root+"/split-info.txt")
opath = os.path.join(target_root, f"group{o+1}",cf)
os.makedirs(opath, exist_ok = True)
for g in grp:
shutil.copy(os.path.join(ipath_dir, g), opath)
return None
## Ensures a given ID is not present in train -val splits within a class
def split_image_folder_byID(source_root, target_root,
ratio= [0.8, 0.2], ):
"""
Written just for two folder logic - train, valid
"""
def parse_study_id(string):
# str_list = [ x for x in string.split("_") if "Study" in x ]
# out = str_list[0]
out = string.split("_")[0] # PateintID
return out
def grouper(name_list):
out_dic = {}
for name in name_list:
id = parse_study_id(name)
if id in out_dic:
out_dic[id]["count"] += 1
out_dic[id]["files"].append(name)
else:
out_dic[id] = {}
out_dic[id]["count"] = 1
out_dic[id]["files"] = [name]
out_dic = dict(sorted(out_dic.items(), key = lambda item: item[1]["count"],
reverse=True ))
return out_dic
random.seed(73)
assert sum(ratio) == 1
assert len(ratio) <= 2
os.makedirs(target_root, exist_ok = True)
class_folder = [n for n in os.listdir(source_root)]
for cf in class_folder:
ipath = os.path.join(source_root, cf, "*.png")
imnames = [os.path.basename(x) for x in glob.glob(ipath)]
total_count = len(imnames)
grouped_dic = grouper(imnames)
#TODO: rewrite logic for N splits when life bestows time
curr_count = 0
valid_dic = {}
while curr_count < total_count * ratio[-1]:
k = random.sample(grouped_dic.keys(), 1)[0]
item = grouped_dic.pop(k)
valid_dic[k] = item
curr_count += item["count"]
print_info(f'Class:{cf},Tot:{total_count} Train:{total_count - curr_count} Valid:{curr_count} ratio:{curr_count/total_count}',
target_root+"/split-info.txt")
ipath_dir = os.path.dirname(ipath)
for o, dic in enumerate([grouped_dic, valid_dic]):
for k in tqdm(dic.keys()):
grp = dic[k]["files"]
opath = os.path.join(target_root, f"group{o+1}",cf)
os.makedirs(opath, exist_ok = True)
for g in grp:
shutil.copy(os.path.join(ipath_dir, g), opath)
return None
if __name__ == "__main__":
split_image_folder_by_nameid(
"/home/joseph.benjamin/WERK/UltraSound/FetalDataImg",
"/home/joseph.benjamin/WERK/data/study-split/",
ratio=[0.8, 0.2]
)
