File size: 2,526 Bytes
857c2e9 | 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 | import os
import shutil
from collections import defaultdict
from pprint import pprint
import matplotlib.pyplot as plt
vid_dir_list = [
"/Users/baizechen/Downloads/2025_10_18_1traj_sr15",
"/Users/baizechen/Downloads/2025_10_19_step29_sr22",
"/Users/baizechen/Downloads/2025_10_19_step59_sr17"
]
# for each vid_dir, split the videos in it into 2 parts
# one is name includes "success=True", the other is name includes "success=False"
# for vid_dir in vid_dir_list:
# for file in os.listdir(vid_dir):
# if "success=True" in file:
# success_True_dir = os.path.join(vid_dir, "success=True")
# if not os.path.exists(success_True_dir):
# os.makedirs(success_True_dir)
# shutil.move(os.path.join(vid_dir, file), os.path.join(success_True_dir, file))
# elif "success=False" in file:
# success_False_dir = os.path.join(vid_dir, "success=False")
# if not os.path.exists(success_False_dir):
# os.makedirs(success_False_dir)
# shutil.move(os.path.join(vid_dir, file), os.path.join(success_False_dir, file))
# for each video dir, after spliting into success=True and success=False, get the success rate for each task
task_success_rate_across_dirs = defaultdict(list)
for vid_dir in vid_dir_list:
success_dict = {}
total_dict = {}
for sub_dir in os.listdir(vid_dir):
if "success=True" in sub_dir:
success_flag = 1
elif "success=False" in sub_dir:
success_flag = 0
else:
continue
for video_file in os.listdir(os.path.join(vid_dir, sub_dir)):
task_name = video_file.split("task=")[1].split(".")[0]
success_dict[task_name] = success_dict.get(task_name, 0) + success_flag
total_dict[task_name] = total_dict.get(task_name, 0) + 1
for task_name in success_dict.keys():
success_rate = success_dict[task_name] / total_dict[task_name]
task_success_rate_across_dirs[task_name].append(success_rate)
# pprint(task_success_rate_across_dirs)
# plot the success rate across trials into sub-plots, each sub-plot is a task
fig, axes = plt.subplots(4, 3, figsize=(16, 12))
for i, task_name in enumerate(task_success_rate_across_dirs.keys()):
plt.subplot(4, 3, i + 1)
plt.plot(task_success_rate_across_dirs[task_name])
plt.title(task_name)
plt.xlabel("Trial")
plt.ylabel("Success Rate")
plt.tight_layout()
plt.savefig("success_rate_across_trials.png")
plt.close()
|