diff --git a/main.py b/main.py index fbfa801..b030988 100644 --- a/main.py +++ b/main.py @@ -214,7 +214,7 @@ def generate_visualizations(csv_path, original_images_dir, output_dir): df = pd.read_csv(csv_path) - body_image = cv2.imread(os.path.join(output_dir, "images", "body", "BodyF.png")) + body_image = cv2.imread("./images/body/BodyF.png") # オリジナルのパスを使用 results_dir = os.path.join(output_dir, "results") os.makedirs(results_dir, exist_ok=True) @@ -226,11 +226,12 @@ if LIGHTGBM_ENABLED: dirs["lightGBM"] = "lightGBM" - # Create all necessary directories - for dir_name in dirs.values(): - os.makedirs(os.path.join(results_dir, dir_name), exist_ok=True) - os.makedirs(os.path.join(results_dir, f"{dir_name}_with_trajectory"), exist_ok=True) - os.makedirs(os.path.join(results_dir, f"{dir_name}_without_trajectory"), exist_ok=True) + # 必要なディレクトリを作成 + os.makedirs(os.path.join(results_dir, "marked_images"), exist_ok=True) + for key in dirs: + if key != "marked": + os.makedirs(os.path.join(results_dir, f"{dirs[key]}_with_trajectory"), exist_ok=True) + os.makedirs(os.path.join(results_dir, f"{dirs[key]}_without_trajectory"), exist_ok=True) points = {key: [] for key in dirs.keys() if key != "marked"} colors = {"conv": CONV_COLOR, "Xgboost": XGBOOST_COLOR, "lightGBM": LIGHTGBM_COLOR} @@ -245,7 +246,7 @@ for point in ["left_shoulder", "right_shoulder", "left_hip", "right_hip", "stethoscope"]: cv2.circle(original_image, (int(row[f"{point}_x"]), int(row[f"{point}_y"])), 10, (255, 255, 0), -1) - cv2.imwrite(os.path.join(results_dir, dirs["marked"], row["image_file_name"]), original_image) + cv2.imwrite(os.path.join(results_dir, "marked_images", row["image_file_name"]), original_image) for key in points: x, y = int(row[f"{key}_stethoscope_x"]), int(row[f"{key}_stethoscope_y"]) @@ -269,6 +270,8 @@ ) # Generate videos + create_video_from_images(os.path.join(results_dir, "marked_images"), os.path.join(results_dir, "marked_video.mp4")) + for key in dirs: if key != "marked": create_video_from_images( @@ -280,9 +283,6 @@ os.path.join(results_dir, f"{key}_video_without_trajectory.mp4"), ) - # Generate video for marked images - create_video_from_images(os.path.join(results_dir, dirs["marked"]), os.path.join(results_dir, "marked_video.mp4")) - def create_video_from_images(image_dir, output_path): images = sorted( @@ -316,15 +316,6 @@ # Create output directory structure os.makedirs(args.output_dir, exist_ok=True) - os.makedirs(os.path.join(args.output_dir, "images", "body"), exist_ok=True) - - # Copy BodyF.png to the output directory - body_image_src = "./images/body/BodyF.png" - body_image_dst = os.path.join(args.output_dir, "images", "body", "BodyF.png") - if not os.path.exists(body_image_dst): - import shutil - - shutil.copy(body_image_src, body_image_dst) frames_dir = os.path.join(args.output_dir, "frames") video_to_frames(args.video_path, frames_dir)