VmafKnapsack.kt

/**
 * Performs a greedy algorithm on the chunk vmaf data to produce a list of cq values that maximizes the average vmaf of
 * the entire encode without going over the [SIZE_LIMIT].
 * 
 * This is a variation of the 0-1 knapsack problem. Due to the extra constraint of needing one of each chunk, the
 * pseudo-polynomial algorithm cannot be applied. Instead, we make the following assumptions to reach a good local
 * maximum.
 * 
 * For any chunk c and cq value q,
 * - vmaf(c, q - 1) > vmaf(c, q) and
 * - size(c, q - 1) > size(c, q) and
 * - vmaf(c, q - 1) / size(c, q - 1) ≈ vmaf(c, q) / size(c, q)
 * 
 * where
 * - vmaf(c, q) is the vmaf of a chunk c at cq-value=q and
 * - size(c, q) is the size of a chunk c at cq-value=q
 * 
 * @param cvmafs A 2d array of chunk data such that `cvmafs[i][q]` gets the data for the ith chunk encoded at cq-level=q
 * @return An IntArray such that for the ith chunk, `return[i]` is the optimal cq value for that chunk
 * */
fun greedyAlgorithm(cvmafs: Array<Array<Chunk>>): IntArray {
	
	// set the cq for all chunks to the highest cq (lowest quality)
	val currentChunkCqs = IntArray(cvmafs.size) { cvmafs.first().size - 1 }
	var currentSize = getTotalSizeFromChoices(cvmafs, currentChunkCqs)
	
	// incrementally increase the quality of a chunk
	while (true) {
		
		// get the chunk that will most benefit the total vmaf if its cq were to be increased
		val availableSize = SIZE_LIMIT - currentSize
		val (bestChunkIndex, sizeIncrease) = getBestChunkToIncrease(cvmafs, currentChunkCqs, availableSize)
		
		// if there isn't a scene we can increase without going over the size limit
		// we have reached the end of our search
		if (bestChunkIndex < 0) break
		
		// increase the quality and update the size
		currentChunkCqs[bestChunkIndex]--
		currentSize += sizeIncrease
		
		// if we hit our target vmaf already we can exit
		if (getAverageVmafFromChoices(cvmafs, currentChunkCqs) > VIDEO_VMAF_EARLY_EXIT) break
		
	}
	
	return currentChunkCqs
	
}

/**
 * Greedily gets the best chunk to increase the cq of that will most increase the worth of the entire encode.
 * 
 * @param cvmafs A 2d array of chunk data such that `cvmafs[i][q]` gets the data for the ith chunk encoded at cq-level=q
 * @param currentChunkCqs An IntArray such that for the ith chunk, `currentChunkCqs[i]` is the current cq value for that
 * chunk
 * @param availableSize The available size left in the encode. Wont return a chunk that would exceed this size
 * @return A pair of the index of the best chunk to increase and the change in size increasing that chunk would require
 * */
fun getBestChunkToIncrease(cvmafs: Array<Array<Chunk>>, currentChunkCqs: IntArray, availableSize: Int): Pair<Int, Int> {
	
	var bestChunk: Int = -1
	var bestWorth: Double = 0.0
	var bestSizeIncrease: Int = 0
	
	// consider each chunk
	for (chunkIndex in cvmafs.indices) {
		
		val currentCq = currentChunkCqs[chunkIndex]
		var proposedCq = currentCq - 1
		
		// we have already maxed out the cq for this chunk
		if (currentCq == 0) continue
		
		// if this chunk's vmaf is already really high, don't include it
		if (cvmafs[chunkIndex][currentCq].averageVmaf > CHUNK_VMAF_EARLY_EXIT) continue
		
		// calculate the worth of increasing this chunk
		var thisChunksWorth = getChunkWorth(cvmafs, chunkIndex, currentCq, proposedCq)
		
		// it is possible for the chunk worth to be negative if a higher cq leads to a lower quality.
		// usually this is a small difference due to error in vmaf. if this is the case, try proposing a lower cq
		// to see if it's just an abnormality. can provide substantial boost to greedyAlgorithm result (~+0.07)
		if (thisChunksWorth < 0 && proposedCq >= 1) {
			proposedCq--
			thisChunksWorth = getChunkWorth(cvmafs, chunkIndex, currentCq, proposedCq)
		}
		
		// if we increased this chunk, how much would be added to the file size
		val sizeIncrease = cvmafs[chunkIndex][proposedCq].size - cvmafs[chunkIndex][currentCq].size
		// make sure increasing this chunk wouldn't be too costly
		if (sizeIncrease > availableSize) continue
		
		// if this chunk's better than the best one we've found, set the best to this
		if (thisChunksWorth > bestWorth) {
			bestChunk = chunkIndex
			bestWorth = thisChunksWorth
			bestSizeIncrease = sizeIncrease
		}
		
	}
	
	return Pair(bestChunk, bestSizeIncrease)
	
}

/**
 * Calculates the worth of decreasing the cq value (increasing quality) for a chunk from [currentCq] to [proposedCq].
 * The worth is calculated by the following equation:
 * 
 * worth = delta vmaf sum / delta size
 * 
 * @param cvmafs A 2d array of chunk data such that `cvmafs[i][q]` gets the data for the ith chunk encoded at cq-level=q
 * @param chunkIndex The index of the chunk to consider decreasing the cq of
 * @param currentCq The current cq of the chunk
 * @param proposedCq The proposed cq of the chunk
 * @return The worth of changing [chunkIndex] from cq=[currentCq] to cq=[proposedCq]
 * */
fun getChunkWorth(cvmafs: Array<Array<Chunk>>, chunkIndex: Int, currentCq: Int, proposedCq: Int): Double {
	
	// get how much size changes if we change to the proposed cq value
	val currentSize = cvmafs[chunkIndex][currentCq].size
	val proposedSize = cvmafs[chunkIndex][proposedCq].size
	val sizeIncrease = proposedSize - currentSize
	
	// get how much vmaf score changes if we change to the proposed cq value
	val currentVmafs = cvmafs[chunkIndex][currentCq].vmafSum
	val proposedVmafs = cvmafs[chunkIndex][proposedCq].vmafSum
	val vmafIncrease = proposedVmafs - currentVmafs
	
	// return the vmaf increase per bytes
	return vmafIncrease / sizeIncrease.toDouble()
	
}