cache-dit 0.2.34__py3-none-any.whl → 0.2.36__py3-none-any.whl

cache_dit/metrics/metrics.py

@@ -5,21 +5,26 @@ import argparse
 import numpy as np
 from tqdm import tqdm
 from functools import partial
+from typing import Callable, Union, Tuple, List
 from skimage.metrics import mean_squared_error
 from skimage.metrics import peak_signal_noise_ratio
 from skimage.metrics import structural_similarity
-from cache_dit.metrics.fid import FrechetInceptionDistance
 from cache_dit.metrics.config import set_metrics_verbose
 from cache_dit.metrics.config import get_metrics_verbose
 from cache_dit.metrics.config import _IMAGE_EXTENSIONS
 from cache_dit.metrics.config import _VIDEO_EXTENSIONS
 from cache_dit.logger import init_logger
+from cache_dit.metrics.fid import compute_fid
+from cache_dit.metrics.fid import compute_video_fid
 from cache_dit.metrics.lpips import compute_lpips_img
+from cache_dit.metrics.clip_score import compute_clip_score
+from cache_dit.metrics.image_reward import compute_reward_score
 
 logger = init_logger(__name__)
 
 
 DISABLE_VERBOSE = not get_metrics_verbose()
+PSNR_TYPE = "custom"
 
 
 def compute_lpips_file(
@@ -51,6 +56,35 @@ def compute_lpips_file(
     )
 
 
+def set_psnr_type(psnr_type: str):
+    global PSNR_TYPE
+    PSNR_TYPE = psnr_type
+    assert PSNR_TYPE in ["skimage", "custom"]
+
+
+def get_psnr_type():
+    global PSNR_TYPE
+    return PSNR_TYPE
+
+
+def calculate_psnr(
+    image_true: np.ndarray,
+    image_test: np.ndarray,
+):
+    """Calculate PSNR (Peak Signal-to-Noise Ratio).
+
+    Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
+
+    Args:
+        image_true (ndarray): Images with range [0, 255].
+        image_test (ndarray): Images with range [0, 255].
+    """
+    mse = np.mean((image_true - image_test) ** 2)
+    if mse == 0:
+        return float("inf")
+    return 20 * np.log10(255.0 / np.sqrt(mse))
+
+
 def compute_psnr_file(
     image_true: np.ndarray | str,
     image_test: np.ndarray | str,
@@ -64,10 +98,13 @@ def compute_psnr_file(
         image_true = cv2.imread(image_true)
     if isinstance(image_test, str):
         image_test = cv2.imread(image_test)
-    return peak_signal_noise_ratio(
-        image_true,
-        image_test,
-    )
+    if get_psnr_type() == "skimage":
+        return peak_signal_noise_ratio(
+            image_true,
+            image_test,
+        )
+    else:
+        return calculate_psnr(image_true, image_test)
 
 
 def compute_mse_file(
@@ -114,7 +151,7 @@ def compute_dir_metric(
     image_true_dir: np.ndarray | str,
     image_test_dir: np.ndarray | str,
     compute_file_func: callable = compute_psnr_file,
-) -> float:
+) -> Union[Tuple[float, int], Tuple[None, None]]:
     # Image
     if isinstance(image_true_dir, np.ndarray) or isinstance(
         image_test_dir, np.ndarray
@@ -235,7 +272,7 @@ def compute_video_metric(
     video_true: str,
     video_test: str,
     compute_frame_func: callable = compute_psnr_file,
-) -> float:
+) -> Union[Tuple[float, int], Tuple[None, None]]:
     """
     video_true = "video_true.mp4"
     video_test = "video_test.mp4"
@@ -335,51 +372,69 @@ def compute_video_metric(
         return None, None
 
 
-compute_lpips = partial(
-    compute_dir_metric,
-    compute_file_func=compute_lpips_file,
+compute_lpips: Callable[..., Union[Tuple[float, int], Tuple[None, None]]] = (
+    partial(
+        compute_dir_metric,
+        compute_file_func=compute_lpips_file,
+    )
 )
 
-compute_psnr = partial(
-    compute_dir_metric,
-    compute_file_func=compute_psnr_file,
+compute_psnr: Callable[..., Union[Tuple[float, int], Tuple[None, None]]] = (
+    partial(
+        compute_dir_metric,
+        compute_file_func=compute_psnr_file,
+    )
 )
 
-compute_ssim = partial(
-    compute_dir_metric,
-    compute_file_func=compute_ssim_file,
+compute_ssim: Callable[..., Union[Tuple[float, int], Tuple[None, None]]] = (
+    partial(
+        compute_dir_metric,
+        compute_file_func=compute_ssim_file,
+    )
 )
 
-compute_mse = partial(
-    compute_dir_metric,
-    compute_file_func=compute_mse_file,
+compute_mse: Callable[..., Union[Tuple[float, int], Tuple[None, None]]] = (
+    partial(
+        compute_dir_metric,
+        compute_file_func=compute_mse_file,
+    )
 )
 
-compute_video_lpips = partial(
+compute_video_lpips: Callable[
+    ..., Union[Tuple[float, int], Tuple[None, None]]
+] = partial(
     compute_video_metric,
     compute_frame_func=compute_lpips_file,
 )
-compute_video_psnr = partial(
+compute_video_psnr: Callable[
+    ..., Union[Tuple[float, int], Tuple[None, None]]
+] = partial(
     compute_video_metric,
     compute_frame_func=compute_psnr_file,
 )
-compute_video_ssim = partial(
+compute_video_ssim: Callable[
+    ..., Union[Tuple[float, int], Tuple[None, None]]
+] = partial(
     compute_video_metric,
     compute_frame_func=compute_ssim_file,
 )
-compute_video_mse = partial(
+compute_video_mse: Callable[
+    ..., Union[Tuple[float, int], Tuple[None, None]]
+] = partial(
     compute_video_metric,
     compute_frame_func=compute_mse_file,
 )
 
 
 METRICS_CHOICES = [
-    "lpips",
-    "psnr",
-    "ssim",
-    "mse",
-    "fid",
-    "all",
+    "lpips",  # img vs img
+    "psnr",  # img vs img
+    "ssim",  # img vs img
+    "mse",  # img vs img
+    "fid",  # img vs img
+    "all",  # img vs img
+    "clip_score",  # img vs prompt
+    "image_reward",  # img vs prompt
 ]
 
 
@@ -405,6 +460,13 @@ def get_args():
         default=None,
         help="Path to ground truth image or Dir to ground truth images",
     )
+    parser.add_argument(
+        "--prompt-true",
+        "-p",
+        type=str,
+        default=None,
+        help="Path to ground truth prompt file for CLIP Score and Image Reward Score.",
+    )
     parser.add_argument(
         "--img-test",
         "-i2",
@@ -442,6 +504,13 @@ def get_args():
         default=None,
         help="Path to ref dir that contains ground truth images",
     )
+    parser.add_argument(
+        "--ref-prompt-true",
+        "-rp",
+        type=str,
+        default=None,
+        help="Path to ground truth prompt file for CLIP Score and Image Reward Score.",
+    )
 
     # Video 1 vs N pattern
     parser.add_argument(
@@ -495,10 +564,11 @@ def get_args():
         help="Path to addtional perf log",
     )
     parser.add_argument(
-        "--perf-tag",
-        "-ptag",
+        "--perf-tags",
+        "-ptags",
+        nargs="+",
         type=str,
-        default=None,
+        default=[],
         help="Tag to parse perf time from perf log",
     )
     parser.add_argument(
@@ -508,6 +578,26 @@ def get_args():
         default=[],
         help="Extra tags to parse perf time from perf log",
     )
+    parser.add_argument(
+        "--psnr-type",
+        type=str,
+        default="custom",
+        choices=["custom", "skimage"],
+        help="The compute type of PSNR, [custom, skimage]",
+    )
+    parser.add_argument(
+        "--cal-speedup",
+        action="store_true",
+        default=False,
+        help="Calculate performance speedup.",
+    )
+    parser.add_argument(
+        "--gen-markdown-table",
+        "-table",
+        action="store_true",
+        default=False,
+        help="Generate performance markdown table",
+    )
     return parser.parse_args()
 
 
@@ -516,16 +606,16 @@ def entrypoint():
     args = get_args()
     logger.debug(args)
 
+    if args.metrics in ["clip_score", "image_reward"]:
+        assert args.prompt_true is not None or args.ref_prompt_true is not None
+        assert args.img_test is not None or args.img_source_dir is not None
+
     if args.enable_verbose:
         global DISABLE_VERBOSE
         set_metrics_verbose(True)
         DISABLE_VERBOSE = not get_metrics_verbose()
 
-    if "all" in args.metrics or "fid" in args.metrics:
-        FID = FrechetInceptionDistance(
-            disable_tqdm=DISABLE_VERBOSE,
-            batch_size=args.fid_batch_size,
-        )
+    set_psnr_type(args.psnr_type)
 
     METRICS_META: dict[str, float] = {}
 
@@ -533,11 +623,11 @@ def entrypoint():
     def _run_metric(
         metric: str,
         img_true: str = None,
+        prompt_true: str = None,
         img_test: str = None,
         video_true: str = None,
         video_test: str = None,
     ) -> None:
-        nonlocal FID
         nonlocal METRICS_META
         metric = metric.lower()
         if img_true is not None and img_test is not None:
@@ -575,9 +665,39 @@ def entrypoint():
                 img_mse, n = compute_mse(img_true, img_test)
                 _logging_msg(img_mse, "mse", n)
             if metric == "fid" or metric == "all":
-                img_fid, n = FID.compute_fid(img_true, img_test)
+                img_fid, n = compute_fid(img_true, img_test)
                 _logging_msg(img_fid, "fid", n)
 
+        if prompt_true is not None and img_test is not None:
+            if any(
+                (
+                    not os.path.exists(prompt_true),  # file
+                    not os.path.exists(img_test),  # dir
+                )
+            ):
+                return
+
+            # img_true and img_test can be files or dirs
+            prompt_true_info = os.path.basename(prompt_true)
+            img_test_info = os.path.basename(img_test)
+
+            def _logging_msg(value: float, name, n: int):
+                if value is None or n is None:
+                    return
+                msg = (
+                    f"{prompt_true_info} vs {img_test_info}, "
+                    f"Num: {n}, {name.upper()}: {value:.5f}"
+                )
+                METRICS_META[msg] = value
+                logger.info(msg)
+
+            if metric == "clip_score":
+                clip_score, n = compute_clip_score(img_test, prompt_true)
+                _logging_msg(clip_score, "clip_score", n)
+            if metric == "image_reward":
+                image_reward, n = compute_reward_score(img_test, prompt_true)
+                _logging_msg(image_reward, "image_reward", n)
+
         if video_true is not None and video_test is not None:
             if any(
                 (
@@ -614,7 +734,7 @@ def entrypoint():
                 video_mse, n = compute_video_mse(video_true, video_test)
                 _logging_msg(video_mse, "mse", n)
             if metric == "fid" or metric == "all":
-                video_fid, n = FID.compute_video_fid(video_true, video_test)
+                video_fid, n = compute_video_fid(video_true, video_test)
                 _logging_msg(video_fid, "fid", n)
 
     # run selected metrics
@@ -627,7 +747,18 @@ def entrypoint():
     def _is_video_1vsN_pattern() -> bool:
         return args.video_source_dir is not None and args.ref_video is not None
 
-    assert not all((_is_image_1vsN_pattern(), _is_video_1vsN_pattern()))
+    def _is_prompt_1vsN_pattern() -> bool:
+        return (
+            args.img_source_dir is not None and args.ref_prompt_true is not None
+        )
+
+    assert not all(
+        (
+            _is_image_1vsN_pattern(),
+            _is_video_1vsN_pattern(),
+            _is_prompt_1vsN_pattern(),
+        )
+    )
 
     if _is_image_1vsN_pattern():
         # Glob Image dirs
@@ -711,11 +842,42 @@ def entrypoint():
                     video_test=video_test,
                 )
 
+    elif _is_prompt_1vsN_pattern():
+        # Glob Image dirs
+        if not os.path.exists(args.img_source_dir):
+            logger.error(f"{args.img_source_dir} not exist!")
+            return
+
+        directories = []
+        for item in os.listdir(args.img_source_dir):
+            item_path = os.path.join(args.img_source_dir, item)
+            if os.path.isdir(item_path):
+                directories.append(item_path)
+
+        if len(directories) == 0:
+            return
+
+        directories = sorted(directories)
+        if not DISABLE_VERBOSE:
+            logger.info(
+                f"Compare {args.ref_prompt_true} vs {directories}, "
+                f"Num compares: {len(directories)}"
+            )
+
+        for metric in args.metrics:
+            for img_test_dir in directories:
+                _run_metric(
+                    metric=metric,
+                    prompt_true=args.ref_prompt_true,
+                    img_test=img_test_dir,
+                )
+
     else:
         for metric in args.metrics:
             _run_metric(
                 metric=metric,
                 img_true=args.img_true,
+                prompt_true=args.prompt_true,
                 img_test=args.img_test,
                 video_true=args.video_true,
                 video_test=args.video_test,
@@ -724,7 +886,7 @@ def entrypoint():
     if args.summary:
 
         def _fetch_perf():
-            if args.perf_log is None or args.perf_tag is None:
+            if args.perf_log is None or len(args.perf_tags) == 0:
                 return []
             if not os.path.exists(args.perf_log):
                 return []
@@ -733,17 +895,20 @@ def entrypoint():
                 perf_lines = file.readlines()
                 for line in perf_lines:
                     line = line.strip()
-                    if args.perf_tag.lower() in line.lower():
-                        if len(args.extra_perf_tags) == 0:
-                            perf_texts.append(line)
-                        else:
-                            has_all_extra_tag = True
-                            for ext_tag in args.extra_perf_tags:
-                                if ext_tag.lower() not in line.lower():
-                                    has_all_extra_tag = False
-                                    break
-                            if has_all_extra_tag:
+                    for perf_tag in args.perf_tags:
+                        if perf_tag.lower() in line.lower():
+                            if len(args.extra_perf_tags) == 0:
                                 perf_texts.append(line)
+                                break
+                            else:
+                                has_all_extra_tag = True
+                                for ext_tag in args.extra_perf_tags:
+                                    if ext_tag.lower() not in line.lower():
+                                        has_all_extra_tag = False
+                                        break
+                                if has_all_extra_tag:
+                                    perf_texts.append(line)
+                                    break
             return perf_texts
 
         PERF_TEXTS: list[str] = _fetch_perf()
@@ -770,8 +935,9 @@ def entrypoint():
             try:
                 if tag.lower() in METRICS_CHOICES:
                     return float(value_str)
-                if args.perf_tag is not None:
-                    if tag.lower() == args.perf_tag.lower():
+                if len(args.perf_tags) > 0:
+                    perf_tags = [tag.lower() for tag in args.perf_tags]
+                    if tag.lower() in perf_tags:
                         return float(value_str)
                 return int(value_str)
             except ValueError:
@@ -779,17 +945,37 @@ def entrypoint():
 
         def _parse_perf(
             compare_tag: str,
+            perf_tag: str,
         ) -> float | None:
             nonlocal PERF_TEXTS
-            perf_times = []
+            perf_values = []
             for line in PERF_TEXTS:
                 if compare_tag in line:
-                    perf_time = _parse_value(line, args.perf_tag)
-                    if perf_time is not None:
-                        perf_times.append(perf_time)
-            if len(perf_times) == 0:
+                    perf_value = _parse_value(line, perf_tag)
+                    if perf_value is not None:
+                        perf_values.append(perf_value)
+            if len(perf_values) == 0:
                 return None
-            return sum(perf_times) / len(perf_times)
+            return sum(perf_values) / len(perf_values)
+
+        def _ref_perf(
+            key: str,
+        ):
+            # U1-Q0-C0-NONE vs U4-Q1-C1-NONE
+            header = key.split(",")[0].strip()
+            reference_tag = None
+            if args.prompt_true is None:
+                reference_tag = header.split("vs")[0].strip()  # U1-Q0-C0-NONE
+
+            if reference_tag is None:
+                return []
+
+            ref_perf_values = []
+            for perf_tag in args.perf_tags:
+                perf_value = _parse_perf(reference_tag, perf_tag)
+                ref_perf_values.append(perf_value)
+
+            return ref_perf_values
 
         def _format_item(
             key: str,
@@ -802,40 +988,129 @@ def entrypoint():
             header = key.split(",")[0].strip()
             compare_tag = header.split("vs")[1].strip()  # U4-Q1-C1-NONE
             has_perf_texts = len(PERF_TEXTS) > 0
+
+            def _perf_msg(perf_tag: str):
+                if "time" in perf_tag.lower():
+                    perf_msg = "Latency(s)"
+                elif "tflops" in perf_tag.lower():
+                    perf_msg = "TFLOPs"
+                elif "flops" in perf_tag.lower():
+                    perf_msg = "FLOPs"
+                else:
+                    perf_msg = perf_tag.upper()
+                return perf_msg
+
             format_str = ""
             # Num / Frames
+            perf_values = []
+            perf_msgs = []
             if n := _parse_value(key, "Num"):
                 if not has_perf_texts:
                     format_str = (
-                        f"{header:<{max_key_len}}  Num: {n}  "
+                        f"{header:<{max_key_len}}, Num: {n}, "
                         f"{metric.upper()}: {value:<7.4f}"
                     )
                 else:
-                    perf_time = _parse_perf(compare_tag)
-                    perf_time = f"{perf_time:<.2f}" if perf_time else None
                     format_str = (
-                        f"{header:<{max_key_len}}  Num: {n}  "
-                        f"{metric.upper()}: {value:<7.4f}  "
-                        f"Perf: {perf_time}"
+                        f"{header:<{max_key_len}}, Num: {n}, "
+                        f"{metric.upper()}: {value:<7.4f}, "
                     )
+                    for perf_tag in args.perf_tags:
+                        perf_value = _parse_perf(compare_tag, perf_tag)
+                        perf_values.append(perf_value)
+
+                        perf_value = (
+                            f"{perf_value:<.2f}" if perf_value else None
+                        )
+                        perf_msg = _perf_msg(perf_tag)
+                        format_str += f"{perf_msg}: {perf_value}, "
+
+                        perf_msgs.append(perf_msg)
+
+                    if not args.cal_speedup:
+                        format_str = format_str.removesuffix(", ")
+
             elif n := _parse_value(key, "Frames"):
                 if not has_perf_texts:
                     format_str = (
-                        f"{header:<{max_key_len}}  Frames: {n}  "
+                        f"{header:<{max_key_len}}, Frames: {n}, "
                         f"{metric.upper()}: {value:<7.4f}"
                     )
                 else:
-                    perf_time = _parse_perf(compare_tag)
-                    perf_time = f"{perf_time:<.2f}" if perf_time else None
                     format_str = (
-                        f"{header:<{max_key_len}}  Frames: {n}  "
-                        f"{metric.upper()}: {value:<7.4f}  "
-                        f"Perf: {perf_time}"
+                        f"{header:<{max_key_len}}, Frames: {n}, "
+                        f"{metric.upper()}: {value:<7.4f}, "
                     )
+                    for perf_tag in args.perf_tags:
+                        perf_value = _parse_perf(compare_tag, perf_tag)
+                        perf_values.append(perf_value)
+
+                        perf_value = (
+                            f"{perf_value:<.2f}" if perf_value else None
+                        )
+                        perf_msg = _perf_msg(perf_tag)
+                        format_str += f"{perf_msg}: {perf_value}, "
+                        perf_msgs.append(perf_msg)
+
+                    if not args.cal_speedup:
+                        format_str = format_str.removesuffix(", ")
             else:
                 raise ValueError("Num or Frames can not be NoneType.")
 
-            return format_str
+            return format_str, perf_values, perf_msgs
+
+        def _format_table(format_strs: List[str], metric: str):
+            if not format_strs:
+                return ""
+
+            metric_upper = metric.upper()
+            all_headers = {"Config", metric_upper}
+            row_data = []
+
+            for line in format_strs:
+                parts = [p.strip() for p in line.split(",")]
+
+                config_part = parts[0].strip()
+                if "vs" in config_part:
+                    config = config_part.split("vs", 1)[1].strip()
+                    if "_DBCACHE_" in config:
+                        config = config.split("_DBCACHE_", 1)[1].strip()
+                else:
+                    config = config_part
+
+                metric_value = next(
+                    p.split(":")[1].strip()
+                    for p in parts
+                    if p.startswith(metric_upper)
+                )
+
+                perf_data = {}
+                for part in parts:
+                    if part.startswith(("Num:", "Frames:", metric_upper)):
+                        continue
+                    if ":" in part:
+                        key, value = part.split(":", 1)
+                        key = key.strip()
+                        value = value.strip()
+                        perf_data[key] = value
+                        all_headers.add(key)
+
+                row_data.append(
+                    {"Config": config, metric_upper: metric_value, **perf_data}
+                )
+
+            sorted_headers = ["Config", metric_upper] + sorted(
+                [h for h in all_headers if h not in ["Config", metric_upper]]
+            )
+
+            table = "| " + " | ".join(sorted_headers) + " |\n"
+            table += "| " + " | ".join(["---"] * len(sorted_headers)) + " |\n"
+
+            for row in row_data:
+                row_values = [row.get(header, "") for header in sorted_headers]
+                table += "| " + " | ".join(row_values) + " |\n"
+
+            return table.strip()
 
         selected_metrics = args.metrics
         if "all" in selected_metrics:
@@ -848,7 +1123,17 @@ def entrypoint():
                 if metric.upper() in key or metric.lower() in key:
                     selected_items[key] = METRICS_META[key]
 
-            reverse = True if metric.lower() in ["psnr", "ssim"] else False
+            reverse = (
+                True
+                if metric.lower()
+                in [
+                    "psnr",
+                    "ssim",
+                    "clip_score",
+                    "image_reward",
+                ]
+                else False
+            )
             sorted_items = sorted(
                 selected_items.items(), key=lambda x: x[1], reverse=reverse
             )
@@ -857,12 +1142,65 @@ def entrypoint():
             ]
             max_key_len = max(len(key) for key in selected_keys)
 
+            ref_perf_values = _ref_perf(key=selected_keys[0])
+            max_perf_values: List[float] = []
+
+            if ref_perf_values and None not in ref_perf_values:
+                max_perf_values = ref_perf_values.copy()
+
+            for key, value in sorted_items:
+                format_str, perf_values, perf_msgs = _format_item(
+                    key, metric, value, max_key_len
+                )
+                # skip 'None' msg but not 'NONE', 'NONE' means w/o cache
+                if "None" in format_str:
+                    continue
+
+                if (
+                    not perf_values
+                    or None in perf_values
+                    or not perf_msgs
+                    or not args.cal_speedup
+                ):
+                    continue
+
+                if not max_perf_values:
+                    max_perf_values = perf_values
+                else:
+                    for i in range(len(max_perf_values)):
+                        max_perf_values[i] = max(
+                            max_perf_values[i], perf_values[i]
+                        )
+
             format_strs = []
             for key, value in sorted_items:
-                format_strs.append(
-                    _format_item(key, metric, value, max_key_len)
+                format_str, perf_values, perf_msgs = _format_item(
+                    key, metric, value, max_key_len
                 )
 
+                # skip 'None' msg but not 'NONE', 'NONE' means w/o cache
+                if "None" in format_str:
+                    continue
+
+                if (
+                    not perf_values
+                    or None in perf_values
+                    or not perf_msgs
+                    or not max_perf_values
+                    or not args.cal_speedup
+                ):
+                    format_strs.append(format_str)
+                    continue
+
+                for perf_value, perf_msg, max_perf_value in zip(
+                    perf_values, perf_msgs, max_perf_values
+                ):
+                    perf_speedup = max_perf_value / perf_value
+                    format_str += f"{perf_msg}(↑): {perf_speedup:<.2f}, "
+
+                format_str = format_str.removesuffix(", ")
+                format_strs.append(format_str)
+
             format_len = max(len(format_str) for format_str in format_strs)
 
             res_len = format_len - len(f"Summary: {metric.upper()}")
@@ -877,6 +1215,12 @@ def entrypoint():
                 print(format_str)
             print("-" * format_len)
 
+            if args.gen_markdown_table:
+                table = _format_table(format_strs, metric)
+                print("-" * format_len)
+                print(f"{table}")
+            print("-" * format_len)
+
 
 if __name__ == "__main__":
     entrypoint()