def test_assembler_with_unique_bodypart(real_assemblies_montblanc):
with open(os.path.join(TEST_DATA_DIR, "montblanc_full.pickle"),
"rb") as file:
data = pickle.load(file)
ass = inferenceutils.Assembler(
data,
max_n_individuals=3,
n_multibodyparts=4,
pcutoff=0.1,
min_affinity=0.1,
)
assert len(ass.metadata["imnames"]) == 180
assert ass.n_keypoints == 5
assert len(ass.graph) == len(ass.paf_inds) == 6
ass.assemble(chunk_size=0)
assert len(ass.assemblies) == len(real_assemblies_montblanc[0])
assert len(ass.unique) == len(real_assemblies_montblanc[1])
assemblies = np.concatenate([
ass.xy for assemblies in ass.assemblies.values() for ass in assemblies
])
assemblies_gt = np.concatenate([
ass.xy for assemblies in real_assemblies_montblanc[0].values()
for ass in assemblies
])
np.testing.assert_equal(assemblies, assemblies_gt)
def test_assembler_with_single_bodypart(real_assemblies):
with open(os.path.join(TEST_DATA_DIR, "trimouse_full.pickle"),
"rb") as file:
temp = pickle.load(file)
data = {"metadata": temp.pop("metadata")}
for k, dict_ in temp.items():
data[k] = {
"coordinates": (dict_["coordinates"][0][:1], ),
"confidence": dict_["confidence"][:1],
}
ass = inferenceutils.Assembler(
data,
max_n_individuals=3,
n_multibodyparts=1,
)
ass.metadata["joint_names"] = ass.metadata["joint_names"][:1]
ass.metadata["num_joints"] = 1
ass.metadata["paf_graph"] = []
ass.metadata["paf"] = []
ass.metadata["bpts"] = [0]
ass.metadata["ibpts"] = [0]
ass.assemble(chunk_size=0)
assert not ass.unique
assert len(ass.assemblies) == len(real_assemblies)
assert all(len(a) == 3 for a in ass.assemblies.values())
def test_assembler_with_identity(tmpdir_factory, real_assemblies):
with open(os.path.join(TEST_DATA_DIR, "trimouse_full.pickle"), "rb") as file:
data = pickle.load(file)
# Generate fake identity predictions
for k, v in data.items():
if k != "metadata":
conf = v["confidence"]
ids = [np.random.rand(c.shape[0], 3) for c in conf]
v["identity"] = ids
ass = inferenceutils.Assembler(data, max_n_individuals=3, n_multibodyparts=12)
assert ass._has_identity
assert len(ass.metadata["imnames"]) == 50
assert ass.n_keypoints == 12
assert len(ass.graph) == len(ass.paf_inds) == 66
# Assemble based on the smallest graph to speed up testing
naive_graph = [
[0, 1],
[7, 8],
[6, 7],
[10, 11],
[4, 5],
[5, 6],
[8, 9],
[9, 10],
[0, 3],
[3, 4],
[0, 2],
]
ass.paf_inds = [ass.graph.index(edge) for edge in naive_graph]
ass.assemble()
assert not ass.unique
assert len(ass.assemblies) == len(real_assemblies)
assert sum(1 for a in ass.assemblies.values() for _ in a) == sum(
1 for a in real_assemblies.values() for _ in a
)
assert all(np.all(_.data[:, -1] != -1) for a in ass.assemblies.values() for _ in a)
# Test now with identity only and ensure assemblies
# contain only parts of a single group ID.
ass.identity_only = True
ass.assemble()
assert len(ass.assemblies) == len(real_assemblies)
eq = []
for a in ass.assemblies.values():
for _ in a:
ids = _.data[:, -1]
ids = ids[~np.isnan(ids)]
eq.append(np.all(ids == ids[0]))
assert all(eq)
output_name = tmpdir_factory.mktemp("data").join("fake.h5")
ass.to_h5(output_name)
ass.to_pickle(str(output_name).replace("h5", "pickle"))
def test_assembler_calibration(real_assemblies):
with open(os.path.join(TEST_DATA_DIR, "trimouse_full.pickle"), "rb") as file:
data = pickle.load(file)
ass = inferenceutils.Assembler(data, max_n_individuals=3, n_multibodyparts=12)
ass.calibrate(os.path.join(TEST_DATA_DIR, "trimouse_calib.h5"))
assert ass._kde is not None
assert ass.safe_edge
assembly = real_assemblies[0][0]
mahal, proba = ass.calc_assembly_mahalanobis_dist(assembly, return_proba=True)
assert np.isclose(mahal, 19.541, atol=1e-3)
assert np.isclose(proba, 1, atol=1e-3)
j1 = inferenceutils.Joint(tuple(assembly.xy[0]), label=0)
j2 = inferenceutils.Joint(tuple(assembly.xy[1]), label=1)
link = inferenceutils.Link(j1, j2)
p = ass.calc_link_probability(link)
assert np.isclose(p, 0.990, atol=1e-3)
def test_assembler(tmpdir_factory, real_assemblies):
with open(os.path.join(TEST_DATA_DIR, "trimouse_full.pickle"),
"rb") as file:
data = pickle.load(file)
with pytest.warns(UserWarning):
ass = inferenceutils.Assembler(
data,
max_n_individuals=3,
n_multibodyparts=12,
identity_only=True, # Test whether warning is properly raised
)
assert len(ass.metadata["imnames"]) == 50
assert ass.n_keypoints == 12
assert len(ass.graph) == len(ass.paf_inds) == 66
# Assemble based on the smallest graph to speed up testing
naive_graph = [
[0, 1],
[7, 8],
[6, 7],
[10, 11],
[4, 5],
[5, 6],
[8, 9],
[9, 10],
[0, 3],
[3, 4],
[0, 2],
]
paf_inds = [ass.graph.index(edge) for edge in naive_graph]
ass.graph = naive_graph
ass.paf_inds = paf_inds
ass.assemble()
assert not ass.unique
assert len(ass.assemblies) == len(real_assemblies)
assert sum(1 for a in ass.assemblies.values()
for _ in a) == sum(1 for a in real_assemblies.values()
for _ in a)
output_name = tmpdir_factory.mktemp("data").join("fake.h5")
ass.to_h5(output_name)
ass.to_pickle(str(output_name).replace("h5", "pickle"))
