diff --git a/sleap/io/format/deeplabcut.py b/sleap/io/format/deeplabcut.py index c625e3b56..6ef428a49 100644 --- a/sleap/io/format/deeplabcut.py +++ b/sleap/io/format/deeplabcut.py @@ -131,7 +131,8 @@ def read_frames( # Create the skeleton from the list of nodes in the csv file. # Note that DeepLabCut doesn't have edges, so these will need to be # added by user later. - node_names = [n[0] for n in list(data)[1::2]] + start_col = 3 if is_new_format else 1 + node_names = [n[0] for n in list(data)[start_col::2]] if skeleton is None: skeleton = Skeleton() @@ -196,14 +197,20 @@ def read_frames( ) else: # Get points for each node. + any_not_missing = False instance_points = dict() for node in node_names: x, y = data[(node, "x")][i], data[(node, "y")][i] instance_points[node] = Point(x, y) - - # Create instance with points assuming there's a single instance per - # frame. - instances.append(Instance(skeleton=skeleton, points=instance_points)) + if ~(np.isnan(x) and np.isnan(y)): + any_not_missing = True + + if any_not_missing: + # Create instance with points assuming there's a single instance per + # frame. + instances.append( + Instance(skeleton=skeleton, points=instance_points) + ) if len(instances) > 0: # Create LabeledFrame and add it to list. diff --git a/tests/data/dlc/dlc_testdata.csv b/tests/data/dlc/dlc_testdata.csv new file mode 100644 index 000000000..0544265e5 --- /dev/null +++ b/tests/data/dlc/dlc_testdata.csv @@ -0,0 +1,7 @@ +scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer +bodyparts,A,A,B,B,C,C +coords,x,y,x,y,x,y +labeled-data/video/img000.png,0,1,2,3,4,5 +labeled-data/video/img001.png,12,13,,,15,16 +labeled-data/video/img002.png,,,,,, +labeled-data/video/img003.png,22,23,24,25,26,27 diff --git a/tests/data/dlc/dlc_testdata_v2.csv b/tests/data/dlc/dlc_testdata_v2.csv new file mode 100644 index 000000000..b273d083f --- /dev/null +++ b/tests/data/dlc/dlc_testdata_v2.csv @@ -0,0 +1,7 @@ +scorer,,,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer +bodyparts,,,A,A,B,B,C,C +coords,,,x,y,x,y,x,y +labeled-data,video,img000.png,0,1,2,3,4,5 +labeled-data,video,img001.png,12,13,,,15,16 +labeled-data,video,img002.png,,,,,, +labeled-data,video,img003.png,22,23,24,25,26,27 diff --git a/tests/data/dlc/img000.png b/tests/data/dlc/img000.png new file mode 100644 index 000000000..eda5c89df Binary files /dev/null and b/tests/data/dlc/img000.png differ diff --git a/tests/data/dlc/img001.png b/tests/data/dlc/img001.png new file mode 100644 index 000000000..eda5c89df Binary files /dev/null and b/tests/data/dlc/img001.png differ diff --git a/tests/data/dlc/img002.png b/tests/data/dlc/img002.png new file mode 100644 index 000000000..eda5c89df Binary files /dev/null and b/tests/data/dlc/img002.png differ diff --git a/tests/data/dlc/img003.png b/tests/data/dlc/img003.png new file mode 100644 index 000000000..eda5c89df Binary files /dev/null and b/tests/data/dlc/img003.png differ diff --git a/tests/data/dlc/madlc_testdata_v2.csv b/tests/data/dlc/madlc_testdata_v2.csv index 3b332c228..702146ee0 100644 --- a/tests/data/dlc/madlc_testdata_v2.csv +++ b/tests/data/dlc/madlc_testdata_v2.csv @@ -2,7 +2,7 @@ scorer,,,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,S individuals,,,Animal1,Animal1,Animal1,Animal1,Animal1,Animal1,Animal2,Animal2,Animal2,Animal2,Animal2,Animal2 bodyparts,,,A,A,B,B,C,C,A,A,B,B,C,C coords,,,x,y,x,y,x,y,x,y,x,y,x,y -labeled-data,species,img000.png,0,1,2,3,4,5,6,7,8,9,10,11 -labeled-data,species,img001.png,12,13,,,15,16,17,18,,,20,21 -labeled-data,species,img002.png,,,,,,,,,,,, -labeled-data,species,img003.png,22,23,24,25,26,27,,,,,, +labeled-data,video,img000.png,0,1,2,3,4,5,6,7,8,9,10,11 +labeled-data,video,img001.png,12,13,,,15,16,17,18,,,20,21 +labeled-data,video,img002.png,,,,,,,,,,,, +labeled-data,video,img003.png,22,23,24,25,26,27,,,,,, diff --git a/tests/io/test_formats.py b/tests/io/test_formats.py index 44cfa1c12..fe5956694 100644 --- a/tests/io/test_formats.py +++ b/tests/io/test_formats.py @@ -161,7 +161,10 @@ def test_matching_adaptor(centered_pair_predictions_hdf5_path): @pytest.mark.parametrize( "test_data", - ["tests/data/dlc/madlc_testdata.csv", "tests/data/dlc/madlc_testdata_v2.csv"], + [ + "tests/data/dlc/madlc_testdata.csv", + "tests/data/dlc/madlc_testdata_v2.csv", + ], ) def test_madlc(test_data): labels = read( @@ -178,7 +181,10 @@ def test_madlc(test_data): assert labels.videos[0].filenames[2].endswith("img002.png") assert labels.videos[0].filenames[3].endswith("img003.png") + # Assert frames without any coor are not labeled assert len(labels) == 3 + + # Assert number of instances per frame is correct assert len(labels[0]) == 2 assert len(labels[1]) == 2 assert len(labels[2]) == 1 @@ -191,6 +197,40 @@ def test_madlc(test_data): assert labels[2].frame_idx == 3 +# TODO: Add test data for old single animal DLC format +@pytest.mark.parametrize( + "test_data", + ["tests/data/dlc/dlc_testdata.csv", "tests/data/dlc/dlc_testdata_v2.csv"], +) +def test_sadlc(test_data): + labels = read( + test_data, + for_object="labels", + as_format="deeplabcut", + ) + + assert labels.skeleton.node_names == ["A", "B", "C"] + assert len(labels.videos) == 1 + assert len(labels.video.filenames) == 4 + assert labels.videos[0].filenames[0].endswith("img000.png") + assert labels.videos[0].filenames[1].endswith("img001.png") + assert labels.videos[0].filenames[2].endswith("img002.png") + assert labels.videos[0].filenames[3].endswith("img003.png") + + # Assert frames without any coor are not labeled + assert len(labels) == 3 + + # Assert number of instances per frame is correct + assert len(labels[0]) == 1 + assert len(labels[1]) == 1 + assert len(labels[2]) == 1 + + assert_array_equal(labels[0][0].numpy(), [[0, 1], [2, 3], [4, 5]]) + assert_array_equal(labels[1][0].numpy(), [[12, 13], [np.nan, np.nan], [15, 16]]) + assert_array_equal(labels[2][0].numpy(), [[22, 23], [24, 25], [26, 27]]) + assert labels[2].frame_idx == 3 + + def test_tracking_scores(tmpdir, centered_pair_predictions_slp_path): # test reading