diff --git a/main.py b/main.py index 69fe303..19e4293 100644 --- a/main.py +++ b/main.py @@ -720,8 +720,6 @@ # --- Conv (Affine) if RTMPOSE_ENABLED and YOLOX_ENABLED and CONV_ENABLED: - # すでに start_conv, end_conv は timings計測用のみ - # → ここでピクセル座標を row に書き込む source_pts = np.array( [ [float(row[f"{pos}_x"]), float(row[f"{pos}_y"])] @@ -739,7 +737,6 @@ ) calc_x, calc_y = calc_position.calc_affine(source_pts, *stetho_pt) - # conv_stethoscope_x/y を row に書き込み row["conv_stethoscope_x"] = calc_x row["conv_stethoscope_y"] = calc_y @@ -773,7 +770,6 @@ row["lightGBM_stethoscope_x"] = lgb_x_pred row["lightGBM_stethoscope_y"] = lgb_y_pred - # row, normalized_row を最終リストに追加 rows.append(row) normalized_rows.append(normalized_row) @@ -803,9 +799,7 @@ print(f"Processed and saved results to: {csvfile_path}") print(f"Processed and saved normalized results to: {normfile_path}") - # ★★★推論結果が出そろった後で描画 (描画時間はFPSに含めない)★★★ generate_visualizations(csvfile_path, base_dir, results_dir) - else: print("No data to write to CSV.") @@ -860,8 +854,10 @@ CSVを読み込み、BodyF.png上や元フレーム上に各手法の結果を描画 → 動画化。 描画時間はFPSに含めず、ここでまとめて行う。 - 手法が True になっているものについては、 - '_with_trajectory' と '_without_trajectory' の両動画を生成。 + もともとの marked_images（姿勢＋聴診器）に加え、 + ・姿勢推定だけ描画した `marked_pose_images` + ・聴診器検出だけ描画した `marked_stethoscope_images` + も生成＆動画化する。 """ df = pd.read_csv(csv_path) body_image_path = "./images/body/BodyF.png" @@ -873,7 +869,7 @@ body_img_pil = Image.open(body_image_path).convert("RGB") body_np_rgb = np.array(body_img_pil) # RGB順 - # 結果格納用ディレクトリの準備 + # 既存と同じのを用意 dirs = {"marked": "marked_images"} if CONV_ENABLED: dirs["conv"] = "conv" @@ -889,11 +885,16 @@ dirs["earsnet"] = "earsnet" if EARSNET_CROP_ENABLED: dirs["earsnet_crop"] = "earsnet_crop" - - # 常に combined も作成 dirs["combined"] = "combined" + # ★追加: 姿勢推定だけ描画するフォルダ & 聴診器だけ描画するフォルダ + pose_only_dir = "marked_pose_images" + stetho_only_dir = "marked_stethoscope_images" + os.makedirs(os.path.join(results_dir, "marked_images"), exist_ok=True) + os.makedirs(os.path.join(results_dir, pose_only_dir), exist_ok=True) + os.makedirs(os.path.join(results_dir, stetho_only_dir), exist_ok=True) + for key in dirs: if key != "marked": os.makedirs( @@ -904,10 +905,9 @@ exist_ok=True, ) - # ★ポイントを各メソッド別に保持(1動画につき1回リセット) + # 各メソッドの軌跡用 points = {key: [] for key in dirs.keys() if key not in ["marked", "combined"]} - # 色設定 colors = { "conv": CONV_COLOR, "Xgboost": XGBOOST_COLOR, @@ -918,7 +918,11 @@ "earsnet_crop": (255, 51, 255), } - # (1) 各フレームにマーキング + BodyF.pngへ軌跡描画 + # ★ 姿勢推定だけの色 (RGB) → (19,80,27) + pose_color_rgb = (19, 80, 27) + # ★ 聴診器検出だけの色 (RGB) → (33,95,154) + stetho_color_rgb = (33, 95, 154) + for _, row in df.iterrows(): original_image_path = os.path.join(original_images_dir, row["image_file_name"]) if not os.path.exists(original_image_path): @@ -927,9 +931,13 @@ if original_image is None: continue - # ---- (1-A) 各フレームへのマーキング (肩・腰・聴診器) + ############################################################################ + # (1) もともとの marked_images: 肩・腰・聴診器をすべて描画 + ############################################################################ pil_marked = Image.fromarray(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)) draw_marked = ImageDraw.Draw(pil_marked) + + # 肩・腰・聴診器などを描画 (既存処理) for point in [ "left_shoulder", "right_shoulder", @@ -946,6 +954,7 @@ and not pd.isna(row[col_y]) ): x, y = int(row[col_x]), int(row[col_y]) + # ★ ここは既存カラー: (255,255,0) など → そのまま pillow_draw_circle(draw_marked, (x, y), 5, fill=(255, 255, 0)) marked_rgb = np.array(pil_marked) @@ -953,7 +962,60 @@ marked_dir = os.path.join(results_dir, "marked_images") cv2.imwrite(os.path.join(marked_dir, row["image_file_name"]), marked_bgr) - # ---- (1-B) BodyF.png へ各メソッドの軌跡を描画 + ############################################################################ + # (2) 姿勢推定だけ描画 → marked_pose_images + ############################################################################ + pil_pose = Image.fromarray(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)) + draw_pose = ImageDraw.Draw(pil_pose) + + # 描画対象: left_shoulder, right_shoulder, left_hip, right_hip + for pose_point in ["left_shoulder", "right_shoulder", "left_hip", "right_hip"]: + col_x = f"{pose_point}_x" + col_y = f"{pose_point}_y" + if ( + col_x in row + and col_y in row + and not pd.isna(row[col_x]) + and not pd.isna(row[col_y]) + ): + x, y = int(row[col_x]), int(row[col_y]) + # ★指定のカラー = RGB(19,80,27) + pillow_draw_circle(draw_pose, (x, y), 5, fill=pose_color_rgb) + + pose_rgb = np.array(pil_pose) + pose_bgr = cv2.cvtColor(pose_rgb, cv2.COLOR_RGB2BGR) + pose_dir_path = os.path.join(results_dir, pose_only_dir) + cv2.imwrite(os.path.join(pose_dir_path, row["image_file_name"]), pose_bgr) + + ############################################################################ + # (3) 聴診器検出だけ描画 → marked_stethoscope_images + ############################################################################ + pil_stetho = Image.fromarray(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)) + draw_stetho = ImageDraw.Draw(pil_stetho) + + # 描画対象: stethoscope_x, stethoscope_y + stx_col = "stethoscope_x" + sty_col = "stethoscope_y" + if ( + stx_col in row + and sty_col in row + and not pd.isna(row[stx_col]) + and not pd.isna(row[sty_col]) + ): + sx, sy = int(row[stx_col]), int(row[sty_col]) + if sx != 0 or sy != 0: # 0の場合は検出されなかったとみなす + # ★指定のカラー = RGB(33,95,154) + pillow_draw_circle(draw_stetho, (sx, sy), 5, fill=stetho_color_rgb) + + stetho_rgb = np.array(pil_stetho) + stetho_bgr = cv2.cvtColor(stetho_rgb, cv2.COLOR_RGB2BGR) + stetho_dir_path = os.path.join(results_dir, stetho_only_dir) + cv2.imwrite(os.path.join(stetho_dir_path, row["image_file_name"]), stetho_bgr) + + ############################################################################ + # (4) 以下は既存処理: combined_with_trajectory / combined_without_trajectory など + ############################################################################ + combined_image_with_traj_rgb = body_np_rgb.copy() combined_image_without_traj_rgb = body_np_rgb.copy() @@ -1015,7 +1077,6 @@ # (D) combined without trajectory pillow_draw_circle(draw_without_traj, (x, y), 10, fill=color) - # 結果 (pil_with_traj / pil_without_traj) を BGR に変換して保存 cwt_np = np.array(pil_with_traj) cwt_bgr = cv2.cvtColor(cwt_np, cv2.COLOR_RGB2BGR) cwd = os.path.join(results_dir, "combined_with_trajectory") @@ -1034,6 +1095,16 @@ os.path.join(results_dir, "marked_video.mp4"), ) + # ★ 新たに姿勢だけの動画, 聴診器だけの動画 も作る + create_video_from_images( + os.path.join(results_dir, pose_only_dir), + os.path.join(results_dir, "pose_video.mp4"), + ) + create_video_from_images( + os.path.join(results_dir, stetho_only_dir), + os.path.join(results_dir, "stethoscope_video.mp4"), + ) + for key in dirs: if key not in ["marked", "combined"]: create_video_from_images(