momentry_core/scripts/find_kids_refined.py

#!/opt/homebrew/bin/python3.11
"""
Refined Kid Detection with stricter filters.
Filters:
1. Ignore tiny faces (background noise).
2. Ignore sitting/squatting poses (Ankle must be lower than Hip).
"""

import json
import math
import os

POSE_JSON_PATH = "output/384b0ff44aaaa1f1/384b0ff44aaaa1f1.pose.json"
# Heuristic Threshold
RATIO_THRESHOLD = 6.2

# Filter constants
MIN_HEAD_WIDTH_PX = 25.0  # Ignore tiny background blobs
STANDING_TOLERANCE = 50.0  # Allow some wiggle room for ankles/hips overlap


def distance(p1, p2):
    return math.sqrt((p1["x"] - p2["x"]) ** 2 + (p1["y"] - p2["y"]) ** 2)


def get_midpoint(p1, p2):
    return {"x": (p1["x"] + p2["x"]) / 2, "y": (p1["y"] + p2["y"]) / 2}


def find_kids():
    if not os.path.exists(POSE_JSON_PATH):
        print(f"Pose JSON not found at {POSE_JSON_PATH}")
        return

    try:
        with open(POSE_JSON_PATH, "r") as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error loading JSON: {e}")
        return

    frames = data.get("frames", {})
    potential_kids = []

    analyzed_poses = 0

    print("Re-scanning pose data with stricter filters...")

    for frame_idx_str, frame_data in frames.items():
        timestamp = frame_data.get("timestamp", 0)
        people_in_frame = frame_data.get("poses", [])

        for person in people_in_frame:
            kps_list = person.get("keypoints", [])
            kp_dict = {kp["name"]: kp for kp in kps_list}

            # Required keypoints for validation
            nose = kp_dict.get("nose")
            l_shoulder = kp_dict.get("left_shoulder")
            r_shoulder = kp_dict.get("right_shoulder")
            l_hip = kp_dict.get("left_hip")
            r_hip = kp_dict.get("right_hip")
            l_ankle = kp_dict.get("left_ankle")
            r_ankle = kp_dict.get("right_ankle")

            # 1. Basic visibility check
            if not nose or not (l_shoulder and r_shoulder):
                continue
            # For strict standing check, we need reliable ankles/hips
            if not (l_ankle and r_ankle):
                continue

            analyzed_poses += 1

            # 2. Check for Sitting/Squatting (Filter 1)
            # If Hips are lower (larger Y) than Ankles, or too close to Ankles vertically,
            # the person is likely sitting or the detection is bad.
            # In Y-down coordinate: Ankle Y should be > Hip Y.

            mid_hip_y = (l_hip["y"] + r_hip["y"]) / 2
            mid_ankle_y = (l_ankle["y"] + r_ankle["y"]) / 2

            if mid_ankle_y < (mid_hip_y + STANDING_TOLERANCE):
                # Ankle is above or level with hip -> Sitting/Crouching
                continue

            # 3. Estimate Head Size
            mid_shoulder = get_midpoint(l_shoulder, r_shoulder)
            dist_nose_shoulder = distance(nose, mid_shoulder)

            # Head Height Estimate
            head_height = dist_nose_shoulder * 2.0

            # 4. Ignore tiny faces (Filter 2)
            # Distance between shoulders (width) is a good proxy for size
            shoulder_width = distance(l_shoulder, r_shoulder)
            if shoulder_width < MIN_HEAD_WIDTH_PX:
                continue

            # 5. Calculate Ratio
            mid_ankle = get_midpoint(l_ankle, r_ankle)
            dist_nose_ankle = distance(nose, mid_ankle)

            # Total Height = Nose to Ankle + half head
            total_height = dist_nose_ankle + (head_height / 2)

            if total_height <= head_height:
                continue

            ratio = total_height / head_height

            if ratio < RATIO_THRESHOLD:
                potential_kids.append(
                    {
                        "frame": frame_idx_str,
                        "timestamp": timestamp,
                        "ratio": round(ratio, 2),
                        "shoulder_width": round(shoulder_width, 1),
                        "confidence": "High" if ratio < 5.5 else "Medium",
                    }
                )

    print(f"Analyzed {analyzed_poses} valid standing poses.")
    print(f"Found {len(potential_kids)} potential kids after filtering.")

    # Group by timestamp
    unique_kids = {}
    for k in potential_kids:
        ts = round(k["timestamp"], 1)
        if ts not in unique_kids:
            unique_kids[ts] = k

    sorted_kids = sorted(unique_kids.values(), key=lambda x: x["timestamp"])

    print("\nRefined Timestamps:")
    for k in sorted_kids:
        print(
            f"  ⏱️ {k['timestamp']:.2f}s | Ratio: {k['ratio']} | Width: {k['shoulder_width']}px | Conf: {k['confidence']}"
        )


if __name__ == "__main__":
    find_kids()