Main function for creating patches and splitting data
(images_dir,
masks_dir,
crops_dir,
patch_size=768)
| 127 | shutil.move(image, os.path.join(dest_crops_dir, "val", label, image_name)) |
| 128 | |
| 129 | def main(images_dir, |
| 130 | masks_dir, |
| 131 | crops_dir, |
| 132 | patch_size=768): |
| 133 | """ |
| 134 | Main function for creating patches and splitting data |
| 135 | """ |
| 136 | dest_images_dir = os.path.join(os.path.abspath(os.path.join(images_dir, os.pardir)), |
| 137 | "patches", "images") |
| 138 | dest_masks_dir = os.path.join(os.path.abspath(os.path.join(masks_dir, os.pardir)), |
| 139 | "patches", "masks") |
| 140 | |
| 141 | images_list = sorted([os.path.join(images_dir, image) for image in os.listdir(images_dir)]) |
| 142 | masks_list = sorted([os.path.join(masks_dir, image) for image in os.listdir(masks_dir)]) |
| 143 | |
| 144 | # create directory if not exist |
| 145 | if not os.path.exists(dest_images_dir): |
| 146 | os.makedirs(dest_images_dir) |
| 147 | if not os.path.exists(dest_masks_dir): |
| 148 | os.makedirs(dest_masks_dir) |
| 149 | |
| 150 | create_patches(images_list, |
| 151 | masks_list, |
| 152 | patch_size, |
| 153 | dest_images_dir, |
| 154 | dest_masks_dir) |
| 155 | |
| 156 | source_images_dir = dest_images_dir |
| 157 | source_masks_dir = dest_masks_dir |
| 158 | source_crops_dir = crops_dir |
| 159 | dest_images_dir = os.path.abspath(os.path.join(images_dir, os.pardir)) |
| 160 | dest_masks_dir = os.path.abspath(os.path.join(masks_dir, os.pardir)) |
| 161 | dest_crops_dir = os.path.abspath(os.path.join(crops_dir, os.pardir)) |
| 162 | |
| 163 | split_data(source_images_dir, |
| 164 | source_masks_dir, |
| 165 | source_crops_dir, |
| 166 | dest_images_dir, |
| 167 | dest_masks_dir, |
| 168 | dest_crops_dir, |
| 169 | split_percentage=0.2) |
| 170 | |
| 171 | if __name__ == "__main__": |
| 172 | parser = argparse.ArgumentParser(prog='PCBSegClassNet') |
no test coverage detected