visualization.py

import cv2
import os
from utils import normalize
import numpy as np

def visualizer(pathes, anomaly_map, img_size, save_path, cls_name):
    for idx, path in enumerate(pathes):
        cls = path.split('/')[-2]
        filename = path.split('/')[-1]
        vis = cv2.cvtColor(cv2.resize(cv2.imread(path), (img_size, img_size)), cv2.COLOR_BGR2RGB)  # RGB
        mask = normalize(anomaly_map[idx])
        vis = apply_ad_scoremap(vis, mask)
        vis = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)  # BGR
        save_vis = os.path.join(save_path, 'imgs', cls_name[idx], cls)
        if not os.path.exists(save_vis):
            os.makedirs(save_vis)
        cv2.imwrite(os.path.join(save_vis, filename), vis)

def apply_ad_scoremap(image, scoremap, alpha=0.5):
    np_image = np.asarray(image, dtype=float)
    scoremap = (scoremap * 255).astype(np.uint8)
    scoremap = cv2.applyColorMap(scoremap, cv2.COLORMAP_JET)
    scoremap = cv2.cvtColor(scoremap, cv2.COLOR_BGR2RGB)
    return (alpha * np_image + (1 - alpha) * scoremap).astype(np.uint8)

def apply_ad_scoremap_batch(images, scoremaps, alpha=0.5):
    """
    批量处理图像和scoremap的叠加
    Args:
        images: np.ndarray 形状 [B, C, H, W] (C=3, RGB)
        scoremaps: np.ndarray 形状 [B, 1, H, W] (值范围[0,1])
        alpha: 混合比例
    Returns:
        np.ndarray 形状 [B, C, H, W] (叠加结果)
    """
    # 确保输入是numpy数组
    images = np.asarray(images, dtype=float)  # [B,3,H,W]
    scoremaps = np.asarray(scoremaps, dtype=float)  # [B,1,H,W]
    
    results = []
    for img, score in zip(images, scoremaps):
        # 处理单个图像
        img = img.transpose(1, 2, 0)  # CHW -> HWC [336,336,3]
        score = score.squeeze(0)  # 移除通道维 [336,336]
        
        # 生成热力图 (与原始函数相同)
        score = (score * 255).astype(np.uint8)
        heatmap = cv2.applyColorMap(score, cv2.COLORMAP_JET)
        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)  # -> [336,336,3]
        
        # 混合并保存结果
        blended = (alpha * img + (1 - alpha) * heatmap).astype(np.uint8)
        blended = blended.transpose(2, 0, 1)  # HWC -> CHW [3,336,336]
        results.append(blended)
    
    return np.stack(results)  # [B,3,H,W]