plot_utils.py

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from keras.utils import load_img, img_to_array
from scipy.ndimage import median_filter as _median_filter
from scipy.signal import wiener as _wiener
from skimage.filters import median as _median_filter
from skimage.metrics import structural_similarity as compare_ssim
from skimage.morphology import disk as _disk
from skimage.restoration import wiener as _wiener_filter
from tensorflow.keras.applications import VGG16
from tensorflow.keras.applications.vgg16 import preprocess_input

import dataset_utils as dataset_utils


def plot_metric(history, metric_suffix):
    import matplotlib.pyplot as plt
    h = history.history
    # try aggregate, then per-output
    train_key = metric_suffix if metric_suffix in h else next(
        (k for k in h if k.endswith(metric_suffix) and not k.startswith('val_')), None)
    val_key = ('val_' + metric_suffix) if ('val_' + metric_suffix) in h else next(
        (k for k in h if k.startswith('val_') and k.endswith(metric_suffix)), None)

    if train_key is None:
        raise ValueError(f"No metric '{metric_suffix}' in history. Available: {list(h.keys())}")

    plt.figure(figsize=(8, 6))
    plt.plot(h[train_key], label=train_key)
    if val_key: plt.plot(h[val_key], label=val_key)
    plt.xlabel('Epochs'); plt.ylabel(metric_suffix); plt.title(metric_suffix)
    plt.legend(); plt.show()

# ==== LPIPS (офіційний). Ініціалізуємо один раз, з прозорими помилками ====
_LPIPS_READY = False
_LPIPS_ERR = None
try:
    import torch
    import lpips  # pip install lpips
    _LPIPS_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    _LPIPS_MODEL = lpips.LPIPS(net='alex').to(_LPIPS_DEVICE)
    _LPIPS_READY = True
except Exception as e:
    _LPIPS_ERR = e
    _LPIPS_READY = False

def _compute_lpips_official(gt_image, pred_image, net='alex'):
    """
    Офіційний LPIPS. Вхід: numpy [H,W,3] у [0,1] RGB. Повертає float.
    Якщо розміри різні — pred ресайзиться до GT bilinear'ом.
    """
    if not _LPIPS_READY:
        raise RuntimeError(f"LPIPS is not available ({_LPIPS_ERR}). "
                           f"Install with: pip install torch torchvision lpips")

    def _to_torch(im):
        t = torch.from_numpy(im.astype(np.float32)).permute(2,0,1).unsqueeze(0)  # [1,3,H,W]
        t = torch.clamp(t, 0.0, 1.0) * 2.0 - 1.0  # -> [-1,1]
        return t.to(_LPIPS_DEVICE)

    gt = np.asarray(gt_image, dtype=np.float32)
    pr = np.asarray(pred_image, dtype=np.float32)
    if gt.shape[:2] != pr.shape[:2]:
        # bilinear resize pred -> gt size
        tpr = torch.from_numpy(pr.transpose(2,0,1)[None]).float()
        tpr = torch.nn.functional.interpolate(tpr, size=gt.shape[:2], mode="bilinear", align_corners=False)
        pr = tpr[0].permute(1,2,0).numpy()

    t_gt = _to_torch(gt)
    t_pr = _to_torch(pr)
    with torch.no_grad():
        d = _LPIPS_MODEL(t_gt, t_pr)  # [1,1,1,1]
    return float(d.mean().item())

def compute_metrics(gt_image, predicted_image, win_size=11, data_range=1.0):
    gt_image = np.clip(gt_image, 0, 1).astype(np.float32)
    predicted_image = np.clip(predicted_image, 0, 1).astype(np.float32)
    ssim = compare_ssim(gt_image, predicted_image, channel_axis=2,
                        win_size=win_size, data_range=data_range)
    psnr = tf.image.psnr(gt_image, predicted_image, max_val=data_range).numpy()
    return float(ssim), float(psnr)

def _apply_wiener(img, balance=0.5):
    out = []
    for c in range(3):
        ch = img[..., c]
        # wiener expects real-valued image; balance ~ regularization
        ch_w = _wiener_filter(ch, psf=None, balance=balance, clip=True)
        out.append(np.clip(ch_w, 0, 1))
    return np.stack(out, axis=-1)

def _apply_median(img, radius=1):
    se = _disk(radius)
    out = []
    for c in range(3):
        out.append(_median_filter(img[..., c], selem=se))
    return np.stack(out, axis=-1)

def _plot_triplet(hazy, pred, gt, title):
    fig, axs = plt.subplots(1, 3, figsize=(12, 4))
    axs[0].imshow(np.clip(hazy, 0, 1)); axs[0].set_title("Hazy"); axs[0].axis("off")
    axs[1].imshow(np.clip(pred, 0, 1)); axs[1].set_title("Predicted"); axs[1].axis("off")
    axs[2].imshow(np.clip(gt,   0, 1)); axs[2].set_title("GT"); axs[2].axis("off")
    fig.suptitle(title)
    plt.tight_layout()
    plt.show()

def predictAndComputeMetrics(model, noisy_img_path, gt_img_path, win_size=11,
                             lpips_net='alex', wiener=False, median=False,
                             show_plot=False):
    # GT
    gt_image = img_to_array(load_img(gt_img_path)) / 255.0

    # Input
    image_batch, hazy_img = dataset_utils.preprocess_test_image(noisy_img_path)

    # Predict
    pred = model.predict(image_batch, verbose=0)
    if isinstance(pred, (list, tuple)):
        pred = pred[0]
    pred = np.clip(np.asarray(pred)[0], 0, 1).astype(np.float32)

    if wiener:
        pred = _apply_wiener(pred)
    if median:
        pred = _apply_median(pred)

    # Metrics
    ssim, psnr = compute_metrics(gt_image, pred, win_size=win_size)
    try:
        lpips = _compute_lpips_official(gt_image, pred, net=lpips_net)
    except Exception as e:
        # Зробимо явний меседж, щоб ви одразу бачили причину, а не NaN.
        print(f"[LPIPS warning] {e}")
        lpips = float("nan")

    if show_plot:
        ttl = f"PSNR={psnr:.2f} | SSIM={ssim:.4f} | LPIPS={lpips if np.isfinite(lpips) else 'NaN'}"
        _plot_triplet(hazy_img, pred, gt_image, ttl)

    return ssim, psnr, lpips


def predictAndComputeMetrics2(model, noisy_img_path, gt_img_path,
                              win_size=11, wiener=False, median=False,
                              lpips_net='alex', show_plot=True):
    # GT
    gt_image = img_to_array(load_img(gt_img_path)) / 255.0
    # Input
    image_batch, hazy_img = dataset_utils.preprocess_test_image(noisy_img_path)
    # Predict
    pred = model.predict(image_batch, verbose=0)
    if isinstance(pred, (list, tuple)):
        pred = pred[0]
    pred = np.clip(np.asarray(pred)[0], 0, 1).astype(np.float32)

    if wiener:
        pred = _apply_wiener(pred)
    if median:
        pred = _apply_median(pred)

    # Metrics
    ssim, psnr = compute_metrics(gt_image, pred, win_size=win_size)
    try:
        lpips = _compute_lpips_official(gt_image, pred, net=lpips_net)
    except Exception as e:
        print(f"[LPIPS warning] {e}")
        lpips = float("nan")

    if show_plot:
        ttl = f"PSNR={psnr:.2f} | SSIM={ssim:.4f} | LPIPS={lpips if np.isfinite(lpips) else 'NaN'}"
        _plot_triplet(hazy_img, pred, gt_image, ttl)

    return ssim, psnr, lpips

def visualize_patches(dataset, num_patches=5):
    """
    Visualize a few image patches from the dataset.

    Args:
        dataset: TensorFlow dataset containing image pairs (noisy_patch, gt_patch).
        num_patches: Number of patches to visualize.
    """
    for noisy_patch, gt_patch in dataset.take(1):  # Take one batch
        for i in range(min(num_patches, noisy_patch.shape[0])):
            plt.figure(figsize=(8, 4))

            # Display the noisy patch (rescaling it back to [0, 255])
            plt.subplot(1, 2, 1)
            plt.imshow(noisy_patch[i].numpy())
            plt.title("Noisy Patch")
            plt.axis("off")

            # Display the ground truth patch (rescaling it back to [0, 255])
            plt.subplot(1, 2, 2)
            plt.imshow(gt_patch[i].numpy())
            plt.title("Ground Truth Patch")
            plt.axis("off")

            plt.show()


def _wiener_color(img, mysize=5, noise=None):
    if img.ndim == 2:
        out = _wiener(img, mysize=mysize, noise=noise)
    else:
        out = np.empty_like(img)
        for c in range(img.shape[-1]):
            out[..., c] = _wiener(img[..., c], mysize=mysize, noise=noise)
    return np.clip(out.astype(np.float32), 0.0, 1.0)

def _median_color(img, k=3):
    if img.ndim == 2:
        out = _median_filter(img, size=k, mode='nearest')
    else:
        out = np.empty_like(img)
        for c in range(img.shape[-1]):
            out[..., c] = _median_filter(img[..., c], size=k, mode='nearest')
    return np.clip(out.astype(np.float32), 0.0, 1.0)

def _plot_series(history, train_key, val_key=None, title="", ylabel=""):
    h = history.history
    x = range(1, len(h.get(train_key, [])) + 1)
    plt.figure(figsize=(6,4))
    if train_key in h: plt.plot(x, h[train_key], label=train_key)
    if val_key and val_key in h: plt.plot(x, h[val_key], label=val_key)
    plt.xlabel("Epoch"); plt.ylabel(ylabel); plt.title(title); plt.legend(); plt.grid(True); plt.tight_layout(); plt.show()

def plot_loss(history):
    _plot_series(history, "loss", "val_loss", title="Training Loss", ylabel="MAE")

def plot_psnr(history):
    _plot_series(history, "clean_psnr_metric", "val_clean_psnr_metric", title="PSNR", ylabel="dB")

def plot_ssim(history):
    _plot_series(history, "clean_ssim_metric", "val_clean_ssim_metric", title="SSIM", ylabel="SSIM")

def plot_training_summary(history):
    plot_loss(history)
    plot_psnr(history)
    plot_ssim(history)