-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathextract_int_mask.py
115 lines (99 loc) · 3.51 KB
/
extract_int_mask.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import os
import numpy as np
import pandas as pd
from PIL import Image
from tqdm import tqdm
# Defining color maps and translation dictionary
# patch_colors_bgr_01_old = {
# "background": [0, 0, 0], # BGR
# "lips": [0, 0, 1], # BGR
# "eyes": [0, 1, 0], # BGR
# "nose": [1, 0, 0], # BGR
# "face": [0.5019607843137255, 0.5019607843137255, 0.5019607843137255], # BGR
# "hair": [0, 1, 1], # BGR
# "eyebrow": [1, 0, 1], # BGR
# "ears": [1, 1, 0], # BGR
# "teeth": [1, 1, 1], # BGR
# "facial_hair": [0.7529411764705882, 0.7529411764705882, 1], # BGR
# "glasses": [0.5019607843137255, 0.5019607843137255, 0] # BGR
# }
patch_colors_bgr_01 = {
"background": [0, 0, 0], # BGR
"generalface": [0.5019607843137255, 0.5019607843137255, 0.5019607843137255], # BGR
"left_eye": [0, 1, 0], # BGR
"right_eye": [0, 0.5019607843137255, 0], # BGR
"nose": [1, 0, 0], # BGR
"left_ear": [1, 1, 0], # BGR
"right_ear": [0.25098039215686274, 0.25098039215686274, 0], # BGR
"lips": [0, 0, 1], # BGR
"left_eyebrow": [1, 0, 1], # BGR
"right_eyebrow": [0.5019607843137255, 0, 0.5019607843137255], # BGR
"hair": [0, 1, 1], # BGR
"teeth": [1, 1, 1], # BGR
"specs": [0.5019607843137255, 0.5019607843137255, 0], # BGR
"beard": [0.7529411764705882, 0.7529411764705882, 1] # BGR
}
patch_colors_rgb_01 = {key: [color[2], color[1], color[0]] for key, color in patch_colors_bgr_01.items()}
# convert patch_colors_rgb_01 to patch_colors_rgb_255
patch_colors_rgb_255 = {}
for key, value in patch_colors_rgb_01.items():
patch_colors_rgb_255[key] = [int(255 * x) for x in value]
# patch_colors_int_old = {
# "background": 0,
# "lips": 1,
# "eyes": 2,
# "nose": 3,
# "face": 4,
# "hair": 5,
# "eyebrow": 6,
# "ears": 7,
# "teeth": 8,
# "facial_hair": 9,
# "glasses": 10
# }
patch_colors_int = {
"background": 0,
"generalface": 1,
"left_eye": 2,
"right_eye": 3,
"nose": 4,
"left_ear": 5,
"right_ear": 6,
"lips": 7,
"left_eyebrow": 8,
"right_eyebrow": 9,
"hair": 10,
"teeth": 11,
"specs": 12,
"beard": 13
}
df_translate = pd.DataFrame({
'Patch': list(patch_colors_int.keys()),
'RGB_255': list(patch_colors_rgb_255.values()),
'Integer': list(patch_colors_int.values())
})
print(df_translate)
def get_int_mask(arr, df_translate):
converted_msk = np.zeros_like(arr, dtype=np.uint8)
# Loop through the DataFrame and replace colors in the image
for index, row in df_translate.iterrows():
patch_color = row['RGB_255']
integer_value = row['Integer']
mask = np.all(arr == patch_color, axis=-1)
converted_msk[mask] = integer_value
return converted_msk
# convert masks
masks_folder_path = "/homeRepo/tanfoni/faceSegmentation/dataset_paid/masks"
masks_output_folder_path = "/homeRepo/tanfoni/faceSegmentation/dataset_paid_integers/masks"
if not os.path.isdir(masks_output_folder_path):
os.makedirs(masks_output_folder_path)
for filename in tqdm(os.listdir(masks_folder_path)):
image_path = os.path.join(masks_folder_path, filename)
save_cmsk_path = os.path.join(masks_output_folder_path, filename)
image = Image.open(image_path)
image = np.asarray(image)
cmsk = get_int_mask(image, df_translate)
# turn cmsk into png image
cmsk = Image.fromarray(cmsk)
cmsk.save(save_cmsk_path)
print("All done!")