seesaw112233/pose-estimation

1import os2import math3import json4import tempfile5from dataclasses import dataclass6from typing import Dict, List, Tuple, Optional7import urllib.request8 9import cv210import numpy as np11import pandas as pd12import gradio as gr13import mediapipe as mp14from mediapipe import solutions15from mediapipe.framework.formats import landmark_pb216from mediapipe.tasks import python17from mediapipe.tasks.python import vision18 19 20# -------------------------21# Model download helper22# -------------------------23def download_models():24    """Download required MediaPipe models if not present"""25    models_dir = "/tmp/mediapipe_models"26    os.makedirs(models_dir, exist_ok=True)27    28    models = {29        "face_landmarker": {30            "url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task",31            "path": os.path.join(models_dir, "face_landmarker.task")32        },33        "pose_landmarker": {34            "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/1/pose_landmarker_heavy.task",35            "path": os.path.join(models_dir, "pose_landmarker_heavy.task")36        }37    }38    39    for model_name, model_info in models.items():40        if not os.path.exists(model_info["path"]):41            print(f"Downloading {model_name}...")42            urllib.request.urlretrieve(model_info["url"], model_info["path"])43            print(f"✓ Downloaded {model_name}")44    45    return models["face_landmarker"]["path"], models["pose_landmarker"]["path"]46 47 48# -------------------------49# Utils: geometry50# -------------------------51def _dist(a: np.ndarray, b: np.ndarray) -> float:52    return float(np.linalg.norm(a - b))53 54def _safe_div(a: float, b: float, eps: float = 1e-8) -> float:55    return a / (b + eps)56 57def eye_aspect_ratio(pts: Dict[int, np.ndarray], idx: List[int]) -> Optional[float]:58    """59    EAR = (||p2-p6|| + ||p3-p5||) / (2*||p1-p4||)60    idx: [p1, p2, p3, p4, p5, p6]61    """62    try:63        p1, p2, p3, p4, p5, p6 = [pts[i] for i in idx]64    except KeyError:65        return None66    A = _dist(p2, p6)67    B = _dist(p3, p5)68    C = _dist(p1, p4)69    return _safe_div((A + B), (2.0 * C))70 71def angle_3pts(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> Optional[float]:72    """73    angle at point b in degrees formed by a-b-c74    """75    ba = a - b76    bc = c - b77    nba = np.linalg.norm(ba)78    nbc = np.linalg.norm(bc)79    if nba < 1e-8 or nbc < 1e-8:80        return None81    cosang = float(np.dot(ba, bc) / (nba * nbc))82    cosang = max(-1.0, min(1.0, cosang))83    return float(np.degrees(np.arccos(cosang)))84 85 86# -------------------------87# MediaPipe indices88# -------------------------89# FaceMesh landmarks for EAR (same indices work for new API)90LEFT_EYE_EAR_IDX  = [33, 160, 158, 133, 153, 144]91RIGHT_EYE_EAR_IDX = [362, 385, 387, 263, 373, 380]92 93# Pose landmark indices for new API94POSE_LANDMARKS = {95    "left_wrist": 15,96    "right_wrist": 16,97    "left_ankle": 27,98    "right_ankle": 28,99    "left_shoulder": 11,100    "right_shoulder": 12,101    "left_elbow": 13,102    "right_elbow": 14,103    "left_hip": 23,104    "right_hip": 24,105    "left_knee": 25,106    "right_knee": 26,107}108 109 110# -------------------------111# Drawing helpers for new API112# -------------------------113mp_drawing = solutions.drawing_utils114mp_drawing_styles = solutions.drawing_styles115 116# Face mesh connections117FACEMESH_TESSELATION = solutions.face_mesh.FACEMESH_TESSELATION118FACEMESH_CONTOURS = solutions.face_mesh.FACEMESH_CONTOURS119 120# Pose connections121POSE_CONNECTIONS = solutions.pose.POSE_CONNECTIONS122 123def draw_face_landmarks(image, face_landmarks):124    """Draw face landmarks on image using new API format - always draw full mesh"""125    if face_landmarks is None:126        return127    128    # Convert to landmark_pb2 format for drawing129    face_landmarks_proto = landmark_pb2.NormalizedLandmarkList()130    face_landmarks_proto.landmark.extend([131        landmark_pb2.NormalizedLandmark(x=lm.x, y=lm.y, z=lm.z)132        for lm in face_landmarks133    ])134    135    # Always draw full tesselation mesh136    mp_drawing.draw_landmarks(137        image=image,138        landmark_list=face_landmarks_proto,139        connections=FACEMESH_TESSELATION,140        landmark_drawing_spec=None,141        connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_tesselation_style()142    )143    144    # Also draw contours for clarity145    mp_drawing.draw_landmarks(146        image=image,147        landmark_list=face_landmarks_proto,148        connections=FACEMESH_CONTOURS,149        landmark_drawing_spec=None,150        connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_contours_style()151    )152 153def draw_pose_landmarks(image, pose_landmarks):154    """Draw pose landmarks on image using new API format"""155    if pose_landmarks is None:156        return157    158    # Convert to landmark_pb2 format for drawing159    pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()160    pose_landmarks_proto.landmark.extend([161        landmark_pb2.NormalizedLandmark(x=lm.x, y=lm.y, z=lm.z)162        for lm in pose_landmarks163    ])164    165    mp_drawing.draw_landmarks(166        image=image,167        landmark_list=pose_landmarks_proto,168        connections=POSE_CONNECTIONS,169        landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style()170    )171 172 173# -------------------------174# Blink detection175# -------------------------176@dataclass177class BlinkState:178    in_blink: bool = False179    blink_count: int = 0180    consec_below: int = 0181 182def update_blink(state: BlinkState, ear: Optional[float], thr: float, min_consec: int) -> BlinkState:183    """184    Basic blink logic:185    - ear below threshold for >= min_consec frames => blink start186    - when ear goes back above => blink end (count once)187    """188    if ear is None:189        return state190 191    if ear < thr:192        state.consec_below += 1193        if (not state.in_blink) and state.consec_below >= min_consec:194            state.in_blink = True195    else:196        if state.in_blink:197            state.blink_count += 1198        state.in_blink = False199        state.consec_below = 0200    return state201 202 203# -------------------------204# Core processing with new API205# -------------------------206def process_video(207    video_path: str,208    min_face_det_conf: float = 0.5,209    min_face_track_conf: float = 0.5,210    min_pose_det_conf: float = 0.5,211    min_pose_track_conf: float = 0.5,212    ear_threshold: float = 0.21,213    blink_min_consec: int = 2,214    max_frames: int = 0,215) -> Tuple[str, str, str, str]:216    """217    Process video using new MediaPipe API with GPU support218    Face mesh is always drawn (not optional)219    """220    # Download models first221    face_model_path, pose_model_path = download_models()222    223    cap = cv2.VideoCapture(video_path)224    if not cap.isOpened():225        raise RuntimeError("Cannot open video. Please upload a valid video file.")226 227    fps = cap.get(cv2.CAP_PROP_FPS)228    if fps <= 1e-6:229        fps = 30.0230    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))231    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))232    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))233 234    # Output paths235    tmpdir = tempfile.mkdtemp(prefix="mp_analysis_")236    out_video = os.path.join(tmpdir, "annotated.mp4")237    out_csv = os.path.join(tmpdir, "per_frame_metrics.csv")238    out_json = os.path.join(tmpdir, "summary.json")239    out_report = os.path.join(tmpdir, "report.md")240 241    fourcc = cv2.VideoWriter_fourcc(*"mp4v")242    writer = cv2.VideoWriter(out_video, fourcc, fps, (width, height))243 244    # Create face landmarker with GPU delegate245    base_options_face = python.BaseOptions(246        model_asset_path=face_model_path,247        delegate=python.BaseOptions.Delegate.GPU248    )249    face_options = vision.FaceLandmarkerOptions(250        base_options=base_options_face,251        running_mode=vision.RunningMode.VIDEO,252        num_faces=1,253        min_face_detection_confidence=min_face_det_conf,254        min_face_presence_confidence=min_face_track_conf,255        min_tracking_confidence=min_face_track_conf,256        output_face_blendshapes=False,257        output_facial_transformation_matrixes=False258    )259    260    # Create pose landmarker with GPU delegate261    base_options_pose = python.BaseOptions(262        model_asset_path=pose_model_path,263        delegate=python.BaseOptions.Delegate.GPU264    )265    pose_options = vision.PoseLandmarkerOptions(266        base_options=base_options_pose,267        running_mode=vision.RunningMode.VIDEO,268        num_poses=1,269        min_pose_detection_confidence=min_pose_det_conf,270        min_pose_presence_confidence=min_pose_track_conf,271        min_tracking_confidence=min_pose_track_conf,272        output_segmentation_masks=False273    )274 275    with vision.FaceLandmarker.create_from_options(face_options) as face_landmarker, \276         vision.PoseLandmarker.create_from_options(pose_options) as pose_landmarker:277 278        rows = []279        prev_pts = {}280        left_blink = BlinkState()281        right_blink = BlinkState()282 283        frame_idx = 0284        while True:285            ok, frame_bgr = cap.read()286            if not ok:287                break288            frame_idx += 1289            if max_frames and frame_idx > max_frames:290                break291 292            # Convert to RGB and create MediaPipe Image293            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)294            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb)295            296            # Timestamp in milliseconds297            timestamp_ms = int((frame_idx - 1) * 1000 / fps)298 299            # Process with new API300            face_result = face_landmarker.detect_for_video(mp_image, timestamp_ms)301            pose_result = pose_landmarker.detect_for_video(mp_image, timestamp_ms)302 303            # Extract face landmarks304            face_pts: Dict[int, np.ndarray] = {}305            face_landmarks = None306            if face_result.face_landmarks:307                face_landmarks = face_result.face_landmarks[0]308                for i, lm in enumerate(face_landmarks):309                    face_pts[i] = np.array([lm.x * width, lm.y * height], dtype=np.float32)310 311            # Calculate EAR312            left_ear = eye_aspect_ratio(face_pts, LEFT_EYE_EAR_IDX)313            right_ear = eye_aspect_ratio(face_pts, RIGHT_EYE_EAR_IDX)314 315            left_blink = update_blink(left_blink, left_ear, ear_threshold, blink_min_consec)316            right_blink = update_blink(right_blink, right_ear, ear_threshold, blink_min_consec)317 318            # Extract pose landmarks319            pose_norm: Dict[str, Optional[np.ndarray]] = {}320            pose_px: Dict[str, Optional[np.ndarray]] = {}321            pose_landmarks = None322            323            if pose_result.pose_landmarks:324                pose_landmarks = pose_result.pose_landmarks[0]325                for name, idx in POSE_LANDMARKS.items():326                    if idx < len(pose_landmarks):327                        lm = pose_landmarks[idx]328                        pose_norm[name] = np.array([lm.x, lm.y], dtype=np.float32)329                        pose_px[name] = np.array([lm.x * width, lm.y * height], dtype=np.float32)330                    else:331                        pose_norm[name] = None332                        pose_px[name] = None333            else:334                for name in POSE_LANDMARKS:335                    pose_norm[name] = None336                    pose_px[name] = None337 338            # Movement metrics339            def movement_metrics(key: str):340                cur = pose_norm.get(key)341                if cur is None:342                    return None, None343                prev = prev_pts.get(key)344                if prev is None:345                    d = 0.0346                else:347                    d = float(np.linalg.norm(cur - prev))348                v = d * fps349                prev_pts[key] = cur350                return d, v351 352            lw_d, lw_v = movement_metrics("left_wrist")353            rw_d, rw_v = movement_metrics("right_wrist")354            la_d, la_v = movement_metrics("left_ankle")355            ra_d, ra_v = movement_metrics("right_ankle")356 357            # Joint angles358            def get_angle(a, b, c):359                if a is None or b is None or c is None:360                    return None361                return angle_3pts(a, b, c)362 363            left_elbow_ang = get_angle(pose_px["left_shoulder"], pose_px["left_elbow"], pose_px["left_wrist"])364            right_elbow_ang = get_angle(pose_px["right_shoulder"], pose_px["right_elbow"], pose_px["right_wrist"])365            left_knee_ang = get_angle(pose_px["left_hip"], pose_px["left_knee"], pose_px["left_ankle"])366            right_knee_ang = get_angle(pose_px["right_hip"], pose_px["right_knee"], pose_px["right_ankle"])367 368            # Draw overlays (face mesh is always drawn, not optional)369            draw_pose_landmarks(frame_bgr, pose_landmarks)370            draw_face_landmarks(frame_bgr, face_landmarks)371 372            # HUD text373            hud_lines = [374                f"Frame: {frame_idx}/{total_frames if total_frames>0 else '?'}  FPS:{fps:.1f}",375                f"EAR L:{left_ear:.3f}" if left_ear is not None else "EAR L:None",376                f"EAR R:{right_ear:.3f}" if right_ear is not None else "EAR R:None",377                f"Blinks L:{left_blink.blink_count}  R:{right_blink.blink_count}",378            ]379            y0 = 24380            for line in hud_lines:381                cv2.putText(frame_bgr, line, (12, y0), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)382                y0 += 22383 384            writer.write(frame_bgr)385 386            rows.append({387                "frame": frame_idx,388                "time_s": (frame_idx - 1) / fps,389                "left_ear": left_ear,390                "right_ear": right_ear,391                "lw_disp": lw_d,392                "rw_disp": rw_d,393                "la_disp": la_d,394                "ra_disp": ra_d,395                "lw_speed": lw_v,396                "rw_speed": rw_v,397                "la_speed": la_v,398                "ra_speed": ra_v,399                "left_elbow_angle": left_elbow_ang,400                "right_elbow_angle": right_elbow_ang,401                "left_knee_angle": left_knee_ang,402                "right_knee_angle": right_knee_ang,403            })404 405    cap.release()406    writer.release()407 408    df = pd.DataFrame(rows)409 410    # Summaries411    def _sum_series(s: pd.Series):412        s2 = s.dropna()413        if len(s2) == 0:414            return {"mean": None, "min": None, "max": None}415        return {"mean": float(s2.mean()), "min": float(s2.min()), "max": float(s2.max())}416 417    summary = {418        "video": {419            "fps": float(fps),420            "width": width,421            "height": height,422            "frames_processed": int(len(df)),423            "duration_s": float(len(df) / fps),424        },425        "blink": {426            "ear_threshold": float(ear_threshold),427            "min_consecutive_frames": int(blink_min_consec),428            "left_blinks": int(left_blink.blink_count),429            "right_blinks": int(right_blink.blink_count),430            "left_blinks_per_min": float(_safe_div(left_blink.blink_count, (len(df)/fps)/60.0)) if len(df) else 0.0,431            "right_blinks_per_min": float(_safe_div(right_blink.blink_count, (len(df)/fps)/60.0)) if len(df) else 0.0,432            "left_ear_stats": _sum_series(df["left_ear"]),433            "right_ear_stats": _sum_series(df["right_ear"]),434        },435        "limb_movement": {436            "total_disp": {437                "left_wrist": float(df["lw_disp"].fillna(0).sum()),438                "right_wrist": float(df["rw_disp"].fillna(0).sum()),439                "left_ankle": float(df["la_disp"].fillna(0).sum()),440                "right_ankle": float(df["ra_disp"].fillna(0).sum()),441            },442            "speed_stats": {443                "left_wrist": _sum_series(df["lw_speed"]),444                "right_wrist": _sum_series(df["rw_speed"]),445                "left_ankle": _sum_series(df["la_speed"]),446                "right_ankle": _sum_series(df["ra_speed"]),447            },448            "angle_stats_deg": {449                "left_elbow": _sum_series(df["left_elbow_angle"]),450                "right_elbow": _sum_series(df["right_elbow_angle"]),451                "left_knee": _sum_series(df["left_knee_angle"]),452                "right_knee": _sum_series(df["right_knee_angle"]),453            }454        }455    }456 457    # Save outputs458    df.to_csv(out_csv, index=False)459    with open(out_json, "w", encoding="utf-8") as f:460        json.dump(summary, f, ensure_ascii=False, indent=2)461 462    report_md = f"""# MediaPipe Face + Pose Analysis Report (GPU Accelerated)463 464## Video Information465- Resolution: {width} x {height}466- FPS: {fps:.2f}467- Frames Processed: {len(df)}468- Duration: {summary["video"]["duration_s"]:.2f} seconds469 470## Blink Analysis (EAR)471- Threshold: {ear_threshold}472- Minimum Consecutive Frames: {blink_min_consec}473- Left Eye Blinks: {summary["blink"]["left_blinks"]} ({summary["blink"]["left_blinks_per_min"]:.2f} blinks/min)474- Right Eye Blinks: {summary["blink"]["right_blinks"]} ({summary["blink"]["right_blinks_per_min"]:.2f} blinks/min)475- Left Eye EAR: mean={summary["blink"]["left_ear_stats"]["mean"]}  min={summary["blink"]["left_ear_stats"]["min"]}  max={summary["blink"]["left_ear_stats"]["max"]}476- Right Eye EAR: mean={summary["blink"]["right_ear_stats"]["mean"]} min={summary["blink"]["right_ear_stats"]["min"]} max={summary["blink"]["right_ear_stats"]["max"]}477 478## Limb Movement (Normalized Units)479> Displacement/speed calculated based on normalized coordinates (0~1), suitable for relative comparison and trend analysis480- Total Displacement (higher = more movement):481  - Left Wrist: {summary["limb_movement"]["total_disp"]["left_wrist"]:.6f}482  - Right Wrist: {summary["limb_movement"]["total_disp"]["right_wrist"]:.6f}483  - Left Ankle: {summary["limb_movement"]["total_disp"]["left_ankle"]:.6f}484  - Right Ankle: {summary["limb_movement"]["total_disp"]["right_ankle"]:.6f}485 486## Output Files487- annotated.mp4: Video with pose skeleton and face mesh overlays488- per_frame_metrics.csv: Frame-by-frame metrics489- summary.json: Statistical summary490 491**Processed with GPU acceleration | New Face Landmarker API | Full Face Mesh Always Enabled**492"""493    with open(out_report, "w", encoding="utf-8") as f:494        f.write(report_md)495 496    return out_video, out_csv, out_json, out_report497 498 499# -------------------------500# Gradio UI501# -------------------------502def ui_process(503    video,504    min_face_det_conf,505    min_face_track_conf,506    min_pose_det_conf,507    min_pose_track_conf,508    ear_threshold,509    blink_min_consec,510    max_frames511):512    if isinstance(video, dict) and "path" in video:513        video_path = video["path"]514    else:515        video_path = video516 517    try:518        out_video, out_csv, out_json, out_report = process_video(519            video_path=str(video_path),520            min_face_det_conf=float(min_face_det_conf),521            min_face_track_conf=float(min_face_track_conf),522            min_pose_det_conf=float(min_pose_det_conf),523            min_pose_track_conf=float(min_pose_track_conf),524            ear_threshold=float(ear_threshold),525            blink_min_consec=int(blink_min_consec),526            max_frames=int(max_frames),527        )528 529        with open(out_report, "r", encoding="utf-8") as f:530            report_text = f.read()531 532        return out_video, out_csv, out_json, report_text533    534    except Exception as e:535        import traceback536        error_msg = f"# Error Processing Video\n\n```\n{traceback.format_exc()}\n```"537        return None, None, None, error_msg538 539 540demo = gr.Blocks(title="Video Pose + Face Analysis (GPU Accelerated)")541 542with demo:543    gr.Markdown("""544    ## Upload Video → MediaPipe GPU Acceleration → Pose + Face Mesh Tracking + Blink/Limb Analysis545    546    **Features:**547    - ✅ GPU Accelerated Processing548    - ✅ New Face Landmarker API (more accurate 478-point face mesh)549    - ✅ Full Face Mesh Always Enabled550    - ✅ Blink Detection (EAR Algorithm)551    - ✅ Limb Movement Quantification552    - ✅ Joint Angle Analysis553    """)554 555    with gr.Row():556        video_in = gr.Video(label="Upload Video")557 558    with gr.Accordion("Parameters (defaults work well for most cases)", open=False):559        gr.Markdown("### Face Detection Parameters")560        min_face_det_conf = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Face Detection Confidence Threshold")561        min_face_track_conf = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Face Tracking Confidence Threshold")562        563        gr.Markdown("### Pose Detection Parameters")564        min_pose_det_conf = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Pose Detection Confidence Threshold")565        min_pose_track_conf = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Pose Tracking Confidence Threshold")566 567        gr.Markdown("### Blink Detection Parameters")568        ear_threshold = gr.Slider(0.10, 0.35, value=0.21, step=0.01, label="Blink Threshold (EAR, lower = stricter)")569        blink_min_consec = gr.Slider(1, 6, value=2, step=1, label="Blink Minimum Consecutive Frames (anti-jitter)")570 571        gr.Markdown("### Processing Options")572        max_frames = gr.Number(value=0, precision=0, label="Maximum Frames to Process (0 = process all, set to 300 for debugging)")573 574    run_btn = gr.Button("🚀 Start Analysis (GPU Accelerated)", variant="primary", size="lg")575 576    with gr.Row():577        video_out = gr.Video(label="Output: Annotated Video")578    with gr.Row():579        csv_out = gr.File(label="Per-Frame Metrics CSV")580        json_out = gr.File(label="Summary JSON")581    report_out = gr.Markdown()582 583    run_btn.click(584        fn=ui_process,585        inputs=[586            video_in,587            min_face_det_conf,588            min_face_track_conf,589            min_pose_det_conf,590            min_pose_track_conf,591            ear_threshold,592            blink_min_consec,593            max_frames,594        ],595        outputs=[video_out, csv_out, json_out, report_out],596    )597 598if __name__ == "__main__":599    demo.launch(server_name="0.0.0.0", server_port=7860)