Spaces:
Running
Running
| from collections import defaultdict | |
| from email.policy import default | |
| import imgui, os | |
| from fairmotion.data import bvh | |
| from IPython import embed | |
| # TODO: list multiple motions (currently, temp fix; viewer.m_id = viewer.m_ids[0]) | |
| def update_motion_by_mid(viewer, mid, ignore_root_skel, ee_as_joint): | |
| viewer.m_id = viewer.m_ids[0] | |
| if mid != viewer.m_id: | |
| viewer.m_id = mid | |
| viewer.worldStepper.reset() | |
| if viewer.all_motions[viewer.m_id] is None: | |
| viewer.all_motions[viewer.m_id] = bvh.load( | |
| viewer.file_names[viewer.m_id], | |
| ignore_root_skel=ignore_root_skel, | |
| ee_as_joint=ee_as_joint, | |
| ) | |
| new_motion = viewer.all_motions[viewer.m_id] | |
| viewer.update_motions([new_motion]) | |
| print(viewer.file_names[viewer.m_id][:-4], "\t", new_motion.num_frames() - 1) | |
| def list_motion_widget(viewer, ignore_root_skel, ee_as_joint): | |
| # motion list | |
| viewer.m_id = viewer.m_ids[0] | |
| if imgui.button(" < "): | |
| new_mid = (viewer.m_id + 1) % len(viewer.all_motions) | |
| update_motion_by_mid(viewer, new_mid, ignore_root_skel, ee_as_joint) | |
| imgui.same_line() | |
| _, new_mid = imgui.combo("", viewer.m_id, viewer.file_names) | |
| update_motion_by_mid(viewer, new_mid, ignore_root_skel, ee_as_joint) | |
| imgui.same_line() | |
| imgui.text(f"{viewer.m_id}/{len(viewer.file_names)}") | |
| imgui.same_line() | |
| if imgui.button(" > "): | |
| new_mid = (viewer.m_id + 1) % len(viewer.all_motions) | |
| update_motion_by_mid(viewer, new_mid, ignore_root_skel, ee_as_joint) | |
| def check_setup(viewer, save_dir="", loaded_check=None): | |
| viewer.all_check_button_check = True | |
| if loaded_check is None: | |
| viewer.checkbox_enabled = [False] * len(viewer.file_names) | |
| else: | |
| assert len(loaded_check) == len(viewer.file_names) | |
| viewer.checkbox_enabled = loaded_check | |
| viewer.save_dir = save_dir | |
| def check_widget(viewer, ee_as_joint): | |
| # select / deselect all (contact solved) | |
| if viewer.all_check_button_check: | |
| if imgui.button("select"): | |
| viewer.all_check_button_check = False | |
| for i in range(len(viewer.checkbox_enabled)): | |
| viewer.checkbox_enabled[i] = True | |
| else: | |
| if imgui.button("deselect"): | |
| viewer.all_check_button_check = True | |
| for i in range(len(viewer.checkbox_enabled)): | |
| viewer.checkbox_enabled[i] = False | |
| # check if contact ok (solved) | |
| for i, fn_i in enumerate(viewer.file_names): | |
| _, viewer.checkbox_enabled[i] = imgui.checkbox( | |
| os.path.basename(fn_i), viewer.checkbox_enabled[i] | |
| ) | |
| changed, text_val = imgui.input_text("save dir", viewer.save_dir, 60) | |
| if changed: | |
| viewer.save_dir = text_val | |
| if imgui.button(" save motions "): | |
| if not os.path.exists(viewer.save_dir): | |
| os.makedirs(viewer.save_dir) | |
| try: | |
| for m_id in range(len(viewer.all_motions)): | |
| if viewer.checkbox_enabled[m_id] and ( | |
| viewer.all_motions[m_id] is not None | |
| ): | |
| file_path = os.path.join(viewer.save_dir, viewer.file_names[m_id]) | |
| bvh.save( | |
| viewer.all_motions[m_id], | |
| file_path, | |
| rot_order="ZXY", | |
| verbose=True, | |
| ee_as_joint=ee_as_joint, | |
| ) | |
| except: | |
| print("Err while saving") | |
| embed() | |
| def tag_setup(viewer, save_dir): | |
| viewer.save_dir = save_dir | |
| viewer.cur_tag = "" | |
| viewer.tag_dicts = defaultdict(set) | |
| def tag_widget(viewer): | |
| _, viewer.cur_tag = imgui.input_text("tag", viewer.cur_tag, 60) | |
| viewer.m_id = viewer.m_ids[0] | |
| if imgui.button("mark tag"): | |
| viewer.tag_dicts[viewer.cur_tag].add(viewer.file_names[viewer.m_id]) | |
| if imgui.button("remove tag"): | |
| viewer.tag_dicts[viewer.cur_tag].remove(viewer.file_names[viewer.m_id]) | |
| imgui.separator() | |
| _, viewer.save_dir = imgui.input_text("save dir", viewer.save_dir, 60) | |
| if imgui.button(" save tags "): | |
| if not os.path.exists(viewer.save_dir): | |
| os.makedirs(viewer.save_dir) | |
| try: | |
| for tag_key in viewer.tag_dicts: | |
| file_path = os.path.join(viewer.save_dir, tag_key + ".txt") | |
| print(file_path) | |
| with open(file_path, "w") as file: | |
| for item in viewer.tag_dicts[tag_key]: | |
| file.write(item + "\n") | |
| except: | |
| print("Err while saving") | |
| embed() | |
| imgui.separator() | |
| for tag_key in viewer.tag_dicts: | |
| if imgui.tree_node(tag_key, imgui.TREE_NODE_DEFAULT_OPEN): | |
| for elem in viewer.tag_dicts[tag_key]: | |
| imgui.text(os.path.basename(elem)) | |
| imgui.tree_pop() | |