-
Notifications
You must be signed in to change notification settings - Fork 0
/
loveda_mask_convert.py
93 lines (71 loc) · 2.83 KB
/
loveda_mask_convert.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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import glob
import os
import numpy as np
import cv2
import multiprocessing.pool as mpp
import multiprocessing as mp
import time
import argparse
import torch
import random
SEED = 42
CLASSES = ('background', 'building', 'road', 'water', 'barren', 'forest',
'agricultural')
PALETTE = [[255, 255, 255], [255, 0, 0], [255, 255, 0], [0, 0, 255],
[159, 129, 183], [0, 255, 0], [255, 195, 128]]
def seed_everything(seed):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--mask-dir", default="data/LoveDA/Train/Rural/masks_png")
parser.add_argument("--output-mask-dir", default="data/LoveDA/Train/Rural/masks_png_convert")
return parser.parse_args()
def convert_label(mask):
mask[mask == 0] = 8
mask -= 1
return mask
def label2rgb(mask):
h, w = mask.shape[0], mask.shape[1]
mask_rgb = np.zeros(shape=(h, w, 3), dtype=np.uint8)
mask_convert = mask[np.newaxis, :, :]
mask_rgb[np.all(mask_convert == 0, axis=0)] = [255, 255, 255]
mask_rgb[np.all(mask_convert == 1, axis=0)] = [255, 0, 0]
mask_rgb[np.all(mask_convert == 2, axis=0)] = [255, 255, 0]
mask_rgb[np.all(mask_convert == 3, axis=0)] = [0, 0, 255]
mask_rgb[np.all(mask_convert == 4, axis=0)] = [159, 129, 183]
mask_rgb[np.all(mask_convert == 5, axis=0)] = [0, 255, 0]
mask_rgb[np.all(mask_convert == 6, axis=0)] = [255, 195, 128]
return mask_rgb
def patch_format(inp):
(mask_path, masks_output_dir) = inp
# print(mask_path, masks_output_dir)
mask_filename = os.path.splitext(os.path.basename(mask_path))[0]
mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
label = convert_label(mask)
rgb_label = label2rgb(label.copy())
rgb_label = cv2.cvtColor(rgb_label, cv2.COLOR_RGB2BGR)
out_mask_path_rgb = os.path.join(masks_output_dir + '_rgb', "{}.png".format(mask_filename))
cv2.imwrite(out_mask_path_rgb, rgb_label)
out_mask_path = os.path.join(masks_output_dir, "{}.png".format(mask_filename))
cv2.imwrite(out_mask_path, label)
if __name__ == "__main__":
seed_everything(SEED)
args = parse_args()
masks_dir = args.mask_dir
masks_output_dir = args.output_mask_dir
mask_paths = glob.glob(os.path.join(masks_dir, "*.png"))
if not os.path.exists(masks_output_dir):
os.makedirs(masks_output_dir)
os.makedirs(masks_output_dir + '_rgb')
inp = [(mask_path, masks_output_dir) for mask_path in mask_paths]
t0 = time.time()
mpp.Pool(processes=mp.cpu_count()).map(patch_format, inp)
t1 = time.time()
split_time = t1 - t0
print('images spliting spends: {} s'.format(split_time))