import cv2
import os
import glob
import numpy as np

def concat_change_detection_images(img1_path, img2_path, label_path, pred_path, output_path):
    img1 = cv2.imread(img1_path)
    img2 = cv2.imread(img2_path)
    label = cv2.imread(label_path) if os.path.exists(label_path) else None
    pred = cv2.imread(pred_path)

    if img1 is None or img2 is None or pred is None:
        print(f"Missing or unreadable image: {img1_path}, {img2_path}, {pred_path}")
        return

    # resize 所有图片为相同大小（以 img1 为基准）
    h, w = img1.shape[:2]
    img2 = cv2.resize(img2, (w, h))
    pred = cv2.resize(pred, (w, h))
    if label is not None:
        label = cv2.resize(label, (w, h))

    # 组合图像（无 label 时跳过）
    if label is not None:
        concat = np.concatenate([img1, img2, label, pred], axis=1)
    else:
        concat = np.concatenate([img1, img2, pred], axis=1)

    cv2.imwrite(output_path, concat)

def batch_process(img1_dir, img2_dir, label_dir, pred_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
        pred_path = os.path.join(pred_dir, filename)
        output_path = os.path.join(output_dir, filename.replace(".png", "_concat.png"))

        print(f"[INFO] img1: {img1_path}, img2: {img2_path}")
        print(f"[INFO] label: {label_path}, pred: {pred_path}")

        concat_change_detection_images(img1_path, img2_path, label_path, pred_path, output_path)
        print(f"Saved: {output_path}")

# 设置路径
img1_dir = "data/WHU_CD/test/image1"
img2_dir = "data/WHU_CD/test/image2"
label_dir = "data/WHU_CD/test/label"            # 如果没有标签图可以设为 None
pred_dir = "work_dirs/CLCD_BS4_epoch200/CDXFormer/version_0/ckpts/test/mask_rgb"
output_dir = "mask_connect_test_dir"

batch_process(img1_dir, img2_dir, label_dir, pred_dir, output_dir)