diff --git a/.gitignore b/.gitignore index 649f51c..fcd616c 100644 --- a/.gitignore +++ b/.gitignore @@ -177,3 +177,7 @@ !modules/posenet/models output/ data/ +*.mp4 +frames/ +*.xml +*.onnx \ No newline at end of file diff --git a/config.py b/config.py index 6c6ab7e..34861a7 100644 --- a/config.py +++ b/config.py @@ -4,7 +4,7 @@ DEVICE = "cuda:0" # Colors for different models (B,G,R format) -CONV_COLOR = (27, 80, 19) # 純粋な緑 +CONV_COLOR = (110, 32, 120) # 純粋な緑 XGBOOST_COLOR = (19, 80, 27) # オレンジ LIGHTGBM_COLOR = (19, 80, 27) # マゼンタ EARSNET_COLOR = (192, 79, 21) # ダークブルー diff --git a/main.py b/main.py index 2ebd93f..6ae653b 100644 --- a/main.py +++ b/main.py @@ -1030,21 +1030,21 @@ }, "lightGBM": { "radius": 8, - "fill_type": "striped", + "fill_type": "fill", "color": LIGHTGBM_COLOR, "outline_color": LIGHTGBM_COLOR, "outline_width": 2, }, "earsnet": { "radius": 8, - "fill_type": "fill", + "fill_type": "outline", "color": EARSNET_COLOR, "outline_color": EARSNET_COLOR, "outline_width": 2, }, "earsnet_crop": { "radius": 8, - "fill_type": "outline", + "fill_type": "fill", "color": EARSNET_CROP_COLOR, "outline_color": EARSNET_CROP_COLOR, "outline_width": 2, @@ -1076,7 +1076,7 @@ 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) + stetho_color_rgb = (110, 32, 120) def draw_glow_marker(draw, center, main_color, radius=5): outer_radius = radius + 3 diff --git a/make-demo-SP.py b/make-demo-SP.py new file mode 100644 index 0000000..19d9f80 --- /dev/null +++ b/make-demo-SP.py @@ -0,0 +1,146 @@ +import os + +import cv2 + + +def mosaic_face_with_yunet( + input_video_path: str, + model_path: str, + output_video_path: str, + frames_output_dir: str, + mosaic_scale: float = 0.1, + score_threshold: float = 0.9, + nms_threshold: float = 0.3, + top_k: int = 5000, +): + """ + YuNet (FaceDetectorYN) を使用して動画の顔のみをモザイク処理し、 + モザイク済みの動画と全フレームを出力する関数。 + + Parameters + ---------- + input_video_path : str + 入力動画ファイルのパス + model_path : str + YuNet の ONNX モデルファイルへのパス + output_video_path : str + モザイクをかけた結果の出力動画ファイルのパス + frames_output_dir : str + 全フレーム画像の出力先ディレクトリパス + mosaic_scale : float + モザイクの縮小率 (例: 0.1 なら幅高さを1/10に縮小して再拡大) + score_threshold : float + 顔検出におけるスコア閾値 (自信度がこれより低い検出結果は無視) + nms_threshold : float + NMS(非最大抑制)の閾値 + top_k : int + NMS 後に保持する最大検出数 + """ + + # 出力用ディレクトリ作成 + os.makedirs(frames_output_dir, exist_ok=True) + + # 動画を開く + cap = cv2.VideoCapture(input_video_path) + if not cap.isOpened(): + print(f"動画ファイルを開けませんでした: {input_video_path}") + return + + # 動画情報を取得 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = cap.get(cv2.CAP_PROP_FPS) + + # 動画出力の設定 (fourcc はmp4用にmp4vを使用) + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height)) + + # YuNet FaceDetectorYN の初期化 + # 入力サイズは推論用に一旦 (320, 320) など固定サイズを与えておき、 + # 各フレーム処理時に setInputSize で動的に変更します。 + face_detector = cv2.FaceDetectorYN.create( + model=model_path, + config="", + input_size=(320, 320), # 後ほどフレームサイズに合わせて上書き + score_threshold=score_threshold, + nms_threshold=nms_threshold, + top_k=top_k, + backend_id=cv2.dnn.DNN_BACKEND_DEFAULT, + target_id=cv2.dnn.DNN_TARGET_CPU, + ) + + frame_count = 0 + + while True: + ret, frame = cap.read() + if not ret: + break + + # フレームサイズが変化する場合は毎回 setInputSize を呼び出す + face_detector.setInputSize((frame.shape[1], frame.shape[0])) + + # 顔検出 + # detect() は (retval, faces) のタプルを返し、faces が検出結果 + # faces の列は [x, y, w, h, confidence, x1, y1, x2, y2, ... x5, y5] (ランドマーク) + _, faces = face_detector.detect(frame) + + if faces is not None: + for face in faces: + x, y, w, h = face[:4].astype(int) + # 必要に応じて座標補正（画像範囲外をはみ出さないようにクリップ） + x = max(0, x) + y = max(0, y) + w = min(w, frame.shape[1] - x) + h = min(h, frame.shape[0] - y) + + # 顔領域を取得 + face_roi = frame[y : y + h, x : x + w] + + # 縮小 + face_roi_small = cv2.resize( + face_roi, + (max(1, int(w * mosaic_scale)), max(1, int(h * mosaic_scale))), + interpolation=cv2.INTER_LINEAR, + ) + # 再拡大 (NEARESTで拡大するとドットが強めのモザイクに) + face_roi_mosaic = cv2.resize( + face_roi_small, (w, h), interpolation=cv2.INTER_NEAREST + ) + + # モザイク部分をフレームに貼り付け + frame[y : y + h, x : x + w] = face_roi_mosaic + + # フレーム画像を保存 + frame_filename = os.path.join(frames_output_dir, f"frame_{frame_count:06d}.jpg") + cv2.imwrite(frame_filename, frame) + + # モザイク処理済みフレームを動画ファイルに書き出し + out.write(frame) + + frame_count += 1 + + cap.release() + out.release() + print("処理が完了しました。") + + +def main(): + input_video = "video/tes.mp4" + model_path = "./face_detection_yunet_2023mar.onnx" + output_video = "output_mosaic_yunet.mp4" + frames_dir = "frames" + + mosaic_face_with_yunet( + input_video_path=input_video, + model_path=model_path, + output_video_path=output_video, + frames_output_dir=frames_dir, + mosaic_scale=0.2, # モザイクの粗さを調整 + score_threshold=0.9, # 検出結果の信頼度が 0.9 未満のものは無視 + nms_threshold=0.3, # NMS 閾値 + top_k=5000, # NMS 後に保持する最大数 + ) + + +if __name__ == "__main__": + main()