diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py index b18dd9a050..e6fa7c913f 100755 --- 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 index 081bbbb755..e5f57b0a44 100644 --- 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 index b85e6adffb..b73036918c 100644 --- 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 index 50905e05a5..c034f0e743 100644 --- 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 index 6021b6f9e6..e4e7da8eb3 100644 --- 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 new file mode 100644 index 0000000000..76fe04ca61 --- /dev/null +++ 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"])