forked from hitufe/attention-unet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
prepare_datasets_STARE.py
90 lines (79 loc) · 3.72 KB
/
prepare_datasets_STARE.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
#==========================================================
#
# This prepare the hdf5 datasets of the DRIVE database
#
#============================================================
import os
import h5py
import numpy as np
from PIL import Image
def write_hdf5(arr,outfile):
with h5py.File(outfile, "w") as f:
f.create_dataset("image", data=arr, dtype=arr.dtype)
#------------Path of the images --------------------------------------------------------------
#train
original_imgs_train = "./STARE/train/images/"
groundTruth_imgs_train = "./STARE/train/1st_manual/"
borderMasks_imgs_train = "./STARE/train/mask/"
#test
original_imgs_test = "./STARE/test/images/"
groundTruth_imgs_test = "./STARE/test/1st_manual/"
borderMasks_imgs_test = "./STARE/test/mask/"
#---------------------------------------------------------------------------------------------
# Nimgs = 20
channels = 3
height = 605
width = 700
dataset_path = "./STARE_datasets_training_testing/"
def get_datasets(imgs_dir,groundTruth_dir,borderMasks_dir,train_test="null"):
if train_test == "train":
Nimgs = 10
else:
Nimgs = 10
imgs = np.empty((Nimgs,height,width,channels))
groundTruth = np.empty((Nimgs,height,width))
border_masks = np.empty((Nimgs,height,width))
for path, subdirs, files in os.walk(imgs_dir): # list all files, directories in the path
print('www')
for i in range(len(files)):
# original
print("original image: " + files[i])
img = Image.open(imgs_dir+files[i])
imgs[i] = np.asarray(img)
# corresponding ground truth
groundTruth_name = files[i][0:-4] + '.ah' + '.ppm'
print("ground truth name: " + groundTruth_name)
g_truth = Image.open(groundTruth_dir + groundTruth_name)
groundTruth[i] = np.asarray(g_truth)
# corresponding border masks
border_masks_name = files[i]
print("border masks name: " + border_masks_name)
b_mask = Image.open(borderMasks_dir + border_masks_name)
border_masks[i] = np.asarray(b_mask)
print("imgs max: " +str(np.max(imgs)))
print("imgs min: " +str(np.min(imgs)))
assert(np.max(groundTruth)==255 and np.max(border_masks)==255)
assert(np.min(groundTruth)==0 and np.min(border_masks)==0)
print("ground truth and border masks are correctly withih pixel value range 0-255 (black-white)")
#reshaping for my standard tensors
imgs = np.transpose(imgs,(0,3,1,2))
assert(imgs.shape == (Nimgs,channels,height,width))
groundTruth = np.reshape(groundTruth,(Nimgs,1,height,width))
border_masks = np.reshape(border_masks,(Nimgs,1,height,width))
assert(groundTruth.shape == (Nimgs,1,height,width))
assert(border_masks.shape == (Nimgs,1,height,width))
return imgs, groundTruth, border_masks
if not os.path.exists(dataset_path):
os.makedirs(dataset_path)
#getting the training datasets
imgs_train, groundTruth_train, border_masks_train = get_datasets(original_imgs_train,groundTruth_imgs_train,borderMasks_imgs_train,"train")
print("saving train datasets")
write_hdf5(imgs_train, dataset_path + "STARE_dataset_imgs_train.hdf5")
write_hdf5(groundTruth_train, dataset_path + "STARE_dataset_groundTruth_train.hdf5")
write_hdf5(border_masks_train,dataset_path + "STARE_dataset_borderMasks_train.hdf5")
#getting the testing datasets
imgs_test, groundTruth_test, border_masks_test = get_datasets(original_imgs_test,groundTruth_imgs_test,borderMasks_imgs_test,"test")
print("saving test datasets")
write_hdf5(imgs_test,dataset_path + "STARE_dataset_imgs_test.hdf5")
write_hdf5(groundTruth_test, dataset_path + "STARE_dataset_groundTruth_test.hdf5")
write_hdf5(border_masks_test,dataset_path + "STARE_dataset_borderMasks_test.hdf5")