def inpaint_or_oom(complnet, image, segmap): ''' If image is too big, return None ''' #mask = binarization(segmap, 127) # 255 / 2 # maybe useless. mask = segmap assert image.shape == mask.shape h, w = image.shape[:2] # 1 image, not batch. image = iu.modulo_padded(image, 8) mask = iu.modulo_padded(mask, 8) image = np.expand_dims(image, 0) # [h,w,c] -> [1,h,w,c] mask = np.expand_dims(mask, 0) input_image = np.concatenate([image, mask], axis=2) result = complnet(input_image) return result[0][:h, :w, ::-1] #---------- remove padding '''
def segment_or_oom(segnet, inp, modulo=16): ''' If image is too big, return None ''' h, w = inp.shape[:2] img = iu.modulo_padded(inp, modulo) img_bat = np.expand_dims(img, 0) segmap = segnet(img_bat) segmap = np.squeeze(segmap[:, :h, :w, :], 0) return segmap '''