-
Notifications
You must be signed in to change notification settings - Fork 1
/
make_masks_fast.py
46 lines (29 loc) · 1.13 KB
/
make_masks_fast.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
# -*- coding: utf-8 -*-
"""
Created on Thu Dec 9 23:41:51 2021
@author: Alex
"""
import argparse
import os
from corescore.masks import CoreImageProcessor
def process_images(image_dir, label):
"""Create masks for labelled images.
For now, merge rock fragment labels from both core boxes"""
coreProcessor = CoreImageProcessor("Images",
labels=label,
merge_fragment_labels=True)
image = coreProcessor.core_types
mask_file = coreProcessor.processImage(image)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--images', help="Directory of images to make masks for")
parser.add_argument('--labels_dir', help="folder of JSON files with labels for images")
args = parser.parse_args()
images = args.images
if not images:
images = './Images'
labels_dir = args.labels_dir
if not labels_dir:
labels_dir = 'train_labels'
for f in os.listdir(labels_dir):
process_images(images, os.path.join(labels_dir,f))