diff --git a/config.py b/config.py index 8a19625..6c6ab7e 100644 --- a/config.py +++ b/config.py @@ -5,9 +5,10 @@ # Colors for different models (B,G,R format) CONV_COLOR = (27, 80, 19) # 純粋な緑 -XGBOOST_COLOR = (0, 165, 255) # オレンジ -LIGHTGBM_COLOR = (255, 0, 255) # マゼンタ -EARSNET_COLOR = (139, 0, 0) # ダークブルー +XGBOOST_COLOR = (19, 80, 27) # オレンジ +LIGHTGBM_COLOR = (19, 80, 27) # マゼンタ +EARSNET_COLOR = (192, 79, 21) # ダークブルー +EARSNET_CROP_COLOR = (192, 79, 21) CATBOOST_COLOR = (0, 0, 255) # 赤 NGBOOST_COLOR = (255, 255, 0) # シアン diff --git a/main.py b/main.py index 7e57b89..1bb4624 100644 --- a/main.py +++ b/main.py @@ -40,6 +40,7 @@ EARSNET_COLOR = config.EARSNET_COLOR CATBOOST_COLOR = config.CATBOOST_COLOR NGBOOST_COLOR = config.NGBOOST_COLOR +EARSNET_CROP_COLOR = config.EARSNET_CROP_COLOR CONV_ENABLED = config.CONV_ENABLED XGBOOST_ENABLED = config.XGBOOST_ENABLED @@ -150,21 +151,72 @@ # Pillow-based drawing helpers ############################################################################### def pillow_draw_circle(draw, center, radius, fill=None, outline=None, width=1): + """PIL.ImageDraw.draw.ellipse 用のラッパ。center=(x,y)指定で円を描く。""" x, y = int(center[0]), int(center[1]) left_up = (x - radius, y - radius) right_down = (x + radius, y + radius) draw.ellipse([left_up, right_down], fill=fill, outline=outline, width=width) -def draw_glow_marker(draw, center, main_color, radius=5): - outer_radius = radius + 3 +### 変更点 1) カスタム描画用の関数を用意 +def draw_custom_circle(draw, center, style): + """ + style辞書により，丸の描画スタイルを変える。 + styleの例: + { + 'radius': 10, + 'fill_type': 'fill' or 'outline' or 'striped', + 'color': (r,g,b), + 'outline_color': (r,g,b), + 'outline_width': 2 + } + """ + radius = style.get("radius", 8) + fill_type = style.get("fill_type", "fill") # 'fill', 'outline', 'striped' + color = style.get("color", (255, 0, 0)) # 中心色 or 枠線色 + outline_color = style.get("outline_color", (255, 255, 255)) + outline_width = style.get("outline_width", 2) + x, y = int(center[0]), int(center[1]) - # 白枠 + outer_radius = radius + outline_width # 外枠込みの半径 + + # まず外枠用の円(太めの枠)を描く（fillはなし） + # fill=None なので枠線のみ pillow_draw_circle( - draw, (x, y), outer_radius, fill=None, outline=(255, 255, 255), width=2 + draw, + (x, y), + outer_radius, + fill=None, + outline=outline_color, + width=outline_width, ) - # 中心塗りつぶし - pillow_draw_circle(draw, (x, y), radius, fill=main_color) + + if fill_type == "fill": + # 中心を塗りつぶす + pillow_draw_circle(draw, (x, y), radius, fill=color, outline=None, width=0) + elif fill_type == "outline": + # 枠線のみ(本体は塗らない) + pillow_draw_circle(draw, (x, y), radius, fill=None, outline=color, width=2) + elif fill_type == "striped": + # 簡易的な縞々を実装 (円内に斜線を引く) + # 円のバウンディングボックスで少しピッチ小さめの線を引く + # fillなし下地に線だけ。 + pillow_draw_circle(draw, (x, y), radius, fill=None, outline=color, width=1) + # 縞々を引く + spacing = 4 # 縞の間隔 + left_up = (x - radius, y - radius) + right_down = (x + radius, y + radius) + + # PILで斜線を引き，円の外にはみ出す部分はあえてそのままに（簡易実装） + # 正確に円内のみ線を引く場合はマスクを使うなど工夫が必要 + lx, ly = left_up + rx, ry = right_down + for i in range(int(rx - lx)): + line_x1 = lx + i * spacing + line_y1 = ly + line_x2 = lx + i * spacing + (ry - ly) + line_y2 = ry + draw.line((line_x1, line_y1, line_x2, line_y2), fill=color, width=1) def pillow_draw_polygon(draw, vertices, outline=(0, 255, 0), width=2): @@ -195,8 +247,15 @@ if stethoscope_x is not None and stethoscope_y is not None: x, y = int(stethoscope_x), int(stethoscope_y) - # 光彩付きマーカー - draw_glow_marker(draw, (x, y), main_color=(255, 0, 0), radius=8) + # ここでは簡易的に fill_type='fill' の赤色マーカーを描画 + style_temp = { + "radius": 8, + "fill_type": "fill", + "color": (255, 0, 0), + "outline_color": (255, 255, 255), + "outline_width": 2, + } + draw_custom_circle(draw, (x, y), style_temp) out_img_rgb = np.array(pil_img) out_img_bgr = cv2.cvtColor(out_img_rgb, cv2.COLOR_RGB2BGR) @@ -567,7 +626,7 @@ timings["rtmpose_single"].append(rtmpose_time) # landmarks = [right_shoulder, left_shoulder, right_hip, left_hip] (例) - # ここは利用環境に合わせて入れ替えてください + # 環境に合わせてインデックスを調整する left_shoulder = (landmarks[0][1], landmarks[0][0]) right_shoulder = (landmarks[1][1], landmarks[1][0]) left_hip = (landmarks[2][1], landmarks[2][0]) @@ -632,7 +691,6 @@ earsnet_time = end_time_earsnet - start_time_earsnet timings["earsnet_single"].append(earsnet_time) timings["pipeline_earsnet"].append(earsnet_time) - else: earsnet_x, earsnet_y = 0, 0 @@ -654,7 +712,6 @@ pipeline_earsnet_cropped_time = rtmpose_time + earsnet_cropped_time timings["pipeline_earsnet_cropped"].append(pipeline_earsnet_cropped_time) - else: earsnet_cropped_x, earsnet_cropped_y = 0, 0 @@ -939,10 +996,8 @@ ・姿勢推定だけ描画した `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" if not os.path.exists(body_image_path): @@ -953,6 +1008,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" @@ -988,23 +1044,70 @@ exist_ok=True, ) - points = {key: [] for key in dirs.keys() if key not in ["marked", "combined"]} - - colors = { - "conv": CONV_COLOR, - "Xgboost": XGBOOST_COLOR, - "lightGBM": LIGHTGBM_COLOR, - "catboost": CATBOOST_COLOR, - "ngboost": NGBOOST_COLOR, - "earsnet": EARSNET_COLOR, - "earsnet_crop": (255, 51, 255), + # 各手法の描画スタイル設定 (色や半径，塗りつぶしの有無など) + ### 変更点 2) ここで各手法ごとの描画スタイルを指定 + marker_styles = { + # 例: conv手法は赤い塗りつぶし + "conv": { + "radius": 8, + "fill_type": "fill", # 'fill', 'outline', 'striped' + "color": CONV_COLOR, # config.CONV_COLOR を使う + "outline_color": CONV_COLOR, + "outline_width": 2, + }, + # 例: Xgboostは青の枠のみ + "Xgboost": { + "radius": 8, + "fill_type": "outline", + "color": XGBOOST_COLOR, + "outline_color": XGBOOST_COLOR, + "outline_width": 2, + }, + # 例: lightGBMは緑の縞々 + "lightGBM": { + "radius": 8, + "fill_type": "striped", + "color": LIGHTGBM_COLOR, + "outline_color": LIGHTGBM_COLOR, + "outline_width": 2, + }, + "earsnet": { + "radius": 8, + "fill_type": "fill", + "color": EARSNET_COLOR, + "outline_color": EARSNET_COLOR, + "outline_width": 2, + }, + "earsnet_crop": { + "radius": 8, + "fill_type": "outline", + "color": EARSNET_CROP_COLOR, # 例としてピンク色 + "outline_color": EARSNET_CROP_COLOR, + "outline_width": 2, + }, } - # Pose color = (33,95,154), Stetho color = (19,80,27) + # デフォルトスタイル(該当keyが無い場合用) + default_style = { + "radius": 8, + "fill_type": "fill", + "color": (255, 0, 0), + "outline_color": (255, 255, 255), + "outline_width": 2, + } + + # 各手法の過去フレーム座標(軌跡描画用) + points = {key: [] for key in dirs.keys() if key not in ["marked", "combined"]} + + # 色固定 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]) # 白枠 @@ -1014,6 +1117,7 @@ # 中心塗りつぶし pillow_draw_circle(draw, (x, y), radius, fill=main_color) + # CSV行ごとに画像を生成 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): @@ -1026,7 +1130,7 @@ pil_marked = Image.fromarray(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)) draw_marked = ImageDraw.Draw(pil_marked) - # 肩・腰・聴診器 + # 肩・腰・聴診器を一括で for point in [ "left_shoulder", "right_shoulder", @@ -1043,6 +1147,7 @@ and not pd.isna(row[col_y]) ): x, y = int(row[col_x]), int(row[col_y]) + # ここはシンプルに塗りつぶし丸などにする draw_glow_marker( draw_marked, (x, y), main_color=(255, 255, 0), radius=5 ) @@ -1105,6 +1210,7 @@ draw_with_traj = ImageDraw.Draw(pil_with_traj) draw_without_traj = ImageDraw.Draw(pil_without_traj) + # 各手法の聴診器推定座標を描画 for key in points.keys(): col_x = f"{key}_stethoscope_x" col_y = f"{key}_stethoscope_y" @@ -1116,16 +1222,20 @@ x, y = int(row[col_x]), int(row[col_y]) points[key].append((x, y)) - color = colors[key] if key in colors else (0, 0, 255) + style = marker_styles.get(key, default_style) # 個別 with trajectory indiv_with_traj_rgb = body_np_rgb.copy() pil_indiv_with = Image.fromarray(indiv_with_traj_rgb) draw_indiv_with = ImageDraw.Draw(pil_indiv_with) + # 軌跡 (過去フレーム分) if len(points[key]) > 1: - pillow_draw_polyline(draw_indiv_with, points[key], color=color, width=2) - draw_glow_marker(draw_indiv_with, (x, y), main_color=color, radius=8) + pillow_draw_polyline( + draw_indiv_with, points[key], color=style["color"], width=2 + ) + # 現フレームのマーカー + draw_custom_circle(draw_indiv_with, (x, y), style) indiv_with_traj_np = np.array(pil_indiv_with) indiv_with_traj_bgr = cv2.cvtColor(indiv_with_traj_np, cv2.COLOR_RGB2BGR) @@ -1138,7 +1248,8 @@ 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) - draw_glow_marker(draw_indiv_without, (x, y), main_color=color, radius=8) + + draw_custom_circle(draw_indiv_without, (x, y), style) indiv_without_traj_np = np.array(pil_indiv_without) indiv_without_traj_bgr = cv2.cvtColor( @@ -1151,11 +1262,13 @@ # combined with trajectory if len(points[key]) > 1: - pillow_draw_polyline(draw_with_traj, points[key], color=color, width=2) - draw_glow_marker(draw_with_traj, (x, y), main_color=color, radius=8) + pillow_draw_polyline( + draw_with_traj, points[key], color=style["color"], width=2 + ) + draw_custom_circle(draw_with_traj, (x, y), style) # combined without trajectory - draw_glow_marker(draw_without_traj, (x, y), main_color=color, radius=8) + draw_custom_circle(draw_without_traj, (x, y), style) cwt_np = np.array(pil_with_traj) cwt_bgr = cv2.cvtColor(cwt_np, cv2.COLOR_RGB2BGR)