Add data merging when there are multiple scorers

Author: jeylau

src/napari_deeplabcut/_reader.py

@@ -182,6 +182,7 @@ def read_hdf(filename: str) -> List[LayerData]:
     layers = []
     for filename in glob.iglob(filename):
         temp = pd.read_hdf(filename)
+        temp = misc.merge_multiple_scorers(temp)
         header = misc.DLCHeader(temp.columns)
         temp = temp.droplevel("scorer", axis=1)
         if "individuals" not in temp.columns.names:

src/napari_deeplabcut/_tests/test_misc.py

@@ -1,3 +1,4 @@
+import numpy as np
 import os
 import pandas as pd
 import pytest
@@ -24,6 +25,31 @@ def test_encode_categories():
     inds = misc.encode_categories(categories, return_map=False)
 
 
+def test_merge_multiple_scorers_no_likelihood(fake_keypoints):
+    temp = fake_keypoints.copy(deep=True)
+    temp.columns = temp.columns.set_levels(["you"], level="scorer")
+    df = fake_keypoints.merge(temp, left_index=True, right_index=True)
+    df = misc.merge_multiple_scorers(df)
+    pd.testing.assert_frame_equal(df, fake_keypoints)
+
+
+def test_merge_multiple_scorers(fake_keypoints):
+    new_columns = pd.MultiIndex.from_product(
+        fake_keypoints.columns.levels[:-1] + [["x", "y", "likelihood"]],
+        names=fake_keypoints.columns.names,
+    )
+    fake_keypoints = fake_keypoints.reindex(new_columns, axis=1)
+    fake_keypoints.loc(axis=1)[:, :, :, "likelihood"] = 1
+    temp = fake_keypoints.copy(deep=True)
+    temp.columns = temp.columns.set_levels(["you"], level="scorer")
+    fake_keypoints.iloc[:5] = np.nan
+    temp.iloc[5:] = np.nan
+    df = fake_keypoints.merge(temp, left_index=True, right_index=True)
+    df = misc.merge_multiple_scorers(df)
+    pd.testing.assert_index_equal(df.columns, fake_keypoints.columns)
+    assert not df.isna().any(axis=None)
+
+
 @pytest.mark.parametrize(
     "path",
     ["/home/to/fake/path", "C:\\Users\\with\\fake\\name"],

src/napari_deeplabcut/misc.py

@@ -27,6 +27,34 @@ def encode_categories(
     return inds
 
 
+def merge_multiple_scorers(
+    df: pd.DataFrame,
+) -> pd.DataFrame:
+    n_frames = df.shape[0]
+    header = DLCHeader(df.columns)
+    n_scorers = len(header._get_unique("scorer"))
+    if n_scorers == 1:
+        return df
+
+    n_bodyparts = len(header.bodyparts)
+    if "likelihood" in header.coords:
+        # Merge annotations from multiple scorers to keep
+        # detections with highest confidence
+        data = df.to_numpy().reshape((n_frames, n_scorers, n_bodyparts, -1))
+        idx = np.nanargmax(data[..., 2], axis=1)
+        data_best = data[
+            np.arange(n_frames)[:, None], idx, np.arange(n_bodyparts)
+        ].reshape((n_frames, -1))
+        df = pd.DataFrame(
+            data_best,
+            index=df.index,
+            columns=header.columns[: data_best.shape[1]],
+        )
+    else:  # Arbitrarily pick data from the first scorer
+        df = df.loc(axis=1)[: header.scorer]
+    return df
+
+
 def to_os_dir_sep(path: str) -> str:
     """
     Replace all directory separators in `path` with `os.path.sep`.