L'électron rare
c49de5c500
Add TAG_SKELETON2D=6, SKEL2D_BYTES=819, decode_skeleton2D() to iphone_usb_bridge.py; handler in IphoneUSBSource writing persons_arkit_2d / persons_arkit_2d_t to state; new state fields in State dataclass. TDD: test_arkit_body.py (round-trip + wrong length → None). Import check verified.
220 lines
9.2 KiB
Python
220 lines
9.2 KiB
Python
"""IphoneUSBSource — a cv2-VideoCapture-shaped frame source backed by the
|
||
iPhone ARBodyTracker USB stream. Decodes the AVLiveWire HEVC video to BGR
|
||
frames for MediaPipe, and (Task 2) writes ARKit skeleton + Vision hands into
|
||
State. Substitutes for cv2.VideoCapture in multi.py under --iphone-usb."""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import socket
|
||
import struct
|
||
import threading
|
||
import time
|
||
|
||
try:
|
||
import av
|
||
except ImportError: # pragma: no cover — av missing is caught at runtime by caller
|
||
av = None # type: ignore[assignment]
|
||
import cv2
|
||
import numpy as np
|
||
|
||
from data_only_viz.scripts.iphone_usb_bridge import (
|
||
connect_device, decode_skeleton, decode_skeleton2D,
|
||
iter_frames, TAG_SKELETON, TAG_SKELETON2D,
|
||
)
|
||
from data_only_viz.state import PoseKp
|
||
|
||
TAG_VIDEO = 2
|
||
TAG_HANDS = 4
|
||
LOG = logging.getLogger("iphone_usb_source")
|
||
|
||
_HAND_BYTES = 1 + 21 * 12 # chirality:u8 + 21 × (x,y,z) × 4 bytes BE f32
|
||
_ARKIT_JOINTS = 91
|
||
|
||
|
||
def _make_point(x: float, y: float, z: float) -> PoseKp:
|
||
"""Wrap raw floats from the wire into a PoseKp landmark."""
|
||
return PoseKp(x=x, y=y, z=z, c=1.0)
|
||
|
||
|
||
def _decode_hands(payload: bytes):
|
||
"""Decode TAG_HANDS payload → list[list[PoseKp]] (may be empty) or None.
|
||
|
||
Wire layout: count:u8, then per hand: chirality:u8 (1=right) followed by
|
||
21 × (x, y, z) big-endian f32. x,y are normalized image coords; z≈confidence.
|
||
Returns None only on empty payload; an empty list when count==0.
|
||
"""
|
||
if not payload:
|
||
return None
|
||
n = payload[0]
|
||
if n == 0:
|
||
return []
|
||
off = 1
|
||
hands: list[list[PoseKp]] = []
|
||
for _ in range(n):
|
||
if off + _HAND_BYTES > len(payload):
|
||
break
|
||
off += 1 # skip chirality byte (not used by HandFeatureExtractor)
|
||
pts = struct.unpack(">" + "f" * 63, payload[off:off + 21 * 12])
|
||
off += 21 * 12
|
||
hand = [_make_point(pts[i * 3], pts[i * 3 + 1], pts[i * 3 + 2])
|
||
for i in range(21)]
|
||
hands.append(hand)
|
||
return hands
|
||
|
||
|
||
def _to_annexb(data: bytes) -> bytes:
|
||
out = bytearray(); i = 0
|
||
while i + 4 <= len(data):
|
||
n = int.from_bytes(data[i:i + 4], "big"); i += 4
|
||
if i + n > len(data):
|
||
break
|
||
out += b"\x00\x00\x00\x01" + data[i:i + n]; i += n
|
||
return bytes(out)
|
||
|
||
|
||
class IphoneUSBSource:
|
||
def __init__(self, state=None, target_size=(640, 480),
|
||
write_hands: bool = True, mirror: bool = True) -> None:
|
||
self.state = state
|
||
self.target_w, self.target_h = target_size
|
||
self._write_hands = write_hands
|
||
# Mirror the video horizontally so the performer facing the camera
|
||
# interacts naturally (move right -> goes right; raise right arm ->
|
||
# the right side responds). Env CONCERT_MIRROR=0 disables it.
|
||
import os as _os
|
||
self._mirror = mirror and (_os.environ.get("CONCERT_MIRROR", "1") != "0")
|
||
self._codec = av.codec.CodecContext.create("hevc", "r") if av is not None else None
|
||
self._lock = threading.Lock()
|
||
self._frame = None # latest BGR np.ndarray
|
||
self._stop = threading.Event()
|
||
self._thread = None
|
||
self._sock = None # stored for shutdown() in release()
|
||
self._opened = False
|
||
|
||
def start(self) -> bool:
|
||
# Try an immediate connect, but do NOT give up if the iPhone app isn't
|
||
# streaming yet: start the read thread regardless. _run retries the
|
||
# connect (with backoff) until the app appears, so a startup race (app
|
||
# not ready / phone locked at launch) recovers the same way a mid-session
|
||
# drop does -- no data_only_viz restart needed.
|
||
sock = connect_device()
|
||
if sock is None:
|
||
LOG.warning("iphone usb: no device yet -- thread will keep retrying")
|
||
else:
|
||
self._sock = sock
|
||
self._opened = True
|
||
self._thread = threading.Thread(
|
||
target=self._run, args=(sock,), name="iphone_usb_src", daemon=True)
|
||
self._thread.start()
|
||
return True
|
||
|
||
def isOpened(self) -> bool:
|
||
return self._opened
|
||
|
||
def _run(self, sock) -> None:
|
||
while not self._stop.is_set():
|
||
if sock is None: # no connection yet (startup or post-drop)
|
||
sock = self._reconnect()
|
||
if sock is None: # release() was called while waiting
|
||
break
|
||
self._sock = sock
|
||
self._opened = True
|
||
LOG.info("iphone usb connected")
|
||
try:
|
||
for tag, pid, payload in iter_frames(sock):
|
||
if self._stop.is_set():
|
||
break
|
||
if tag == TAG_VIDEO and len(payload) > 1:
|
||
annexb = _to_annexb(payload[1:])
|
||
if self._codec is not None and av is not None:
|
||
try:
|
||
for fr in self._codec.decode(av.Packet(annexb)):
|
||
img = fr.to_ndarray(format="bgr24")
|
||
img = cv2.resize(img, (self.target_w, self.target_h))
|
||
if self._mirror:
|
||
img = cv2.flip(img, 1)
|
||
with self._lock:
|
||
self._frame = img
|
||
except Exception as e: # av.AVError is a removal-prone alias
|
||
LOG.debug("hevc decode: %s", e)
|
||
elif tag == TAG_SKELETON and self.state is not None:
|
||
joints = decode_skeleton(payload)
|
||
if joints is not None:
|
||
with self.state.lock():
|
||
arr = self.state.persons_arkit_joints.get(pid)
|
||
if arr is None or arr.shape != (_ARKIT_JOINTS, 3):
|
||
arr = np.zeros((_ARKIT_JOINTS, 3), dtype=np.float32)
|
||
for i, (x, y, z, valid) in enumerate(joints):
|
||
if valid:
|
||
arr[i] = (x, y, z)
|
||
self.state.persons_arkit_joints[pid] = arr
|
||
self.state.persons_arkit_last_t[pid] = time.perf_counter()
|
||
elif tag == TAG_SKELETON2D and self.state is not None:
|
||
pts = decode_skeleton2D(payload)
|
||
if pts is not None:
|
||
arr = np.array([[x, y] for (x, y, _v) in pts], dtype=np.float32)
|
||
with self.state.lock():
|
||
self.state.persons_arkit_2d[pid] = arr
|
||
self.state.persons_arkit_2d_t[pid] = time.perf_counter()
|
||
elif tag == TAG_HANDS:
|
||
hands = _decode_hands(payload)
|
||
if hands is not None and self.state is not None:
|
||
with self.state.lock():
|
||
# Always expose iPhone Vision hands (stable,
|
||
# rotation-invariant) for the air-piano.
|
||
self.state.persons_hands_iphone = hands
|
||
self.state.persons_hands_iphone_t = \
|
||
time.perf_counter()
|
||
if self._write_hands:
|
||
self.state.persons_hands = hands
|
||
except OSError as e:
|
||
LOG.warning("iphone usb stream error: %s", e)
|
||
finally:
|
||
try:
|
||
sock.close()
|
||
except OSError:
|
||
pass
|
||
# Stream ended. Drop the stale frame so read() reports no-frame,
|
||
# then reconnect (unless stopping) so a transient USB/app drop
|
||
# doesn't freeze the pose pipeline (the iPhone backgrounding or a
|
||
# USB hiccup ends iter_frames; we re-establish :7000).
|
||
with self._lock:
|
||
self._frame = None
|
||
self._opened = False
|
||
sock = None # forces the reconnect at the loop top
|
||
if not self._stop.is_set():
|
||
LOG.info("iphone usb stream ended -- reconnecting...")
|
||
self._opened = False
|
||
|
||
def _reconnect(self):
|
||
"""Retry connect_device() with capped backoff until it succeeds or
|
||
release() is called. Returns the new socket, or None if stopping."""
|
||
delay = 0.5
|
||
while not self._stop.is_set():
|
||
sock = connect_device()
|
||
if sock is not None:
|
||
return sock
|
||
self._stop.wait(delay)
|
||
delay = min(delay * 2, 3.0)
|
||
return None
|
||
|
||
def read(self):
|
||
with self._lock:
|
||
if self._frame is None or not self._opened:
|
||
return False, None
|
||
return True, self._frame.copy()
|
||
|
||
def set(self, *args) -> bool:
|
||
return True # cv2 CAP_PROP_* no-op
|
||
|
||
def release(self) -> None:
|
||
self._stop.set()
|
||
if self._sock is not None:
|
||
try:
|
||
self._sock.shutdown(socket.SHUT_RDWR)
|
||
except OSError:
|
||
pass
|
||
if self._thread is not None:
|
||
self._thread.join(timeout=2.0)
|
||
self._opened = False
|