Created
September 4, 2024 00:22
-
-
Save rdev/8ca146ebc9e44a293bd8f4e3a307b177 to your computer and use it in GitHub Desktop.
This file contains hidden or 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 AVFoundation | |
| import CoreImage | |
| import SwiftUI | |
| import UniformTypeIdentifiers | |
| let context = CIContext() | |
| struct ContentView: View { | |
| @State var nsImage: NSImage? | |
| @State var ciImage = CIImage(contentsOf: Bundle.main.url(forResource: "1080p", withExtension: "jpg")!)! | |
| let TOTAL_FRAMES = 1_500 | |
| var body: some View { | |
| VStack { | |
| if let nsImage { | |
| Image(nsImage: nsImage) | |
| .resizable() | |
| .scaledToFit() | |
| } | |
| Button("CROP ME BABY, WITH A VIDEO") { | |
| let outputURL = FileManager.default.temporaryDirectory.appending(component: "output.mp4") // Set the output file path | |
| let started = Date() | |
| self.createVideo( | |
| outputURL: outputURL, | |
| frameCount: self.TOTAL_FRAMES, | |
| fps: 60, | |
| generateCIImage: { frameIndex in | |
| self.generateFrame(frame: frameIndex) | |
| }) { success in | |
| if success { | |
| let ended = Date() | |
| print("Video created in \(ended.timeIntervalSince(started).toString()) ms") | |
| NSWorkspace.shared.open(outputURL) | |
| } else { | |
| print("Failed to create video.") | |
| } | |
| } | |
| } | |
| } | |
| .padding() | |
| .onAppear { | |
| // Set preview | |
| self.nsImage = NSImage.fromCIImage(self.ciImage) | |
| } | |
| } | |
| func generateFrame(frame: Int) -> CIImage { | |
| // Dimension stuff | |
| let width = self.ciImage.extent.width | |
| let height = self.ciImage.extent.height | |
| let frameStarted = Date() | |
| let scale = 1 + Double(frame) * 0.001 | |
| let adjustedWidth = width * scale | |
| let adjustedHeight = height * scale | |
| let offsetX = (width - adjustedWidth) / 2 | |
| let offsetY = (height - adjustedHeight) / 2 | |
| // Affine transform | |
| let affineTransform = CIFilter( | |
| name: "CIAffineTransform", | |
| parameters: [ | |
| "inputImage": ciImage as Any, | |
| "inputTransform": AffineTransform(m11: scale, m12: 0.0, m21: 0.0, m22: scale, tX: offsetX, tY: offsetY) | |
| ]) | |
| // Bail if unsuccessful | |
| guard let transformedImage = affineTransform?.outputImage else { | |
| return CIImage(color: .red) // return broken frame | |
| } | |
| // Crop transformed image to original bounds | |
| let cropRect = CGRect(x: 0, y: 0, width: width, height: height) | |
| let cropFilter = CIFilter( | |
| name: "CICrop", | |
| parameters: [ | |
| "inputImage": transformedImage, | |
| "inputRectangle": CIVector(cgRect: cropRect) | |
| ]) | |
| // Bail if unsuccessful x2 | |
| guard let outputImage = cropFilter?.outputImage else { | |
| return CIImage(color: .red) // return broken frame | |
| } | |
| let textFilter = CIFilter( | |
| name: "CITextImageGenerator", | |
| parameters: [ | |
| "inputText": "Frame: \(String(format: "%04d", frame))", | |
| "inputFontName": "Helvetica", | |
| "inputFontSize": 36, | |
| "inputScaleFactor": 1.0, | |
| "inputPadding": 50 | |
| ]) | |
| let blendFilter = CIFilter( | |
| name: "CISourceOverCompositing", | |
| parameters: [ | |
| "inputImage" : textFilter?.outputImage as Any, | |
| "inputBackgroundImage" : outputImage as Any | |
| ]) | |
| let frameEnded = Date() | |
| print("Frame rendered in \(frameEnded.timeIntervalSince(frameStarted).toString()) ms") | |
| // Literally have to delay to make intermediate crops visible | |
| DispatchQueue.main.asyncAfter(deadline: .now() + 0.001 * Double(frame)) { | |
| self.nsImage = NSImage.fromCIImage(outputImage) | |
| } | |
| return blendFilter!.outputImage! | |
| } | |
| func createVideo(outputURL: URL, frameCount: Int, fps: Int32 = 30, generateCIImage: (Int) -> CIImage, completion: @escaping (Bool) -> Void) { | |
| // Video settings | |
| let exampleImage = generateCIImage(0) // Generate the first image to get the size | |
| let width = Int(exampleImage.extent.width) | |
| let height = Int(exampleImage.extent.height) | |
| let videoSettings: [String: Any] = [ | |
| AVVideoCodecKey: AVVideoCodecType.h264, | |
| AVVideoWidthKey: width, | |
| AVVideoHeightKey: height | |
| ] | |
| // Create AVAssetWriter | |
| guard let assetWriter = try? AVAssetWriter(outputURL: outputURL, fileType: .mp4) else { | |
| completion(false) | |
| return | |
| } | |
| // Create AVAssetWriterInput for video | |
| let writerInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoSettings) | |
| writerInput.expectsMediaDataInRealTime = false | |
| // Create pixel buffer adapter | |
| let sourcePixelBufferAttributes: [String: Any] = [ | |
| kCVPixelBufferPixelFormatTypeKey as String: Int(kCVPixelFormatType_32ARGB), | |
| kCVPixelBufferWidthKey as String: width, | |
| kCVPixelBufferHeightKey as String: height | |
| ] | |
| let pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(assetWriterInput: writerInput, sourcePixelBufferAttributes: sourcePixelBufferAttributes) | |
| assetWriter.add(writerInput) | |
| // Start writing | |
| assetWriter.startWriting() | |
| assetWriter.startSession(atSourceTime: .zero) | |
| // Frame timing | |
| let frameDuration = CMTimeMake(value: 1, timescale: fps) | |
| // Write each image frame in the loop | |
| for i in 0 ..< frameCount { | |
| let presentationTime = CMTimeMultiply(frameDuration, multiplier: Int32(i)) | |
| // Generate the CIImage for this frame | |
| let ciImage = generateCIImage(i) | |
| // Create pixel buffer | |
| var pixelBuffer: CVPixelBuffer? | |
| let status = CVPixelBufferPoolCreatePixelBuffer(nil, pixelBufferAdaptor.pixelBufferPool!, &pixelBuffer) | |
| if status == kCVReturnSuccess, let pixelBuffer = pixelBuffer { | |
| context.render(ciImage, to: pixelBuffer) | |
| // Wait until ready and append the pixel buffer to the video | |
| while !writerInput.isReadyForMoreMediaData { | |
| Thread.sleep(forTimeInterval: 0.1) | |
| } | |
| pixelBufferAdaptor.append(pixelBuffer, withPresentationTime: presentationTime) | |
| } | |
| } | |
| // Finish writing | |
| writerInput.markAsFinished() | |
| assetWriter.finishWriting { | |
| completion(assetWriter.status == .completed) | |
| } | |
| } | |
| } | |
| // For measurement | |
| extension Double { | |
| func toString(decimal: Int = 9) -> String { | |
| let value = decimal < 0 ? 0 : decimal | |
| var string = String(format: "%.\(value)f", self * 1_000) | |
| while string.last == "0" || string.last == "." { | |
| if string.last == "." { string = String(string.dropLast()); break } | |
| string = String(string.dropLast()) | |
| } | |
| return string | |
| } | |
| } | |
| // For convenience | |
| extension NSImage { | |
| // Generate an NSImage from a CIImage. | |
| static func fromCIImage(_ ciImage: CIImage) -> NSImage { | |
| let rep = NSCIImageRep(ciImage: ciImage) | |
| let nsImage = NSImage(size: rep.size) | |
| nsImage.addRepresentation(rep) | |
| return nsImage | |
| } | |
| } | |
| func createVideo(outputURL: URL, frameCount: Int, fps: Int32 = 30, generateCIImage: (Int) -> CIImage, completion: @escaping (Bool) -> Void) { | |
| // Generate the first image to get the size | |
| let exampleImage = generateCIImage(0) | |
| let width = Int(exampleImage.extent.width) | |
| let height = Int(exampleImage.extent.height) | |
| let videoSettings: [String: Any] = [ | |
| AVVideoCodecKey: AVVideoCodecType.h264, | |
| AVVideoWidthKey: width, | |
| AVVideoHeightKey: height | |
| ] | |
| // Create AVAssetWriter | |
| guard let assetWriter = try? AVAssetWriter(outputURL: outputURL, fileType: .mp4) else { | |
| print("Failed to create AVAssetWriter") | |
| completion(false) | |
| return | |
| } | |
| // Create AVAssetWriterInput for video | |
| let writerInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoSettings) | |
| writerInput.expectsMediaDataInRealTime = false | |
| // Create pixel buffer adapter | |
| let sourcePixelBufferAttributes: [String: Any] = [ | |
| kCVPixelBufferPixelFormatTypeKey as String: Int(kCVPixelFormatType_32ARGB), | |
| kCVPixelBufferWidthKey as String: width, | |
| kCVPixelBufferHeightKey as String: height | |
| ] | |
| let pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(assetWriterInput: writerInput, sourcePixelBufferAttributes: sourcePixelBufferAttributes) | |
| guard let pixelBufferPool = pixelBufferAdaptor.pixelBufferPool else { | |
| print("Failed to create pixel buffer pool") | |
| completion(false) | |
| return | |
| } | |
| assetWriter.add(writerInput) | |
| // Start writing | |
| assetWriter.startWriting() | |
| assetWriter.startSession(atSourceTime: .zero) | |
| let context = CIContext() | |
| let frameDuration = CMTimeMake(value: 1, timescale: fps) | |
| for i in 0 ..< frameCount { | |
| let presentationTime = CMTimeMultiply(frameDuration, multiplier: Int32(i)) | |
| // Generate the CIImage for this frame | |
| let ciImage = generateCIImage(i) | |
| // Create pixel buffer | |
| var pixelBuffer: CVPixelBuffer? | |
| let status = CVPixelBufferPoolCreatePixelBuffer(nil, pixelBufferPool, &pixelBuffer) | |
| if status == kCVReturnSuccess, let pixelBuffer = pixelBuffer { | |
| context.render(ciImage, to: pixelBuffer) | |
| // Wait until ready and append the pixel buffer to the video | |
| while !writerInput.isReadyForMoreMediaData { | |
| Thread.sleep(forTimeInterval: 0.1) | |
| } | |
| pixelBufferAdaptor.append(pixelBuffer, withPresentationTime: presentationTime) | |
| } else { | |
| print("Failed to create pixel buffer at frame \(i)") | |
| } | |
| } | |
| // Finish writing | |
| writerInput.markAsFinished() | |
| assetWriter.finishWriting { | |
| completion(assetWriter.status == .completed) | |
| } | |
| } | |
| #Preview { | |
| ContentView() | |
| .frame(width: 600, height: 460) | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment