add removesuffix func (#3)

src/python/pose_format/bin/directory.py

@@ -5,14 +5,21 @@
 from tqdm import tqdm
 
 
+def removesuffix(text: str, suffix: str):
+	if text.endswith(suffix):
+		return text[:-len(suffix)]
+	else:
+		return text
+
+
 def find_missing_pose_files(directory: str):
     all_files = os.listdir(directory)
     mp4_files = [f for f in all_files if f.endswith(".mp4")]
-    pose_files = {os.path.splitext(f)[0] for f in all_files if f.endswith(".pose")}
+    pose_files = {removesuffix(f, ".pose") for f in all_files if f.endswith(".pose")}
     missing_pose_files = []
 
     for mp4_file in mp4_files:
-        base_name = os.path.splitext(mp4_file)[0]
+        base_name = removesuffix(mp4_file, ".mp4")
         if base_name not in pose_files:
             missing_pose_files.append(os.path.join(directory, mp4_file))
 
@@ -32,5 +39,5 @@ def main():
     missing_pose_files = find_missing_pose_files(args.directory)
 
     for mp4_path in tqdm(missing_pose_files):
-        pose_file_name = os.path.splitext(mp4_path)[0] + ".pose"
+        pose_file_name = removesuffix(mp4_path, ".mp4") + ".pose"
         pose_video(mp4_path, pose_file_name, 'mediapipe')