| import os |
| import time |
| import pickle |
| import glob |
| import folder_paths |
|
|
|
|
| def _ensure_output_dir(): |
| out_dir = folder_paths.get_output_directory() |
| os.makedirs(out_dir, exist_ok=True) |
| return out_dir |
|
|
|
|
| |
| |
| |
| def _list_all_pkl_under_input(): |
| inp = folder_paths.get_input_directory() |
| exts = (".pkl", ".pickle", ".pt") |
|
|
| files = [] |
| for ext in exts: |
| pattern = os.path.join(inp, "**", f"*{ext}") |
| files.extend(glob.glob(pattern, recursive=True)) |
|
|
| rel = [] |
| for f in files: |
| if os.path.isfile(f): |
| r = os.path.relpath(f, inp).replace("\\", "/") |
| rel.append(r) |
|
|
| rel = sorted(set(rel)) |
| return rel if rel else [""] |
|
|
|
|
| def _abs_from_input(rel_path: str) -> str: |
| inp = folder_paths.get_input_directory() |
| return os.path.join(inp, rel_path).replace("\\", "/") |
|
|
|
|
| def _make_unique_path(base_path: str) -> str: |
| """ |
| If file exists, append incremental suffix: |
| pose_data.pkl |
| pose_data_0001.pkl |
| pose_data_0002.pkl |
| """ |
| if not os.path.exists(base_path): |
| return base_path |
|
|
| directory = os.path.dirname(base_path) |
| name = os.path.basename(base_path) |
| base, ext = os.path.splitext(name) |
|
|
| idx = 1 |
| while True: |
| new_name = f"{base}_{idx:04d}{ext}" |
| new_path = os.path.join(directory, new_name) |
| if not os.path.exists(new_path): |
| return new_path |
| idx += 1 |
|
|
|
|
| def _default_filename(prefix: str, ext: str): |
| ts = time.strftime("%Y%m%d_%H%M%S") |
| return f"{prefix}_{ts}{ext}" |
|
|
|
|
| class TSSavePoseDataAsPickle: |
| OUTPUT_NODE = True |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "pose_data": ("POSEDATA",), |
| "filename": ("STRING", {"default": ""}), |
| } |
| } |
|
|
| RETURN_TYPES = ("STRING",) |
| RETURN_NAMES = ("saved_path",) |
| FUNCTION = "save" |
| CATEGORY = "save" |
|
|
| def save(self, pose_data, filename=""): |
| out_dir = _ensure_output_dir() |
| filename = (filename or "").strip() |
| if not filename: |
| filename = _default_filename("pose_data", ".pkl") |
| if not filename.lower().endswith((".pkl", ".pickle")): |
| filename += ".pkl" |
|
|
| abs_path = _make_unique_path(os.path.join(out_dir, filename)) |
|
|
| with open(abs_path, "wb") as f: |
| pickle.dump(pose_data, f, protocol=pickle.HIGHEST_PROTOCOL) |
|
|
| return (abs_path,) |
|
|
|
|
| class TSLoadPoseDataPickle: |
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| |
| "file": (_list_all_pkl_under_input(),), |
| } |
| } |
|
|
| RETURN_TYPES = ("POSEDATA",) |
| RETURN_NAMES = ("pose_data",) |
| FUNCTION = "load" |
| CATEGORY = "load" |
|
|
| def load(self, file): |
| if not isinstance(file, str) or not file.strip(): |
| raise ValueError("TS PoseData Pickle: Please select a .pkl/.pt file.") |
|
|
| abs_path = _abs_from_input(file) |
| if not os.path.isfile(abs_path): |
| raise ValueError(f"TS PoseData Pickle: File not found: {abs_path}") |
|
|
| with open(abs_path, "rb") as f: |
| pose_data = pickle.load(f) |
|
|
| return (pose_data,) |
|
|
