-
Notifications
You must be signed in to change notification settings - Fork 11
/
data.py
107 lines (82 loc) · 3.33 KB
/
data.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
import os
import numpy as np
import cv2
from glob import glob
from tqdm import tqdm
import imageio
from albumentations import HorizontalFlip, VerticalFlip, ElasticTransform, GridDistortion, OpticalDistortion, CoarseDropout
def create_dir(path):
if not os.path.exists(path):
os.makedirs(path)
def load_data(path):
""" X = Images and Y = masks """
train_x = sorted(glob(os.path.join(path, "training", "images", "*.tif")))
train_y = sorted(glob(os.path.join(path, "training", "1st_manual", "*.gif")))
test_x = sorted(glob(os.path.join(path, "test", "images", "*.tif")))
test_y = sorted(glob(os.path.join(path, "test", "1st_manual", "*.gif")))
return (train_x, train_y), (test_x, test_y)
def augment_data(images, masks, save_path, augment=True):
H = 512
W = 512
for idx, (x, y) in tqdm(enumerate(zip(images, masks)), total=len(images)):
""" Extracting names """
name = x.split("/")[-1].split(".")[0]
""" Reading image and mask """
x = cv2.imread(x, cv2.IMREAD_COLOR)
y = imageio.mimread(y)[0]
if augment == True:
aug = HorizontalFlip(p=1.0)
augmented = aug(image=x, mask=y)
x1 = augmented["image"]
y1 = augmented["mask"]
aug = VerticalFlip(p=1.0)
augmented = aug(image=x, mask=y)
x2 = augmented["image"]
y2 = augmented["mask"]
aug = ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03)
augmented = aug(image=x, mask=y)
x3 = augmented['image']
y3 = augmented['mask']
aug = GridDistortion(p=1)
augmented = aug(image=x, mask=y)
x4 = augmented['image']
y4 = augmented['mask']
aug = OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5)
augmented = aug(image=x, mask=y)
x5 = augmented['image']
y5 = augmented['mask']
X = [x, x1, x2, x3, x4, x5]
Y = [y, y1, y2, y3, y4, y5]
else:
X = [x]
Y = [y]
index = 0
for i, m in zip(X, Y):
i = cv2.resize(i, (W, H))
m = cv2.resize(m, (W, H))
if len(X) == 1:
tmp_image_name = f"{name}.jpg"
tmp_mask_name = f"{name}.jpg"
else:
tmp_image_name = f"{name}_{index}.jpg"
tmp_mask_name = f"{name}_{index}.jpg"
image_path = os.path.join(save_path, "image", tmp_image_name)
mask_path = os.path.join(save_path, "mask", tmp_mask_name)
cv2.imwrite(image_path, i)
cv2.imwrite(mask_path, m)
index += 1
if __name__ == "__main__":
""" Seeding """
np.random.seed(42)
""" Load the data """
data_path = "/media/nikhil/ML/ml_dataset/Retina blood vessel segmentation/"
(train_x, train_y), (test_x, test_y) = load_data(data_path)
print(f"Train: {len(train_x)} - {len(train_y)}")
print(f"Test: {len(test_x)} - {len(test_y)}")
""" Creating directories """
create_dir("new_data/train/image")
create_dir("new_data/train/mask")
create_dir("new_data/test/image")
create_dir("new_data/test/mask")
augment_data(train_x, train_y, "new_data/train/", augment=False)
augment_data(test_x, test_y, "new_data/test/", augment=False)