feat(ios): HEVC video capture, drop OSC

Add VideoEncoder (VideoToolbox HEVC) and stream encoded frames over USB as .video AVLiveWire frames alongside the skeleton. Remove the legacy OSC/UDP fanout and its host/port config UI — the iPhone link is now USB-only.
2026-05-18 17:27:34 +02:00
parent 39457dde41
commit 1af056dbf1
3 changed files with 217 additions and 117 deletions
@@ -1,7 +1,7 @@
 import ARKit
+import AVLiveWire
 import Combine
 import Foundation
-import Network
 import RealityKit
 import SwiftUI

@@ -31,15 +31,15 @@ final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {
    @Published var jointsPerSec: Double = 0
    @Published var bodyCount: Int = 0
    @Published var skeleton2D: SkeletonSnapshot?
+    @Published var usbState: USBServer.State = .idle
    /// Set by the SwiftUI view via GeometryReader so the projection
    /// matches the on-screen ARView size.
    var viewportSize: CGSize = .zero
-    private var host: String = "192.168.0.159"
-    private var pythonPort: UInt16 = 57128
-    private var swiftPort: UInt16 = 57129
    private var sendEnvMesh: Bool = false
    private let session = ARSession()
-    private var conns: [NWConnection] = []
+    private let usb = USBServer()
+    private let videoEncoder = VideoEncoder()
+    private var videoStarted = false
    private var lastFrameTime: TimeInterval = 0
    private var jointsInSecond: Int = 0
    private var lastSecond: TimeInterval = 0
@@ -54,13 +54,22 @@ final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {
        arView.session.delegate = self
        arView.environment.background = .color(.black)
        arView.debugOptions = []
+        usb.onState = { [weak self] s in
+            Task { @MainActor in self?.usbState = s }
+        }
+        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())
+            }
+        }
    }

-    func configure(host: String, pythonPort: UInt16, swiftPort: UInt16,
-                   sendEnvMesh: Bool) {
-        self.host = host
-        self.pythonPort = pythonPort
-        self.swiftPort = swiftPort
+    func configure(sendEnvMesh: Bool) {
        self.sendEnvMesh = sendEnvMesh
    }

@@ -83,7 +92,7 @@ final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {
            feats.append("env-mesh: requires separate session (TODO)")
        }
        cfg.automaticImageScaleEstimationEnabled = true
-        openUDP()
+        usb.start()
        session.run(cfg, options: [.resetTracking, .removeExistingAnchors])
        status = feats.isEmpty
            ? "running (RGB only)"
@@ -93,32 +102,13 @@ final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {

    func stop() {
        session.pause()
-        for c in conns { c.cancel() }
-        conns.removeAll()
+        usb.stop()
+        videoEncoder.stop()
+        videoStarted = false
        running = false
        status = "stopped"
    }

-    // MARK: - UDP fanout
-
-    private func openUDP() {
-        let ports: [UInt16] = [pythonPort, swiftPort]
-        for p in ports where p != 0 {
-            guard let nwPort = NWEndpoint.Port(rawValue: p) else { continue }
-            let conn = NWConnection(
-                to: .hostPort(host: NWEndpoint.Host(host), port: nwPort),
-                using: .udp)
-            conn.start(queue: .global(qos: .userInitiated))
-            conns.append(conn)
-        }
-    }
-
-    private func sendDatagram(_ data: Data) {
-        for c in conns {
-            c.send(content: data, completion: .idempotent)
-        }
-    }
-
    // MARK: - ARSessionDelegate

    nonisolated func session(_ s: ARSession, didUpdate frame: ARFrame) {
@@ -128,16 +118,26 @@ final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {
            if t - self.lastFrameTime < 1.0 / 30.0 { return }
            self.lastFrameTime = t

+            // 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)
+            }
+
            var count: Int = 0
            var firstBody: ARBodyAnchor?
            for anchor in frame.anchors {
                guard let body = anchor as? ARBodyAnchor else { continue }
-                self.publishJoints(pid: count, body: body)
+                self.publishUSB(pid: count, timestamp: t, body: body)
                if count == 0 { firstBody = body }
                count += 1
            }
-            self.sendOSC(addr: "/body3d/count",
-                         args: [.int32(Int32(count))])
            self.framesSent &+= 1
            self.bodyCount = count
            self.updateSkeleton2D(body: firstBody, camera: frame.camera)
@@ -191,64 +191,27 @@ final class ARBodySession: NSObject, ObservableObject, ARSessionDelegate {
    /// indices without re-reading the ARKit skeleton definition.
    var bodyParentIndices: [Int] { bodyParents }

-    private func publishJoints(pid: Int, body: ARBodyAnchor) {
+    private func publishUSB(pid: Int, timestamp: TimeInterval,
+                             body: ARBodyAnchor) {
+        guard usbState == .connected else { return }
        let skeleton = body.skeleton
        let transforms = skeleton.jointModelTransforms
        let root = body.transform
-        for (idx, m) in transforms.enumerated() {
-            let world = root * m
-            sendOSC(addr: "/body3d/kp",
-                    args: [.int32(Int32(pid)),
-                           .int32(Int32(idx)),
-                           .float32(world.columns.3.x),
-                           .float32(world.columns.3.y),
-                           .float32(world.columns.3.z)])
+        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())
    }

-    // MARK: - OSC minimal encoder
-
-    enum OSCArg {
-        case int32(Int32)
-        case float32(Float)
-        case string(String)
-    }
-
-    private func sendOSC(addr: String, args: [OSCArg]) {
-        var data = Data()
-        appendOSCString(addr, into: &data)
-        var types = ","
-        for a in args {
-            switch a {
-            case .int32: types.append("i")
-            case .float32: types.append("f")
-            case .string: types.append("s")
-            }
-        }
-        appendOSCString(types, into: &data)
-        for a in args {
-            switch a {
-            case .int32(let v):
-                var be = v.bigEndian
-                withUnsafeBytes(of: &be) { data.append(contentsOf: $0) }
-            case .float32(let v):
-                var be = v.bitPattern.bigEndian
-                withUnsafeBytes(of: &be) { data.append(contentsOf: $0) }
-            case .string(let s):
-                appendOSCString(s, into: &data)
-            }
-        }
-        sendDatagram(data)
-    }
-
-    private func appendOSCString(_ s: String, into data: inout Data) {
-        let bytes = Array(s.utf8) + [0]
-        data.append(contentsOf: bytes)
-        let pad = (4 - data.count % 4) % 4
-        if pad > 0 {
-            data.append(contentsOf: [UInt8](repeating: 0, count: pad))
-        }
-    }
 }

 struct ARViewContainer: UIViewRepresentable {
@@ -4,9 +4,6 @@ import RealityKit

 struct ContentView: View {
    @StateObject private var session = ARBodySession()
-    @State private var host: String = "192.168.0.159"
-    @State private var pythonPort: String = "57128"   // -> data_only_viz IphoneOSCListener
-    @State private var swiftPort: String = "57129"    // -> AVLiveBody ArkitOSCListener (diagnostic)
    @State private var sendEnvMesh: Bool = false

    /// Replace the live ARView with a gradient placeholder. Camera is
@@ -40,6 +37,22 @@ struct ContentView: View {
        }
    }

+    private var usbDotColor: Color {
+        switch session.usbState {
+        case .idle:       return .gray
+        case .listening:  return .yellow
+        case .connected:  return .green
+        }
+    }
+
+    private var usbStateLabel: String {
+        switch session.usbState {
+        case .idle:       return "idle"
+        case .listening:  return "listening :\(USBServer.port)"
+        case .connected:  return "connected"
+        }
+    }
+
    @ViewBuilder
    private var cameraBackground: some View {
        if useMockBackground {
@@ -66,24 +79,6 @@ struct ContentView: View {
            Text("AR Body → AV-Live")
                .font(.headline)
                .foregroundColor(.white)
-            HStack {
-                Text("Host").foregroundColor(.white)
-                TextField("GrosMac IP", text: $host)
-                    .keyboardType(.numbersAndPunctuation)
-                    .textFieldStyle(.roundedBorder)
-            }
-            HStack {
-                Text("Py").foregroundColor(.white)
-                TextField("57128", text: $pythonPort)
-                    .keyboardType(.numberPad)
-                    .textFieldStyle(.roundedBorder)
-                    .frame(width: 70)
-                Text("Swift").foregroundColor(.white)
-                TextField("57129", text: $swiftPort)
-                    .keyboardType(.numberPad)
-                    .textFieldStyle(.roundedBorder)
-                    .frame(width: 70)
-            }
            Toggle(isOn: $sendEnvMesh) {
                Text("Env mesh (LiDAR)").foregroundColor(.white)
            }
@@ -92,11 +87,7 @@ struct ContentView: View {
                    if session.running {
                        session.stop()
                    } else {
-                        session.configure(
-                            host: host,
-                            pythonPort: UInt16(pythonPort) ?? 57128,
-                            swiftPort: UInt16(swiftPort) ?? 57129,
-                            sendEnvMesh: sendEnvMesh)
+                        session.configure(sendEnvMesh: sendEnvMesh)
                        session.start()
                    }
                }
@@ -109,9 +100,18 @@ struct ContentView: View {
                    .background(.black.opacity(0.5))
                    .cornerRadius(6)
            }
-            Text("bodies: \(session.bodyCount)  frames: \(session.framesSent)  joints/s: \(Int(session.jointsPerSec))")
-                .font(.caption2)
-                .foregroundColor(.white)
+            HStack(spacing: 6) {
+                Circle()
+                    .fill(usbDotColor)
+                    .frame(width: 8, height: 8)
+                Text("USB \(usbStateLabel)")
+                    .font(.caption2)
+                    .foregroundColor(.white.opacity(0.8))
+                Spacer(minLength: 8)
+                Text("bodies: \(session.bodyCount)  frames: \(session.framesSent)  j/s: \(Int(session.jointsPerSec))")
+                    .font(.caption2)
+                    .foregroundColor(.white)
+            }
        }
        .padding(12)
        .background(.black.opacity(0.5))
@@ -0,0 +1,137 @@
+import AVLiveWire
+import CoreMedia
+import CoreVideo
+import Foundation
+import VideoToolbox
+
+/// Hardware HEVC encoder. Feed `CVPixelBuffer`s from ARKit frames in;
+/// receive one `VideoPayload` per encoded access unit via `onPayload`.
+/// Keyframe payloads carry the VPS/SPS/PPS parameter sets prepended,
+/// each as a 4-byte big-endian length prefix followed by the NAL
+/// bytes, so the Mac decoder can build its format description without
+/// a side channel.
+final class VideoEncoder {
+    var onPayload: ((VideoPayload) -> Void)?
+
+    private var session: VTCompressionSession?
+    private let lock = NSLock()
+
+    /// Create the compression session for a given frame size.
+    func start(width: Int32, height: Int32) {
+        stop()
+        var s: VTCompressionSession?
+        let status = VTCompressionSessionCreate(
+            allocator: kCFAllocatorDefault,
+            width: width, height: height,
+            codecType: kCMVideoCodecType_HEVC,
+            encoderSpecification: nil,
+            imageBufferAttributes: nil,
+            compressedDataAllocator: nil,
+            outputCallback: nil,
+            refcon: nil,
+            compressionSessionOut: &s)
+        guard status == noErr, let s else {
+            NSLog("VideoEncoder: VTCompressionSessionCreate failed %d",
+                  status)
+            return
+        }
+        VTSessionSetProperty(s, key: kVTCompressionPropertyKey_RealTime,
+                             value: kCFBooleanTrue)
+        VTSessionSetProperty(s,
+            key: kVTCompressionPropertyKey_AllowFrameReordering,
+            value: kCFBooleanFalse)
+        VTSessionSetProperty(s,
+            key: kVTCompressionPropertyKey_MaxKeyFrameInterval,
+            value: 30 as CFNumber)
+        VTCompressionSessionPrepareToEncodeFrames(s)
+        lock.lock(); session = s; lock.unlock()
+    }
+
+    /// Encode one frame. `pts` is the capture timestamp in seconds.
+    func encode(_ pixelBuffer: CVPixelBuffer, pts: Double) {
+        lock.lock(); let s = session; lock.unlock()
+        guard let s else { return }
+        let time = CMTime(seconds: pts, preferredTimescale: 1_000_000)
+        VTCompressionSessionEncodeFrame(
+            s, imageBuffer: pixelBuffer, presentationTimeStamp: time,
+            duration: .invalid, frameProperties: nil,
+            infoFlagsOut: nil) { [weak self] status, _, sample in
+                guard status == noErr, let sample else { return }
+                self?.handle(sample)
+            }
+    }
+
+    func stop() {
+        lock.lock(); let s = session; session = nil; lock.unlock()
+        if let s {
+            VTCompressionSessionInvalidate(s)
+        }
+    }
+
+    deinit { stop() }
+
+    // MARK: - Sample -> VideoPayload
+
+    private func handle(_ sample: CMSampleBuffer) {
+        let isKeyframe = !Self.notSync(sample)
+        var out = Data()
+        if isKeyframe,
+           let fmt = CMSampleBufferGetFormatDescription(sample) {
+            out.append(Self.parameterSets(fmt))
+        }
+        if let block = CMSampleBufferGetDataBuffer(sample) {
+            var lengthOut = 0
+            var ptr: UnsafeMutablePointer<Int8>?
+            if CMBlockBufferGetDataPointer(
+                block, atOffset: 0, lengthAtOffsetOut: nil,
+                totalLengthOut: &lengthOut,
+                dataPointerOut: &ptr) == noErr, let ptr {
+                out.append(UnsafeBufferPointer(
+                    start: UnsafeRawPointer(ptr)
+                        .assumingMemoryBound(to: UInt8.self),
+                    count: lengthOut))
+            }
+        }
+        guard !out.isEmpty else { return }
+        onPayload?(VideoPayload(isKeyframe: isKeyframe, data: out))
+    }
+
+    /// True if the sample is NOT a sync (key) frame.
+    private static func notSync(_ sample: CMSampleBuffer) -> Bool {
+        guard let arr = CMSampleBufferGetSampleAttachmentsArray(
+            sample, createIfNecessary: false),
+            CFArrayGetCount(arr) > 0 else { return false }
+        let dict = unsafeBitCast(CFArrayGetValueAtIndex(arr, 0),
+                                 to: CFDictionary.self)
+        let key = Unmanaged.passUnretained(
+            kCMSampleAttachmentKey_NotSync).toOpaque()
+        return CFDictionaryContainsKey(dict, key)
+    }
+
+    /// Concatenate the HEVC VPS/SPS/PPS parameter sets, each as a
+    /// 4-byte big-endian length prefix followed by the NAL bytes.
+    private static func parameterSets(
+        _ fmt: CMFormatDescription) -> Data {
+        var count = 0
+        CMVideoFormatDescriptionGetHEVCParameterSetAtIndex(
+            fmt, parameterSetIndex: 0, parameterSetPointerOut: nil,
+            parameterSetSizeOut: nil, parameterSetCountOut: &count,
+            nalUnitHeaderLengthOut: nil)
+        var data = Data()
+        for i in 0..<count {
+            var ptr: UnsafePointer<UInt8>?
+            var size = 0
+            guard CMVideoFormatDescriptionGetHEVCParameterSetAtIndex(
+                fmt, parameterSetIndex: i,
+                parameterSetPointerOut: &ptr,
+                parameterSetSizeOut: &size,
+                parameterSetCountOut: nil,
+                nalUnitHeaderLengthOut: nil) == noErr,
+                let ptr else { continue }
+            var be = UInt32(size).bigEndian
+            withUnsafeBytes(of: &be) { data.append(contentsOf: $0) }
+            data.append(UnsafeBufferPointer(start: ptr, count: size))
+        }
+        return data
+    }
+}