Add additional drop_likelihood_columns guards (#3333)

C-Achard · web-flow · commit 02c78da321cc · 2026-05-20T14:16:49.000+02:00
* Expose drop_likelihood_columns and apply it

Rename the private _drop_likelihood_columns to drop_likelihood_columns and use it across the codebase to consistently sanitize annotation DataFrames (drop any coord level named 'likelihood'). Updated imports/usages in modelzoo (ma_dlc, ma_dlc_dataframe), pose_estimation_pytorch (dlcloader), and utils (skeleton). Added a docstring note about centralizing project I/O and maintenance concerns. Added a unit test (tests/pose_estimation_pytorch/data/test_dlc_dataloader.py) to ensure DLCLoader.to_coco ignores likelihood columns.

* Warn and drop likelihood columns; normalize data var

Change logging from info to warning (with stacklevel=2) when likelihood columns are detected and drop them. Call drop_likelihood_columns immediately after reading the HDF file. Normalize variable naming by replacing Data with data throughout mergeandsplit and update subsequent references (index ranges, scorer extraction, and enumeration) to use the lower-case variable.
diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
@@ -493,11 +493,16 @@ def parse_video_filenames(videos: list[str]) -> list[str]:
     return filenames
 
 
-def _drop_likelihood_columns(df: pd.DataFrame) -> pd.DataFrame:
+def drop_likelihood_columns(df: pd.DataFrame) -> pd.DataFrame:
     """Drop any columns whose coord level is named 'likelihood'.
 
     This sanitizes annotation DataFrames coming from h5/csv files before they are
     used for training dataset generation.
+
+    # NOTE @C-Achard 2026-05-18: This is used in several places as a guard
+        Most call sites using this should instead go through a canonical, validated project loading function
+        AND THEN do any custom local processing they require. The current design is hard to maintain and error prone,
+        and lacks a clearly documented, centralized project I/O interface.
     """
     if not isinstance(df.columns, pd.MultiIndex):
         return df
@@ -507,7 +512,7 @@ def _drop_likelihood_columns(df: pd.DataFrame) -> pd.DataFrame:
 
     likelihood_mask = coord_values == "likelihood"
     if likelihood_mask.any():
-        logging.info("Detected likelihood columns in annotation data; dropping them.")
+        logging.warning("Detected likelihood columns in annotation data; dropping them.", stacklevel=2)
         df = df.drop(columns=df.columns[likelihood_mask])
 
     return df
@@ -569,7 +574,7 @@ def merge_annotateddatasets(cfg, trainingsetfolder_full):
     AnnotationData = AnnotationData.reindex(bodyparts, axis=1, level=AnnotationData.columns.names.index("bodyparts"))
     # Filter out any stray likelihood columns that may have been concatenated in
     # see napari-deeplabcut #204 and DeepLabCut #3319
-    AnnotationData = _drop_likelihood_columns(AnnotationData)
+    AnnotationData = drop_likelihood_columns(AnnotationData)
 
     if AnnotationData.empty:
         logging.warning(
@@ -701,23 +706,24 @@ def mergeandsplit(config, trainindex=0, uniform=True):
     fn = os.path.join(project_path, trainingsetfolder, "CollectedData_" + cfg["scorer"])
 
     try:
-        Data = pd.read_hdf(fn + ".h5")
+        data = pd.read_hdf(fn + ".h5")
+        data = drop_likelihood_columns(data)
     except FileNotFoundError:
-        Data = merge_annotateddatasets(
+        data = merge_annotateddatasets(
             cfg,
             Path(os.path.join(project_path, trainingsetfolder)),
         )
-        if Data is None:
+        if data is None:
             return [], []
 
-    conversioncode.guarantee_multiindex_rows(Data)
-    Data = Data[scorer]  # extract labeled data
+    conversioncode.guarantee_multiindex_rows(data)
+    data = data[scorer]  # extract labeled data
 
     if uniform:
         TrainingFraction = cfg["TrainingFraction"]
         trainFraction = TrainingFraction[trainindex]
         trainIndices, testIndices = SplitTrials(
-            range(len(Data.index)),
+            range(len(data.index)),
             trainFraction,
             True,
         )
@@ -726,7 +732,7 @@ def mergeandsplit(config, trainindex=0, uniform=True):
         test_video_name = [Path(i).stem for i in videos][trainindex]
         print("Excluding the following folder (from training):", test_video_name)
         trainIndices, testIndices = [], []
-        for index, name in enumerate(Data.index):
+        for index, name in enumerate(data.index):
             if test_video_name == name[1]:  # this is the video name
                 # print(name,test_video_name)
                 testIndices.append(index)
@@ -754,7 +760,7 @@ def to_matlab_cell(array):
         return outer
 
     # Again, remove likelihood if present
-    df = _drop_likelihood_columns(df)
+    df = drop_likelihood_columns(df)
 
     if isinstance(df.columns, pd.MultiIndex):
         coord_level = "coords" if "coords" in df.columns.names else df.columns.names[-1]
diff --git a/deeplabcut/modelzoo/generalized_data_converter/datasets/ma_dlc.py b/deeplabcut/modelzoo/generalized_data_converter/datasets/ma_dlc.py
@@ -13,6 +13,7 @@
 import numpy as np
 import pandas as pd
 
+from deeplabcut.generate_training_dataset.trainingsetmanipulation import drop_likelihood_columns
 from deeplabcut.modelzoo.generalized_data_converter.datasets.base_dlc import (
     BaseDLCPoseDataset,
 )
@@ -27,7 +28,7 @@ def __init__(self, proj_root, dataset_name, shuffle=1, modelprefix=""):
         super().__init__(proj_root, dataset_name, shuffle=shuffle, modelprefix=modelprefix)
 
     def _df2generic(self, df, image_id_offset=0):
-
+        df = drop_likelihood_columns(df)
         individuals = df.columns.get_level_values("individuals").unique().tolist()
 
         unique_bpts = []
diff --git a/deeplabcut/modelzoo/generalized_data_converter/datasets/ma_dlc_dataframe.py b/deeplabcut/modelzoo/generalized_data_converter/datasets/ma_dlc_dataframe.py
@@ -15,6 +15,7 @@
 import pandas as pd
 
 from deeplabcut.generate_training_dataset.trainingsetmanipulation import (
+    drop_likelihood_columns,
     parse_video_filenames,
 )
 from deeplabcut.modelzoo.generalized_data_converter.datasets.base import BasePoseDataset
@@ -79,6 +80,7 @@ def merge_annotateddatasets(cfg):
     else:
         bodyparts = cfg["bodyparts"]
     AnnotationData = AnnotationData.reindex(bodyparts, axis=1, level=AnnotationData.columns.names.index("bodyparts"))
+    AnnotationData = drop_likelihood_columns(AnnotationData)
 
     return AnnotationData
 
@@ -140,7 +142,7 @@ def populate_generic(self):
         self.whether_anno_image_match(self.generic_test_images, self.generic_test_annotations)
 
     def _df2generic(self, df, image_id_offset=0):
-
+        df = drop_likelihood_columns(df)
         individuals = df.columns.get_level_values("individuals").unique().tolist()
 
         unique_bpts = []
diff --git a/deeplabcut/pose_estimation_pytorch/data/dlcloader.py b/deeplabcut/pose_estimation_pytorch/data/dlcloader.py
@@ -23,6 +23,7 @@
 
 import deeplabcut.utils.auxiliaryfunctions as af
 from deeplabcut.core.engine import Engine
+from deeplabcut.generate_training_dataset.trainingsetmanipulation import drop_likelihood_columns
 from deeplabcut.pose_estimation_pytorch.data.base import Loader
 from deeplabcut.pose_estimation_pytorch.data.dataset import PoseDatasetParameters
 from deeplabcut.pose_estimation_pytorch.data.snapshots import Snapshot
@@ -373,6 +374,8 @@ def to_coco(
         Returns:
             the coco format data
         """
+        df = drop_likelihood_columns(df)
+
         with_individuals = "individuals" in df.columns.names
         if not with_individuals and (len(parameters.individuals) > 1 or len(parameters.unique_bpts) > 0):
             raise ValueError(
diff --git a/deeplabcut/utils/skeleton.py b/deeplabcut/utils/skeleton.py
@@ -31,6 +31,8 @@
 from scipy.spatial import KDTree
 from skimage import io
 
+from deeplabcut.generate_training_dataset.trainingsetmanipulation import drop_likelihood_columns
+
 
 # NOTE @C-Achard 2026-03-26 duplicate config read/write functions
 # should be addressed in config refactor
@@ -60,6 +62,7 @@ def __init__(self, config_path):
             folder = os.path.join(root, dir_)
             if os.path.isdir(folder) and not any(folder.endswith(s) for s in ("cropped", "labeled")):
                 self.df = pd.read_hdf(os.path.join(folder, f"CollectedData_{self.cfg['scorer']}.h5"))
+                self.df = drop_likelihood_columns(self.df)
                 row, col = self.pick_labeled_frame()
                 if "individuals" in self.df.columns.names:
                     self.df = self.df.xs(col, axis=1, level="individuals")
diff --git a/tests/pose_estimation_pytorch/data/test_dlc_dataloader.py b/tests/pose_estimation_pytorch/data/test_dlc_dataloader.py
@@ -0,0 +1,68 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pandas as pd
+
+import deeplabcut.pose_estimation_pytorch.data.dlcloader as dlcloader_mod
+from deeplabcut.pose_estimation_pytorch.data.dlcloader import DLCLoader
+
+
+def test_to_coco_ignores_likelihood_columns(monkeypatch, tmp_path):
+    fake_shape = (3, 480, 640)
+    monkeypatch.setattr(
+        dlcloader_mod,
+        "read_image_shape_fast",
+        lambda _: fake_shape,
+    )
+
+    scorer = "testscorer"
+    bodyparts = ["nose", "tail"]
+
+    index = pd.MultiIndex.from_tuples(
+        [("labeled-data", "video1", "img0001.png")],
+        names=["set", "video", "image"],
+    )
+
+    # Baseline dataframe: x/y only
+    columns_xy = pd.MultiIndex.from_product(
+        [[scorer], bodyparts, ["x", "y"]],
+        names=["scorer", "bodyparts", "coords"],
+    )
+    df_xy = pd.DataFrame(
+        [[10.0, 20.0, 30.0, 40.0]],
+        index=index,
+        columns=columns_xy,
+    )
+
+    # Same data, but with likelihood columns added
+    columns_xyl = pd.MultiIndex.from_product(
+        [[scorer], bodyparts, ["x", "y", "likelihood"]],
+        names=["scorer", "bodyparts", "coords"],
+    )
+    df_xyl = pd.DataFrame(
+        [[10.0, 20.0, 0.9, 30.0, 40.0, 0.8]],
+        index=index,
+        columns=columns_xyl,
+    )
+
+    # to_coco only needs these attributes from parameters
+    params = SimpleNamespace(
+        bodyparts=bodyparts,
+        unique_bpts=[],
+        individuals=["animal"],
+    )
+
+    baseline = DLCLoader.to_coco(tmp_path, df_xy, params)
+    got = DLCLoader.to_coco(tmp_path, df_xyl, params)
+
+    assert len(got["images"]) == len(baseline["images"]) == 1
+    assert len(got["annotations"]) == len(baseline["annotations"]) == 1
+
+    got_ann = got["annotations"][0]
+    expected_ann = baseline["annotations"][0]
+
+    assert got_ann["image_id"] == expected_ann["image_id"]
+    assert got_ann["category_id"] == expected_ann["category_id"]
+    assert got_ann["num_keypoints"] == expected_ann["num_keypoints"] == 2
+    assert np.array_equal(got_ann["keypoints"], expected_ann["keypoints"])
+    assert np.allclose(got_ann["bbox"], expected_ann["bbox"])
