File size: 2,719 Bytes
cd47a59 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import os
import subprocess
def main():
base_dir = r"D:\Users\roy\CV_final_project\champ-master"
transferd_dir = os.path.join(base_dir, "transferd_result")
if not os.path.exists(transferd_dir):
print(f"錯誤：找不到轉換結果資料夾 {transferd_dir}")
return
# 取得所有的影片資料夾名稱
video_names = [d for d in os.listdir(transferd_dir) if os.path.isdir(os.path.join(transferd_dir, d))]
video_names.sort()
print(f"總共找到 {len(video_names)} 支影片需要進行 DWPose 骨架特徵提取...")
for video_name in video_names:
input_normal_dir = os.path.join("transferd_result", video_name, "normal")
output_dwpose_dir = os.path.join("transferd_result", video_name, "dwpose")
# 檢查該影片的 normal 是否存在 (因為 generate_dwpose 是吃 normal 的影像)
full_input_dir = os.path.join(base_dir, input_normal_dir)
if not os.path.exists(full_input_dir):
print(f"[{video_name}] 找不到 normal 影像資料夾 (跳過): {input_normal_dir}")
continue
normal_imgs = [f for f in os.listdir(full_input_dir) if f.endswith('.png') or f.endswith('.jpg')]
if len(normal_imgs) == 0:
print(f"[{video_name}] normal 資料夾內沒有圖片 (跳過): {input_normal_dir}")
continue
# 斷點續傳檢查：檢查 dwpose 是否已完成
full_output_dir = os.path.join(base_dir, output_dwpose_dir)
if os.path.exists(full_output_dir):
dwpose_imgs = [f for f in os.listdir(full_output_dir) if f.endswith('.png') or f.endswith('.jpg')]
if len(dwpose_imgs) >= len(normal_imgs):
print(f"[{video_name}] 所有影格皆已完成 DWPose 提取，自動跳過...")
continue
elif len(dwpose_imgs) > 0:
print(f"[{video_name}] 發現未完成的進度，腳本將自動略過已存在的圖片繼續執行...")
print(f"\n========================================")
print(f"正在提取骨架特徵圖: {video_name}")
print(f"========================================")
cmd = [
"python", "-m", "scripts.data_processors.dwpose.generate_dwpose",
"--input", input_normal_dir,
"--output", output_dwpose_dir
]
result = subprocess.run(cmd, cwd=base_dir)
if result.returncode != 0:
print(f"[{video_name}] 執行 DWPose 時發生錯誤！")
else:
print(f"[{video_name}] DWPose 提取完成！\n")
if __name__ == "__main__":
main()
|