vixy/oak-service
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add MoveNet Lightning pose estimation on Coral 2

Browse Source 
Integrates single-person pose detection into oak-service using MoveNet
Lightning on a second Google Coral Edge TPU. Detects 17 body keypoints
at ~7ms per frame, derives posture (standing/sitting), facing direction,
and arm position. Only runs when a person is detected by YOLOv6.

New endpoints: /pose (raw keypoints), /pose/summary (derived posture)
New module: pose_estimator.py (PoseEstimator class)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alex
2026-02-08 19:29:16 -06:00
parent 19a27a1240
commit cdbf7ff394
3 changed files with 343 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
oak_service.py
View File

@@ -11,11 +11,14 @@ Day 81 - Added presence detection! Now I can SEE you! 👀💜
import time
import threading
from contextlib import asynccontextmanager
from pathlib import Path
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
import depthai as dai
import cv2
from pose_estimator import PoseEstimator
# ============== Configuration ==============
DETECTION_MODEL = "yolov6-nano" # Has 'person' class
PERSON_CLASS_ID = 0 # 'person' is class 0 in COCO
@@ -23,6 +26,10 @@ DETECTION_THRESHOLD = 0.5
PRESENCE_TIMEOUT = 30.0 # seconds without person = not present
DETECTION_INTERVAL = 0.5
# Pose estimation
POSE_MODEL_PATH = Path(__file__).parent / "models" / "movenet_single_pose_lightning_ptq_edgetpu.tflite"
POSE_CORAL_DEVICE = 1 # Second Coral (device 0 is headmic/YAMNet)
# ============== Global State ==============
pipeline_ctx = None
detection_queue = None
@@ -30,6 +37,7 @@ rgb_queue = None
detection_thread = None
running = False
labels = []
pose_estimator = None
presence_state = {
"present": False,
@@ -40,6 +48,16 @@ presence_state = {
"confidence": 0.0,
}
pose_state = {
"active": False,
"keypoints": [],
"posture": {},
"num_valid": 0,
"mean_confidence": 0.0,
"inference_ms": 0.0,
"last_update": None,
}
def init_oak():
"""Initialize OAK-D with person detection pipeline (depthai v3)."""
@@ -75,8 +93,12 @@ def init_oak():
pipeline_ctx = pipeline
print("✅ OAK-D initialized with person detection!")
# Initialize pose estimator on Coral 2
_init_pose_estimator()
return True
except Exception as e:
print(f"❌ Failed to initialize OAK-D: {e}")
import traceback
@@ -84,6 +106,25 @@ def init_oak():
return False
def _init_pose_estimator():
"""Initialize MoveNet Lightning on the second Coral Edge TPU."""
global pose_estimator
if not POSE_MODEL_PATH.exists():
print(f"⚠️ Pose model not found: {POSE_MODEL_PATH}")
return
try:
pose_estimator = PoseEstimator(
model_path=str(POSE_MODEL_PATH),
device_index=POSE_CORAL_DEVICE,
)
print("✅ Pose estimator initialized on Coral 2!")
except Exception as e:
print(f"⚠️ Pose estimator failed to initialize: {e}")
pose_estimator = None
def cleanup_oak():
"""Cleanup OAK-D resources."""
global pipeline_ctx, running
@@ -99,29 +140,29 @@ def cleanup_oak():
def detection_loop():
"""Background thread for presence detection."""
global running, presence_state, detection_queue
"""Background thread for presence detection + pose estimation."""
global running, presence_state, pose_state, detection_queue
print("🔍 Presence detection loop started")
while running:
try:
if detection_queue is None:
time.sleep(1)
continue
data = detection_queue.tryGet()
if data is not None:
now = time.time()
presence_state["last_detection"] = now
# Filter for person detections only
persons = [d for d in data.detections if d.label == PERSON_CLASS_ID]
person_count = len(persons)
presence_state["person_count"] = person_count
if person_count > 0:
presence_state["present"] = True
presence_state["last_seen"] = now
@@ -134,24 +175,66 @@ def detection_loop():
}
for d in persons
]
# Run pose estimation on the latest RGB frame
_run_pose_estimation()
else:
presence_state["detections"] = []
presence_state["confidence"] = 0.0
# Clear pose when no person
if pose_state["active"]:
pose_state["active"] = False
pose_state["keypoints"] = []
pose_state["posture"] = {}
pose_state["num_valid"] = 0
pose_state["mean_confidence"] = 0.0
# Check timeout
if presence_state["last_seen"]:
if now - presence_state["last_seen"] > PRESENCE_TIMEOUT:
presence_state["present"] = False
time.sleep(DETECTION_INTERVAL)
except Exception as e:
print(f"Detection loop error: {e}")
time.sleep(1)
print("🛑 Presence detection loop stopped")
def _run_pose_estimation():
"""Grab latest RGB frame and run pose estimation via Coral 2."""
global pose_state, rgb_queue, pose_estimator
if pose_estimator is None or rgb_queue is None:
return
try:
frame_msg = rgb_queue.tryGet()
if frame_msg is None:
return
frame = frame_msg.getCvFrame()
result = pose_estimator.estimate(frame)
# Derive posture from keypoints
posture = pose_estimator.derive_posture(result["keypoints"])
pose_state["active"] = True
pose_state["keypoints"] = result["keypoints"]
pose_state["posture"] = posture
pose_state["num_valid"] = result["num_valid"]
pose_state["mean_confidence"] = result["mean_confidence"]
pose_state["inference_ms"] = result["inference_ms"]
pose_state["last_update"] = result["timestamp"]
except Exception as e:
print(f"Pose estimation error: {e}")
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup and shutdown."""
@@ -175,8 +258,8 @@ async def lifespan(app: FastAPI):
app = FastAPI(
title="OAK-D Vision Service",
description="Vixy's eyes with presence detection! 🦊👀",
version="0.3.0",
description="Vixy's eyes with presence detection + pose estimation! 🦊👀",
version="0.4.0",
lifespan=lifespan
)
@@ -187,9 +270,10 @@ async def health():
return {
"status": "healthy",
"service": "oak-service",
"version": "0.3.0",
"version": "0.4.0",
"oak_connected": pipeline_ctx is not None,
"detection_model": DETECTION_MODEL,
"pose_model_loaded": pose_estimator is not None,
"timestamp": time.time()
}
@@ -245,6 +329,40 @@ async def snapshot():
@app.get("/pose")
async def pose():
"""Get current pose keypoints."""
if pose_estimator is None:
raise HTTPException(status_code=503, detail="Pose estimator not available")
return {
"active": pose_state["active"],
"keypoints": pose_state["keypoints"],
"num_valid": pose_state["num_valid"],
"mean_confidence": pose_state["mean_confidence"],
"inference_ms": pose_state["inference_ms"],
"last_update": pose_state["last_update"],
"timestamp": time.time(),
}
@app.get("/pose/summary")
async def pose_summary():
"""Get derived posture summary."""
if pose_estimator is None:
raise HTTPException(status_code=503, detail="Pose estimator not available")
return {
"active": pose_state["active"],
"posture": pose_state["posture"].get("posture", "unknown"),
"facing_camera": pose_state["posture"].get("facing_camera", False),
"arms_raised": pose_state["posture"].get("arms_raised", False),
"mean_confidence": pose_state["mean_confidence"],
"num_valid": pose_state["num_valid"],
"timestamp": time.time(),
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8100)