import os import sys import cv2 import numpy as np import gradio as gr import tempfile import torch from torchvision import transforms import tensorflow as tf import tensorflow_hub as hub import mediapipe as mp def setup_environment(): """環境セットアップ関数""" print("環境セットアップを開始します...") # 必要なパッケージをインストール os.system("pip install torch torchvision opencv-python numpy gradio tensorflow tensorflow-hub mediapipe") print("環境セットアップが完了しました！") # セットアップ実行 setup_environment() # モデル初期化 def initialize_models(): """各モデルを初期化""" models = {} # MediaPipe Pose mp_pose = mp.solutions.pose models['mediapipe'] = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5) # MoveNet MultiPose movenet_model = hub.load("https://tfhub.dev/google/movenet/multipose/lightning/1") models['movenet'] = movenet_model.signatures['serving_default'] return models # 共通設定 KEYPOINT_EDGES = [ (0, 1), (1, 3), (0, 2), (2, 4), (5, 7), (7, 9), (6, 8), (8, 10), (5, 6), (5, 11), (6, 12), (11, 13), (13, 15), (12, 14), (14, 16), (11, 12) ] # 各モデル用のポーズ検出関数 def detect_pose_mediapipe(frame, pose): rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) results = pose.process(rgb) return results.pose_landmarks if results.pose_landmarks else None def detect_pose_movenet(frame, movenet): image = tf.image.resize_with_pad(tf.expand_dims(frame, axis=0), 256, 256) input_image = tf.cast(image, dtype=tf.int32) outputs = movenet(input_image) return outputs['output_0'].numpy()[0] # 棒人間描画関数 def draw_stick_figure(frame, landmarks, model_type, width, height, head_size_factor, line_thickness, fixed_size=None): blank = np.ones((height, width, 3), dtype=np.uint8) * 255 black = (0, 0, 0) if model_type == 'mediapipe': def get_point(index): lm = landmarks.landmark[index] return int(lm.x * width), int(lm.y * height) head_x, head_y = get_point(0) if fixed_size is not None: shoulder_hip_dist = fixed_size else: def get_distance(a, b): xa, ya = get_point(a) xb, yb = get_point(b) return ((xa - xb)**2 + (ya - yb)**2)**0.5 shoulder_hip_dist = get_distance(11, 13) radius = int(shoulder_hip_dist * head_size_factor) cv2.circle(blank, (head_x, head_y), radius, black, thickness=line_thickness) connections = [ (11, 12), (11, 13), (13, 15), (12, 14), (14, 16), (11, 23), (12, 24), (23, 24), (23, 25), (24, 26), (25, 27), (26, 28) ] for a, b in connections: pt1 = get_point(a) pt2 = get_point(b) cv2.line(blank, pt1, pt2, black, line_thickness) elif model_type == 'movenet': for person in landmarks: overall_score = person[-1] if overall_score < 0.2: continue keypoints = [] for i in range(17): y, x, score = person[i * 3:(i + 1) * 3] if score < 0.2: keypoints.append(None) else: px, py = int(x * width), int(y * height) keypoints.append((px, py)) if keypoints[0] and keypoints[5] and keypoints[6]: head_x, head_y = keypoints[0] d1 = np.linalg.norm(np.array(keypoints[5]) - np.array(keypoints[11])) if keypoints[11] else 0 d2 = np.linalg.norm(np.array(keypoints[6]) - np.array(keypoints[12])) if keypoints[12] else 0 shoulder_hip_dist = (d1 + d2) / 2 if d1 and d2 else 50 radius = int(shoulder_hip_dist * head_size_factor) cv2.circle(blank, (head_x, head_y), radius, black, thickness=line_thickness) for idx, pt in enumerate(keypoints): if pt and idx not in [0, 1, 2, 3, 4]: cv2.circle(blank, pt, line_thickness, black, -1) for a, b in KEYPOINT_EDGES: if keypoints[a] and keypoints[b]: cv2.line(blank, keypoints[a], keypoints[b], black, line_thickness) return blank # 動画処理メイン関数 def process_video(video_path, model_type, head_size_factor, line_thickness, use_average_head_size): models = initialize_models() model = models[model_type] cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError("動画ファイルを開けません") total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) 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) average_shoulder_hip_dist = None if use_average_head_size and model_type == 'mediapipe': distances = [] while True: ret, frame = cap.read() if not ret: break landmarks = detect_pose_mediapipe(frame, model) if landmarks: def get_point(index): lm = landmarks.landmark[index] return int(lm.x * width), int(lm.y * height) def get_distance(a, b): xa, ya = get_point(a) xb, yb = get_point(b) return ((xa - xb)**2 + (ya - yb)**2)**0.5 dist = get_distance(11, 13) distances.append(dist) average_shoulder_hip_dist = np.mean(distances) if distances else 50 cap.set(cv2.CAP_PROP_POS_FRAMES, 0) output_path = tempfile.mktemp(suffix=".mp4") fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) frame_idx = 0 while True: ret, frame = cap.read() if not ret: break frame_idx += 1 print(f"処理中: フレーム {frame_idx}/{total_frames} ({(frame_idx/total_frames)*100:.1f}%)") if model_type == 'mediapipe': landmarks = detect_pose_mediapipe(frame, model) if landmarks: drawn = draw_stick_figure( frame, landmarks, model_type, width, height, head_size_factor, line_thickness, average_shoulder_hip_dist ) else: drawn = np.ones((height, width, 3), dtype=np.uint8) * 255 elif model_type == 'movenet': landmarks = detect_pose_movenet(frame, model) drawn = draw_stick_figure( frame, landmarks, model_type, width, height, head_size_factor, line_thickness, None ) out.write(drawn) cap.release() out.release() print("処理完了!") return output_path # Gradioインターフェース def gradio_interface(video_file, model_type, head_size_factor, line_thickness, use_avg): return process_video( video_file, model_type, head_size_factor, line_thickness, use_avg == "全体平均で固定" ) model_info = """ - **MediaPipe Pose**: 単独人物向け、高精度 - **MoveNet MultiPose**: 複数人物検出可能 """ demo = gr.Interface( fn=gradio_interface, inputs=[ gr.Video(label="動画をアップロード"), gr.Radio( ["mediapipe", "movenet"], label="モデル選択", value="mediapipe", info=model_info ), gr.Slider(minimum=0.1, maximum=1.0, value=0.4, label="頭の大きさ（肩〜腰の比率）"), gr.Slider(minimum=1, maximum=10, step=1, value=2, label="線の太さ"), gr.Radio( ["フレームごとに計算", "全体平均で固定"], value="フレームごとに計算", label="頭サイズの計算方法 (MediaPipeのみ有効)" ) ], outputs=gr.Video(label="棒人間動画"), title="統合版 棒人間モーショントラッキング", description="MediaPipe または MoveNet による棒人間動画生成ツール" ) if __name__ == "__main__": demo.launch()