Created
August 26, 2024 00:10
-
-
Save Matt54/c24e9163d70905fe3b14cb5fc209c9f8 to your computer and use it in GitHub Desktop.
RealityView loading a USDZ model, extracting a LowLevelMesh from it, and then animating a warp reveal of it's vertices
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import SwiftUI | |
| import RealityKit | |
| import Metal | |
| struct LowLevelMeshRevealView: View { | |
| @State var entity: ModelEntity? | |
| @State var lowLevelMesh: LowLevelMesh? | |
| @State var originalVertices: [VertexData] = [] | |
| @State var timer: Timer? | |
| @State var isForward: Bool = true | |
| @State var revealProgress: Float = -0.55 // Start slightly below the model | |
| let timerUpdateDuration: TimeInterval = 1/120.0 | |
| let bottomValue: Float = -0.375 | |
| let topValue: Float = 0.55 | |
| let device: MTLDevice | |
| let commandQueue: MTLCommandQueue | |
| let computePipelineState: MTLComputePipelineState | |
| init() { | |
| self.device = MTLCreateSystemDefaultDevice()! | |
| self.commandQueue = device.makeCommandQueue()! | |
| let library = device.makeDefaultLibrary()! | |
| let kernelFunction = library.makeFunction(name: "displacementRevealKernel")! | |
| self.computePipelineState = try! device.makeComputePipelineState(function: kernelFunction) | |
| } | |
| var body: some View { | |
| RealityView { content in | |
| let model = try! await loadModelEntity() | |
| content.add(model) | |
| let lowLevelMesh = createMesh(from: model)! | |
| // Store original vertex data | |
| lowLevelMesh.withUnsafeBytes(bufferIndex: 0) { buffer in | |
| let vertices = buffer.bindMemory(to: VertexData.self) | |
| self.originalVertices = Array(vertices) | |
| } | |
| model.model?.mesh = try! await MeshResource(from: lowLevelMesh) | |
| self.entity = model | |
| self.lowLevelMesh = lowLevelMesh | |
| } | |
| .onAppear { startTimer() } | |
| .onDisappear { stopTimer() } | |
| } | |
| func startTimer() { | |
| timer = Timer.scheduledTimer(withTimeInterval: timerUpdateDuration, repeats: true) { timer in | |
| let stepValue: Float = 0.0025 | |
| if isForward { | |
| revealProgress += stepValue | |
| } else { | |
| revealProgress -= stepValue | |
| } | |
| if revealProgress >= topValue { | |
| revealProgress = topValue | |
| isForward = false | |
| } else if revealProgress < bottomValue { | |
| revealProgress = bottomValue | |
| isForward = true | |
| } | |
| print(revealProgress) | |
| updateMesh() | |
| } | |
| } | |
| func stopTimer() { | |
| timer?.invalidate() | |
| timer = nil | |
| } | |
| func createMesh(from modelEntity: ModelEntity) -> LowLevelMesh? { | |
| guard let meshResource = modelEntity.model?.mesh, | |
| let meshPart = meshResource.contents.models.first?.parts.first else { | |
| return nil | |
| } | |
| let positions = meshPart[MeshBuffers.positions]?.elements ?? [] | |
| let normals = meshPart[MeshBuffers.normals]?.elements ?? [] | |
| let textureCoordinates = meshPart[MeshBuffers.textureCoordinates]?.elements ?? [] | |
| let triangleIndices = meshPart.triangleIndices?.elements ?? [] | |
| print("positions.count \(positions.count)") | |
| print("normals.count \(normals.count)") | |
| print("textureCoordinates.count \(textureCoordinates.count)") | |
| print("triangleIndices.count \(triangleIndices.count)") | |
| var descriptor = VertexData.descriptor | |
| descriptor.vertexCapacity = positions.count | |
| descriptor.indexCapacity = triangleIndices.count | |
| guard let lowLevelMesh = try? LowLevelMesh(descriptor: descriptor) else { | |
| return nil | |
| } | |
| // Copy vertex data | |
| lowLevelMesh.withUnsafeMutableBytes(bufferIndex: 0) { buffer in | |
| let vertices = buffer.bindMemory(to: (SIMD3<Float>, SIMD3<Float>, SIMD2<Float>).self) | |
| for i in 0..<positions.count { | |
| vertices[i] = (positions[i], normals[i], textureCoordinates[i]) | |
| } | |
| } | |
| // Copy index data | |
| lowLevelMesh.withUnsafeMutableIndices { buffer in | |
| let indices = buffer.bindMemory(to: UInt32.self) | |
| for (index, triangleIndex) in triangleIndices.enumerated() { | |
| indices[index] = UInt32(triangleIndex) | |
| } | |
| } | |
| // Set up parts | |
| let bounds = meshResource.bounds | |
| lowLevelMesh.parts.replaceAll([ | |
| LowLevelMesh.Part( | |
| indexCount: triangleIndices.count, | |
| topology: .triangle, | |
| bounds: bounds | |
| ) | |
| ]) | |
| return lowLevelMesh | |
| } | |
| func updateMesh() { | |
| guard let mesh = lowLevelMesh, | |
| let commandBuffer = commandQueue.makeCommandBuffer(), | |
| let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return } | |
| // Reset mesh to original state | |
| mesh.withUnsafeMutableBytes(bufferIndex: 0) { buffer in | |
| let vertices = buffer.bindMemory(to: VertexData.self) | |
| for i in 0..<originalVertices.count { | |
| vertices[i] = originalVertices[i] | |
| } | |
| } | |
| let vertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer) | |
| computeEncoder.setComputePipelineState(computePipelineState) | |
| computeEncoder.setBuffer(vertexBuffer, offset: 0, index: 0) | |
| computeEncoder.setBytes(&revealProgress, length: MemoryLayout<Float>.size, index: 1) | |
| let threadsPerGrid = MTLSize(width: mesh.vertexCapacity, height: 1, depth: 1) | |
| let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1) | |
| computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup) | |
| computeEncoder.endEncoding() | |
| commandBuffer.commit() | |
| } | |
| func loadModelEntity(url: URL = URL(string: "https://matt54.github.io/Resources/Anubis_Statue_1.usdz")!) async throws -> ModelEntity { | |
| let (downloadedURL, _) = try await URLSession.shared.download(from: url) | |
| let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first! | |
| let destinationURL = documentsDirectory.appendingPathComponent("downloadedModel.usdz") | |
| if FileManager.default.fileExists(atPath: destinationURL.path) { | |
| try FileManager.default.removeItem(at: destinationURL) | |
| } | |
| try FileManager.default.moveItem(at: downloadedURL, to: destinationURL) | |
| let entity = try await ModelEntity.init(contentsOf: destinationURL) | |
| try FileManager.default.removeItem(at: destinationURL) | |
| return entity | |
| } | |
| struct VertexData { | |
| var position: SIMD3<Float> = .zero | |
| var normal: SIMD3<Float> = .zero | |
| var uv: SIMD2<Float> = .zero | |
| static var vertexAttributes: [LowLevelMesh.Attribute] = [ | |
| .init(semantic: .position, format: .float3, offset: MemoryLayout<Self>.offset(of: \.position)!), | |
| .init(semantic: .normal, format: .float3, offset: MemoryLayout<Self>.offset(of: \.normal)!), | |
| .init(semantic: .uv0, format: .float2, offset: MemoryLayout<Self>.offset(of: \.uv)!) | |
| ] | |
| static var vertexLayouts: [LowLevelMesh.Layout] = [ | |
| .init(bufferIndex: 0, bufferStride: MemoryLayout<Self>.stride) | |
| ] | |
| static var descriptor: LowLevelMesh.Descriptor { | |
| var desc = LowLevelMesh.Descriptor() | |
| desc.vertexAttributes = VertexData.vertexAttributes | |
| desc.vertexLayouts = VertexData.vertexLayouts | |
| desc.indexType = .uint32 | |
| return desc | |
| } | |
| } | |
| } | |
| #Preview { | |
| LowLevelMeshRevealView() | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment