import cv2
import numpy as np
from scipy import ndimage
def fast_hist(im, gt, n=9):
"""
n is num_of_classes
"""
k = (gt >= 0) & (gt < n)
return np.bincount(n * gt[k].astype(int) + im[k], minlength=n**2).reshape(n, n)
def flood_fill(test_array, h_max=255):
"""
fill in the hole
"""
input_array = np.copy(test_array)
el = ndimage.generate_binary_structure(2,2).astype(np.int)
inside_mask = ndimage.binary_erosion(~np.isnan(input_array), structure=el)
output_array = np.copy(input_array)
output_array[inside_mask]=h_max
output_old_array = np.copy(input_array)
output_old_array.fill(0)
el = ndimage.generate_binary_structure(2,1).astype(np.int)
while not np.array_equal(output_old_array, output_array):
output_old_array = np.copy(output_array)
output_array = np.maximum(input_array,ndimage.grey_erosion(output_array, size=(3,3), footprint=el))
return output_array
def fill_break_line(cw_mask):
broken_line_h = np.array([[0,0,0,0,0],
[0,0,0,0,0],
[1,0,0,0,1],
[0,0,0,0,0],
[0,0,0,0,0]], dtype=np.uint8)
broken_line_h2 = np.array([[0,0,0,0,0],
[0,0,0,0,0],
[1,1,0,1,1],
[0,0,0,0,0],
[0,0,0,0,0]], dtype=np.uint8)
broken_line_v = np.transpose(broken_line_h)
broken_line_v2 = np.transpose(broken_line_h2)
cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_h)
cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_v)
cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_h2)
cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_v2)
return cw_mask
def refine_room_region(cw_mask, rm_ind):
label_rm, num_label = ndimage.label((1-cw_mask))
new_rm_ind = np.zeros(rm_ind.shape)
for j in xrange(1, num_label+1):
mask = (label_rm == j).astype(np.uint8)
ys, xs = np.where(mask!=0)
area = (np.amax(xs)-np.amin(xs))*(np.amax(ys)-np.amin(ys))
if area < 100:
continue
else:
room_types, type_counts = np.unique(mask*rm_ind, return_counts=True)
if len(room_types) > 1:
room_types = room_types[1:] # ignore background type which is zero
type_counts = type_counts[1:] # ignore background count
new_rm_ind += mask*room_types[np.argmax(type_counts)]
return new_rm_ind
