From c0834ed7ab9441c8a3ad01bbf6c556e6c0bfce1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=27=C3=A9lectron=20rare?= <108685187+electron-rare@users.noreply.github.com> Date: Mon, 18 May 2026 21:33:12 +0200 Subject: [PATCH] feat(avlivebody-mac): multi-hmr and body fusion Context: Task 4 of the macOS rewrite needs the dense-mesh half of the pipeline alongside the USB skeleton consumer landed in task 3. Approach: Add a CoreML wrapper that mirrors the validated Python reference (data_only_viz/multihmr_coreml.py) and a pure-logic fusion stage that corrects the mesh pelvis depth using the LiDAR-precise USB skeleton. Changes: - MultiHMRCoreML.swift: 1x3x672x672 ImageNet-normalized image input, 1x3x3 cam_K input, K=4 SMPL-X person outputs at 10475 vertices, det threshold 0.3. - BodyFusion.swift: stateless fuse(persons, skeletons) overrides the highest-score mesh translation.z with the skeleton pelvis Z when available, passes through otherwise. - BodyFusionTests.swift: pelvis override and pass-through cases. Impact: Unlocks the mesh renderer wiring in later tasks and gives the macOS app metrically-correct depth in front of the camera. --- .../Sources/AVLiveBody/usb/BodyFusion.swift | 26 +++ .../AVLiveBody/usb/MultiHMRCoreML.swift | 156 ++++++++++++++++++ .../AVLiveBodyTests/BodyFusionTests.swift | 29 ++++ 3 files changed, 211 insertions(+) create mode 100644 avlivebody-mac/Sources/AVLiveBody/usb/BodyFusion.swift create mode 100644 avlivebody-mac/Sources/AVLiveBody/usb/MultiHMRCoreML.swift create mode 100644 avlivebody-mac/Tests/AVLiveBodyTests/BodyFusionTests.swift diff --git a/avlivebody-mac/Sources/AVLiveBody/usb/BodyFusion.swift b/avlivebody-mac/Sources/AVLiveBody/usb/BodyFusion.swift new file mode 100644 index 0000000..43de1b2 --- /dev/null +++ b/avlivebody-mac/Sources/AVLiveBody/usb/BodyFusion.swift @@ -0,0 +1,26 @@ +import AVLiveWire +import Foundation +import simd + +/// Associates Multi-HMR meshes with USB skeletons and corrects the +/// mesh pelvis depth. Pure, stateless — unit-testable. +enum BodyFusion { + static let pelvisJoint = 0 + + static func fuse(persons: [MultiHMRPerson], + skeletons: [Int: SkeletonPayload]) + -> [MultiHMRPerson] { + let pelvisZs: [Float] = skeletons.values.compactMap { s in + guard pelvisJoint < s.valid.count, + s.valid[pelvisJoint] else { return nil } + return s.joints[pelvisJoint].z + } + guard !pelvisZs.isEmpty, + let primaryIdx = persons.indices.max(by: { + persons[$0].score < persons[$1].score + }) else { return persons } + var out = persons + out[primaryIdx].translation.z = pelvisZs[0] + return out + } +} diff --git a/avlivebody-mac/Sources/AVLiveBody/usb/MultiHMRCoreML.swift b/avlivebody-mac/Sources/AVLiveBody/usb/MultiHMRCoreML.swift new file mode 100644 index 0000000..4472b40 --- /dev/null +++ b/avlivebody-mac/Sources/AVLiveBody/usb/MultiHMRCoreML.swift @@ -0,0 +1,156 @@ +import CoreML +import CoreVideo +import CoreImage +import Foundation + +/// One detected SMPL-X body from Multi-HMR. +struct MultiHMRPerson { + var vertices: [SIMD3] // 10475 SMPL-X verts, model space + var translation: SIMD3 // pelvis translation + var score: Float +} + +/// CoreML wrapper around the bundled `multihmr_full_672_s.mlpackage`. +/// Mirrors `data_only_viz/multihmr_coreml.py`: two MLMultiArray inputs +/// (`image` 1x3x672x672 ImageNet-normalized, `cam_K` 1x3x3), fixed +/// K=4 person outputs. +final class MultiHMRCoreML { + static let inputSize = 672 + static let vertexCount = 10475 + static let maxPersons = 4 + private static let detThreshold: Float = 0.3 + private static let normMean: [Float] = [0.485, 0.456, 0.406] + private static let normStd: [Float] = [0.229, 0.224, 0.225] + + private let model: MLModel + private let ciContext = CIContext() + + /// Loads the bundled model. Returns nil if the resource or load + /// fails — callers fall back to skeleton-only rendering. + init?() { + guard let url = Bundle.main.url( + forResource: "multihmr_full_672_s", + withExtension: "mlpackage") else { + NSLog("MultiHMRCoreML: mlpackage resource missing") + return nil + } + let cfg = MLModelConfiguration() + cfg.computeUnits = .cpuAndGPU + do { + let compiled = try MLModel.compileModel(at: url) + model = try MLModel(contentsOf: compiled, configuration: cfg) + } catch { + NSLog("MultiHMRCoreML: load failed %@", + String(describing: error)) + return nil + } + } + + /// Run inference on one camera frame. `cameraK` is the 3x3 camera + /// intrinsics row-major. + func infer(_ pixelBuffer: CVPixelBuffer, + cameraK: [Float]) -> [MultiHMRPerson] { + guard let image = makeImageInput(pixelBuffer), + let k = makeKInput(cameraK) else { return [] } + let inputs: [String: MLFeatureValue] = [ + "image": MLFeatureValue(multiArray: image), + "cam_K": MLFeatureValue(multiArray: k), + ] + guard let provider = try? MLDictionaryFeatureProvider( + dictionary: inputs), + let out = try? model.prediction(from: provider) else { + return [] + } + return parse(out) + } + + // MARK: - Input preprocessing + + /// `CVPixelBuffer` -> [1,3,672,672] Float32, RGB, ImageNet-normed. + private func makeImageInput(_ pb: CVPixelBuffer) -> MLMultiArray? { + let n = Self.inputSize + // Resize to n x n BGRA via CoreImage. + let ci = CIImage(cvPixelBuffer: pb) + let sx = CGFloat(n) / ci.extent.width + let sy = CGFloat(n) / ci.extent.height + let scaled = ci.transformed( + by: CGAffineTransform(scaleX: sx, y: sy)) + var dst: CVPixelBuffer? + CVPixelBufferCreate(kCFAllocatorDefault, n, n, + kCVPixelFormatType_32BGRA, nil, &dst) + guard let dst else { return nil } + ciContext.render(scaled, to: dst) + CVPixelBufferLockBaseAddress(dst, .readOnly) + defer { CVPixelBufferUnlockBaseAddress(dst, .readOnly) } + guard let base = CVPixelBufferGetBaseAddress(dst) else { + return nil + } + let rowBytes = CVPixelBufferGetBytesPerRow(dst) + let px = base.assumingMemoryBound(to: UInt8.self) + guard let arr = try? MLMultiArray( + shape: [1, 3, NSNumber(value: n), NSNumber(value: n)], + dataType: .float32) else { return nil } + let ptr = arr.dataPointer.assumingMemoryBound(to: Float.self) + let plane = n * n + for y in 0.. [1,3,3] Float32. + private func makeKInput(_ k: [Float]) -> MLMultiArray? { + guard k.count == 9, + let arr = try? MLMultiArray( + shape: [1, 3, 3], dataType: .float32) else { return nil } + let ptr = arr.dataPointer.assumingMemoryBound(to: Float.self) + for i in 0..<9 { ptr[i] = k[i] } + return arr + } + + // MARK: - Output parsing + + private func parse(_ out: MLFeatureProvider) -> [MultiHMRPerson] { + guard let v3d = out.featureValue(for: "var_2420")? + .multiArrayValue, + let transl = out.featureValue(for: "var_2423")? + .multiArrayValue, + let scores = out.featureValue(for: "var_2436")? + .multiArrayValue else { return [] } + var persons: [MultiHMRPerson] = [] + let vc = Self.vertexCount + for k in 0..]( + repeating: .zero, count: vc) + let base = k * vc * 3 + for i in 0.. SkeletonPayload { + var p = SkeletonPayload() + p.joints[0] = SIMD3(0, 0, pelvisZ) + p.valid[0] = true + return p + } + + func testPelvisDepthOverride() { + let mesh = MultiHMRPerson( + vertices: [SIMD3](repeating: .zero, count: 1), + translation: SIMD3(0, 0, -1.0), score: 0.9) + let fused = BodyFusion.fuse( + persons: [mesh], skeletons: [0: skeleton(pelvisZ: -2.5)]) + XCTAssertEqual(fused[0].translation.z, -2.5, accuracy: 1e-4) + } + + func testPassthroughWhenNoSkeleton() { + let mesh = MultiHMRPerson( + vertices: [SIMD3](repeating: .zero, count: 1), + translation: SIMD3(0, 0, -1.0), score: 0.9) + let fused = BodyFusion.fuse(persons: [mesh], skeletons: [:]) + XCTAssertEqual(fused[0].translation.z, -1.0, accuracy: 1e-4) + } +}