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