def create_angle(data_dir, gt_dir, save_dir):
if not os.path.exists(save_dir):
os.makedirs(save_dir)
gt_files = natsorted(glob(os.path.join(gt_dir, '*.csv')))
for gt_file in tqdm(gt_files):
# parse ground truth
file_name = os.path.splitext(os.path.basename(gt_file))[0]
boxes, lines = data_utils.csv_to_annotation(gt_file)
# read rgb image
rgb_file = os.path.join(data_dir, '{}.tif'.format(file_name))
rgb = imageio.imread(rgb_file)
# make line map
line_map, line_pts = data_utils.render_line_graph(
rgb.shape[:2], boxes, lines)
# make angle map
orient_map, angle_map = data_utils.render_angle_map(
line_map, boxes, lines, line_pts)
# save image
save_name = os.path.join(save_dir, '{}_angle.png'.format(file_name))
imageio.imsave(save_name, angle_map.astype(np.uint8))
save_name = os.path.join(save_dir, '{}_vecmap.png'.format(file_name))
imageio.imsave(save_name, orient_map.astype(np.uint8))
def compare_results():
results_dir = os.path.join(r'/hdd/Results/line_mtl_cust',
os.path.basename(MODEL_DIR))
image_dir = r'~/Documents/bohao/data/transmission_line/raw2'
pred_files = natsorted(glob(os.path.join(results_dir, '*.png')))
conf_files = natsorted(glob(os.path.join(results_dir, '*.npy')))
for pred_file, conf_file in zip(pred_files, conf_files):
file_name = os.path.splitext(os.path.basename(pred_file))[0]
rgb_file = os.path.join(image_dir, '{}.tif'.format(file_name))
rgb = misc_utils.load_file(rgb_file)[..., :3]
pred = misc_utils.load_file(pred_file)
# make line map
lbl_file = os.path.join(
r'/media/ei-edl01/data/remote_sensing_data/transmission_line/parsed_annotation',
'{}.csv'.format(file_name))
boxes, lines = data_utils.csv_to_annotation(lbl_file)
lbl, _ = data_utils.render_line_graph(rgb.shape[:2], boxes, lines)
plt.figure(figsize=(15, 8))
ax1 = plt.subplot(131)
plt.imshow(rgb)
plt.axis('off')
plt.subplot(132, sharex=ax1, sharey=ax1)
plt.title(file_name.replace('_resize', ''))
plt.imshow(lbl)
plt.axis('off')
plt.subplot(133, sharex=ax1, sharey=ax1)
plt.imshow(pred)
plt.axis('off')
plt.tight_layout()
plt.show()
def evaluate(self,
model,
patch_size,
overlap,
pred_dir=None,
report_dir=None,
save_conf=False):
iou_a, iou_b = 0, 0
report = []
if pred_dir:
misc_utils.make_dir_if_not_exist(pred_dir)
for lbl_file in self.lbl_files:
# only do NZ image
if 'NZ' not in lbl_file:
continue
# parse ground truth
file_name = os.path.splitext(os.path.basename(lbl_file))[0]
boxes, lines = data_utils.csv_to_annotation(lbl_file)
tile_id = int(''.join([a for a in file_name if a.isdigit()]))
if tile_id > 3:
continue
# read rgb image
rgb_file = os.path.join(self.image_dir, '{}.tif'.format(file_name))
rgb = misc_utils.load_file(rgb_file)[..., :3]
# make line map
lbl, _ = data_utils.render_line_graph(rgb.shape[:2], boxes, lines)
# evaluate on tiles
tile_dim = rgb.shape[:2]
grid_list = patch_extractor.make_grid(tile_dim, patch_size,
overlap)
tile_preds = []
for patch in patch_extractor.patch_block(rgb, 0, grid_list,
patch_size, False):
for tsfm in self.tsfm:
tsfm_image = tsfm(image=patch)
patch = tsfm_image['image']
patch = torch.unsqueeze(patch, 0).to(self.device)
pred = model.inference(patch).detach().cpu().numpy()
tile_preds.append(change_channel_order(pred, True)[0, :, :, :])
# stitch back to tiles
tile_preds = patch_extractor.unpatch_block(np.array(tile_preds),
tile_dim,
patch_size,
tile_dim,
patch_size,
overlap=0)
if save_conf:
misc_utils.save_file(
os.path.join(pred_dir, '{}.npy'.format(file_name)),
tile_preds)
tile_preds = np.argmax(tile_preds, -1)
a, b = metric_utils.iou_metric(lbl, tile_preds)
file_name = os.path.splitext(os.path.basename(lbl_file))[0]
print('{}: IoU={:.2f}'.format(file_name, a / b * 100))
report.append('{},{},{},{}\n'.format(file_name, a, b, a / b * 100))
iou_a += a
iou_b += b
if pred_dir:
misc_utils.save_file(
os.path.join(pred_dir, '{}.png'.format(file_name)),
tile_preds)
print('Overall: IoU={:.2f}'.format(iou_a / iou_b * 100))
report.append('Overall,{},{},{}\n'.format(iou_a, iou_b,
iou_a / iou_b * 100))
if report_dir:
misc_utils.make_dir_if_not_exist(report_dir)
misc_utils.save_file(os.path.join(report_dir, 'result.txt'),
report)
def make_patches(data_dir, gt_dir, vec_dir, save_dir, patch_size, overlap=0):
"""
Preprocess the standard inria dataset
:param data_dir: path to the original inria dataset
:param save_dir: directory to save the extracted patches
:param patch_size: size of the patches, should be a tuple of (h, w)
:param pad: #pixels to be padded around each tile, should be either one element or four elements
:param overlap: #overlapping pixels between two patches in both vertical and horizontal direction
:return:
"""
# create folders and files
patch_dir = os.path.join(save_dir, 'patches')
if not os.path.exists(patch_dir):
os.makedirs(patch_dir)
record_file_train = open(os.path.join(save_dir, 'file_list_train.txt'),
'w+')
record_file_valid = open(os.path.join(save_dir, 'file_list_valid.txt'),
'w+')
# get rgb and gt files
gt_files = natsorted(glob(os.path.join(gt_dir, '*.csv')))
for gt_file in tqdm(gt_files):
# only do NZ image
if 'NZ' not in gt_file:
continue
# parse ground truth
file_name = os.path.splitext(os.path.basename(gt_file))[0]
boxes, lines = data_utils.csv_to_annotation(gt_file)
tile_id = int(''.join([a for a in file_name if a.isdigit()]))
# read rgb image
rgb_file = os.path.join(data_dir, '{}.tif'.format(file_name))
rgb = imageio.imread(rgb_file)[..., :3]
# make line map
line_map, line_pts = data_utils.render_line_graph(
rgb.shape[:2], boxes, lines)
# read vec image
vec_file = os.path.join(vec_dir, '{}_angle.png'.format(file_name))
vec = imageio.imread(vec_file)
for rgb_patch, gt_patch, vec_patch, y, x, in patch_tile(
rgb, line_map, vec, patch_size, overlap):
rgb_patchname = '{}_y{}x{}.jpg'.format(file_name.replace(' ', '_'),
int(y), int(x))
gt_patchname = '{}_y{}x{}.png'.format(file_name.replace(' ', '_'),
int(y), int(x))
vec_patchname = '{}_y{}x{}_angle.png'.format(
file_name.replace(' ', '_'), int(y), int(x))
imageio.imsave(os.path.join(patch_dir, rgb_patchname),
rgb_patch.astype(np.uint8))
imageio.imsave(os.path.join(patch_dir, gt_patchname),
gt_patch.astype(np.uint8))
imageio.imsave(os.path.join(patch_dir, vec_patchname),
vec_patch.astype(np.uint8))
if tile_id <= 3:
record_file_valid.write('{} {} {}\n'.format(
rgb_patchname, gt_patchname, vec_patchname))
else:
record_file_train.write('{} {} {}\n'.format(
rgb_patchname, gt_patchname, vec_patchname))
record_file_train.close()
record_file_valid.close()
if __name__ == '__main__':
'''create_angle(
data_dir=r'~/Documents/bohao/data/transmission_line/raw2',
gt_dir=r'/media/ei-edl01/data/remote_sensing_data/transmission_line/parsed_annotation',
save_dir=r'/hdd/pgm/angle',
)'''
img = imageio.imread(
os.path.join(r'~/Documents/bohao/data/transmission_line/raw2',
'NZ_Gisborne_3_resize.tif'))[..., :3]
boxes, lines = data_utils.csv_to_annotation(
os.path.join(
r'/media/ei-edl01/data/remote_sensing_data/transmission_line/parsed_annotation',
'NZ_Gisborne_3_resize.csv'))
line_map, line_pts = data_utils.render_line_graph(img.shape[:2], boxes,
lines)
orient_map, angle_map = data_utils.render_angle_map(
line_map, boxes, lines, line_pts)
def plotVecMap(vectorMap, image, figsize_=(6, 6), name=''):
from numpy import ma
# fig = plt.figure(figsize=figsize_)
plt.axis('off')
plt.imshow(image)
U = vectorMap[:, :, 0] * -1
V = vectorMap[:, :, 1]
X, Y = np.meshgrid(np.arange(U.shape[1]), np.arange(U.shape[0]))
print(U.shape, vectorMap.shape)
print(X.shape, Y.shape)