From 1b2c725bab80542f067c50cc4dcefa46a71f2a99 Mon Sep 17 00:00:00 2001
From: Daniel Perry <dperry@danieltperry.me>
Date: Sat, 29 Jun 2024 12:41:21 -0400
Subject: [PATCH] Add a bitrate command line argument to control bitrate of
 output video

---
 Sources/Upscaling/UpscalingExportSession.swift | 14 ++++++++++++--
 Sources/fx-upscale/FXUpscale.swift             |  2 ++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/Sources/Upscaling/UpscalingExportSession.swift b/Sources/Upscaling/UpscalingExportSession.swift
index 8bb0f5b..8f8200c 100644
--- a/Sources/Upscaling/UpscalingExportSession.swift
+++ b/Sources/Upscaling/UpscalingExportSession.swift
@@ -11,6 +11,7 @@ public class UpscalingExportSession {
         outputCodec: AVVideoCodecType? = nil,
         preferredOutputURL: URL,
         outputSize: CGSize,
+        outputBitRate: Int? = nil,
         creator: String? = nil
     ) {
         self.asset = asset
@@ -23,6 +24,7 @@ public class UpscalingExportSession {
             outputURL = preferredOutputURL
         }
         self.outputSize = outputSize
+        self.outputBitRate = outputBitRate
         self.creator = creator
         progress = Progress(parent: nil, userInfo: [
             .fileURLKey: outputURL
@@ -42,6 +44,7 @@ public class UpscalingExportSession {
     public let outputCodec: AVVideoCodecType?
     public let outputURL: URL
     public let outputSize: CGSize
+    public let outputBitRate: Int?
     public let creator: String?
 
     public let progress: Progress
@@ -80,7 +83,8 @@ public class UpscalingExportSession {
                     for: track,
                     formatDescription: formatDescription,
                     outputSize: outputSize,
-                    outputCodec: outputCodec
+                    outputCodec: outputCodec,
+                    outputBitRate: outputBitRate
                 ) else { continue }
 
             if assetReader.canAdd(assetReaderOutput) {
@@ -275,7 +279,8 @@ public class UpscalingExportSession {
         for track: AVAssetTrack,
         formatDescription: CMFormatDescription?,
         outputSize: CGSize,
-        outputCodec: AVVideoCodecType?
+        outputCodec: AVVideoCodecType?,
+        outputBitRate: Int?
     ) async throws -> AVAssetWriterInput? {
         switch track.mediaType {
         case .video:
@@ -297,6 +302,11 @@ public class UpscalingExportSession {
                formatDescription?.hasLeftAndRightEye ?? false {
                 var compressionProperties: [CFString: Any] = [:]
                 compressionProperties[kVTCompressionPropertyKey_MVHEVCVideoLayerIDs] = [0, 1]
+                
+                if let outputBitRate {
+                    compressionProperties[AVVideoAverageBitRateKey as CFString] = outputBitRate
+                }
+                
                 if let extensions = formatDescription?.extensions {
                     for key in [
                         kVTCompressionPropertyKey_HeroEye,
diff --git a/Sources/fx-upscale/FXUpscale.swift b/Sources/fx-upscale/FXUpscale.swift
index 8b745cd..aeb5a11 100644
--- a/Sources/fx-upscale/FXUpscale.swift
+++ b/Sources/fx-upscale/FXUpscale.swift
@@ -11,6 +11,7 @@ import Upscaling
 
     @Option(name: .shortAndLong, help: "The output file width") var width: Int?
     @Option(name: .shortAndLong, help: "The output file height") var height: Int?
+    @Option(name: .shortAndLong, help: "The output file bitrate") var bitrate: Int?
 
     mutating func run() async throws {
         guard ["mov", "m4v", "mp4"].contains(url.pathExtension.lowercased()) else {
@@ -60,6 +61,7 @@ import Upscaling
             outputCodec: convertToProRes ? .proRes422 : nil,
             preferredOutputURL: url.renamed { "\($0) Upscaled" },
             outputSize: outputSize,
+            outputBitRate: bitrate,
             creator: ProcessInfo.processInfo.processName
         )