diff --git a/src/cli.ts b/src/cli.ts index c8b488d..76b7459 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -104,7 +104,7 @@ async function run() { type: 'string', describe: 'Folder with vmaf measurement results', demandOption: true - }); + }) }, runSuggestLadder ) @@ -123,6 +123,10 @@ async function run() { type: 'boolean', description: 'Read bitrate of transcoded file with ffprobe', default: false + }, + variables: { + type: 'string', + description: 'List of variables to include as columns in csv' } }); }, @@ -186,20 +190,30 @@ async function exportWmafResultToCsv(argv) { argv.probeBitrate ) ).flatMap((result) => { - return result[1].map((resolutionVmaf) => ({ - folder, - filename: resolutionVmaf.vmafFile, - resolution: `${resolutionVmaf.resolution.width}X${resolutionVmaf.resolution.height}`, - qvbr: resolutionVmaf.qvbr, - vmaf: resolutionVmaf.vmaf, - vmafHd: resolutionVmaf.vmafHd, - vmafHdPhone: resolutionVmaf.vmafHdPhone, - bitrate: result[0], - realTime: resolutionVmaf.cpuTime?.realTime, - cpuTime: resolutionVmaf.cpuTime?.cpuTime - })); + return result[1].map((resolutionVmaf) => { + const obj = { + folder, + filename: resolutionVmaf.vmafFile, + resolution: `${resolutionVmaf.resolution.width}X${resolutionVmaf.resolution.height}`, + vmaf: resolutionVmaf.vmaf, + vmafHd: resolutionVmaf.vmafHd, + vmafHdPhone: resolutionVmaf.vmafHdPhone, + bitrate: result[0], + realTime: resolutionVmaf.cpuTime?.realTime, + cpuTime: resolutionVmaf.cpuTime?.cpuTime, + variables: Object.keys(resolutionVmaf.variables).map((k) => `${k}=${resolutionVmaf.variables[k]}`).join(':') + } + if (argv.variables) { + for (const v of argv.variables.split(',')) { + obj[v.toLowerCase()] = resolutionVmaf.variables[v] || ''; + } + } + return obj; + }); }); + console.log(pairs); + await new ObjectsToCsv(pairs).toDisk(`${argv.folder}/results.csv`, { allColumns: true, append: true @@ -221,9 +235,9 @@ async function transcodeAndAnalyse(argv) { logger.info( `saveAsCsv: ${job.saveAsCsv}, ` + - (job.saveAsCsv - ? `also saving results as a .csv file.` - : `will not save results as a .csv file.`) + (job.saveAsCsv + ? `also saving results as a .csv file.` + : `will not save results as a .csv file.`) ); if (job.saveAsCsv) { models.forEach((model) => diff --git a/src/models/vmaf-bitrate-pair.ts b/src/models/vmaf-bitrate-pair.ts index d1e8254..87796d8 100644 --- a/src/models/vmaf-bitrate-pair.ts +++ b/src/models/vmaf-bitrate-pair.ts @@ -5,7 +5,7 @@ export type VmafBitratePair = { vmaf?: number; vmafHd?: number; vmafHdPhone?: number; - qvbr: number | null; + variables: Record; cpuTime?: { realTime: number; cpuTime: number; diff --git a/src/pairVmaf.test.ts b/src/pairVmaf.test.ts new file mode 100644 index 0000000..03c2b4c --- /dev/null +++ b/src/pairVmaf.test.ts @@ -0,0 +1,29 @@ +import { parseVmafFilename } from './pairVmaf'; + +describe('parseFilename', () => { + it('basename only, should return resolution and bitrate', () => { + expect(parseVmafFilename('1920x1080_0.mp4')) + .toEqual({width: 1920, height: 1080, bitrate: 0, variables: {}}); + }); + + it('full path, should return resolution and bitrate', () => { + expect(parseVmafFilename('/apa/bepa/1920x1080_0.mp4')) + .toEqual({width: 1920, height: 1080, bitrate: 0, variables: {}}); + }); + + it('with variables, should return resolution, bitrate, and variables', () => { + expect(parseVmafFilename('1920x1080_0_VAR1_value1_VAR2_value2_vmaf.json')) + .toEqual({width: 1920, height: 1080, bitrate: 0, variables: { + VAR1: 'value1', + VAR2: 'value2' + }}); + }); + + it('with variables and generic file extension, should return resolution, bitrate, and variables', () => { + expect(parseVmafFilename('1920x1080_0_VAR1_value1_VAR2_value2.mp4')) + .toEqual({width: 1920, height: 1080, bitrate: 0, variables: { + VAR1: 'value1', + VAR2: 'value2' + }}); + }); +}) \ No newline at end of file diff --git a/src/pairVmaf.ts b/src/pairVmaf.ts index 8cd1b5d..12c6553 100644 --- a/src/pairVmaf.ts +++ b/src/pairVmaf.ts @@ -7,11 +7,6 @@ import * as fs from 'fs'; import * as path from 'path'; import { VmafBitratePair } from './models/vmaf-bitrate-pair'; -function extractQVBRNumberFromFilename(filename: string): number | null { - const match = filename.match(/_QVBR_(\d+)_/); - return match ? parseInt(match[1], 10) : null; -} - export async function pairVmafWithResolutionAndBitrate( directoryWithVmafFiles: string, filterFunction: ( @@ -60,14 +55,14 @@ export async function pairVmafWithResolutionAndBitrate( `Loaded VMAF data from ${JSON.stringify(analysisData, null, 2)}.` ); analysisData.vmafList.forEach(({ filename, vmafScores }) => { - const [resolutionStr, bitrateStr] = filename.split('_'); - const [widthStr, heightStr] = resolutionStr.split('x'); - - const width = parseInt(widthStr); - const height = parseInt(heightStr); - const bitrate = bitrates ? bitrates[filename] : parseInt(bitrateStr); + const dataFromFilename = parseVmafFilename(filename); + if (!dataFromFilename) { + logger.error('Unable to parse data from filename: ', filename); + return; + } + const { width, height, bitrate: bitrateFromFile, variables } = dataFromFilename; + const bitrate = bitrates ? bitrates[filename] : bitrateFromFile; const cpuTime = cpuTimes ? cpuTimes[filename] : undefined; - const qvbr = extractQVBRNumberFromFilename(filename); const vmaf = vmafScores.vmaf; const vmafHd = vmafScores.vmafHd; const vmafHdPhone = vmafScores.vmafHdPhone; @@ -76,7 +71,7 @@ export async function pairVmafWithResolutionAndBitrate( if (pairs.has(bitrate)) { pairs.get(bitrate)?.push({ resolution: { width, height }, - qvbr, + variables, vmaf, vmafHd, vmafHdPhone, @@ -87,7 +82,7 @@ export async function pairVmafWithResolutionAndBitrate( pairs.set(bitrate, [ { resolution: { width, height }, - qvbr, + variables, vmaf, vmafHd, vmafHdPhone, @@ -188,3 +183,31 @@ export async function runFfprobe(file) { }); }); } + +const vmafFilenameRegex: RegExp = /(.*\/)?(?\d+)x(?\d+)_(?\d+)(?(_([A-Za-z0-9-]+)_([A-Za-z0-9.]+))*)?(_vmaf\.json|\.[A-Za-z0-9]+)/; + +export function parseVmafFilename(file): { + width: number, + height: number, + bitrate: number, + variables: Record +} | undefined { + const result = vmafFilenameRegex.exec(file); + if (!result) { + return undefined; + } + const groups = result.groups as { width: string, height: string, bitrate: string, variables: string }; + const variables = {} + if (groups.variables) { + const variablesList = groups.variables.split('_').slice(1); + for (let i = 0; i < variablesList.length - 1; i += 2) { + variables[variablesList[i]] = variablesList[i + 1]; + } + } + return { + width: parseInt(groups.width), + height: parseInt(groups.height), + bitrate: parseInt(groups.bitrate), + variables + }; +}