ffmpeg_quality_metrics

 84class FfmpegQualityMetrics:
 85    """
 86    A class to calculate quality metrics with FFmpeg
 87    """
 88
 89    ALLOWED_SCALERS = [
 90        "fast_bilinear",
 91        "bilinear",
 92        "bicubic",
 93        "experimental",
 94        "neighbor",
 95        "area",
 96        "bicublin",
 97        "gauss",
 98        "sinc",
 99        "lanczos",
100        "spline",
101    ]
102    DEFAULT_SCALER = "bicubic"
103    DEFAULT_THREADS = 0
104
105    DEFAULT_VMAF_THREADS = 0  # used to be os.cpu_count(), now auto
106    DEFAULT_VMAF_SUBSAMPLE = 1  # sample every frame
107    DEFAULT_VMAF_MODEL_DIRECTORY = os.path.join(
108        os.path.dirname(__file__), "vmaf_models"
109    )
110    DEFAULT_VMAF_OPTIONS: VmafOptions = {
111        "model_path": None,
112        "model_params": [],
113        "n_threads": DEFAULT_VMAF_THREADS,
114        "n_subsample": DEFAULT_VMAF_SUBSAMPLE,
115        "features": [],
116    }
117    POSSIBLE_FILTERS: List[FilterName] = [
118        "libvmaf",
119        "psnr",
120        "ssim",
121        "vif",
122        "msad",
123    ]
124    METRIC_TO_FILTER_MAP: Dict[MetricName, FilterName] = {
125        "vmaf": "libvmaf",
126        "psnr": "psnr",
127        "ssim": "ssim",
128        "vif": "vif",
129        "msad": "msad",
130    }
131
132    def __init__(
133        self,
134        ref: str,
135        dist: str,
136        scaling_algorithm: str = DEFAULT_SCALER,
137        framerate: Union[float, None] = None,
138        dist_delay: float = 0,
139        dry_run: Union[bool, None] = False,
140        verbose: Union[bool, None] = False,
141        threads: int = DEFAULT_THREADS,
142        progress: Union[bool, None] = False,
143        keep_tmp_files: Union[bool, None] = False,
144        tmp_dir: Union[str, None] = None,
145        num_frames: Union[int, None] = None,
146        start_offset: Union[str, None] = None,
147        ffmpeg_path: str = "ffmpeg",
148    ):
149        """Instantiate a new FfmpegQualityMetrics
150
151        Args:
152            ref (str): reference file
153            dist (str): distorted file
154            scaling_algorithm (str, optional): A scaling algorithm. Must be one of the following: ["fast_bilinear", "bilinear", "bicubic", "experimental", "neighbor", "area", "bicublin", "gauss", "sinc", "lanczos", "spline"]. Defaults to "bicubic"
155            framerate (float, optional): Force a frame rate. Defaults to None.
156            dist_delay (float): Delay the distorted file against the reference by this amount of seconds. Defaults to 0.
157            dry_run (bool, optional): Don't run anything, just print commands. Defaults to False.
158            verbose (bool, optional): Show more output. Defaults to False.
159            threads (int, optional): Number of ffmpeg threads. Defaults to 0 (auto).
160            progress (bool, optional): Show a progress bar. Defaults to False.
161            keep_tmp_files (bool, optional): Keep temporary files for debugging purposes. Defaults to False.
162            tmp_dir (str, optional): Directory to store temporary files. Will use system default if not specified. Defaults to None.
163            num_frames (int, optional): Number of frames to analyze from the input files. Defaults to None (all frames).
164            start_offset (str, optional): Seek to this position before analyzing. Accepts timestamp (e.g., '00:00:10' or '10.5') or frame number with 'f:' prefix (e.g., 'f:100'). Defaults to None.
165            ffmpeg_path (str, optional): Path to ffmpeg executable. Defaults to "ffmpeg".
166
167        Raises:
168            FfmpegQualityMetricsError: A generic error
169        """
170        self.ref = str(ref)
171        self.dist = str(dist)
172        self.scaling_algorithm = str(scaling_algorithm)
173        self.framerate = float(framerate) if framerate is not None else None
174        self.dist_delay = float(dist_delay)
175        self.dry_run = bool(dry_run)
176        self.verbose = bool(verbose)
177        self.threads = int(threads)
178        self.progress = bool(progress)
179        self.keep_tmp_files = bool(keep_tmp_files)
180        self.tmp_dir = str(tmp_dir) if tmp_dir is not None else tempfile.gettempdir()
181        self.num_frames = int(num_frames) if num_frames is not None else None
182        self.start_offset = str(start_offset) if start_offset is not None else None
183        self.ffmpeg_path = ffmpeg_path
184
185        if not os.path.isfile(self.ref):
186            raise FfmpegQualityMetricsError(f"Reference file not found: {self.ref}")
187        if not os.path.isfile(self.dist):
188            raise FfmpegQualityMetricsError(f"Distorted file not found: {self.dist}")
189
190        if self.ref == self.dist:
191            logger.warning(
192                "Reference and distorted files are the same! This may lead to unexpected results or numerical issues."
193            )
194
195        if ref.endswith(".yuv") or dist.endswith(".yuv"):
196            raise FfmpegQualityMetricsError(
197                "YUV files are not supported, please convert to a format that ffmpeg can read natively, such as Y4M or FFV1."
198            )
199
200        self.data: MetricData = {
201            "vmaf": [],
202            "psnr": [],
203            "ssim": [],
204            "vif": [],
205            "msad": [],
206        }
207
208        self.available_filters: List[str] = []
209
210        self.global_stats: GlobalStats = {}
211
212        if not os.path.isdir(self.tmp_dir):
213            logger.debug(f"Creating temporary directory: {self.tmp_dir}")
214            os.makedirs(self.tmp_dir)
215        self.temp_files: Dict[FilterName, str] = {}
216
217        for filter_name in self.POSSIBLE_FILTERS:
218            suffix = "txt" if filter_name != "libvmaf" else "json"
219
220            self.temp_files[filter_name] = os.path.join(
221                self.tmp_dir,
222                f"ffmpeg_quality_metrics_{filter_name}_{os.path.basename(self.ref)}_{os.path.basename(self.dist)}.{suffix}",
223            )
224            logger.debug(
225                f"Writing temporary {filter_name.upper()} information to: {self.temp_files[filter_name]}"
226            )
227
228        if scaling_algorithm not in self.ALLOWED_SCALERS:
229            raise FfmpegQualityMetricsError(
230                f"Allowed scaling algorithms: {self.ALLOWED_SCALERS}"
231            )
232
233        self._check_available_filters()
234
235    def _check_available_filters(self):
236        """
237        Check which filters are available
238        """
239        cmd = [self.ffmpeg_path, "-filters"]
240        stdout, _ = run_command(cmd)
241        filter_list = []
242        for line in stdout.split("\n"):
243            line = line.strip()
244            if line == "":
245                continue
246            cols = line.split(" ")
247            if len(cols) > 1:
248                filter_name = cols[1]
249                filter_list.append(filter_name)
250
251        for key in FfmpegQualityMetrics.POSSIBLE_FILTERS:
252            if key in filter_list:
253                self.available_filters.append(key)
254
255        logger.debug(f"Available filters: {self.available_filters}")
256
257    @staticmethod
258    def get_framerate(input_file: str, ffmpeg_path: str = "ffmpeg") -> float:
259        """Parse the FPS from the input file.
260
261        Args:
262            input_file (str): Input file path
263            ffmpeg_path (str, optional): Path to ffmpeg executable. Defaults to "ffmpeg".
264
265        Raises:
266            FfmpegQualityMetricsError: A generic error
267
268        Returns:
269            float: The FPS parsed
270        """
271        cmd = [ffmpeg_path, "-nostdin", "-y", "-i", input_file]
272
273        output = run_command(cmd, allow_error=True)
274        pattern = re.compile(r"(\d+(\.\d+)?) fps")
275        try:
276            if pattern_ret := pattern.search(str(output)):
277                match = pattern_ret.groups()[0]
278                return float(match)
279        except Exception:
280            pass
281
282        raise FfmpegQualityMetricsError(f"could not parse FPS from file {input_file}!")
283
284    def _get_framerates(self) -> Tuple[float, float]:
285        """
286        Get the framerates of the reference and distorted files.
287
288        Returns:
289            Tuple[float, float]: The framerates of the reference and distorted files
290        """
291        ref_framerate = FfmpegQualityMetrics.get_framerate(self.ref, self.ffmpeg_path)
292        dist_framerate = FfmpegQualityMetrics.get_framerate(self.dist, self.ffmpeg_path)
293
294        if ref_framerate != dist_framerate:
295            logger.warning(
296                f"ref, dist framerates differ: {ref_framerate}, {dist_framerate}. "
297                "This may result in inaccurate quality metrics. Force an input framerate via the -r option."
298            )
299
300        return ref_framerate, dist_framerate
301
302    def _parse_start_offset(self, framerate: float) -> Union[str, None]:
303        """
304        Parse the start_offset parameter and convert frame numbers to timestamps if needed.
305
306        Args:
307            framerate (float): The framerate to use for frame-to-timestamp conversion
308
309        Returns:
310            Union[str, None]: The timestamp string for ffmpeg's -ss option, or None if no offset
311        """
312        if self.start_offset is None:
313            return None
314
315        # Check if it's a frame number (format: "f:100" or "f:100.5")
316        if self.start_offset.startswith("f:"):
317            try:
318                frame_num = float(self.start_offset[2:])
319                timestamp = frame_num / framerate
320                return str(timestamp)
321            except ValueError:
322                raise FfmpegQualityMetricsError(
323                    f"Invalid frame number in start_offset: {self.start_offset}"
324                )
325
326        # Otherwise, assume it's a timestamp string (e.g., "00:00:10" or "10.5")
327        return self.start_offset
328
329    def _get_filter_opts(self, filter_name: FilterName) -> str:
330        """
331        Returns:
332            str: Specific ffmpeg filter options for a chosen metric filter.
333        """
334        if filter_name in ["ssim", "psnr"]:
335            return f"{filter_name}='{win_path_check(self.temp_files[filter_name])}'"
336        elif filter_name == "libvmaf":
337            return f"libvmaf='{self._get_libvmaf_filter_opts()}'"
338        elif filter_name == "vif":
339            return "vif,metadata=mode=print"
340        elif filter_name == "msad":
341            return "msad,metadata=mode=print"
342        else:
343            raise FfmpegQualityMetricsError(f"Unknown filter {filter_name}!")
344
345    def calculate(
346        self,
347        metrics: List[MetricName] = ["ssim", "psnr"],
348        vmaf_options: Union[VmafOptions, None] = None,
349    ) -> Dict[MetricName, SingleMetricData]:
350        """Calculate one or more metrics.
351
352        Args:
353            metrics (list, optional): A list of metrics to calculate.
354                Possible values are ["ssim", "psnr", "vmaf"].
355                Defaults to ["ssim", "psnr"].
356            vmaf_options (dict, optional): VMAF-specific options. Uses defaults if not specified.
357
358        Raises:
359            FfmpegQualityMetricsError: In case of an error
360            e: A generic error
361
362        Returns:
363            dict: A dictionary of per-frame info, with the key being the metric name and the value being a dict of frame numbers ('n') and metric values.
364        """
365        if not metrics:
366            raise FfmpegQualityMetricsError("No metrics specified!")
367
368        # check available metrics
369        for metric_name in metrics:
370            filter_name = self.METRIC_TO_FILTER_MAP.get(metric_name, None)
371            if filter_name not in self.POSSIBLE_FILTERS:
372                raise FfmpegQualityMetricsError(f"No such metric '{metric_name}'")
373            if filter_name not in self.available_filters:
374                raise FfmpegQualityMetricsError(
375                    f"Your ffmpeg version does not have the filter '{filter_name}'"
376                )
377
378        # set VMAF options specifically
379        if "vmaf" in metrics:
380            self._check_libvmaf_availability()
381            self.vmaf_options = self.DEFAULT_VMAF_OPTIONS
382            # override with user-supplied options
383            if vmaf_options:
384                for key, value in vmaf_options.items():
385                    if value is not None:
386                        self.vmaf_options[key] = value  # type: ignore
387            self._set_vmaf_model_path(self.vmaf_options["model_path"])
388
389        # ffmpeg 7.1 or higher: scale2ref filter is deprecated
390        # input 0: ref, input 1: dist --> swapped for scale filter
391
392        # Apply select filter if num_frames is specified
393        select_filter = ""
394        if self.num_frames is not None:
395            select_filter = f"select='lt(n\\,{self.num_frames})',"
396
397        filter_chains = [
398            f"[1][0]scale=rw:rh:flags={self.scaling_algorithm}[dist]",
399            f"[dist]{select_filter}settb=AVTB,setpts=PTS-STARTPTS[distpts]",
400            f"[0]{select_filter}settb=AVTB,setpts=PTS-STARTPTS[refpts]",
401        ]
402
403        # generate split filters depending on the number of models
404        n_splits = len(metrics)
405        if n_splits > 1:
406            for source in ["dist", "ref"]:
407                suffixes = "".join([f"[{source}{n}]" for n in range(1, n_splits + 1)])
408                filter_chains.extend(
409                    [
410                        f"[{source}pts]split={n_splits}{suffixes}",
411                    ]
412                )
413
414        # special case, only one metric:
415        if n_splits == 1:
416            metric_name = metrics[0]
417            filter_chains.extend(
418                [
419                    f"[distpts][refpts]{self._get_filter_opts(self.METRIC_TO_FILTER_MAP[metric_name])}"
420                ]
421            )
422        # all other cases:
423        else:
424            for n, metric_name in zip(range(1, n_splits + 1), metrics):
425                filter_chains.extend(
426                    [
427                        f"[dist{n}][ref{n}]{self._get_filter_opts(self.METRIC_TO_FILTER_MAP[metric_name])}"
428                    ]
429                )
430
431        try:
432            output = self._run_ffmpeg_command(filter_chains, desc=", ".join(metrics))
433            self._read_temp_files(metrics)
434            if output:
435                self._read_ffmpeg_output(output, metrics)
436            else:
437                raise FfmpegQualityMetricsError("ffmpeg output is empty!")
438        except Exception as e:
439            raise e
440        finally:
441            self._cleanup_temp_files()
442
443        # return only those data entries containing values
444        return {k: v for k, v in self.data.items() if v}
445
446    def _get_libvmaf_filter_opts(self) -> str:
447        """
448        Returns:
449
450            str: A string to use for VMAF in ffmpeg filter chain
451        """
452        # we only have one model, and its path parameter is not optional
453        all_model_params: Dict[str, str] = {
454            "path": win_vmaf_model_path_check(self.vmaf_model_path)
455        }
456
457        # add further model parameters
458        for model_param in self.vmaf_options["model_params"]:
459            key, value = model_param.split("=")
460            all_model_params[key] = value
461
462        all_model_params_str = "\\:".join(
463            f"{k}={v}" for k, v in all_model_params.items()
464        )
465
466        vmaf_opts: Dict[str, str] = {
467            "model": all_model_params_str,
468            "log_path": win_path_check(self.temp_files["libvmaf"]),
469            "log_fmt": "json",
470            "n_threads": str(self.vmaf_options["n_threads"]),
471            "n_subsample": str(self.vmaf_options["n_subsample"]),
472        }
473
474        if self.vmaf_options["features"]:
475            features = []
476            for feature in self.vmaf_options["features"]:
477                if not feature.startswith("name"):
478                    feature = f"name={feature}"
479                features.append(feature.replace(":", "\\:"))
480            vmaf_opts["feature"] = "|".join(features)
481
482        vmaf_opts_string = ":".join(
483            f"{k}={v}" for k, v in vmaf_opts.items() if v is not None
484        )
485
486        return vmaf_opts_string
487
488    def _check_libvmaf_availability(self) -> None:
489        if "libvmaf" not in self.available_filters:
490            raise FfmpegQualityMetricsError(
491                "Your ffmpeg build does not have support for VMAF. Make sure you download or build a version compiled with --enable-libvmaf!"
492            )
493
494    def _set_vmaf_model_path(self, model_path: Union[str, None] = None) -> None:
495        """
496        Logic to set the model path depending on the default or the user-supplied string
497        """
498        if model_path is None:
499            self.vmaf_model_path = FfmpegQualityMetrics.get_default_vmaf_model_path()
500        else:
501            self.vmaf_model_path = str(model_path)
502
503        supplied_models = FfmpegQualityMetrics.get_supplied_vmaf_models()
504
505        if not os.path.isfile(self.vmaf_model_path):
506            # check if this is one of the supplied ones? e.g. user passed only a filename
507            if self.vmaf_model_path in supplied_models:
508                self.vmaf_model_path = os.path.join(
509                    FfmpegQualityMetrics.DEFAULT_VMAF_MODEL_DIRECTORY,
510                    self.vmaf_model_path,
511                )
512            else:
513                raise FfmpegQualityMetricsError(
514                    f"Could not find model at {self.vmaf_model_path}. Please set --model-path to a valid VMAF .json model file."
515                )
516
517    def _read_vmaf_temp_file(self) -> None:
518        """
519        Read the VMAF temp file and append the data to the data dict.
520        """
521        with open(self.temp_files["libvmaf"], "r") as in_vmaf:
522            vmaf_log = json.load(in_vmaf)
523            logger.debug(f"VMAF log: {json.dumps(vmaf_log, indent=4)}")
524            for frame_data in vmaf_log["frames"]:
525                # append frame number, increase +1
526                frame_data["metrics"]["n"] = int(frame_data["frameNum"]) + 1
527                self.data["vmaf"].append(frame_data["metrics"])
528
529    def _read_ffmpeg_output(self, ffmpeg_output: str, metrics=[]) -> None:
530        """
531        Read the metric values from ffmpeg's stderr, for those that don't output
532        to a file.
533        """
534        if self.dry_run:
535            return
536        if "vif" in metrics:
537            self._parse_ffmpeg_metadata_output(ffmpeg_output, "vif")
538        if "msad" in metrics:
539            self._parse_ffmpeg_metadata_output(ffmpeg_output, "msad")
540
541    def _parse_ffmpeg_metadata_output(
542        self, ffmpeg_output: str, metric_name: Literal["vif", "msad"]
543    ) -> None:
544        """
545        Parse the filter output written to ffmpeg's metadata output
546
547        Args:
548            ffmpeg_output (str): The output of ffmpeg's stderr
549            metric_name (Literal["vif", "msad"]): The name of the metric to parse
550        """
551        # Example for VIF:
552        #
553        # [Parsed_metadata_4 @ 0x7f995cd08640] frame:1    pts:1       pts_time:0.0401x
554        # [Parsed_metadata_4 @ 0x7f995cd08640] lavfi.vif.scale.0=0.263582
555        # [Parsed_metadata_4 @ 0x7f995cd08640] lavfi.vif.scale.1=0.560129
556        # [Parsed_metadata_4 @ 0x7f995cd08640] lavfi.vif.scale.2=0.626596
557        # [Parsed_metadata_4 @ 0x7f995cd08640] lavfi.vif.scale.3=0.682183
558        #
559        # Example for MSAD:
560        #
561        # [Parsed_metadata_6 @ 0x10ad04ea0] lavfi.msad.msad.Y=0.029998
562        # [Parsed_metadata_6 @ 0x10ad04ea0] lavfi.msad.msad.U=0.019501
563        # [Parsed_metadata_6 @ 0x10ad04ea0] lavfi.msad.msad.V=0.026455
564        # [Parsed_metadata_6 @ 0x10ad04ea0] lavfi.msad.msad_avg=0.025318
565
566        lines = [line.strip() for line in ffmpeg_output.split("\n")]
567        current_frame = None
568        frame_data: Dict[str, float] = {}
569
570        for line in lines:
571            if not line.startswith("[Parsed_metadata"):
572                continue
573
574            fields = line.split(" ")
575
576            # a new frame appears
577            if fields[3].startswith("frame"):
578                # if we have data already
579                if frame_data:
580                    self.data[metric_name].append(frame_data)
581
582                # get the frame number and reset the frame data
583                current_frame = int(fields[3].split(":")[1])
584                frame_data = {"n": current_frame}
585                continue
586
587            # no frame was set, or no VIF info present
588            if current_frame is None or not fields[3].startswith(
589                f"lavfi.{metric_name}"
590            ):
591                continue
592
593            # we have a frame
594            key, value = fields[3].split("=")
595            key = key.replace(f"lavfi.{metric_name}.", "").replace(".", "_").lower()
596            frame_data[key] = round(float(value), 3)
597
598        # append final frame data
599        if frame_data:
600            self.data[metric_name].append(frame_data)
601
602    def _read_temp_files(self, metrics=[]):
603        """
604        Read the data from multiple temp files
605        """
606        if self.dry_run:
607            return
608        if "vmaf" in metrics:
609            self._read_vmaf_temp_file()
610        if "ssim" in metrics:
611            self._read_ssim_temp_file()
612        if "psnr" in metrics:
613            self._read_psnr_temp_file()
614
615    def _run_ffmpeg_command(
616        self, filter_chains: List[str] = [], desc: str = ""
617    ) -> Union[str, None]:
618        """
619        Run the ffmpeg command to get the quality metrics.
620        The filter chains must be specified manually.
621        'desc' can be a human readable description for the progress bar.
622
623        Returns:
624            Union[str, None]: The output of ffmpeg's stderr
625        """
626        if not self.framerate:
627            ref_framerate, dist_framerate = self._get_framerates()
628        else:
629            ref_framerate = self.framerate
630            dist_framerate = self.framerate
631
632        # Parse start_offset
633        start_offset_timestamp = self._parse_start_offset(ref_framerate)
634
635        cmd = [
636            self.ffmpeg_path,
637            "-nostdin",
638            "-nostats",
639            "-y",
640            "-threads",
641            str(self.threads),
642        ]
643
644        # Add -ss before reference input if start_offset is specified
645        if start_offset_timestamp is not None:
646            cmd.extend(["-ss", start_offset_timestamp])
647
648        # Add -r before -i only if no start_offset (to avoid seeking issues)
649        # When seeking is used, -r before -i can interfere with frame-accurate seeking
650        if start_offset_timestamp is None:
651            cmd.extend(["-r", str(ref_framerate)])
652
653        cmd.extend(["-i", self.ref, "-itsoffset", str(self.dist_delay)])
654
655        # Add -ss before distorted input if start_offset is specified
656        if start_offset_timestamp is not None:
657            cmd.extend(["-ss", start_offset_timestamp])
658
659        # Add -r before -i only if no start_offset
660        if start_offset_timestamp is None:
661            cmd.extend(["-r", str(dist_framerate)])
662
663        cmd.extend(
664            [
665                "-i",
666                self.dist,
667                "-filter_complex",
668                ";".join(filter_chains),
669                "-an",
670                "-f",
671                "null",
672                NUL,
673            ]
674        )
675
676        if self.progress:
677            logger.debug(quoted_cmd(cmd))
678            with FfmpegProgress(cmd, self.dry_run) as ff:
679                with tqdm(total=100, position=1, desc=desc) as pbar:
680                    for progress in ff.run_command_with_progress():
681                        pbar.update(progress - pbar.n)
682                return ff.stderr
683        else:
684            _, stderr = run_command(cmd, dry_run=self.dry_run)
685            return stderr
686
687    def _cleanup_temp_files(self) -> None:
688        """
689        Remove the temporary files
690        """
691        for temp_file in self.temp_files.values():
692            if os.path.isfile(temp_file):
693                if self.keep_tmp_files:
694                    logger.debug(f"Keeping temp file {temp_file}")
695                else:
696                    os.remove(temp_file)
697
698    def _read_psnr_temp_file(self) -> None:
699        """
700        Parse the PSNR generated logfile
701        """
702        with open(self.temp_files["psnr"], "r") as in_psnr:
703            # n:1 mse_avg:529.52 mse_y:887.00 mse_u:233.33 mse_v:468.25 psnr_avg:20.89 psnr_y:18.65 psnr_u:24.45 psnr_v:21.43
704            lines = in_psnr.readlines()
705            for line in lines:
706                line = line.strip()
707                fields = line.split(" ")
708                frame_data = {}
709                for field in fields:
710                    k, v = field.split(":")
711                    frame_data[k] = round(float(v), 3) if k != "n" else int(v)
712                self.data["psnr"].append(frame_data)
713
714    def _read_ssim_temp_file(self) -> None:
715        """
716        Parse the SSIM generated logfile
717        """
718        with open(self.temp_files["ssim"], "r") as in_ssim:
719            # n:1 Y:0.937213 U:0.961733 V:0.945788 All:0.948245 (12.860441)\n
720            lines = in_ssim.readlines()
721            for line in lines:
722                line = line.strip().split(" (")[0]  # remove excess
723                fields = line.split(" ")
724                frame_data = {}
725                for field in fields:
726                    k, v = field.split(":")
727                    if k != "n":
728                        # make psnr and ssim keys the same
729                        k = "ssim_" + k.lower()
730                        k = k.replace("all", "avg")
731                    frame_data[k] = round(float(v), 3) if k != "n" else int(v)
732                self.data["ssim"].append(frame_data)
733
734    @staticmethod
735    def get_brewed_vmaf_model_path() -> Union[str, None]:
736        """
737        Hack to get path for VMAF model from Homebrew or Linuxbrew.
738        This works for libvmaf 2.x
739
740        Returns:
741            str or None: the path or None if not found
742        """
743        stdout, _ = run_command(["brew", "--prefix", "libvmaf"])
744        cellar_path = stdout.strip()
745
746        model_path = os.path.join(cellar_path, "share", "libvmaf", "model")
747
748        if not os.path.isdir(model_path):
749            logger.warning(
750                f"{model_path} does not exist. Are you sure you have installed the most recent version of libvmaf with Homebrew?"
751            )
752            return None
753
754        return model_path
755
756    @staticmethod
757    def get_default_vmaf_model_path() -> str:
758        """
759        Return the default model path depending on whether the user is running Homebrew
760        or has a static build.
761
762        Returns:
763            str: the path
764        """
765        if has_brew() and ffmpeg_is_from_brew():
766            # If the user installed ffmpeg using homebrew
767            model_path = FfmpegQualityMetrics.get_brewed_vmaf_model_path()
768            if model_path is not None:
769                return os.path.join(
770                    model_path,
771                    "vmaf_v0.6.1.json",
772                )
773
774        share_path = os.path.join("/usr", "local", "share", "model")
775        if os.path.isdir(share_path):
776            return os.path.join(share_path, "vmaf_v0.6.1.json")
777        else:
778            # return the bundled file as a fallback
779            return os.path.join(
780                FfmpegQualityMetrics.DEFAULT_VMAF_MODEL_DIRECTORY, "vmaf_v0.6.1.json"
781            )
782
783    @staticmethod
784    def get_supplied_vmaf_models() -> List[str]:
785        """
786        Return a list of VMAF models supplied with the software.
787
788        Returns:
789            List[str]: A list of VMAF model names
790        """
791        return [
792            f
793            for f in os.listdir(FfmpegQualityMetrics.DEFAULT_VMAF_MODEL_DIRECTORY)
794            if f.endswith(".json")
795        ]
796
797    def get_global_stats(self) -> GlobalStats:
798        """
799        Return a dictionary for each calculated metric, with different statstics
800
801        Returns:
802            dict: A dictionary with stats, each key being a metric name and each value being a dictionary with the stats for every submetric. The stats are: 'average', 'median', 'stdev', 'min', 'max'.
803        """
804        for metric_name in self.data:
805            logger.debug(f"Aggregating stats for {metric_name}")
806            metric_data = self.data[metric_name]
807            if len(metric_data) == 0:
808                continue
809            submetric_keys = [k for k in metric_data[0].keys() if k != "n"]
810
811            stats: Dict[str, GlobalStatsData] = {}
812            for submetric_key in submetric_keys:
813                values = [float(frame[submetric_key]) for frame in metric_data]
814                # Filter out non-finite values (inf, -inf, nan) for robust statistics
815                finite = [v for v in values if math.isfinite(v)]
816                # Fallback to [0.0] if all values are non-finite to prevent crashes
817                all_values = finite if finite else [0.0]
818
819                stats[submetric_key] = {
820                    "average": round(float(mean(all_values)), 3),
821                    "median": round(float(median(all_values)), 3),
822                    "stdev": round(
823                        float(pstdev(finite)) if len(finite) > 1 else 0.0, 3
824                    ),
825                    "min": round(min(all_values), 3),
826                    "max": round(max(all_values), 3),
827                }
828            self.global_stats[metric_name] = stats
829
830        return self.global_stats
831
832    def get_results_csv(self) -> str:
833        """
834        Return a CSV string with the data
835
836        Returns:
837            str: The CSV string
838        """
839        # Check if we have any data
840        has_data = any(metric_data for metric_data in self.data.values())
841        if not has_data:
842            raise FfmpegQualityMetricsError("No data calculated!")
843
844        # Collect all frames and merge data by frame number
845        frames_data: Dict[int, Dict[str, Union[str, float, int]]] = {}
846
847        # Process each metric's data
848        for metric_data in self.data.values():
849            if not metric_data:
850                continue
851
852            for frame_info in cast(SingleMetricData, metric_data):
853                frame_num = int(frame_info["n"])
854                if frame_num not in frames_data:
855                    frames_data[frame_num] = {"n": frame_num}
856
857                # Add all metric properties for this frame
858                for key, value in frame_info.items():
859                    if key != "n":  # Skip frame number as it's already added
860                        frames_data[frame_num][key] = value
861
862        if not frames_data:
863            raise FfmpegQualityMetricsError("No frame data found!")
864
865        # Sort frames by frame number
866        sorted_frames = sorted(frames_data.keys())
867
868        # Collect all unique column names (excluding 'n' which we'll put first)
869        all_columns: set[str] = set()
870        for frame_data in frames_data.values():
871            all_columns.update(frame_data.keys())
872        all_columns.discard("n")
873
874        # Create column order: n first, then sorted metric columns, then input files
875        columns = ["n"] + sorted(all_columns) + ["input_file_dist", "input_file_ref"]
876
877        # Generate CSV using StringIO and csv module
878        output = StringIO()
879        writer = csv.writer(output)
880
881        # Write header
882        writer.writerow(columns)
883
884        # Write data rows
885        for frame_num in sorted_frames:
886            frame_data = frames_data[frame_num]
887            row = []
888            for col in columns:
889                if col == "input_file_dist":
890                    row.append(self.dist)
891                elif col == "input_file_ref":
892                    row.append(self.ref)
893                else:
894                    # Use the frame data value or empty string if not present
895                    row.append(str(frame_data.get(col, "")))
896            writer.writerow(row)
897
898        return output.getvalue()
899
900    def get_results_json(self) -> str:
901        """
902        Return the results as JSON string
903
904        Returns:
905            str: The JSON string
906        """
907        ret: Dict = {}
908        for key in self.data:
909            metric_data = self.data[key]
910            if len(metric_data) == 0:
911                continue
912            ret[key] = metric_data
913        ret["global"] = self.get_global_stats()
914        ret["input_file_dist"] = self.dist
915        ret["input_file_ref"] = self.ref
916
917        return json.dumps(ret, indent=4)