diff --git a/train_test_indices.py b/train_test_indices.py new file mode 100644 index 0000000..055c943 --- /dev/null +++ b/train_test_indices.py @@ -0,0 +1,39 @@ +# In the future we might need to define exact train and test sets to compare different models. This script has initiated for this purpose. + +import os +import cv2 + +def load_images_from_folder(folder): + images = [] + for filename in os.listdir(folder): + if filename.endswith(".png"): + img = cv2.imread(os.path.join(folder, filename)) + if img is not None: + images.append(img) + return images +root_folder = '/home/sabrina/stinkbugs/labeled-data/' + +directory_list = list() +for root, dirs, files in os.walk("/home/sabrina/stinkbugs/labeled-data", topdown=False): + for name in dirs: + directory_list.append(os.path.join(root, name)) + +folders = [os.path.join(root_folder, x) for x in directory_list] +all_images = [img for folder in folders for img in load_images_from_folder(folder)] +# %% +n_frames = len(all_images) +frames = list(range(0,n_frames)) +random.Random(4).shuffle(frames) +split = 95 +percent = int(split * len(frames) / 100) +train1 = frames[:percent] +test1 = frames[percent:] + +list1 = deeplabcut.create_training_dataset( + config_path, + num_shuffles= suffle, + net_type="resnet_50", + augmenter_type="imgaug", + trainIndices = [train1], + testIndices = [test1], +)