def test_frame_merge_between_predicted_and_user(skeleton, centered_pair_vid):
user_inst = Instance(skeleton=skeleton, points={skeleton.nodes[0]: Point(1, 2)},)
user_labels = Labels(
[LabeledFrame(video=centered_pair_vid, frame_idx=0, instances=[user_inst],)]
)
pred_inst = PredictedInstance(
skeleton=skeleton,
points={skeleton.nodes[0]: PredictedPoint(1, 2, score=1.0)},
score=1.0,
)
pred_labels = Labels(
[LabeledFrame(video=centered_pair_vid, frame_idx=0, instances=[pred_inst],)]
)
# Merge predictions into current labels dataset
_, _, new_conflicts = Labels.complex_merge_between(
user_labels,
new_labels=pred_labels,
unify=False, # since we used match_to when loading predictions file
)
# new predictions should replace old ones
Labels.finish_complex_merge(user_labels, new_conflicts)
# We should be able to cleanly merge the user and the predicted instance,
# and we want to retain both even though they perfectly match.
assert user_inst in user_labels[0].instances
assert pred_inst in user_labels[0].instances
assert len(user_labels[0].instances) == 2
def test_frame_merge_predicted_and_user(skeleton, centered_pair_vid):
user_inst = Instance(
skeleton=skeleton,
points={skeleton.nodes[0]: Point(1, 2)},
)
user_frame = LabeledFrame(
video=centered_pair_vid,
frame_idx=0,
instances=[user_inst],
)
pred_inst = PredictedInstance(
skeleton=skeleton,
points={skeleton.nodes[0]: PredictedPoint(1, 2, score=1.0)},
score=1.0,
)
pred_frame = LabeledFrame(
video=centered_pair_vid,
frame_idx=0,
instances=[pred_inst],
)
LabeledFrame.complex_frame_merge(user_frame, pred_frame)
# We should be able to cleanly merge the user and the predicted instance,
# and we want to retain both even though they perfectly match.
assert user_inst in user_frame.instances
assert pred_inst in user_frame.instances
assert len(user_frame.instances) == 2
def test_instance_access():
labels = Labels()
dummy_skeleton = Skeleton()
dummy_video = Video(backend=MediaVideo)
dummy_video2 = Video(backend=MediaVideo)
for i in range(10):
labels.append(
LabeledFrame(
dummy_video,
frame_idx=i,
instances=[Instance(dummy_skeleton),
Instance(dummy_skeleton)],
))
for i in range(10):
labels.append(
LabeledFrame(
dummy_video2,
frame_idx=i,
instances=[
Instance(dummy_skeleton),
Instance(dummy_skeleton),
Instance(dummy_skeleton),
],
))
assert len(labels.all_instances) == 50
assert len(list(labels.instances(video=dummy_video))) == 20
assert len(list(labels.instances(video=dummy_video2))) == 30
def test_label_mutability():
dummy_video = Video(backend=MediaVideo)
dummy_skeleton = Skeleton()
dummy_instance = Instance(dummy_skeleton)
dummy_frame = LabeledFrame(dummy_video, frame_idx=0, instances=[dummy_instance])
labels = Labels()
labels.append(dummy_frame)
assert dummy_video in labels.videos
assert dummy_video in labels
assert dummy_skeleton in labels.skeletons
assert dummy_skeleton in labels
assert dummy_frame in labels.labeled_frames
assert dummy_frame in labels
assert (dummy_video, 0) in labels
assert (dummy_video, 1) not in labels
dummy_video2 = Video(backend=MediaVideo)
dummy_skeleton2 = Skeleton(name="dummy2")
dummy_instance2 = Instance(dummy_skeleton2)
dummy_frame2 = LabeledFrame(dummy_video2, frame_idx=0, instances=[dummy_instance2])
assert dummy_video2 not in labels
assert dummy_skeleton2 not in labels
assert dummy_frame2 not in labels
labels.append(dummy_frame2)
assert dummy_video2 in labels
assert dummy_frame2 in labels
labels.remove_video(dummy_video2)
assert dummy_video2 not in labels
assert dummy_frame2 not in labels
assert len(labels.find(dummy_video2)) == 0
assert len(labels) == 1
labels.append(LabeledFrame(dummy_video, frame_idx=0))
assert len(labels) == 1
dummy_frames = [LabeledFrame(dummy_video, frame_idx=i) for i in range(10)]
dummy_frames2 = [LabeledFrame(dummy_video2, frame_idx=i) for i in range(10)]
for f in dummy_frames + dummy_frames2:
labels.append(f)
assert len(labels) == 20
labels.remove_video(dummy_video2)
assert len(labels) == 10
assert len(labels.find(dummy_video)) == 10
assert dummy_frame in labels
assert all([label in labels for label in dummy_frames[1:]])
assert dummy_video2 not in labels
assert len(labels.find(dummy_video2)) == 0
assert all([label not in labels for label in dummy_frames2])
labels.remove_video(dummy_video)
assert len(labels.find(dummy_video)) == 0
def simple_predictions():
video = Video.from_filename("video.mp4")
skeleton = Skeleton()
skeleton.add_node("a")
skeleton.add_node("b")
track_a = Track(0, "a")
track_b = Track(0, "b")
labels = Labels()
instances = []
instances.append(
PredictedInstance(
skeleton=skeleton,
score=2,
track=track_a,
points=dict(a=PredictedPoint(1, 1, score=0.5),
b=PredictedPoint(1, 1, score=0.5)),
))
instances.append(
PredictedInstance(
skeleton=skeleton,
score=5,
track=track_b,
points=dict(a=PredictedPoint(1, 1, score=0.7),
b=PredictedPoint(1, 1, score=0.7)),
))
labeled_frame = LabeledFrame(video, frame_idx=0, instances=instances)
labels.append(labeled_frame)
instances = []
instances.append(
PredictedInstance(
skeleton=skeleton,
score=3,
track=track_a,
points=dict(a=PredictedPoint(4, 5, score=1.5),
b=PredictedPoint(1, 1, score=1.0)),
))
instances.append(
PredictedInstance(
skeleton=skeleton,
score=6,
track=track_b,
points=dict(a=PredictedPoint(6, 13, score=1.7),
b=PredictedPoint(1, 1, score=1.0)),
))
labeled_frame = LabeledFrame(video, frame_idx=1, instances=instances)
labels.append(labeled_frame)
return labels
def test_duplicate_skeletons_serializing():
vid = Video.from_filename("foo.mp4")
skeleton_a = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
skeleton_b = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
lf_a = LabeledFrame(vid, frame_idx=2, instances=[Instance(skeleton_a)])
lf_b = LabeledFrame(vid, frame_idx=3, instances=[Instance(skeleton_b)])
new_labels = Labels(labeled_frames=[lf_a, lf_b])
new_labels_json = new_labels.to_dict()
def test_scalar_properties():
# Scalar
dummy_video = Video(backend=MediaVideo)
dummy_skeleton = Skeleton()
dummy_instance = Instance(dummy_skeleton)
dummy_frame = LabeledFrame(dummy_video,
frame_idx=0,
instances=[dummy_instance])
labels = Labels()
labels.append(dummy_frame)
assert labels.video == dummy_video
assert labels.skeleton == dummy_skeleton
# Empty
labels = Labels()
with pytest.raises(ValueError):
labels.video
with pytest.raises(ValueError):
labels.skeleton
# More than one video
dummy_skeleton = Skeleton()
labels = Labels()
labels.append(
LabeledFrame(Video(backend=MediaVideo),
frame_idx=0,
instances=[Instance(dummy_skeleton)]))
labels.append(
LabeledFrame(Video(backend=MediaVideo),
frame_idx=0,
instances=[Instance(dummy_skeleton)]))
assert labels.skeleton == dummy_skeleton
with pytest.raises(ValueError):
labels.video
# More than one skeleton
dummy_video = Video(backend=MediaVideo)
labels = Labels()
labels.append(
LabeledFrame(dummy_video,
frame_idx=0,
instances=[Instance(Skeleton())]))
labels.append(
LabeledFrame(dummy_video,
frame_idx=1,
instances=[Instance(Skeleton())]))
assert labels.video == dummy_video
with pytest.raises(ValueError):
labels.skeleton
def test_distinct_skeletons_serializing():
vid = Video.from_filename("foo.mp4")
skeleton_a = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
skeleton_b = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
skeleton_b.add_node("foo")
lf_a = LabeledFrame(vid, frame_idx=2, instances=[Instance(skeleton_a)])
lf_b = LabeledFrame(vid, frame_idx=3, instances=[Instance(skeleton_b)])
new_labels = Labels(labeled_frames=[lf_a, lf_b])
# Make sure we can serialize this
new_labels_json = new_labels.to_dict()
def multi_skel_vid_labels(hdf5_vid, small_robot_mp4_vid, skeleton, stickman):
"""
Build a big list of LabeledFrame objects and wrap it in Labels class.
Args:
hdf5_vid: An HDF5 video fixture
small_robot_mp4_vid: An MP4 video fixture
skeleton: A fly skeleton.
stickman: A stickman skeleton
Returns:
The Labels object containing all the labeled frames
"""
labels = []
stick_tracks = [
Track(spawned_on=0, name=f"Stickman {i}") for i in range(6)
]
fly_tracks = [Track(spawned_on=0, name=f"Fly {i}") for i in range(6)]
# Make some tracks None to test that
fly_tracks[3] = None
stick_tracks[2] = None
for f in range(500):
vid = [hdf5_vid, small_robot_mp4_vid][f % 2]
label = LabeledFrame(video=vid, frame_idx=f % vid.frames)
fly_instances = []
for i in range(6):
fly_instances.append(
Instance(skeleton=skeleton, track=fly_tracks[i]))
for node in skeleton.nodes:
fly_instances[i][node] = Point(x=i % vid.width,
y=i % vid.height)
stickman_instances = []
for i in range(6):
stickman_instances.append(
Instance(skeleton=stickman, track=stick_tracks[i]))
for node in stickman.nodes:
stickman_instances[i][node] = Point(x=i % vid.width,
y=i % vid.height)
label.instances = stickman_instances + fly_instances
labels.append(label)
labels = Labels(labels)
return labels
def test_save_labels_with_images(min_labels_slp, tmpdir):
labels = Labels(min_labels_slp.labeled_frames)
labels.append(LabeledFrame(video=labels.video, frame_idx=1))
labels.suggestions.append(SuggestionFrame(video=labels.video, frame_idx=2))
fn = os.path.join(tmpdir, "test_user_only.slp")
labels.save(
fn, with_images=True, embed_all_labeled=False, embed_suggested=False,
)
assert Labels.load_file(fn).video.embedded_frame_inds == [0]
fn = os.path.join(tmpdir, "test_all_labeled.slp")
labels.save(
fn, with_images=True, embed_all_labeled=True, embed_suggested=False,
)
assert Labels.load_file(fn).video.embedded_frame_inds == [0, 1]
fn = os.path.join(tmpdir, "test_suggested.slp")
labels.save(
fn, with_images=True, embed_all_labeled=False, embed_suggested=True,
)
assert Labels.load_file(fn).video.embedded_frame_inds == [0, 2]
fn = os.path.join(tmpdir, "test_all.slp")
labels.save(
fn, with_images=True, embed_all_labeled=True, embed_suggested=True,
)
assert Labels.load_file(fn).video.embedded_frame_inds == [0, 1, 2]
def test_save_frame_data_hdf5(min_labels_slp, tmpdir):
labels = Labels(min_labels_slp.labeled_frames)
labels.append(LabeledFrame(video=labels.video, frame_idx=1))
labels.suggestions.append(SuggestionFrame(video=labels.video, frame_idx=2))
fn = os.path.join(tmpdir, "test_user_only.slp")
labels.save_frame_data_hdf5(
fn, format="png", user_labeled=True, all_labeled=False, suggested=False,
)
assert Video.from_filename(fn, dataset="video0").embedded_frame_inds == [0]
fn = os.path.join(tmpdir, "test_all_labeled.slp")
labels.save_frame_data_hdf5(
fn, format="png", user_labeled=False, all_labeled=True, suggested=False,
)
assert Video.from_filename(fn, dataset="video0").embedded_frame_inds == [0, 1]
fn = os.path.join(tmpdir, "test_suggested.slp")
labels.save_frame_data_hdf5(
fn, format="png", user_labeled=False, all_labeled=False, suggested=True,
)
assert Video.from_filename(fn, dataset="video0").embedded_frame_inds == [2]
fn = os.path.join(tmpdir, "test_all.slp")
labels.save_frame_data_hdf5(
fn, format="png", user_labeled=False, all_labeled=True, suggested=True,
)
assert Video.from_filename(fn, dataset="video0").embedded_frame_inds == [0, 1, 2]
def removal_test_labels():
skeleton = Skeleton()
video = Video(backend=MediaVideo)
lf_user_only = LabeledFrame(
video=video, frame_idx=0, instances=[Instance(skeleton=skeleton)]
)
lf_pred_only = LabeledFrame(
video=video, frame_idx=1, instances=[PredictedInstance(skeleton=skeleton)]
)
lf_both = LabeledFrame(
video=video,
frame_idx=2,
instances=[Instance(skeleton=skeleton), PredictedInstance(skeleton=skeleton)],
)
labels = Labels([lf_user_only, lf_pred_only, lf_both])
return labels
def test_merge_predictions():
dummy_video_a = Video.from_filename("foo.mp4")
dummy_video_b = Video.from_filename("foo.mp4")
dummy_skeleton_a = Skeleton()
dummy_skeleton_a.add_node("node")
dummy_skeleton_b = Skeleton()
dummy_skeleton_b.add_node("node")
dummy_instances_a = []
dummy_instances_a.append(
Instance(skeleton=dummy_skeleton_a, points=dict(node=Point(1, 1)))
)
dummy_instances_a.append(
Instance(skeleton=dummy_skeleton_a, points=dict(node=Point(2, 2)))
)
labels_a = Labels()
labels_a.append(
LabeledFrame(dummy_video_a, frame_idx=0, instances=dummy_instances_a)
)
dummy_instances_b = []
dummy_instances_b.append(
Instance(skeleton=dummy_skeleton_b, points=dict(node=Point(1, 1)))
)
dummy_instances_b.append(
PredictedInstance(
skeleton=dummy_skeleton_b, points=dict(node=Point(3, 3)), score=1
)
)
labels_b = Labels()
labels_b.append(
LabeledFrame(dummy_video_b, frame_idx=0, instances=dummy_instances_b)
)
# Frames have one redundant instance (perfect match) and all the
# non-matching instances are different types (one predicted, one not).
merged, extra_a, extra_b = Labels.complex_merge_between(labels_a, labels_b)
assert len(merged[dummy_video_a]) == 1
assert len(merged[dummy_video_a][0]) == 1 # the predicted instance was merged
assert not extra_a
assert not extra_b
def test_has_frame():
video = Video(backend=MediaVideo)
labels = Labels([LabeledFrame(video=video, frame_idx=0)])
assert labels.has_frame(labels[0])
assert labels.has_frame(labels[0], use_cache=False)
assert labels.has_frame(LabeledFrame(video=video, frame_idx=0))
assert labels.has_frame(video=video, frame_idx=0)
assert labels.has_frame(video=video, frame_idx=0, use_cache=False)
assert not labels.has_frame(LabeledFrame(video=video, frame_idx=1))
assert not labels.has_frame(LabeledFrame(video=video, frame_idx=1), use_cache=False)
assert not labels.has_frame(video=video, frame_idx=1)
with pytest.raises(ValueError):
labels.has_frame()
with pytest.raises(ValueError):
labels.has_frame(video=video)
with pytest.raises(ValueError):
labels.has_frame(frame_idx=1)
def test_dont_unify_skeletons():
vid = Video.from_filename("foo.mp4")
skeleton_a = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
skeleton_b = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
lf_a = LabeledFrame(vid, frame_idx=2, instances=[Instance(skeleton_a)])
lf_b = LabeledFrame(vid, frame_idx=3, instances=[Instance(skeleton_b)])
labels = Labels(labeled_frames=[lf_a])
labels.extend_from([lf_b], unify=False)
ids = skeleton_ids_from_label_instances(labels)
# Make sure we still have two distinct skeleton objects
assert len(set(ids)) == 2
# Make sure we can serialize this
labels.to_dict()
def test_instance_labeled_frame_ref(skeleton, centered_pair_vid):
"""
Test whether links between labeled frames and instances are kept
"""
instances = [Instance(skeleton=skeleton) for i in range(3)]
frame = LabeledFrame(video=centered_pair_vid, frame_idx=0, instances=instances)
assert frame.instances[0].frame == frame
assert frame[0].frame == frame
assert frame[0].frame_idx == 0
def test_unify_skeletons():
vid = Video.from_filename("foo.mp4")
skeleton_a = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
skeleton_b = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
lf_a = LabeledFrame(vid, frame_idx=2, instances=[Instance(skeleton_a)])
lf_b = LabeledFrame(vid, frame_idx=3, instances=[Instance(skeleton_b)])
labels = Labels()
labels.extend_from([lf_a], unify=True)
labels.extend_from([lf_b], unify=True)
ids = skeleton_ids_from_label_instances(labels)
# Make sure that skeleton_b got replaced with skeleton_a when we
# added the frame with "unify" set
assert len(set(ids)) == 1
# Make sure we can serialize this
labels.to_dict()
def test_multivideo_tracks():
vid_a = Video.from_filename("foo.mp4")
vid_b = Video.from_filename("bar.mp4")
skeleton = Skeleton.load_json("tests/data/skeleton/fly_skeleton_legs.json")
track_a = Track(spawned_on=2, name="A")
track_b = Track(spawned_on=3, name="B")
inst_a = Instance(track=track_a, skeleton=skeleton)
inst_b = Instance(track=track_b, skeleton=skeleton)
lf_a = LabeledFrame(vid_a, frame_idx=2, instances=[inst_a])
lf_b = LabeledFrame(vid_b, frame_idx=3, instances=[inst_b])
labels = Labels(labeled_frames=[lf_a, lf_b])
# Try setting video B instance to track used in video A
labels.track_swap(vid_b, new_track=track_a, old_track=track_b, frame_range=(3, 4))
assert inst_b.track == track_a
def test_basic_suggestions(small_robot_mp4_vid):
dummy_video = small_robot_mp4_vid
dummy_skeleton = Skeleton()
dummy_instance = Instance(dummy_skeleton)
dummy_frame = LabeledFrame(dummy_video, frame_idx=0, instances=[dummy_instance])
labels = Labels()
labels.append(dummy_frame)
suggestions = VideoFrameSuggestions.suggest(
labels=labels, params=dict(method="sample", per_video=13)
)
labels.set_suggestions(suggestions)
assert len(labels.get_video_suggestions(dummy_video)) == 13
def read(
cls,
file: FileHandle,
video_path: str,
skeleton_path: str,
*args,
**kwargs,
) -> Labels:
f = file.file
video = Video.from_filename(video_path)
skeleton_data = pd.read_csv(skeleton_path, header=0)
skeleton = Skeleton()
skeleton.add_nodes(skeleton_data["name"])
nodes = skeleton.nodes
for name, parent, swap in skeleton_data.itertuples(index=False,
name=None):
if parent is not np.nan:
skeleton.add_edge(parent, name)
lfs = []
pose_matrix = f["pose"][:]
track_count, frame_count, node_count, _ = pose_matrix.shape
tracks = [Track(0, f"Track {i}") for i in range(track_count)]
for frame_idx in range(frame_count):
lf_instances = []
for track_idx in range(track_count):
points_array = pose_matrix[track_idx, frame_idx, :, :]
points = dict()
for p in range(len(points_array)):
x, y, score = points_array[p]
points[nodes[p]] = Point(x, y) # TODO: score
inst = Instance(skeleton=skeleton,
track=tracks[track_idx],
points=points)
lf_instances.append(inst)
lfs.append(
LabeledFrame(video,
frame_idx=frame_idx,
instances=lf_instances))
return Labels(labeled_frames=lfs)
def test_remove_predictions_with_new_labels(removal_test_labels):
labels = removal_test_labels
assert len(labels) == 3
new_labels = Labels([
LabeledFrame(
video=labels.video,
frame_idx=1,
instances=[PredictedInstance(skeleton=labels.skeleton)],
)
])
labels.remove_predictions(new_labels=new_labels)
assert len(labels) == 2
assert labels[0].frame_idx == 0
assert labels[0].has_user_instances
assert not labels[0].has_predicted_instances
assert labels[1].frame_idx == 2
assert labels[1].has_user_instances
assert labels[1].has_predicted_instances
def test_labels_merge():
dummy_video = Video(backend=MediaVideo)
dummy_skeleton = Skeleton()
dummy_skeleton.add_node("node")
labels = Labels()
dummy_frames = []
# Add 10 instances with different points (so they aren't "redundant")
for i in range(10):
instance = Instance(skeleton=dummy_skeleton, points=dict(node=Point(i, i)))
dummy_frame = LabeledFrame(dummy_video, frame_idx=0, instances=[instance])
dummy_frames.append(dummy_frame)
labels.labeled_frames.extend(dummy_frames)
assert len(labels) == 10
assert len(labels.labeled_frames[0].instances) == 1
labels.merge_matching_frames()
assert len(labels) == 1
assert len(labels.labeled_frames[0].instances) == 10
def test_deserialize_suggestions(small_robot_mp4_vid, tmpdir):
dummy_video = small_robot_mp4_vid
dummy_skeleton = Skeleton()
dummy_instance = Instance(dummy_skeleton)
dummy_frame = LabeledFrame(dummy_video, frame_idx=0, instances=[dummy_instance])
labels = Labels()
labels.append(dummy_frame)
suggestions = VideoFrameSuggestions.suggest(
labels=labels, params=dict(method="sample", per_video=13)
)
labels.set_suggestions(suggestions)
filename = os.path.join(tmpdir, "new_suggestions.h5")
Labels.save_file(filename=filename, labels=labels)
new_suggestion_labels = Labels.load_file(filename)
assert len(suggestions) == len(new_suggestion_labels.suggestions)
assert [frame.frame_idx for frame in suggestions] == [
frame.frame_idx for frame in new_suggestion_labels.suggestions
]
def read_frames(
cls,
file: FileHandle,
skeleton: Optional[Skeleton] = None,
full_video: Optional[Video] = None,
*args,
**kwargs,
) -> List[LabeledFrame]:
filename = file.filename
data = pd.read_csv(filename, header=[1, 2])
# 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]]
if skeleton is None:
skeleton = Skeleton()
skeleton.add_nodes(node_names)
img_files = data.iloc[:, 0] # get list of all images
if full_video:
video = full_video
index_frames_by_original_index = True
else:
# Create the Video object
img_dir = os.path.dirname(filename)
video = cls.make_video_for_image_list(img_dir, img_files)
# The frames in the video we created will be indexed from 0 to N
# rather than having their index from the original source video.
index_frames_by_original_index = False
frames = []
for i in range(len(data)):
# get points for each node
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.
# For DeepLabCut we're assuming there's a single instance per frame.
instance = Instance(skeleton=skeleton, points=instance_points)
if index_frames_by_original_index:
# extract "0123" from "path/img0123.png" as original frame index
frame_idx_match = re.search("(?<=img)(\\d+)(?=\.png)", img_files[i])
if frame_idx_match is not None:
frame_idx = int(frame_idx_match.group(0))
else:
raise ValueError(
f"Unable to determine frame index for image {img_files[i]}"
)
else:
frame_idx = i
# create labeledframe and add it to list
frames.append(
LabeledFrame(video=video, frame_idx=frame_idx, instances=[instance])
)
return frames
def test_complex_merge():
dummy_video_a = Video.from_filename("foo.mp4")
dummy_video_b = Video.from_filename("foo.mp4")
dummy_skeleton_a = Skeleton()
dummy_skeleton_a.add_node("node")
dummy_skeleton_b = Skeleton()
dummy_skeleton_b.add_node("node")
dummy_instances_a = []
dummy_instances_a.append(
Instance(skeleton=dummy_skeleton_a, points=dict(node=Point(1, 1))))
dummy_instances_a.append(
Instance(skeleton=dummy_skeleton_a, points=dict(node=Point(2, 2))))
labels_a = Labels()
labels_a.append(
LabeledFrame(dummy_video_a, frame_idx=0, instances=dummy_instances_a))
dummy_instances_b = []
dummy_instances_b.append(
Instance(skeleton=dummy_skeleton_b, points=dict(node=Point(1, 1))))
dummy_instances_b.append(
Instance(skeleton=dummy_skeleton_b, points=dict(node=Point(3, 3))))
labels_b = Labels()
labels_b.append(
LabeledFrame(dummy_video_b, frame_idx=0,
instances=dummy_instances_b)) # conflict
labels_b.append(
LabeledFrame(dummy_video_b, frame_idx=1,
instances=dummy_instances_b)) # clean
merged, extra_a, extra_b = Labels.complex_merge_between(labels_a, labels_b)
# Check that we have the cleanly merged frame
assert dummy_video_a in merged
assert len(merged[dummy_video_a]) == 1 # one merged frame
assert len(merged[dummy_video_a][1]) == 2 # with two instances
# Check that labels_a includes redundant and clean
assert len(labels_a.labeled_frames) == 2
assert len(labels_a.labeled_frames[0].instances) == 1
assert labels_a.labeled_frames[0].instances[0].points[0].x == 1
assert len(labels_a.labeled_frames[1].instances) == 2
assert labels_a.labeled_frames[1].instances[0].points[0].x == 1
assert labels_a.labeled_frames[1].instances[1].points[0].x == 3
# Check that extra_a/b includes the appropriate conflicting instance
assert len(extra_a) == 1
assert len(extra_b) == 1
assert len(extra_a[0].instances) == 1
assert len(extra_b[0].instances) == 1
assert extra_a[0].instances[0].points[0].x == 2
assert extra_b[0].instances[0].points[0].x == 3
# Check that objects were unified
assert extra_a[0].video == extra_b[0].video
# Check resolving the conflict using new
Labels.finish_complex_merge(labels_a, extra_b)
assert len(labels_a.labeled_frames) == 2
assert len(labels_a.labeled_frames[0].instances) == 2
assert labels_a.labeled_frames[0].instances[1].points[0].x == 3
def read(
cls, file: FileHandle, gui: bool = True, *args, **kwargs,
):
filename = file.filename
mat_contents = sio.loadmat(filename)
box_path = cls._unwrap_mat_scalar(mat_contents["boxPath"])
# If the video file isn't found, try in the same dir as the mat file
if not os.path.exists(box_path):
file_dir = os.path.dirname(filename)
box_path_name = box_path.split("\\")[-1] # assume windows path
box_path = os.path.join(file_dir, box_path_name)
if not os.path.exists(box_path):
if gui:
video_paths = [box_path]
missing = [True]
okay = MissingFilesDialog(video_paths, missing).exec_()
if not okay or missing[0]:
return
box_path = video_paths[0]
else:
# Ignore missing videos if not loading from gui
box_path = ""
if os.path.exists(box_path):
vid = Video.from_hdf5(
dataset="box", filename=box_path, input_format="channels_first"
)
else:
vid = None
nodes_ = mat_contents["skeleton"]["nodes"]
edges_ = mat_contents["skeleton"]["edges"]
points_ = mat_contents["positions"]
edges_ = edges_ - 1 # convert matlab 1-indexing to python 0-indexing
nodes = cls._unwrap_mat_array(nodes_)
edges = cls._unwrap_mat_array(edges_)
nodes = list(map(str, nodes)) # convert np._str to str
sk = Skeleton(name=filename)
sk.add_nodes(nodes)
for edge in edges:
sk.add_edge(source=nodes[edge[0]], destination=nodes[edge[1]])
labeled_frames = []
node_count, _, frame_count = points_.shape
for i in range(frame_count):
new_inst = Instance(skeleton=sk)
for node_idx, node in enumerate(nodes):
x = points_[node_idx][0][i]
y = points_[node_idx][1][i]
new_inst[node] = Point(x, y)
if len(new_inst.points):
new_frame = LabeledFrame(video=vid, frame_idx=i)
new_frame.instances = (new_inst,)
labeled_frames.append(new_frame)
labels = Labels(labeled_frames=labeled_frames, videos=[vid], skeletons=[sk])
return labels
def read(
cls,
file: FileHandle,
video: Union[Video, str],
*args,
**kwargs,
) -> Labels:
connect_adj_nodes = False
if video is None:
raise ValueError(
"Cannot read analysis hdf5 if no video specified.")
if not isinstance(video, Video):
video = Video.from_filename(video)
f = file.file
tracks_matrix = f["tracks"][:].T
track_names_list = f["track_names"][:].T
node_names_list = f["node_names"][:].T
# shape: frames * nodes * 2 * tracks
frame_count, node_count, _, track_count = tracks_matrix.shape
tracks = [
Track(0, track_name.decode()) for track_name in track_names_list
]
skeleton = Skeleton()
last_node_name = None
for node_name in node_names_list:
node_name = node_name.decode()
skeleton.add_node(node_name)
if connect_adj_nodes and last_node_name:
skeleton.add_edge(last_node_name, node_name)
last_node_name = node_name
frames = []
for frame_idx in range(frame_count):
instances = []
for track_idx in range(track_count):
points = tracks_matrix[frame_idx, ..., track_idx]
if not np.all(np.isnan(points)):
point_scores = np.ones(len(points))
# make everything a PredictedInstance since the usual use
# case is to export predictions for analysis
instances.append(
PredictedInstance.from_arrays(
points=points,
point_confidences=point_scores,
skeleton=skeleton,
track=tracks[track_idx],
instance_score=1,
))
if instances:
frames.append(
LabeledFrame(video=video,
frame_idx=frame_idx,
instances=instances))
return Labels(labeled_frames=frames)
def read(
cls,
file: format.filehandle.FileHandle,
video_search: Union[Callable, List[Text], None] = None,
match_to: Optional[Labels] = None,
*args,
**kwargs,
):
f = file.file
labels = cls.read_headers(file, video_search, match_to)
frames_dset = f["frames"][:]
instances_dset = f["instances"][:]
points_dset = f["points"][:]
pred_points_dset = f["pred_points"][:]
# Shift the *non-predicted* points since these used to be saved with a gridline
# coordinate system.
if (file.format_id or 0) < 1.1:
points_dset[:]["x"] -= 0.5
points_dset[:]["y"] -= 0.5
# Rather than instantiate a bunch of Point\PredictedPoint objects, we will use
# inplace numpy recarrays. This will save a lot of time and memory when reading
# things in.
points = PointArray(buf=points_dset, shape=len(points_dset))
pred_points = PredictedPointArray(
buf=pred_points_dset, shape=len(pred_points_dset)
)
# Extend the tracks list with a None track. We will signify this with a -1 in
# the data which will map to last element of tracks
tracks = labels.tracks.copy()
tracks.extend([None])
# A dict to keep track of instances that have a from_predicted link. The key is
# the instance and the value is the index of the instance.
from_predicted_lookup = {}
# Create the instances
instances = []
for i in instances_dset:
track = tracks[i["track"]]
skeleton = labels.skeletons[i["skeleton"]]
if i["instance_type"] == 0: # Instance
instance = Instance(
skeleton=skeleton,
track=track,
points=points[i["point_id_start"] : i["point_id_end"]],
)
else: # PredictedInstance
instance = PredictedInstance(
skeleton=skeleton,
track=track,
points=pred_points[i["point_id_start"] : i["point_id_end"]],
score=i["score"],
)
instances.append(instance)
if i["from_predicted"] != -1:
from_predicted_lookup[instance] = i["from_predicted"]
# Make a second pass to add any from_predicted links
for instance, from_predicted_idx in from_predicted_lookup.items():
instance.from_predicted = instances[from_predicted_idx]
# Create the labeled frames
frames = [
LabeledFrame(
video=labels.videos[frame["video"]],
frame_idx=frame["frame_idx"],
instances=instances[
frame["instance_id_start"] : frame["instance_id_end"]
],
)
for i, frame in enumerate(frames_dset)
]
labels.labeled_frames = frames
# Do the stuff that should happen after we have labeled frames
labels.update_cache()
return labels
def read(
cls,
file: FileHandle,
img_dir: str,
use_missing_gui: bool = False,
*args,
**kwargs,
) -> Labels:
dicts = file.json
# Make skeletons from "categories"
skeleton_map = dict()
for category in dicts["categories"]:
skeleton = Skeleton(name=category["name"])
skeleton_id = category["id"]
node_names = category["keypoints"]
skeleton.add_nodes(node_names)
try:
for src_idx, dst_idx in category["skeleton"]:
skeleton.add_edge(node_names[src_idx], node_names[dst_idx])
except IndexError as e:
# According to the COCO data format specifications[^1], the edges
# are supposed to be 1-indexed. But in some of their own
# dataset the edges are 1-indexed! So we'll try.
# [1]: http://cocodataset.org/#format-data
# Clear any edges we already created using 0-indexing
skeleton.clear_edges()
# Add edges
for src_idx, dst_idx in category["skeleton"]:
skeleton.add_edge(node_names[src_idx - 1], node_names[dst_idx - 1])
skeleton_map[skeleton_id] = skeleton
# Make videos from "images"
# Remove images that aren't referenced in the annotations
img_refs = [annotation["image_id"] for annotation in dicts["annotations"]]
dicts["images"] = list(filter(lambda im: im["id"] in img_refs, dicts["images"]))
# Key in JSON file should be "file_name", but sometimes it's "filename",
# so we have to check both.
img_filename_key = "file_name"
if img_filename_key not in dicts["images"][0].keys():
img_filename_key = "filename"
# First add the img_dir to each image filename
img_paths = [
os.path.join(img_dir, image[img_filename_key]) for image in dicts["images"]
]
# See if there are any missing files
img_missing = [not os.path.exists(path) for path in img_paths]
if sum(img_missing):
if use_missing_gui:
okay = MissingFilesDialog(img_paths, img_missing).exec_()
if not okay:
return None
else:
raise FileNotFoundError(
f"Images for COCO dataset could not be found in {img_dir}."
)
# Update the image paths (with img_dir or user selected path)
for image, path in zip(dicts["images"], img_paths):
image[img_filename_key] = path
# Create the video objects for the image files
image_video_map = dict()
vid_id_video_map = dict()
for image in dicts["images"]:
image_id = image["id"]
image_filename = image[img_filename_key]
# Sometimes images have a vid_id which links multiple images
# together as one video. If so, we'll use that as the video key.
# But if there isn't a vid_id, we'll treat each images as a
# distinct video and use the image id as the video id.
vid_id = image.get("vid_id", image_id)
if vid_id not in vid_id_video_map:
kwargs = dict(filenames=[image_filename])
for key in ("width", "height"):
if key in image:
kwargs[key] = image[key]
video = Video.from_image_filenames(**kwargs)
vid_id_video_map[vid_id] = video
frame_idx = 0
else:
video = vid_id_video_map[vid_id]
frame_idx = video.num_frames
video.backend.filenames.append(image_filename)
image_video_map[image_id] = (video, frame_idx)
# Make instances from "annotations"
lf_map = dict()
track_map = dict()
for annotation in dicts["annotations"]:
skeleton = skeleton_map[annotation["category_id"]]
image_id = annotation["image_id"]
video, frame_idx = image_video_map[image_id]
keypoints = np.array(annotation["keypoints"], dtype="int").reshape(-1, 3)
track = None
if "track_id" in annotation:
track_id = annotation["track_id"]
if track_id not in track_map:
track_map[track_id] = Track(frame_idx, str(track_id))
track = track_map[track_id]
points = dict()
any_visible = False
for i in range(len(keypoints)):
node = skeleton.nodes[i]
x, y, flag = keypoints[i]
if flag == 0:
# node not labeled for this instance
continue
is_visible = flag == 2
any_visible = any_visible or is_visible
points[node] = Point(x, y, is_visible)
if points:
# If none of the points had 2 has the "visible" flag, we'll
# assume this incorrect and just mark all as visible.
if not any_visible:
for point in points.values():
point.visible = True
inst = Instance(skeleton=skeleton, points=points, track=track)
if image_id not in lf_map:
lf_map[image_id] = LabeledFrame(video, frame_idx)
lf_map[image_id].insert(0, inst)
return Labels(labeled_frames=list(lf_map.values()))
def read_frames(
cls,
file: FileHandle,
skeleton: Optional[Skeleton] = None,
full_video: Optional[Video] = None,
*args,
**kwargs,
) -> List[LabeledFrame]:
filename = file.filename
# Read CSV file.
data = pd.read_csv(filename, header=[1, 2])
# Check if this is in the new multi-animal format.
is_multianimal = data.columns[0][0] == "individuals"
if is_multianimal:
# Reload with additional header rows if using new format.
data = pd.read_csv(filename, header=[1, 2, 3])
# Pull out animal and node names from the columns.
animal_names = []
node_names = []
for animal_name, node_name, _ in data.columns[1:][::2]:
if animal_name not in animal_names:
animal_names.append(animal_name)
if node_name not in node_names:
node_names.append(node_name)
else:
# 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]]
if skeleton is None:
skeleton = Skeleton()
skeleton.add_nodes(node_names)
# Get list of all images filenames.
img_files = data.iloc[:, 0]
if full_video:
video = full_video
index_frames_by_original_index = True
else:
# Create the Video object
img_dir = os.path.dirname(filename)
video = cls.make_video_for_image_list(img_dir, img_files)
# The frames in the video we created will be indexed from 0 to N
# rather than having their index from the original source video.
index_frames_by_original_index = False
lfs = []
for i in range(len(data)):
# Figure out frame index to use.
if index_frames_by_original_index:
# Extract "0123" from "path/img0123.png" as original frame index.
frame_idx_match = re.search("(?<=img)(\\d+)(?=\\.png)", img_files[i])
if frame_idx_match is not None:
frame_idx = int(frame_idx_match.group(0))
else:
raise ValueError(
f"Unable to determine frame index for image {img_files[i]}"
)
else:
frame_idx = i
instances = []
if is_multianimal:
for animal_name in animal_names:
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],
)
instance_points[node] = Point(x, y)
if ~(np.isnan(x) and np.isnan(y)):
any_not_missing = True
if any_not_missing:
# Create instance with points.
instances.append(
Instance(skeleton=skeleton, points=instance_points)
)
else:
# Get points for each node.
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))
# Create LabeledFrame and add it to list.
lfs.append(
LabeledFrame(video=video, frame_idx=frame_idx, instances=instances)
)
return lfs
