ffmpeg_quality_metrics

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