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 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
