clement
364 lines
14 KiB
Swift
364 lines
14 KiB
Swift
import ARKit
|
|
import AVLiveWire
|
|
import Combine
|
|
import Foundation
|
|
import SwiftUI
|
|
import Vision
|
|
|
|
/// Drives the ARKit body-tracking session and streams the 91-joint
|
|
/// skeleton plus HEVC-encoded camera video to the tethered Mac over
|
|
/// USB (AVLiveWire frames via usbmuxd). No network involved.
|
|
///
|
|
/// Lightweight 2D snapshot of the tracked skeleton, ready for SwiftUI
|
|
/// Canvas. Joint indices follow `ARSkeletonDefinition.defaultBody3D`.
|
|
struct SkeletonSnapshot: Equatable {
|
|
/// Projected joint positions in viewport coordinates, or nil if the
|
|
/// joint falls outside the view / is not finite.
|
|
let points: [CGPoint?]
|
|
/// Per-joint tracking flag (`ARSkeleton.isJointTracked`).
|
|
let tracked: [Bool]
|
|
}
|
|
|
|
@MainActor
|
|
final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {
|
|
@Published var running: Bool = false
|
|
@Published var status: String = "idle"
|
|
@Published var framesSent: Int = 0
|
|
@Published var jointsPerSec: Double = 0
|
|
@Published var bodyCount: Int = 0
|
|
@Published var skeleton2D: SkeletonSnapshot?
|
|
@Published var usbState: USBServer.State = .idle
|
|
@Published var arkitFps: Double = 0
|
|
@Published var visionMs: Double = 0
|
|
@Published var visionDropRate: Double = 0
|
|
@Published var thermalLabel: String = "nominal"
|
|
/// Set by the SwiftUI view via GeometryReader so the projection
|
|
/// matches the on-screen ARView size.
|
|
var viewportSize: CGSize = .zero
|
|
private var sendEnvMesh: Bool = false
|
|
private let session = ARSession()
|
|
private let usb = USBServer()
|
|
private let videoEncoder = VideoEncoder()
|
|
private var videoStarted = false
|
|
private var lastFrameTime: TimeInterval = 0
|
|
private let visionQueue = DispatchQueue(
|
|
label: "cc.saillant.arbody.vision", qos: .userInitiated)
|
|
private var visionBusy = false
|
|
private var lastVisionTime: TimeInterval = 0
|
|
private var visionAttempts: Int = 0
|
|
private var visionDrops: Int = 0
|
|
private var jointsInSecond: Int = 0
|
|
private var lastSecond: TimeInterval = 0
|
|
private let bodyParents: [Int] =
|
|
ARSkeletonDefinition.defaultBody3D.parentIndices
|
|
|
|
nonisolated static func thermalString(
|
|
_ s: ProcessInfo.ThermalState) -> String {
|
|
switch s {
|
|
case .nominal: return "nominal"
|
|
case .fair: return "fair"
|
|
case .serious: return "serious"
|
|
case .critical: return "critical"
|
|
@unknown default: return "?"
|
|
}
|
|
}
|
|
|
|
override init() {
|
|
super.init()
|
|
session.delegate = self
|
|
usb.onState = { [weak self] s in
|
|
Task { @MainActor in
|
|
guard let self else { return }
|
|
let was = self.usbState
|
|
self.usbState = s
|
|
if s == .connected && was != .connected {
|
|
self.publishTopology()
|
|
}
|
|
}
|
|
}
|
|
videoEncoder.onPayload = { [weak self] payload in
|
|
Task { @MainActor in
|
|
guard let self, self.usbState == .connected else {
|
|
return
|
|
}
|
|
self.usb.send(tag: .video, pid: -1,
|
|
timestamp: self.lastFrameTime,
|
|
payload: payload.encoded())
|
|
}
|
|
}
|
|
thermalLabel = Self.thermalString(
|
|
ProcessInfo.processInfo.thermalState)
|
|
NotificationCenter.default.addObserver(
|
|
forName: ProcessInfo.thermalStateDidChangeNotification,
|
|
object: nil, queue: .main) { [weak self] _ in
|
|
Task { @MainActor in
|
|
self?.thermalLabel = Self.thermalString(
|
|
ProcessInfo.processInfo.thermalState)
|
|
}
|
|
}
|
|
}
|
|
|
|
func configure(sendEnvMesh: Bool) {
|
|
self.sendEnvMesh = sendEnvMesh
|
|
}
|
|
|
|
func start() {
|
|
guard ARBodyTrackingConfiguration.isSupported else {
|
|
status = "ARBodyTracking unsupported (need A12+, iPhone XR/XS+)"
|
|
return
|
|
}
|
|
let cfg = ARBodyTrackingConfiguration()
|
|
var feats: [String] = []
|
|
// No extra frame semantics: `.sceneDepth` is reserved to
|
|
// ARWorldTracking, and `.personSegmentationWithDepth` is
|
|
// rejected per-frame by ABPKPersonIDTracker in this config
|
|
// (spams the console without producing usable depth).
|
|
// NOTE: ARBodyTrackingConfiguration does not expose
|
|
// sceneReconstruction (that's ARWorldTrackingConfiguration
|
|
// territory). Env mesh capture requires a separate ARSession
|
|
// with body tracking off — out of scope for this scaffold.
|
|
if sendEnvMesh {
|
|
feats.append("env-mesh: requires separate session (TODO)")
|
|
}
|
|
cfg.automaticImageScaleEstimationEnabled = true
|
|
usb.start()
|
|
session.run(cfg, options: [.resetTracking, .removeExistingAnchors])
|
|
status = feats.isEmpty
|
|
? "running (RGB only)"
|
|
: "running (\(feats.joined(separator: ", ")))"
|
|
running = true
|
|
}
|
|
|
|
func stop() {
|
|
session.pause()
|
|
usb.stop()
|
|
videoEncoder.stop()
|
|
videoStarted = false
|
|
running = false
|
|
status = "stopped"
|
|
}
|
|
|
|
// MARK: - ARSessionDelegate
|
|
|
|
nonisolated func session(_ s: ARSession, didUpdate frame: ARFrame) {
|
|
let t = frame.timestamp
|
|
Task { @MainActor in
|
|
// Throttle to 30 fps max.
|
|
let dt = t - self.lastFrameTime
|
|
if dt < 1.0 / 30.0 { return }
|
|
self.lastFrameTime = t
|
|
if dt > 0, dt < 1.0 {
|
|
let inst = 1.0 / dt
|
|
self.arkitFps = self.arkitFps == 0
|
|
? inst : self.arkitFps * 0.9 + inst * 0.1
|
|
}
|
|
|
|
// Encode the camera frame to HEVC and stream it over USB.
|
|
let img = frame.capturedImage
|
|
let w = Int32(CVPixelBufferGetWidth(img))
|
|
let h = Int32(CVPixelBufferGetHeight(img))
|
|
if !self.videoStarted, w > 0, h > 0 {
|
|
self.videoEncoder.start(width: w, height: h)
|
|
self.videoStarted = true
|
|
}
|
|
if self.videoStarted {
|
|
self.videoEncoder.encode(img, pts: t)
|
|
}
|
|
|
|
// Vision hands at ~15 fps, off-main, drop if busy.
|
|
if t - self.lastVisionTime >= 1.0 / 15.0 {
|
|
self.visionAttempts += 1
|
|
if self.visionBusy {
|
|
self.visionDrops += 1
|
|
} else {
|
|
self.lastVisionTime = t
|
|
self.visionBusy = true
|
|
// CVPixelBuffer is non-Sendable; this is a deliberate
|
|
// single-owner handoff to the Vision worker queue.
|
|
nonisolated(unsafe) let buf = img
|
|
self.visionQueue.async {
|
|
self.runVision(on: buf, timestamp: t)
|
|
Task { @MainActor in self.visionBusy = false }
|
|
}
|
|
}
|
|
if self.visionAttempts >= 60 {
|
|
self.visionDropRate =
|
|
Double(self.visionDrops) / Double(self.visionAttempts)
|
|
self.visionAttempts = 0
|
|
self.visionDrops = 0
|
|
}
|
|
}
|
|
|
|
var count: Int = 0
|
|
var firstBody: ARBodyAnchor?
|
|
for anchor in frame.anchors {
|
|
guard let body = anchor as? ARBodyAnchor else { continue }
|
|
self.publishUSB(pid: count, timestamp: t, body: body,
|
|
camera: frame.camera)
|
|
if count == 0 { firstBody = body }
|
|
count += 1
|
|
}
|
|
self.framesSent &+= 1
|
|
self.bodyCount = count
|
|
self.updateSkeleton2D(body: firstBody, camera: frame.camera)
|
|
|
|
let now = Date().timeIntervalSinceReferenceDate
|
|
self.jointsInSecond &+= count * 91
|
|
if now - self.lastSecond >= 1.0 {
|
|
self.jointsPerSec = Double(self.jointsInSecond)
|
|
/ max(0.001, now - self.lastSecond)
|
|
self.jointsInSecond = 0
|
|
self.lastSecond = now
|
|
}
|
|
}
|
|
}
|
|
|
|
private func currentInterfaceOrientation() -> UIInterfaceOrientation {
|
|
for scene in UIApplication.shared.connectedScenes {
|
|
if let ws = scene as? UIWindowScene {
|
|
return ws.interfaceOrientation
|
|
}
|
|
}
|
|
return .portrait
|
|
}
|
|
|
|
private func updateSkeleton2D(body: ARBodyAnchor?, camera: ARCamera) {
|
|
guard let body, viewportSize.width > 1, viewportSize.height > 1
|
|
else {
|
|
if skeleton2D != nil { skeleton2D = nil }
|
|
return
|
|
}
|
|
let xforms = body.skeleton.jointModelTransforms
|
|
let root = body.transform
|
|
let orient = currentInterfaceOrientation()
|
|
var pts: [CGPoint?] = Array(repeating: nil, count: xforms.count)
|
|
var tracked: [Bool] = Array(repeating: false, count: xforms.count)
|
|
for (i, m) in xforms.enumerated() {
|
|
let w = root * m
|
|
let p3 = simd_make_float3(w.columns.3.x,
|
|
w.columns.3.y,
|
|
w.columns.3.z)
|
|
let p2 = camera.projectPoint(p3,
|
|
orientation: orient,
|
|
viewportSize: viewportSize)
|
|
if p2.x.isFinite && p2.y.isFinite { pts[i] = p2 }
|
|
tracked[i] = body.skeleton.isJointTracked(i)
|
|
}
|
|
skeleton2D = SkeletonSnapshot(points: pts, tracked: tracked)
|
|
}
|
|
|
|
/// Exposed for SwiftUI overlays that need to wire bone parent
|
|
/// indices without re-reading the ARKit skeleton definition.
|
|
var bodyParentIndices: [Int] { bodyParents }
|
|
|
|
// MARK: - Vision hand capture
|
|
|
|
nonisolated func runVision(on pixelBuffer: CVPixelBuffer,
|
|
timestamp t: TimeInterval) {
|
|
let handReq = VNDetectHumanHandPoseRequest()
|
|
handReq.maximumHandCount = 2
|
|
let handler = VNImageRequestHandler(
|
|
cvPixelBuffer: pixelBuffer, orientation: .right, options: [:])
|
|
let t0 = Date()
|
|
try? handler.perform([handReq])
|
|
let ms = Date().timeIntervalSince(t0) * 1000.0
|
|
Task { @MainActor in
|
|
self.visionMs = self.visionMs == 0
|
|
? ms : self.visionMs * 0.8 + ms * 0.2
|
|
}
|
|
|
|
// ---- hands ----
|
|
// Fixed 21-joint order matching MediaPipe-style indexing.
|
|
let order: [VNHumanHandPoseObservation.JointName] = [
|
|
.wrist,
|
|
.thumbCMC, .thumbMP, .thumbIP, .thumbTip,
|
|
.indexMCP, .indexPIP, .indexDIP, .indexTip,
|
|
.middleMCP, .middlePIP, .middleDIP, .middleTip,
|
|
.ringMCP, .ringPIP, .ringDIP, .ringTip,
|
|
.littleMCP, .littlePIP, .littleDIP, .littleTip,
|
|
]
|
|
var handsOut: [HandsPayload.Hand] = []
|
|
for obs in (handReq.results ?? []).prefix(2) {
|
|
guard let pts = try? obs.recognizedPoints(.all) else { continue }
|
|
let p = order.map { name -> SIMD3<Float> in
|
|
if let rp = pts[name] {
|
|
// Vision: normalized, bottom-left origin -> flip y.
|
|
return SIMD3(Float(rp.location.x),
|
|
Float(1.0 - rp.location.y),
|
|
Float(rp.confidence))
|
|
}
|
|
return .zero
|
|
}
|
|
handsOut.append(HandsPayload.Hand(
|
|
isRight: obs.chirality == .right, points: p))
|
|
}
|
|
|
|
let handsData = HandsPayload(hands: handsOut).encoded()
|
|
Task { @MainActor in
|
|
guard self.usbState == .connected else { return }
|
|
self.usb.send(tag: .hands, pid: -1,
|
|
timestamp: t, payload: handsData)
|
|
}
|
|
}
|
|
|
|
// MARK: - USB skeleton
|
|
|
|
private func publishUSB(pid: Int, timestamp: TimeInterval,
|
|
body: ARBodyAnchor, camera: ARCamera) {
|
|
guard usbState == .connected else { return }
|
|
let skeleton = body.skeleton
|
|
let transforms = skeleton.jointModelTransforms
|
|
let root = body.transform
|
|
var payload = SkeletonPayload()
|
|
let n = min(SkeletonPayload.jointCount, transforms.count)
|
|
for i in 0..<n {
|
|
let w = root * transforms[i]
|
|
payload.joints[i] = SIMD3(w.columns.3.x,
|
|
w.columns.3.y,
|
|
w.columns.3.z)
|
|
payload.valid[i] = skeleton.isJointTracked(i)
|
|
}
|
|
usb.send(tag: .skeleton,
|
|
pid: Int16(clamping: pid),
|
|
timestamp: timestamp,
|
|
payload: payload.encoded())
|
|
|
|
// 2D projected skeleton: viewport-normalized [0,1] coords per joint.
|
|
var p2 = Skeleton2DPayload()
|
|
let orient = currentInterfaceOrientation()
|
|
for i in 0..<min(Skeleton2DPayload.jointCount, transforms.count) {
|
|
let w = root * transforms[i]
|
|
let p3 = SIMD3<Float>(w.columns.3.x, w.columns.3.y, w.columns.3.z)
|
|
let proj = camera.projectPoint(p3,
|
|
orientation: orient,
|
|
viewportSize: viewportSize)
|
|
let nx = Float(proj.x / max(viewportSize.width, 1))
|
|
let ny = Float(proj.y / max(viewportSize.height, 1))
|
|
if nx.isFinite && ny.isFinite {
|
|
p2.points[i] = SIMD2<Float>(nx, ny)
|
|
p2.valid[i] = skeleton.isJointTracked(i)
|
|
}
|
|
}
|
|
usb.send(tag: .skeleton2D,
|
|
pid: Int16(clamping: pid),
|
|
timestamp: timestamp,
|
|
payload: p2.encoded())
|
|
}
|
|
|
|
/// Send the ARKit body joint names + parent indices once, so the Mac
|
|
/// renders bones generically. Static data — available before any body
|
|
/// is tracked, so it is safe to emit right after connect.
|
|
@MainActor
|
|
private func publishTopology() {
|
|
guard usbState == .connected else { return }
|
|
let def = ARSkeletonDefinition.defaultBody3D
|
|
let names = def.jointNames
|
|
let parents = def.parentIndices.map { Int16(clamping: $0) }
|
|
let payload = TopologyPayload(jointNames: names, parents: parents)
|
|
usb.send(tag: .topology, pid: -1,
|
|
timestamp: lastFrameTime,
|
|
payload: payload.encoded())
|
|
}
|
|
|
|
}
|