feat(export-csv): parse all variables from filename and include in csv

Signed-off-by: Gustav Grusell <[email protected]>

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) =>

src/models/vmaf-bitrate-pair.ts

@@ -5,7 +5,7 @@ export type VmafBitratePair = {
   vmaf?: number;
   vmafHd?: number;
   vmafHdPhone?: number;
-  qvbr: number | null;
+  variables: Record<string,string>;
   cpuTime?: {
     realTime: number;
     cpuTime: number;

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'
+        }});
+  });
+})

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 = /(.*\/)?(?<width>\d+)x(?<height>\d+)_(?<bitrate>\d+)(?<variables>(_([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<string,string>
+} | 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
+  };
+}