def test_read_image_shape_fast(tmp_path):
path_rgb_image = os.path.join(TEST_DATA_DIR, "image.png")
img = io.imread(path_rgb_image)
shape = img.shape
assert read_image_shape_fast(path_rgb_image) == (shape[2], shape[0], shape[1])
path_gray_image = str(tmp_path / "gray.png")
io.imsave(path_gray_image, color.rgb2gray(img).astype(np.uint8))
assert read_image_shape_fast(path_gray_image) == (1, shape[0], shape[1])
def format_multianimal_training_data(
df,
train_inds,
project_path,
n_decimals=2,
):
train_data = []
nrows = df.shape[0]
filenames = df.index.to_list()
n_bodyparts = df.columns.get_level_values("bodyparts").unique().size
individuals = df.columns.get_level_values("individuals")
n_individuals = individuals.unique().size
mask_single = individuals.str.contains("single")
n_animals = n_individuals - 1 if np.any(mask_single) else n_individuals
array = np.full((nrows, n_individuals, n_bodyparts, 3),
fill_value=np.nan,
dtype=np.float32)
array[..., 0] = np.arange(n_bodyparts)
temp = df.to_numpy()
temp_multi = temp[:, ~mask_single].reshape((nrows, n_animals, -1, 2))
n_multibodyparts = temp_multi.shape[2]
array[:, :n_animals, :n_multibodyparts, 1:] = temp_multi
if n_animals != n_individuals: # There is a unique individual
n_uniquebodyparts = n_bodyparts - n_multibodyparts
temp_single = np.reshape(temp[:, mask_single],
(nrows, 1, n_uniquebodyparts, 2))
array[:, -1:, -n_uniquebodyparts:, 1:] = temp_single
array = np.round(array, decimals=n_decimals)
for i in tqdm(train_inds):
filename = filenames[i]
img_shape = read_image_shape_fast(os.path.join(project_path,
*filename))
joints = dict()
has_data = False
for n, xy in enumerate(array[i]):
# Drop missing body parts
xy = xy[~np.isnan(xy).any(axis=1)]
# Drop points lying outside the image
inside = np.logical_and.reduce((
xy[:, 1] < img_shape[2],
xy[:, 1] > 0,
xy[:, 2] < img_shape[1],
xy[:, 2] > 0,
))
xy = xy[inside]
if xy.size:
has_data = True
joints[n] = xy
if has_data:
data = {
"image": filename,
"size": np.asarray(img_shape),
"joints": joints,
}
train_data.append(data)
return train_data
