feat: extract qvbr for csv export

Eyevinn · Aug 20, 2024 · 52fb9bf · 52fb9bf
1 parent 12c8785
commit 52fb9bf
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/cli.ts b/src/cli.ts
@@ -153,6 +153,7 @@ async function exportWmafResultToCsv(argv) {
       folder,
       filename: resolutionVmaf.vmafFile,
       resolution: `${resolutionVmaf.resolution.width}X${resolutionVmaf.resolution.height}`,
+      qvbr: resolutionVmaf.qvbr,
       vmaf: resolutionVmaf.vmaf,
       bitrate: result[0],
       realTime: resolutionVmaf.cpuTime?.realTime,

diff --git a/src/pairVmaf.ts b/src/pairVmaf.ts
@@ -5,6 +5,11 @@ import ffmpeg from 'fluent-ffmpeg';
 import * as fs from 'fs';
 import * as path from 'path';
 
+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: (
@@ -24,6 +29,7 @@ export async function pairVmafWithResolutionAndBitrate(
     {
       resolution: Resolution;
       vmaf: number;
+      qvbr: number | null;
       cpuTime?: { realTime: number; cpuTime: number };
       vmafFile: string;
     }[]
@@ -66,18 +72,20 @@ export async function pairVmafWithResolutionAndBitrate(
     const height = parseInt(heightStr);
     const bitrate = bitrates ? bitrates[filename] : parseInt(bitrateStr);
     const cpuTime = cpuTimes ? cpuTimes[filename] : undefined;
+    const qvbr = extractQVBRNumberFromFilename(filename);
 
     if (filterFunction(bitrate, { width, height }, vmaf)) {
       if (pairs.has(bitrate)) {
         pairs.get(bitrate)?.push({
           resolution: { width, height },
+          qvbr,
           vmaf,
           cpuTime,
           vmafFile: filename
         });
       } else {
         pairs.set(bitrate, [
-          { resolution: { width, height }, vmaf, cpuTime, vmafFile: filename }
+          { resolution: { width, height }, qvbr, vmaf, cpuTime, vmafFile: filename }
         ]);
       }
     }