five_connect.py

import cv2
import os
import glob
import numpy as np

def concat_image_heatmap(img1_path, img2_path, label_path, mask_path, heatmap_path, output_path):
    img1 = cv2.imread(img1_path)
    img2 = cv2.imread(img2_path)
    mask = cv2.imread(mask_path)
    heatmap = cv2.imread(heatmap_path)
    label = cv2.imread(label_path) if label_path and os.path.exists(label_path) else None

    if img1 is None or img2 is None or mask is None or heatmap is None:
        print(f"❌ Missing image: {img1_path}, {img2_path}, {mask_path}, {heatmap_path}")
        return

    h, w = img1.shape[:2]
    img2 = cv2.resize(img2, (w, h))
    mask = cv2.resize(mask, (w, h))
    heatmap = cv2.resize(heatmap, (w, h))
    label = cv2.resize(label, (w, h)) if label is not None else np.zeros_like(img1)

    top_row = np.concatenate([img1, img2, label], axis=1)
    bottom_row = np.concatenate([mask, heatmap], axis=1)

    # 补齐对齐
    max_width = max(top_row.shape[1], bottom_row.shape[1])
    if top_row.shape[1] < max_width:
        pad = max_width - top_row.shape[1]
        top_row = cv2.copyMakeBorder(top_row, 0, 0, 0, pad, cv2.BORDER_CONSTANT, value=0)
    if bottom_row.shape[1] < max_width:
        pad = max_width - bottom_row.shape[1]
        bottom_row = cv2.copyMakeBorder(bottom_row, 0, 0, 0, pad, cv2.BORDER_CONSTANT, value=0)

    full_image = np.concatenate([top_row, bottom_row], axis=0)
    cv2.imwrite(output_path, full_image)
    print(f"✅ Saved: {output_path}")

def batch_process(img1_dir, img2_dir, label_dir, mask_dir, heatmap_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    img1_paths = glob.glob(os.path.join(img1_dir, "*.png"))

    for img1_path in img1_paths:
        filename = os.path.basename(img1_path)
        img2_path = os.path.join(img2_dir, filename)
        label_path = os.path.join(label_dir, filename) if label_dir else None
        mask_path = os.path.join(mask_dir, filename)
        heatmap_path = os.path.join(heatmap_dir, filename)
        output_path = os.path.join(output_dir, filename.replace(".png", "_full.png"))

        concat_image_heatmap(img1_path, img2_path, label_path, mask_path, heatmap_path, output_path)

# 设置路径
img1_dir = "data/WHU_CD/test/image1"
img2_dir = "data/WHU_CD/test/image2"
label_dir = "data/WHU_CD/test/label"  # 可设为 None
mask_dir = "mask_connect_test_dir/mask_rgb"
heatmap_dir = "mask_connect_test_dir/grad_cam/model.net.decoderhead.LHBlock2"
output_dir = "mask_heatmap_concat_dir"

batch_process(img1_dir, img2_dir, label_dir, mask_dir, heatmap_dir, output_dir)