diff --git a/sleap/io/format/deeplabcut.py b/sleap/io/format/deeplabcut.py index bb5dc3410..5892dba1a 100644 --- a/sleap/io/format/deeplabcut.py +++ b/sleap/io/format/deeplabcut.py @@ -19,10 +19,10 @@ import numpy as np import pandas as pd -from typing import List, Optional +from typing import List, Optional, Dict from sleap import Labels, Video, Skeleton -from sleap.instance import Instance, LabeledFrame, Point +from sleap.instance import Instance, LabeledFrame, Point, Track from sleap.util import find_files_by_suffix from .adaptor import Adaptor, SleapObjectType @@ -119,11 +119,12 @@ def read_frames( # Pull out animal and node names from the columns. start_col = 3 if is_new_format else 1 - animal_names = [] + tracks: Dict[str, Optional[Track]] = {} node_names = [] for animal_name, node_name, _ in data.columns[start_col:][::2]: - if animal_name not in animal_names: - animal_names.append(animal_name) + # Keep the starting frame index for each individual/track + if animal_name not in tracks.keys(): + tracks[animal_name] = None if node_name not in node_names: node_names.append(node_name) @@ -177,23 +178,33 @@ def read_frames( instances = [] if is_multianimal: - for animal_name in animal_names: + for animal_name in tracks.keys(): any_not_missing = False # Get points for each node. instance_points = dict() for node in node_names: - x, y = ( - data[(animal_name, node, "x")][i], - data[(animal_name, node, "y")][i], - ) + if (animal_name, node) in data.columns: + x, y = ( + data[(animal_name, node, "x")][i], + data[(animal_name, node, "y")][i], + ) + else: + x, y = np.nan, np.nan instance_points[node] = Point(x, y) if ~(np.isnan(x) and np.isnan(y)): any_not_missing = True if any_not_missing: + # Create track + if tracks[animal_name] is None: + tracks[animal_name] = Track(spawned_on=i, name=animal_name) # Create instance with points. instances.append( - Instance(skeleton=skeleton, points=instance_points) + Instance( + skeleton=skeleton, + points=instance_points, + track=tracks[animal_name], + ) ) else: # Get points for each node. @@ -270,6 +281,8 @@ def read( skeleton = Skeleton() if project_data.get("multianimalbodyparts", False): skeleton.add_nodes(project_data["multianimalbodyparts"]) + if "uniquebodyparts" in project_data: + skeleton.add_nodes(project_data["uniquebodyparts"]) else: skeleton.add_nodes(project_data["bodyparts"]) @@ -298,13 +311,24 @@ def read( # If subdirectory is foo, we look for foo.mp4 in videos dir. shortname = os.path.split(data_subdir)[-1] - video_path = os.path.join(videos_dir, f"{shortname}.mp4") - - if os.path.exists(video_path): + video_path = None + if os.path.exists(videos_dir): + with os.scandir(videos_dir) as file_iterator: + for file in file_iterator: + if not file.is_file(): + continue + if os.path.splitext(file.name)[0] != shortname: + continue + video_path = os.path.join(videos_dir, file.name) + break + + if video_path is not None and os.path.exists(video_path): video = Video.from_filename(video_path) else: # When no video is found, the individual frame images # stored in the labeled data subdir will be used. + if video_path is None: + video_path = os.path.join(videos_dir, f"{shortname}.mp4") print( f"Unable to find {video_path} so using individual frame images." ) diff --git a/sleap/nn/system.py b/sleap/nn/system.py index eeb3f3ca4..f73dc8457 100644 --- a/sleap/nn/system.py +++ b/sleap/nn/system.py @@ -159,7 +159,7 @@ def summary(): for gpu in all_gpus: print(f" Device: {gpu.name}") print(f" Available: {gpu in gpus}") - print(f" Initalized: {is_initialized(gpu)}") + print(f" Initialized: {is_initialized(gpu)}") print( f" Memory growth: {tf.config.experimental.get_memory_growth(gpu)}" ) diff --git a/tests/data/dlc/labeled-data/video/CollectedData_LM.csv b/tests/data/dlc/labeled-data/video/CollectedData_LM.csv index f57b667f4..27c86f8af 100644 --- a/tests/data/dlc/labeled-data/video/CollectedData_LM.csv +++ b/tests/data/dlc/labeled-data/video/CollectedData_LM.csv @@ -1,8 +1,8 @@ -scorer,,,LM,LM,LM,LM,LM,LM,LM,LM,LM,LM,LM,LM -individuals,,,individual1,individual1,individual1,individual1,individual1,individual1,individual2,individual2,individual2,individual2,individual2,individual2 -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,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 +scorer,,,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer +individuals,,,Animal1,Animal1,Animal1,Animal1,Animal1,Animal1,Animal2,Animal2,Animal2,Animal2,Animal2,Animal2,single,single,single,single +bodyparts,,,A,A,B,B,C,C,A,A,B,B,C,C,D,D,E,E +coords,,,x,y,x,y,x,y,x,y,x,y,x,y,x,y,x,y +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,22,23,24,25 labeled-data,video,img002.png,,,,,,,,,,,, -labeled-data,video,img003.png,22,23,24,25,26,27,,,,,, +labeled-data,video,img003.png,26,27,28,29,30,31,,,,,,,32,33,34,35 diff --git a/tests/data/dlc/labeled-data/video/maudlc_testdata.csv b/tests/data/dlc/labeled-data/video/maudlc_testdata.csv new file mode 100644 index 000000000..4e3e3c28c --- /dev/null +++ b/tests/data/dlc/labeled-data/video/maudlc_testdata.csv @@ -0,0 +1,8 @@ +scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer +individuals,Animal1,Animal1,Animal1,Animal1,Animal1,Animal1,Animal2,Animal2,Animal2,Animal2,Animal2,Animal2,single,single,single,single +bodyparts,A,A,B,B,C,C,A,A,B,B,C,C,D,D,E,E +coords,x,y,x,y,x,y,x,y,x,y,x,y,x,y,x,y +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,22,23,24,25 +labeled-data/video/img002.png,,,,,,,,,,,, +labeled-data/video/img003.png,26,27,28,29,30,31,,,,,,,32,33,34,35 diff --git a/tests/data/dlc/labeled-data/video/maudlc_testdata_v2.csv b/tests/data/dlc/labeled-data/video/maudlc_testdata_v2.csv new file mode 100644 index 000000000..27c86f8af --- /dev/null +++ b/tests/data/dlc/labeled-data/video/maudlc_testdata_v2.csv @@ -0,0 +1,8 @@ +scorer,,,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer +individuals,,,Animal1,Animal1,Animal1,Animal1,Animal1,Animal1,Animal2,Animal2,Animal2,Animal2,Animal2,Animal2,single,single,single,single +bodyparts,,,A,A,B,B,C,C,A,A,B,B,C,C,D,D,E,E +coords,,,x,y,x,y,x,y,x,y,x,y,x,y,x,y,x,y +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,22,23,24,25 +labeled-data,video,img002.png,,,,,,,,,,,, +labeled-data,video,img003.png,26,27,28,29,30,31,,,,,,,32,33,34,35 diff --git a/tests/data/dlc/madlc_230_config.yaml b/tests/data/dlc/madlc_230_config.yaml index ae2cbb44b..01e1d32c1 100644 --- a/tests/data/dlc/madlc_230_config.yaml +++ b/tests/data/dlc/madlc_230_config.yaml @@ -1,12 +1,12 @@ # Project definitions (do not edit) -Task: madlc_2.3.0 +Task: maudlc_2.3.0 scorer: LM date: Mar1 multianimalproject: true identity: false # Project path (change when moving around) -project_path: D:\social-leap-estimates-animal-poses\pull-requests\sleap\tests\data\dlc\madlc_testdata_v3 +project_path: D:\social-leap-estimates-animal-poses\pull-requests\sleap\tests\data\dlc\maudlc_testdata_v3 # Annotation data set configuration (and individual video cropping parameters) video_sets: @@ -16,7 +16,9 @@ individuals: - individual1 - individual2 - individual3 -uniquebodyparts: [] +uniquebodyparts: +- D +- E multianimalbodyparts: - A - B diff --git a/tests/gui/test_commands.py b/tests/gui/test_commands.py index 13aa60e6b..899b1f4a0 100644 --- a/tests/gui/test_commands.py +++ b/tests/gui/test_commands.py @@ -65,7 +65,7 @@ def test_import_labels_from_dlc_folder(): assert len(labels.videos) == 2 assert len(labels.skeletons) == 1 assert len(labels.nodes) == 3 - assert len(labels.tracks) == 0 + assert len(labels.tracks) == 3 assert set( [fix_path_separator(l.video.backend.filename) for l in labels.labeled_frames] diff --git a/tests/io/test_formats.py b/tests/io/test_formats.py index a89bf60d7..725a00ed3 100644 --- a/tests/io/test_formats.py +++ b/tests/io/test_formats.py @@ -198,7 +198,6 @@ def test_matching_adaptor(centered_pair_predictions_hdf5_path): [ "tests/data/dlc/labeled-data/video/madlc_testdata.csv", "tests/data/dlc/labeled-data/video/madlc_testdata_v2.csv", - "tests/data/dlc/madlc_230_config.yaml", ], ) def test_madlc(test_data): @@ -232,6 +231,78 @@ def test_madlc(test_data): assert labels[2].frame_idx == 3 +@pytest.mark.parametrize( + "test_data", + [ + "tests/data/dlc/labeled-data/video/maudlc_testdata.csv", + "tests/data/dlc/labeled-data/video/maudlc_testdata_v2.csv", + "tests/data/dlc/madlc_230_config.yaml", + ], +) +def test_maudlc(test_data): + labels = read( + test_data, + for_object="labels", + as_format="deeplabcut", + ) + + assert labels.skeleton.node_names == ["A", "B", "C", "D", "E"] + 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]) == 2 + assert len(labels[1]) == 3 + assert len(labels[2]) == 2 + + assert_array_equal( + labels[0][0].numpy(), + [[0, 1], [2, 3], [4, 5], [np.nan, np.nan], [np.nan, np.nan]], + ) + assert_array_equal( + labels[0][1].numpy(), + [[6, 7], [8, 9], [10, 11], [np.nan, np.nan], [np.nan, np.nan]], + ) + assert_array_equal( + labels[1][0].numpy(), + [[12, 13], [np.nan, np.nan], [15, 16], [np.nan, np.nan], [np.nan, np.nan]], + ) + assert_array_equal( + labels[1][1].numpy(), + [[17, 18], [np.nan, np.nan], [20, 21], [np.nan, np.nan], [np.nan, np.nan]], + ) + assert_array_equal( + labels[1][2].numpy(), + [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan], [22, 23], [24, 25]], + ) + assert_array_equal( + labels[2][0].numpy(), + [[26, 27], [28, 29], [30, 31], [np.nan, np.nan], [np.nan, np.nan]], + ) + assert_array_equal( + labels[2][1].numpy(), + [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan], [32, 33], [34, 35]], + ) + assert labels[2].frame_idx == 3 + + # Assert tracks are correct + assert len(labels.tracks) == 3 + sorted_animals = sorted(["Animal1", "Animal2", "single"]) + assert sorted([t.name for t in labels.tracks]) == sorted_animals + for t in labels.tracks: + if t.name == "single": + assert t.spawned_on == 1 + else: + assert t.spawned_on == 0 + + @pytest.mark.parametrize( "test_data", [