| 133 | |
| 134 | |
| 135 | def test_assembler(tmpdir_factory, real_assemblies): |
| 136 | with open(os.path.join(TEST_DATA_DIR, "trimouse_full.pickle"), "rb") as file: |
| 137 | data = pickle.load(file) |
| 138 | with pytest.warns(UserWarning): |
| 139 | ass = inferenceutils.Assembler( |
| 140 | data, |
| 141 | max_n_individuals=3, |
| 142 | n_multibodyparts=12, |
| 143 | identity_only=True, # Test whether warning is properly raised |
| 144 | ) |
| 145 | assert len(ass.metadata["imnames"]) == 50 |
| 146 | assert ass.n_keypoints == 12 |
| 147 | assert len(ass.graph) == len(ass.paf_inds) == 66 |
| 148 | # Assemble based on the smallest graph to speed up testing |
| 149 | naive_graph = [ |
| 150 | [0, 1], |
| 151 | [7, 8], |
| 152 | [6, 7], |
| 153 | [10, 11], |
| 154 | [4, 5], |
| 155 | [5, 6], |
| 156 | [8, 9], |
| 157 | [9, 10], |
| 158 | [0, 3], |
| 159 | [3, 4], |
| 160 | [0, 2], |
| 161 | ] |
| 162 | ass.paf_inds = [ass.graph.index(edge) for edge in naive_graph] |
| 163 | ass.assemble() |
| 164 | assert not ass.unique |
| 165 | assert len(ass.assemblies) == len(real_assemblies) |
| 166 | assert sum(1 for a in ass.assemblies.values() for _ in a) == sum(1 for a in real_assemblies.values() for _ in a) |
| 167 | |
| 168 | output_dir = tmpdir_factory.mktemp("data") |
| 169 | ass.to_h5(output_dir.join("fake.h5")) |
| 170 | ass.to_pickle(output_dir.join("fake.pickle")) |
| 171 | |
| 172 | |
| 173 | def test_assembler_with_single_bodypart(real_assemblies): |