-
Notifications
You must be signed in to change notification settings - Fork 11
/
data.py
102 lines (85 loc) · 3.35 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
# -*- coding: utf-8 -*-
"""
Created on Fri Dec 18 09:53:51 2020
@author: Asma Baccouche
"""
from __future__ import print_function
import os
import numpy as np
import cv2
image_rows = 256
image_cols = 256
def create_train_data(data_path, name):
train_data_path = os.path.join(data_path, 'img')
train_mask_path = os.path.join(data_path, 'msk')
images = os.listdir(train_data_path)
masks = os.listdir(train_mask_path)
total = len(images)
imgs = np.ndarray((total, image_rows, image_cols, 3), dtype=np.uint8)
imgs_mask = np.ndarray((total, image_rows, image_cols), dtype=np.uint8)
i = 0
print('-'*30)
print('Creating training images...')
print('-'*30)
for j in range(len(images)):
img = cv2.imread(os.path.join(train_data_path, images[j]))
img = cv2.resize(img, (256,256))
#enhancement
gray_img=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
gray_img_eqhist=cv2.equalizeHist(gray_img)
img = cv2.cvtColor(gray_img_eqhist, cv2.COLOR_GRAY2BGR)
#end enhancement
img_mask = cv2.imread(os.path.join(train_mask_path, masks[j]), 0)
img_mask = cv2.resize(img_mask, (256,256))
img = np.array([img])
img_mask = np.array([img_mask])
imgs[i] = img
imgs_mask[i] = img_mask
if i % 100 == 0:
print('Done: {0}/{1} images'.format(i, total))
i += 1
print('Loading done.')
np.save('imgs_train_'+name+'.npy', imgs)
np.save('imgs_mask_train_'+name+'.npy', imgs_mask)
print('Saving to .npy files done.')
def load_train_data(name):
imgs_train = np.load('UNets files/imgs_train_'+name+'.npy')
imgs_mask_train = np.load('UNets files/imgs_mask_train_'+name+'.npy')
return imgs_train, imgs_mask_train
def create_test_data(data_path2, name):
test_data_path = os.path.join(data_path2, 'img')
test_mask_path = os.path.join(data_path2, 'msk')
images = os.listdir(test_data_path)
masks = os.listdir(test_mask_path)
total = len(images)
imgs = np.ndarray((total, image_rows, image_cols, 3), dtype=np.uint8)
imgs_mask = np.ndarray((total, image_rows, image_cols), dtype=np.uint8)
i = 0
print('-'*30)
print('Creating testing images...')
print('-'*30)
for j in range(len(images)):
img = cv2.imread(os.path.join(test_data_path, images[j]))
img = cv2.resize(img, (256,256))
#enhancement
gray_img=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
gray_img_eqhist=cv2.equalizeHist(gray_img)
img = cv2.cvtColor(gray_img_eqhist, cv2.COLOR_GRAY2BGR)
#end enhancement
img_mask = cv2.imread(os.path.join(test_mask_path, masks[j]), 0)
img_mask = cv2.resize(img_mask, (256,256))
img = np.array([img])
img_mask = np.array([img_mask])
imgs[i] = img
imgs_mask[i] = img_mask
if i % 100 == 0:
print('Done: {0}/{1} images'.format(i, total))
i += 1
print('Loading done.')
np.save('imgs_test_'+name+'.npy', imgs)
np.save('imgs_id_test_'+name+'.npy', imgs_mask)
print('Saving to .npy files done.')
def load_test_data(name):
imgs_test = np.load('UNets files/imgs_test_'+name+'.npy')
imgs_id = np.load('UNets files/imgs_id_test_'+name+'.npy')
return imgs_test, imgs_id