user
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
from __future__ import annotations
|
|
from dataclasses import dataclass
|
|
from typing import Dict, List
|
|
import numpy as np
|
|
|
|
from entities.rgb_image512 import RgbImage512
|
|
from entities.color_space import ColorSpace
|
|
from entities.psnr_metric import PsnrMetric
|
|
|
|
|
|
@dataclass
|
|
class AnalyzeQuality:
|
|
"""CLI-режим расчёта метрик качества."""
|
|
|
|
def run(
|
|
self,
|
|
original_path: str,
|
|
stego_path: str,
|
|
space: str = "ycbcr",
|
|
metrics: List[str] = None,
|
|
) -> Dict[str, float]:
|
|
"""Посчитать метрики между изображениями."""
|
|
if metrics is None:
|
|
metrics = ["psnr"]
|
|
metrics = [m.lower() for m in metrics]
|
|
if any(m != "psnr" for m in metrics):
|
|
raise ValueError("Поддерживается только метрика PSNR.")
|
|
|
|
orig = RgbImage512.from_file(original_path).to_array()
|
|
steg = RgbImage512.from_file(stego_path).to_array()
|
|
|
|
if space.lower() == "ycbcr":
|
|
a = ColorSpace.rgb_to_ycbcr(orig)
|
|
b = ColorSpace.rgb_to_ycbcr(steg)
|
|
elif space.lower() == "rgb":
|
|
a, b = orig, steg
|
|
else:
|
|
raise ValueError("space должен быть 'rgb' или 'ycbcr'.")
|
|
|
|
psnr = PsnrMetric().psnr(a, b)
|
|
return {"psnr": float(psnr)}
|