diff --git a/VHR/segmentation.py b/VHR/segmentation.py index b3c7142ccbb2e12a9f8efd26f227744813da6586..43346a9a233a41b9dfa96b0effb95fdafaaa3260 100644 --- a/VHR/segmentation.py +++ b/VHR/segmentation.py @@ -2,7 +2,7 @@ import sys import numpy as np -from Common.otb_numpy_proc import to_otb_pipeline +from Common.otb_numpy_proc import to_otb_pipeline, get_region import os.path import otbApplication as otb import multiprocessing as mp @@ -189,41 +189,21 @@ def write_qgis_seg_style(out_file, line_color='255,255,0,255', line_width=0.46): ) return out_file -def vectorize_tile(obj, startx, starty, sizex, sizey, to_skip, out): - print('Using startx {} and starty {}'.format(startx, starty)) +def vectorize_tile(obj, region, to_keep, out): r = otb.itkRegion() - r['index'][0], r['index'][1] = startx, starty - r['size'][0], r['size'][1] = sizex, sizey + r['index'][0], r['index'][1] = region[0], region[1] + r['size'][0], r['size'][1] = region[2], region[3] obj.PropagateRequestedRegion('out', r) clip = obj.ExportImage('out') - min_bb_x = 0 - min_bb_y = 0 - - rp = regionprops(np.squeeze(clip['array'].astype(np.int))) - - for o in rp: - if o.bbox[2] > sizey - 1 or o.bbox[3] > sizex - 1: - clip['array'][clip['array'] == o.label] = 0 - if o.bbox[3] > sizex - 1 and sizex - o.bbox[1] > min_bb_x: - min_bb_x = sizex - o.bbox[1] - if o.bbox[2] > sizey - 1 and sizey - o.bbox[0] > min_bb_y: - min_bb_y = sizey - o.bbox[0] - elif o.label in to_skip: - clip['array'][clip['array'] == o.label] = 0 - - for o in rp: - if (o.bbox[2] >= sizey - min_bb_y and o.bbox[2] <= sizey - 1) or (o.bbox[3] >= sizex - min_bb_x and o.bbox[3] <= sizex - 1) : - to_skip.append(o.label) + clip['array'] *= np.isin(clip['array'], to_keep) vec = otb.Registry.CreateApplication('SimpleVectorization') vec.ImportImage('in', clip) vec.SetParameterString('out', out) vec.ExecuteAndWriteOutput() - print('Returning min_bb_x {} and min_bb_y {}'.format(min_bb_x, min_bb_y)) - - return min_bb_x, min_bb_y, to_skip + return out def tiled_vectorization(input_segm, nominal_tile_size, output_template): in_seg = to_otb_pipeline(input_segm) @@ -231,23 +211,29 @@ def tiled_vectorization(input_segm, nominal_tile_size, output_template): rp = regionprops(arr) tx, ty = int(arr.shape[1] / nominal_tile_size[0]) + 1, int(arr.shape[0] / nominal_tile_size[1]) + 1 + arr = None obj_to_tile = dict.fromkeys(range(tx*ty)) + tiles = dict.fromkeys(range(tx*ty)) for i in range(tx*ty): obj_to_tile[i] = [] + tiles[i] = [np.inf, np.inf, 0, 0] for o in rp: ix, iy = int(o.bbox[1] / nominal_tile_size[0]), int(o.bbox[0] / nominal_tile_size[1]) idx = ix * ty + iy obj_to_tile[idx].append(o.label) + tiles[idx][0] = min(o.bbox[1], tiles[idx][0]) + tiles[idx][1] = min(o.bbox[0], tiles[idx][1]) + tiles[idx][2] = max(o.bbox[3] - tiles[idx][0], tiles[idx][2]) + tiles[idx][3] = max(o.bbox[2] - tiles[idx][1], tiles[idx][3]) - import matplotlib.pyplot as plt for i in range(tx*ty): - test = np.isin(arr, np.array(obj_to_tile[i])) - plt.figure() - plt.imshow(test) + in_seg = to_otb_pipeline(input_segm) + if len(obj_to_tile[i]) > 0: + vectorize_tile(in_seg, tiles[i], obj_to_tile[i], output_template.format(i)) - return obj_to_tile + return obj_to_tile, tiles def get_bounding_boxes(input_segm): in_seg = to_otb_pipeline(input_segm)