Helper function to move patches and crop data into train and val
(source_images_dir,
source_masks_dir,
source_crops_dir,
dest_images_dir,
dest_masks_dir,
dest_crops_dir,
split_percentage)
| 61 | pbar.update(1) |
| 62 | |
| 63 | def split_data(source_images_dir, |
| 64 | source_masks_dir, |
| 65 | source_crops_dir, |
| 66 | dest_images_dir, |
| 67 | dest_masks_dir, |
| 68 | dest_crops_dir, |
| 69 | split_percentage): |
| 70 | """ |
| 71 | Helper function to move patches and crop data into train and val |
| 72 | """ |
| 73 | images_list = sorted([os.path.join(source_images_dir, image) |
| 74 | for image in os.listdir(source_images_dir)]) |
| 75 | masks_list = sorted([os.path.join(source_masks_dir, image) |
| 76 | for image in os.listdir(source_masks_dir)]) |
| 77 | crops_list = sorted(glob(os.path.join(source_crops_dir, "*/*"))) |
| 78 | crops_labels_list = [crop.split(os.path.sep)[-2] for crop in crops_list] |
| 79 | |
| 80 | train_images, test_images, train_masks, test_masks = train_test_split(images_list, |
| 81 | masks_list, |
| 82 | test_size=split_percentage, |
| 83 | random_state=42) |
| 84 | |
| 85 | train_crops, test_crops, train_labels, test_labels = train_test_split(crops_list, |
| 86 | crops_labels_list, |
| 87 | test_size=split_percentage, |
| 88 | random_state=42) |
| 89 | |
| 90 | if not os.path.exists(os.path.join(dest_images_dir, "train", "images")): |
| 91 | os.makedirs(os.path.join(dest_images_dir, "train", "images")) |
| 92 | if not os.path.exists(os.path.join(dest_images_dir, "train", "masks")): |
| 93 | os.makedirs(os.path.join(dest_images_dir, "train", "masks")) |
| 94 | |
| 95 | if not os.path.exists(os.path.join(dest_images_dir, "val", "images")): |
| 96 | os.makedirs(os.path.join(dest_images_dir, "val", "images")) |
| 97 | if not os.path.exists(os.path.join(dest_images_dir, "val", "masks")): |
| 98 | os.makedirs(os.path.join(dest_images_dir, "val", "masks")) |
| 99 | |
| 100 | if not os.path.exists(os.path.join(dest_crops_dir, "train")): |
| 101 | os.makedirs(os.path.join(dest_crops_dir, "train")) |
| 102 | if not os.path.exists(os.path.join(dest_crops_dir, "val")): |
| 103 | os.makedirs(os.path.join(dest_crops_dir, "val")) |
| 104 | |
| 105 | for image, mask in zip(train_images, train_masks): |
| 106 | image_name = os.path.basename(image) |
| 107 | mask_name = os.path.basename(mask) |
| 108 | shutil.move(image, os.path.join(dest_images_dir, "train", "images", image_name)) |
| 109 | shutil.move(mask, os.path.join(dest_masks_dir, "train", "masks", mask_name)) |
| 110 | |
| 111 | for image, mask in zip(test_images, test_masks): |
| 112 | image_name = os.path.basename(image) |
| 113 | mask_name = os.path.basename(mask) |
| 114 | shutil.move(image, os.path.join(dest_images_dir, "val", "images", image_name)) |
| 115 | shutil.move(mask, os.path.join(dest_masks_dir, "val", "masks", mask_name)) |
| 116 | |
| 117 | for image, label in zip(train_crops, train_labels): |
| 118 | image_name = os.path.basename(image) |
| 119 | if not os.path.exists(os.path.join(dest_crops_dir, "train", label)): |
| 120 | os.makedirs(os.path.join(dest_crops_dir, "train", label)) |