FFMpegProcessor.kt (9453B) - raw

      1 package me.rhunk.snapenhance.download
      2 
      3 import android.media.AudioFormat
      4 import android.media.MediaMetadataRetriever
      5 import com.arthenica.ffmpegkit.FFmpegKit
      6 import com.arthenica.ffmpegkit.FFmpegSession
      7 import com.arthenica.ffmpegkit.Level
      8 import com.arthenica.ffmpegkit.Statistics
      9 import kotlinx.coroutines.suspendCancellableCoroutine
     10 import me.rhunk.snapenhance.LogManager
     11 import me.rhunk.snapenhance.RemoteSideContext
     12 import me.rhunk.snapenhance.common.config.impl.DownloaderConfig
     13 import me.rhunk.snapenhance.common.data.download.AudioStreamFormat
     14 import me.rhunk.snapenhance.common.logger.LogLevel
     15 import me.rhunk.snapenhance.task.PendingTask
     16 import java.io.File
     17 import java.util.concurrent.Executors
     18 
     19 
     20 class ArgumentList  {
     21     private val arguments = mutableListOf<Pair<String, String>>()
     22 
     23     operator fun plusAssign(stringPair: Pair<String, String>) {
     24         arguments += stringPair
     25     }
     26 
     27     operator fun plusAssign(key: String) {
     28         arguments += key to ""
     29     }
     30 
     31     operator fun minusAssign(key: String) {
     32         arguments.removeIf { it.first == key }
     33     }
     34 
     35     operator fun get(key: String) = arguments.find { it.first == key }?.second
     36 
     37     fun forEach(action: (Pair<String, String>) -> Unit) {
     38         arguments.forEach(action)
     39     }
     40 
     41     fun clear() {
     42         arguments.clear()
     43     }
     44 }
     45 
     46 
     47 class FFMpegProcessor(
     48     private val logManager: LogManager,
     49     private val ffmpegOptions: DownloaderConfig.FFMpegOptions,
     50     private val onStatistics: (Statistics) -> Unit = {}
     51 ) {
     52     companion object {
     53         private const val TAG = "ffmpeg-processor"
     54 
     55         fun newFFMpegProcessor(context: RemoteSideContext, pendingTask: PendingTask) = FFMpegProcessor(
     56             logManager = context.log,
     57             ffmpegOptions = context.config.root.downloader.ffmpegOptions,
     58             onStatistics = {
     59                 pendingTask.updateProgress("Processing (frames=${it.videoFrameNumber}, fps=${it.videoFps}, time=${it.time}, bitrate=${it.bitrate}, speed=${it.speed})")
     60             }
     61         )
     62     }
     63     enum class Action {
     64         DOWNLOAD_DASH,
     65         MERGE_OVERLAY,
     66         CONVERSION,
     67         MERGE_MEDIA,
     68         DOWNLOAD_AUDIO_STREAM,
     69     }
     70 
     71     data class Request(
     72         val action: Action,
     73         val inputs: List<String>,
     74         val output: File,
     75         val overlay: File? = null, //only for MERGE_OVERLAY
     76         val startTime: Long? = null, //only for DOWNLOAD_DASH
     77         val duration: Long? = null, //only for DOWNLOAD_DASH
     78         val audioStreamFormat: AudioStreamFormat? = null, //only for DOWNLOAD_AUDIO_STREAM
     79 
     80         var videoCodec: String? = null,
     81         var audioCodec: String? = null,
     82     )
     83 
     84 
     85     private suspend fun newFFMpegTask(globalArguments: ArgumentList, inputArguments: ArgumentList, outputArguments: ArgumentList) = suspendCancellableCoroutine<FFmpegSession> {
     86         val stringBuilder = StringBuilder()
     87         arrayOf(globalArguments, inputArguments, outputArguments).forEach { argumentList ->
     88             argumentList.forEach { (key, value) ->
     89                 stringBuilder.append("$key ${value.takeIf { it.isNotEmpty() }?.plus(" ") ?: ""}")
     90             }
     91         }
     92 
     93         logManager.debug("arguments: $stringBuilder", "FFMpegProcessor")
     94 
     95         FFmpegKit.executeAsync(stringBuilder.toString(),
     96             { session ->
     97                 it.resumeWith(
     98                     if (session.returnCode.isValueSuccess) {
     99                         Result.success(session)
    100                     } else {
    101                         Result.failure(Exception(session.output))
    102                     }
    103                 )
    104             }, logFunction@{ log ->
    105                 logManager.internalLog(TAG, when (log.level) {
    106                     Level.AV_LOG_ERROR, Level.AV_LOG_FATAL -> LogLevel.ERROR
    107                     Level.AV_LOG_WARNING -> LogLevel.WARN
    108                     Level.AV_LOG_VERBOSE -> LogLevel.VERBOSE
    109                     else -> return@logFunction
    110                 }, log.message)
    111             }, { onStatistics(it) }, Executors.newSingleThreadExecutor())
    112     }
    113 
    114     suspend fun execute(args: Request) {
    115         // load ffmpeg native sync to avoid native crash
    116         synchronized(this) { FFmpegKit.listSessions() }
    117         val globalArguments = ArgumentList().apply {
    118             this += "-y"
    119             this += "-threads" to ffmpegOptions.threads.get().toString()
    120         }
    121 
    122         val inputArguments = ArgumentList().apply {
    123             args.inputs.forEach { path ->
    124                 this += "-i" to path
    125             }
    126         }
    127 
    128         val outputArguments = ArgumentList().apply {
    129             this += "-preset" to (ffmpegOptions.preset.getNullable() ?: "ultrafast")
    130             this += "-c:v" to (ffmpegOptions.customVideoCodec.get().takeIf { it.isNotEmpty() } ?: "h264_mediacodec")
    131             this += "-c:a" to (ffmpegOptions.customAudioCodec.get().takeIf { it.isNotEmpty() } ?: "copy")
    132             this += "-crf" to ffmpegOptions.constantRateFactor.get().let { "\"$it\"" }
    133             this += "-b:v" to ffmpegOptions.videoBitrate.get().toString() + "K"
    134             this += "-b:a" to ffmpegOptions.audioBitrate.get().toString() + "K"
    135         }
    136 
    137         when (args.action) {
    138             Action.DOWNLOAD_DASH -> {
    139                 outputArguments += "-ss" to "'${args.startTime}ms'"
    140                 if (args.duration != null) {
    141                     outputArguments += "-t" to "'${args.duration}ms'"
    142                 }
    143             }
    144             Action.MERGE_OVERLAY -> {
    145                 inputArguments += "-i" to args.overlay!!.absolutePath
    146                 outputArguments += "-filter_complex" to "\"[0]scale2ref[img][vid];[img]setsar=1[img];[vid]nullsink;[img][1]overlay=(W-w)/2:(H-h)/2,scale=2*trunc(iw*sar/2):2*trunc(ih/2)\""
    147             }
    148             Action.CONVERSION -> {
    149                 if (ffmpegOptions.customAudioCodec.isEmpty()) {
    150                     outputArguments -= "-c:a"
    151                 }
    152                 outputArguments -= "-c:v"
    153                 args.videoCodec?.let {
    154                     outputArguments += "-c:v" to it
    155                 } ?: run {
    156                     outputArguments += "-vn"
    157                 }
    158                 args.audioCodec?.let {
    159                     outputArguments -= "-c:a"
    160                     outputArguments += "-c:a" to it
    161                 }
    162             }
    163             Action.MERGE_MEDIA -> {
    164                 inputArguments.clear()
    165                 val filesInfo = args.inputs.mapNotNull { file ->
    166                     runCatching {
    167                         MediaMetadataRetriever().apply { setDataSource(file) }
    168                     }.getOrNull()?.let { file to it }
    169                 }
    170 
    171                 val (maxWidth, maxHeight) = filesInfo.maxByOrNull { (_, r) ->
    172                     r.extractMetadata(MediaMetadataRetriever.METADATA_KEY_VIDEO_WIDTH)?.toIntOrNull() ?: 0
    173                 }?.let { (_, r) ->
    174                     r.extractMetadata(MediaMetadataRetriever.METADATA_KEY_VIDEO_WIDTH)?.toIntOrNull() to
    175                     r.extractMetadata(MediaMetadataRetriever.METADATA_KEY_VIDEO_HEIGHT)?.toIntOrNull()
    176                 } ?: throw Exception("Failed to get video size")
    177 
    178                 val filterFirstPart = StringBuilder()
    179                 val filterSecondPart = StringBuilder()
    180                 var containsNoSound = false
    181 
    182                 filesInfo.forEachIndexed { index, (file, retriever) ->
    183                     filterFirstPart.append("[$index:v]scale=$maxWidth:$maxHeight,setsar=1[v$index];")
    184                     if (retriever.extractMetadata(MediaMetadataRetriever.METADATA_KEY_HAS_AUDIO) == "yes") {
    185                         filterSecondPart.append("[v$index][$index:a]")
    186                     } else {
    187                         containsNoSound = true
    188                         filterSecondPart.append("[v$index][${filesInfo.size}]")
    189                     }
    190                     inputArguments += "-i" to file
    191                 }
    192 
    193                 if (containsNoSound) {
    194                     inputArguments += "-f" to "lavfi"
    195                     inputArguments += "-t" to "0.1"
    196                     inputArguments += "-i" to "anullsrc=channel_layout=stereo:sample_rate=44100"
    197                 }
    198 
    199                 if (outputArguments["-c:a"] == "copy") {
    200                     outputArguments -= "-c:a"
    201                 }
    202 
    203                 outputArguments += "-fps_mode" to "vfr"
    204 
    205                 outputArguments += "-filter_complex" to "\"$filterFirstPart ${filterSecondPart}concat=n=${args.inputs.size}:v=1:a=1[vout][aout]\""
    206                 outputArguments += "-map" to "\"[aout]\""
    207                 outputArguments += "-map" to "\"[vout]\""
    208 
    209                 filesInfo.forEach { it.second.close() }
    210             }
    211             Action.DOWNLOAD_AUDIO_STREAM -> {
    212                 outputArguments.clear()
    213                 globalArguments += "-f" to when (args.audioStreamFormat!!.encoding) {
    214                     AudioFormat.ENCODING_PCM_8BIT -> "u8"
    215                     AudioFormat.ENCODING_PCM_16BIT -> "s16le"
    216                     AudioFormat.ENCODING_PCM_FLOAT -> "f32le"
    217                     AudioFormat.ENCODING_PCM_32BIT -> "s32le"
    218                     else -> throw IllegalArgumentException("Unsupported audio encoding")
    219                 }
    220                 globalArguments += "-ar" to args.audioStreamFormat.sampleRate.toString()
    221                 globalArguments += "-ac" to args.audioStreamFormat.channels.toString()
    222             }
    223         }
    224         outputArguments += args.output.absolutePath
    225         newFFMpegTask(globalArguments, inputArguments, outputArguments)
    226     }
    227 }