forked from anilbatra2185/road_connectivity
-
Notifications
You must be signed in to change notification settings - Fork 0
/
create_crops.py
137 lines (110 loc) · 4.37 KB
/
create_crops.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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#!/usr/bin/env python2
"""
create_crops.py: script to crops for training and validation.
It will save the crop images and mask in the format:
<image_name>_<x>_<y>.<suffix>
where x = [0, (Image_Height - crop_size) / stride]
y = [0, (Image_Width - crop_size) / stride]
It will create following directory structure:
base_dir
| train_crops.txt # created by script
| val_crops.txt # created by script
|
└───train_crops # created by script
│ └───gt
│ └───images
└───val_crops # created by script
│ └───gt
│ └───images
"""
from __future__ import print_function
import argparse
import os
import mmap
import cv2
import time
import numpy as np
from skimage import io
from tqdm import tqdm
tqdm.monitor_interval = 0
def verify_image(img_file):
try:
img = io.imread(img_file)
except:
return False
return True
def CreatCrops(base_dir, crop_type, size, stride, image_suffix, gt_suffix):
crops = os.path.join(base_dir, '{}_crops'.format(crop_type))
crops_file = open(os.path.join(base_dir,'{}_crops.txt'.format(crop_type)),'w')
full_file_path = os.path.join(base_dir,'{}.txt'.format(crop_type))
full_file = open(full_file_path,'r')
def get_num_lines(file_path):
fp = open(file_path, "r+")
buf = mmap.mmap(fp.fileno(), 0)
lines = 0
while buf.readline():
lines += 1
return lines
failure_images = []
for name in tqdm(full_file, ncols=100, desc="{}_crops".format(crop_type),
total=get_num_lines(full_file_path)):
name = name.strip("\n")
image_file = os.path.join(base_dir,'{}/images'.format(crop_type),name+image_suffix)
gt_file = os.path.join(base_dir,'{}/gt'.format(crop_type),name+gt_suffix)
if not verify_image(image_file):
failure_images.append(image_file)
continue
image = cv2.imread(image_file)
gt = cv2.imread(gt_file,0)
if image is None:
failure_images.append(image_file)
continue
if gt is None:
failure_images.append(image_file)
continue
H,W,C = image.shape
maxx = (H-size)/stride
maxy = (W-size)/stride
for x in range(maxx+1):
for y in range(maxy+1):
im_ = image[x*stride:x*stride + size,y*stride:y*stride + size,:]
gt_ = gt[x*stride:x*stride + size,y*stride:y*stride + size]
crops_file.write('{}_{}_{}\n'.format(name,x,y))
cv2.imwrite(crops+'/images/{}_{}_{}.png'.format(name,x,y), im_)
cv2.imwrite(crops+'/gt/{}_{}_{}.png'.format(name,x,y), gt_)
crops_file.close()
full_file.close()
if len(failure_images) > 0:
print("Unable to process {} images : {}".format(len(failure_images), failure_images))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-d', '--base_dir', type=str, required=True,
help='Base directory for Spacenent Dataset.')
parser.add_argument('--crop_size', type=int, required=True,
help='Crop Size of Image')
parser.add_argument('--crop_overlap', type=int, required=True,
help='Crop overlap Size of Image')
parser.add_argument('--im_suffix', type=str, required=True,
help='Dataset specific image suffix.')
parser.add_argument('--gt_suffix', type=str, required=True,
help='Dataset specific gt suffix.')
args = parser.parse_args()
start = time.clock()
## Create overlapping Crops for training
CreatCrops(args.base_dir,
crop_type='train',
size=args.crop_size,
stride=args.crop_overlap,
image_suffix=args.im_suffix,
gt_suffix=args.gt_suffix)
## Create non-overlapping Crops for validation
CreatCrops(args.base_dir,
crop_type='val',
size=args.crop_size,
stride=args.crop_size, ## Non-overlapping
image_suffix=args.im_suffix,
gt_suffix=args.gt_suffix)
end = time.clock()
print('Finished Creating crops, time {0}s'.format(end - start))
if __name__ == "__main__":
main()