diff --git a/config.py b/config.py index 8a19625..83cfdf0 100644 --- a/config.py +++ b/config.py @@ -18,10 +18,10 @@ CATBOOST_ENABLED = False NGBOOST_ENABLED = False NORMALIZE_ENABLED = True -POSENET_ENABLED = False -RTMPOSE_ENABLED = True -MOBILENETV1SSD_ENABLED = False -YOLOX_ENABLED = True +POSENET_ENABLED = True +RTMPOSE_ENABLED = False +MOBILENETV1SSD_ENABLED = True +YOLOX_ENABLED = False EARSNET_CROP_ENABLED = True # Neural network model settings diff --git a/main.py b/main.py index 19e4293..4ef530c 100644 --- a/main.py +++ b/main.py @@ -132,13 +132,32 @@ # Pillow-based drawing helpers ############################################################################### def pillow_draw_circle(draw, center, radius, fill=None, outline=None, width=1): + """ + 円を描画するユーティリティ。 + fill で塗りつぶし色、outline + width で外枠を指定可能。 + """ x, y = int(center[0]), int(center[1]) left_up = (x - radius, y - radius) right_down = (x + radius, y + radius) - if fill is not None: - draw.ellipse([left_up, right_down], fill=fill, outline=outline, width=width) - else: - draw.ellipse([left_up, right_down], outline=outline, width=width) + draw.ellipse([left_up, right_down], fill=fill, outline=outline, width=width) + + +def draw_glow_marker(draw, center, main_color, radius=5): + """ + 光彩風に見えるように、やや大きめの枠 + 中心塗りつぶし を描画。 + main_color: (R, G, B) + radius: 中心の塗りつぶし半径 + """ + # 外側の枠（光彩用）：radius+3くらいにして薄い色 or 白枠など + outer_radius = radius + 3 + x, y = int(center[0]), int(center[1]) + + # 1) 白枠で大きめの円 + pillow_draw_circle( + draw, (x, y), outer_radius, fill=None, outline=(255, 255, 255), width=2 + ) + # 2) 中心を main_color で塗りつぶし + pillow_draw_circle(draw, (x, y), radius, fill=main_color, outline=None, width=0) def pillow_draw_polygon(draw, vertices, outline=(0, 255, 0), width=2): @@ -169,10 +188,8 @@ if stethoscope_x is not None and stethoscope_y is not None: x, y = int(stethoscope_x), int(stethoscope_y) - pillow_draw_circle(draw, (x, y), 10, fill=(255, 0, 0)) - pillow_draw_circle( - draw, (x, y), 12, fill=None, outline=(255, 255, 255), width=2 - ) + # 光彩付きマーカー + draw_glow_marker(draw, (x, y), main_color=(255, 0, 0), radius=8) out_img_rgb = np.array(pil_img) out_img_bgr = cv2.cvtColor(out_img_rgb, cv2.COLOR_RGB2BGR) @@ -402,12 +419,10 @@ rows = [] normalized_rows = [] - # YOLOX初期化 yolox_inferencer = None if YOLOX_ENABLED: yolox_inferencer = init_yolox() - # 時間計測用 dict (描画の時間は含まない) timings = { "rtmpose_single": [], "yolox_single": [], @@ -423,7 +438,7 @@ "pipeline_earsnet_cropped": [], } - # 各モデルの事前ロード + # モデルロード if LIGHTGBM_ENABLED: lgb_model_x = load_model("./models/LightGBM/stethoscope_calc_x_best_model.pkl") lgb_model_y = load_model("./models/LightGBM/stethoscope_calc_y_best_model.pkl") @@ -481,9 +496,7 @@ "stethoscope_y", ] - # ----------------------------- - # メインループ：推論 & 座標計算のみ - # ----------------------------- + # メインループ for image_file_name in png_files: image_path = os.path.join(base_dir, image_file_name) frame = cv2.imread(image_path) @@ -491,7 +504,7 @@ print(f"Failed to load image: {image_path}") continue - # (A) RTMPose or PoseNet + # (A) RTMPose rtmpose_time = 0.0 if RTMPOSE_ENABLED: start_time_rtmpose = time.time() @@ -502,7 +515,6 @@ bboxes = np.concatenate( (pred_instance.bboxes, pred_instance.scores[:, None]), axis=1 ) - # 人物のみ bboxes = bboxes[ np.logical_and(pred_instance.labels == 0, pred_instance.scores > 0.3) ] @@ -521,7 +533,6 @@ processed_frames += 1 continue - # PoseOverlay(可視化) → 時間計測には含めない if visualizer is not None: visualizer.add_datasample( "result", @@ -554,10 +565,10 @@ rtmpose_time = end_time_rtmpose - start_time_rtmpose timings["rtmpose_single"].append(rtmpose_time) - left_shoulder = landmarks[0] - right_shoulder = landmarks[1] - left_hip = landmarks[2] - right_hip = landmarks[3] + left_shoulder = landmarks[1] + right_shoulder = landmarks[0] + left_hip = landmarks[3] + right_hip = landmarks[2] cv2.imwrite( os.path.join(pose_overlay_dir, image_file_name), pose_overlay_img @@ -691,7 +702,7 @@ "stethoscope_y": normalized_points[4, 1], } - # --- EARSNet の Normalized + # EARSNet の Normalized if EARSNET_ENABLED: stetho_point_earsnet = np.array( [ @@ -718,7 +729,7 @@ normalized_row["earsnet_crop_stethoscope_x"] = norm_earsnet_crop[4, 0] normalized_row["earsnet_crop_stethoscope_y"] = norm_earsnet_crop[4, 1] - # --- Conv (Affine) + # Conv (Affine) if RTMPOSE_ENABLED and YOLOX_ENABLED and CONV_ENABLED: source_pts = np.array( [ @@ -736,37 +747,32 @@ [float(row["stethoscope_x"]), float(row["stethoscope_y"])] ) calc_x, calc_y = calc_position.calc_affine(source_pts, *stetho_pt) - row["conv_stethoscope_x"] = calc_x row["conv_stethoscope_y"] = calc_y - # --- XGBoost + # XGBoost if RTMPOSE_ENABLED and YOLOX_ENABLED and XGBOOST_ENABLED: if NORMALIZE_ENABLED: input_data_xg = pd.DataFrame([normalized_row]) else: input_data_xg = pd.DataFrame([row]) - X_scaled_x = xg_scaler_x.transform(input_data_xg[input_columns]) x_pred = xg_model_x.predict(X_scaled_x)[0] X_scaled_y = xg_scaler_y.transform(input_data_xg[input_columns]) y_pred = xg_model_y.predict(X_scaled_y)[0] - row["Xgboost_stethoscope_x"] = x_pred row["Xgboost_stethoscope_y"] = y_pred - # --- LightGBM + # LightGBM if RTMPOSE_ENABLED and YOLOX_ENABLED and LIGHTGBM_ENABLED: if NORMALIZE_ENABLED: input_data_lgb = pd.DataFrame([normalized_row]) else: input_data_lgb = pd.DataFrame([row]) - X_scaled_x = lgb_scaler_x.transform(input_data_lgb[input_columns]) lgb_x_pred = lgb_model_x.predict(X_scaled_x)[0] X_scaled_y = lgb_scaler_y.transform(input_data_lgb[input_columns]) lgb_y_pred = lgb_model_y.predict(X_scaled_y)[0] - row["lightGBM_stethoscope_x"] = lgb_x_pred row["lightGBM_stethoscope_y"] = lgb_y_pred @@ -775,7 +781,7 @@ processed_frames += 1 - # (G) CSV書き込み + # CSV書き込み if rows: fieldnames = list(rows[0].keys()) csvfile_path = os.path.join(results_dir, "results.csv") @@ -803,7 +809,7 @@ else: print("No data to write to CSV.") - # (H) FPS計算結果をCSV保存 + # FPS計測結果をCSV保存 fps_data = [] for method_name, time_list in timings.items(): if not time_list: @@ -858,6 +864,9 @@ ・姿勢推定だけ描画した `marked_pose_images` ・聴診器検出だけ描画した `marked_stethoscope_images` も生成＆動画化する。 + + 姿勢推定の色 = (33,95,154), 聴診器検出の色 = (19,80,27) + 各マーカーには光彩風の枠をつけて視認性を上げる。 """ df = pd.read_csv(csv_path) body_image_path = "./images/body/BodyF.png" @@ -869,7 +878,6 @@ 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" @@ -887,7 +895,7 @@ dirs["earsnet_crop"] = "earsnet_crop" dirs["combined"] = "combined" - # ★追加: 姿勢推定だけ描画するフォルダ & 聴診器だけ描画するフォルダ + # 追加フォルダ pose_only_dir = "marked_pose_images" stetho_only_dir = "marked_stethoscope_images" @@ -905,7 +913,6 @@ exist_ok=True, ) - # 各メソッドの軌跡用 points = {key: [] for key in dirs.keys() if key not in ["marked", "combined"]} colors = { @@ -918,10 +925,20 @@ "earsnet_crop": (255, 51, 255), } - # ★ 姿勢推定だけの色 (RGB) → (19,80,27) - pose_color_rgb = (19, 80, 27) - # ★ 聴診器検出だけの色 (RGB) → (33,95,154) - stetho_color_rgb = (33, 95, 154) + # Pose color = (33,95,154), Stetho color = (19,80,27) + pose_color_rgb = (33, 95, 154) + stetho_color_rgb = (19, 80, 27) + + def draw_glow_marker(draw, center, main_color, radius=5): + # 光彩用外枠を白などにして少し大きめ + outer_radius = radius + 3 + x, y = int(center[0]), int(center[1]) + # 白枠 + pillow_draw_circle( + draw, (x, y), outer_radius, fill=None, outline=(255, 255, 255), width=2 + ) + # 中心塗りつぶし + pillow_draw_circle(draw, (x, y), radius, fill=main_color) for _, row in df.iterrows(): original_image_path = os.path.join(original_images_dir, row["image_file_name"]) @@ -931,13 +948,9 @@ if original_image is None: continue - ############################################################################ - # (1) もともとの marked_images: 肩・腰・聴診器をすべて描画 - ############################################################################ + # 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", @@ -954,21 +967,20 @@ 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)) + # 既存デフォルト色 (255,255,0) → 小さいマーカー + draw_glow_marker( + draw_marked, (x, y), main_color=(255, 255, 0), radius=5 + ) marked_rgb = np.array(pil_marked) marked_bgr = cv2.cvtColor(marked_rgb, cv2.COLOR_RGB2BGR) marked_dir = os.path.join(results_dir, "marked_images") cv2.imwrite(os.path.join(marked_dir, row["image_file_name"]), marked_bgr) - ############################################################################ - # (2) 姿勢推定だけ描画 → marked_pose_images - ############################################################################ + # 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" @@ -979,43 +991,38 @@ 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) + # 光彩付き, main_color = (33,95,154) + draw_glow_marker( + draw_pose, (x, y), main_color=pose_color_rgb, radius=15 + ) 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 - ############################################################################ + # 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]) + "stethoscope_x" in row + and "stethoscope_y" in row + and not pd.isna(row["stethoscope_x"]) + and not pd.isna(row["stethoscope_y"]) ): - 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) + sx, sy = int(row["stethoscope_x"]), int(row["stethoscope_y"]) + if sx != 0 or sy != 0: # 0なら未検出とみなす + draw_glow_marker( + draw_stetho, (sx, sy), main_color=stetho_color_rgb, radius=15 + ) 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 など - ############################################################################ - + # 4) combined系（従来） combined_image_with_traj_rgb = body_np_rgb.copy() combined_image_without_traj_rgb = body_np_rgb.copy() @@ -1024,7 +1031,6 @@ draw_with_traj = ImageDraw.Draw(pil_with_traj) draw_without_traj = ImageDraw.Draw(pil_without_traj) - # 各メソッドの推定結果(earsnet, conv, xgboost, lightGBM等)を取得 for key in points.keys(): col_x = f"{key}_stethoscope_x" col_y = f"{key}_stethoscope_y" @@ -1045,7 +1051,7 @@ if len(points[key]) > 1: pillow_draw_polyline(draw_indiv_with, points[key], color=color, width=2) - pillow_draw_circle(draw_indiv_with, (x, y), 10, fill=color) + draw_glow_marker(draw_indiv_with, (x, y), main_color=color, radius=8) indiv_with_traj_np = np.array(pil_indiv_with) indiv_with_traj_bgr = cv2.cvtColor(indiv_with_traj_np, cv2.COLOR_RGB2BGR) @@ -1058,7 +1064,7 @@ indiv_without_traj_rgb = body_np_rgb.copy() pil_indiv_without = Image.fromarray(indiv_without_traj_rgb) draw_indiv_without = ImageDraw.Draw(pil_indiv_without) - pillow_draw_circle(draw_indiv_without, (x, y), 10, fill=color) + draw_glow_marker(draw_indiv_without, (x, y), main_color=color, radius=8) indiv_without_traj_np = np.array(pil_indiv_without) indiv_without_traj_bgr = cv2.cvtColor( @@ -1072,10 +1078,10 @@ # (C) combined with trajectory if len(points[key]) > 1: pillow_draw_polyline(draw_with_traj, points[key], color=color, width=2) - pillow_draw_circle(draw_with_traj, (x, y), 10, fill=color) + draw_glow_marker(draw_with_traj, (x, y), main_color=color, radius=8) # (D) combined without trajectory - pillow_draw_circle(draw_without_traj, (x, y), 10, fill=color) + draw_glow_marker(draw_without_traj, (x, y), main_color=color, radius=8) cwt_np = np.array(pil_with_traj) cwt_bgr = cv2.cvtColor(cwt_np, cv2.COLOR_RGB2BGR) @@ -1095,7 +1101,7 @@ 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"), @@ -1165,7 +1171,6 @@ help="Directory to save output images and results", ) - # RTMpose用 det_config = "modules/rtmpose/mmdetection_cfg/rtmdet_m_640-8xb32_coco-person.py" det_checkpoint = ( "models/rtmpose/rtmdet_m_8xb32-100e_coco-obj365-person-235e8209.pth" @@ -1200,7 +1205,6 @@ ) process_images(args, detector, pose_estimator, visualizer) else: - # PoseNet or no keypoints usage process_images(args, None, None, None) # (4) スレッド終了