from multiprocessing import cpu_count
import numpy as np
from tqdm import tqdm
from multiprocessing.dummy import Pool
def compute_patch_indices(image_shape, patch_size, overlap, start=None):
if isinstance(overlap, int):
overlap = np.asarray([overlap] * len(image_shape))
if start is None:
n_patches = np.ceil(image_shape / (patch_size - overlap))
overflow = (patch_size - overlap) * n_patches - image_shape + overlap
start = -np.ceil(overflow / 2)
elif isinstance(start, int):
start = np.asarray([start] * len(image_shape))
stop = image_shape + start
step = patch_size - overlap
return get_set_of_patch_indices(start, stop, step)
# def compute_patch_indices_full(image_shape, patch_size, overlap, start=None):
# if isinstance(overlap, int):
# overlap = np.asarray([overlap] * len(image_shape))
# step = patch_size - overlap
# max_index = image_shape - patch_size
# if start is None:
# n_patches = np.ceil(max_index / step)
# overflow = step * n_patches - image_shape
# start = -np.ceil(overflow / 2)
# elif isinstance(start, int):
# start = np.asarray([start] * len(image_shape))
# stop = max_index - start
# return get_set_of_patch_indices(start, stop, step)
def get_set_of_patch_indices(start, stop, step):
return np.asarray(np.mgrid[start[0]:stop[0]:step[0], start[1]:stop[1]:step[1],
start[2]:stop[2]:step[2]].reshape(3, -1).T, dtype=np.int)
def get_random_patch_index(image_shape, patch_shape):
"""
Returns a random corner index for a patch. If this is used during training, the middle pixels will be seen by
the model way more often than the edge pixels (which is probably a bad thing).
:param image_shape: Shape of the image
:param patch_shape: Shape of the patch
:return: a tuple containing the corner index which can be used to get a patch from an image
"""
return get_random_nd_index(np.subtract(image_shape, patch_shape))
def get_random_nd_index(index_max):
return tuple([np.random.choice(index_max[index] + 1) for index in range(len(index_max))])
def get_patch_from_3d_data(data, patch_shape, patch_index):
"""
Returns a patch from a numpy array.
:param data: numpy array from which to get the patch.
:param patch_shape: shape/size of the patch.
:param patch_index: corner index of the patch.
:return: numpy array take from the data with the patch shape specified.
"""
patch_index = np.asarray(patch_index, dtype=np.int16)
patch_shape = np.asarray(patch_shape)
image_shape = data.shape[-3:]
if np.any(patch_index < 0) or np.any((patch_index + patch_shape) > image_shape):
data, patch_index = fix_out_of_bound_patch_attempt(data, patch_shape, patch_index)
return data[..., patch_index[0]:patch_index[0] + patch_shape[0],
patch_index[1]:patch_index[1] + patch_shape[1],
patch_index[2]:patch_index[2] + patch_shape[2]]
def fix_out_of_bound_patch_attempt(data, patch_shape, patch_index, ndim=3):
"""
Pads the data and alters the patch index so that a patch will be correct.
:param data:
:param patch_shape:
:param patch_index:
:return: padded data, fixed patch index
"""
image_shape = data.shape[-ndim:]
pad_before = np.abs((patch_index < 0) * patch_index)
pad_after = np.abs(((patch_index + patch_shape) > image_shape) * ((patch_index + patch_shape) - image_shape))
pad_args = np.stack([pad_before, pad_after], axis=1)
if pad_args.shape[0] < len(data.shape):
pad_args = [[0, 0]] * (len(data.shape) - pad_args.shape[0]) + pad_args.tolist()
data = np.pad(data, pad_args, mode="edge")
patch_index += pad_before
return data, patch_index
def reconstruct_from_patches(patches, patch_indices, data_shape, default_value=0):
"""
Reconstructs an array of the original shape from the lists of patches and corresponding patch indices. Overlapping
patches are averaged.
:param patches: List of numpy array patches.
:param patch_indices: List of indices that corresponds to the list of patches.
:param data_shape: Shape of the array from which the patches were extracted.
:param default_value: The default value of the resulting data. if the patch coverage is complete, this value will
be overwritten.
:return: numpy array containing the data reconstructed by the patches.
"""
def compute_full_picture(patches, patch_indices, data_shape):
data = np.zeros(data_shape)
image_shape = data_shape[-4:-1]
count = np.zeros(data_shape, dtype=np.int)
for patch, index in zip(patches, patch_indices):
image_patch_shape = patch.shape[-4:-1]
if np.any(index < 0):
fix_patch = np.asarray((index < 0) * np.abs(index), dtype=np.int)
patch = patch[fix_patch[0]:, fix_patch[1]:, fix_patch[2]:, ...]
index[index < 0] = 0
if np.any((index + image_patch_shape) >= image_shape):
fix_patch = np.asarray(image_patch_shape - (((index + image_patch_shape) >= image_shape)
* ((index + image_patch_shape) - image_shape)),
dtype=np.int)
patch = patch[:fix_patch[0], :fix_patch[1], :fix_patch[2], ...]
patch_index = np.zeros(data_shape, dtype=np.bool)
patch_index[index[0]:index[0] + patch.shape[-4],
index[1]:index[1] + patch.shape[-3],
index[2]:index[2] + patch.shape[-2], ...] = True
# patch_data = np.zeros(data_shape)
# patch_data[patch_index] = patch.flatten()
data[patch_index] += patch.flatten()
# new_data_index = np.logical_and(patch_index, np.logical_not(count > 0))
# data[new_data_index] = patch_data[new_data_index]
#
# averaged_data_index = np.logical_and(patch_index, count > 0)
# #if np.any(averaged_data_index):
# data[averaged_data_index] = (data[averaged_data_index] * count[averaged_data_index] + patch_data[
# averaged_data_index]) / (count[averaged_data_index] + 1)
count[patch_index] += 1
return data, count
workers = cpu_count()
pool = Pool(workers)
results = []
for i in range(workers):
patches_i, patch_indices_i = patches[i::workers], \
patch_indices[i::workers]
results.append(pool.apply_async(compute_full_picture,
args=(patches_i, patch_indices_i, data_shape)))
data = np.zeros(data_shape)
count = np.zeros(data_shape, dtype=np.int)
for i in range(len(results)):
data_i, count_i = results[i].get()
data += data_i
count += count_i
assert np.all(count > 0)
return data / count
