context_cluster.py

# MIT License
# Copyright (c) 2023 Jiewen Yang
# Please refer to the paper : https://arxiv.org/abs/2309.11145.

import os
import importlib
import torch
import torch.nn.functional as F
import json
from monai.data import DataLoader
import matplotlib
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np
import cv2
from sklearn.cluster import KMeans
torch.autograd.set_detect_anomaly(True)
matplotlib.use('Agg')


def plot_cluster_center_profiles(cluster_id, H_profile, V_profile, save_dir):
    """save cluster Horizontal & Vertical Profile"""
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    axes[0].plot(H_profile)
    axes[0].set_title(f"Cluster {cluster_id} Horizontal Profile")
    axes[1].plot(V_profile)
    axes[1].set_title(f"Cluster {cluster_id} Vertical Profile")
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, f"cluster_{cluster_id}_profiles.png"), dpi=300)
    plt.close()


def combined_threshold_weighted(img):
    """
    Binarization using a weighted average of maximum entropy thresholding and Otsu thresholding
    Args:
        img: one-channel image (uint8, range [0,255])
    Returns:
        binary: Binary image after weighted average thresholding
    """
    _, binary = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    return binary


def compute_histogram_features(img_batch):
    """
    Extracting horizontal/vertical projection histogram features from batch images
    img_batch: Tensor [B, C, F, H, W], Byte
    return: [B, feature_dim] numpy array
    """
    B, C, F, H, W = img_batch.shape
    img_batch = img_batch.float().mean(dim=2)  # [B, C, H, W]
    img_batch = img_batch.squeeze(1)           # [B, H, W]

    features = []
    for img in img_batch:
        img_np = img.detach().cpu().numpy()
        img_np = img_np.astype(np.uint8)  

        binary = combined_threshold_weighted(img_np)

        # Horizontal/Vertical Projection Histogram
        vertical_sum = np.sum(binary, axis=1)
        horizontal_sum = np.sum(binary, axis=0)

        # Max Normalization
        vertical_norm = vertical_sum / (np.max(vertical_sum) + 1e-5)
        horizontal_norm = horizontal_sum / (np.max(horizontal_sum) + 1e-5)

        # concatenate feat

        feature_vec = np.concatenate([horizontal_norm, vertical_norm])

        features.append(feature_vec)

    return np.array(features)


import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import os

def joint_cluster_and_save(src_loader, tgt_loader, output_dir, device, k=2, samples_per_cluster=30):
    """
    Combines cluster source and target domains, and saves samples, statistics, and visualizes cluster centers.
    """
    all_features, all_imgs, all_domains = [], [], []
    vis_save_dir = "/home/fangxinyan/StaffEcho_jn/cluster_samples/binary_samples"
    from itertools import islice


    src_count = len(src_loader.dataset)
    tgt_count = len(tgt_loader.dataset)
    min_count = min(src_count, tgt_count)

    print(f"src count: {src_count}, tgt count: {tgt_count}")
    print(f"sample count: {min_count}")

    
    print("Extract source domain features...")
    for img_batch, *_ in tqdm(islice(src_loader, 0, min_count // src_loader.batch_size), desc="Loading SRC"):
        img_batch = img_batch.to(device)
        feats = compute_histogram_features(img_batch)
        all_features.append(feats)
        all_imgs.append(img_batch.cpu())
        all_domains.extend(["src"] * img_batch.shape[0])


    print("Extract target domain features...")
    for img_batch, *_ in tqdm(islice(tgt_loader, 0, min_count // tgt_loader.batch_size), desc="Loading TGT"):
        img_batch = img_batch.to(device)
        feats = compute_histogram_features(img_batch)
        all_features.append(feats)
        all_imgs.append(img_batch.cpu())
        all_domains.extend(["tgt"] * img_batch.shape[0])


    all_features = np.vstack(all_features)
    all_imgs = torch.cat(all_imgs, dim=0)  # [N, C, F, H, W]
    B, C, F, H, W = all_imgs.shape


    print("Kmeans...")
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    all_labels = kmeans.fit_predict(all_features)

   
    print("Save clustered samples...")
    for cluster_id in range(k):
        for domain in ["src", "tgt"]:
            cluster_dir = os.path.join(output_dir, f"{domain}_cluster_{cluster_id}")
            os.makedirs(cluster_dir, exist_ok=True)

            indices = [
                idx for idx, (label, dom) in enumerate(zip(all_labels, all_domains))
                if label == cluster_id and dom == domain
            ]
            selected_indices = indices[:samples_per_cluster]  # Save only the first N images

            for idx in tqdm(selected_indices, desc=f"Saving {domain} cluster {cluster_id}"):
                img = all_imgs[idx].float().mean(dim=1).squeeze(0).numpy()
                img_uint8 = np.clip(img, 0, 255).astype(np.uint8)
                save_path = os.path.join(cluster_dir, f"{domain}_{idx}.png")
                cv2.imwrite(save_path, img_uint8)

        # === Visualizing cluster centers + samples ===
        cluster_center_dir = os.path.join(output_dir, "cluster_centers")
        os.makedirs(cluster_center_dir, exist_ok=True)

        cluster_center = kmeans.cluster_centers_[cluster_id]
        cluster_indices = np.where(all_labels == cluster_id)[0]

        # Sort by distance from cluster center
        distances = np.linalg.norm(all_features[cluster_indices] - cluster_center, axis=1)
        sorted_indices = cluster_indices[np.argsort(distances)]

        # === Visualizing tSNE + Clustering Representative Samples ===
        sorted_indices_per_cluster = {}
        for cluster_id in range(k):
            cluster_center = kmeans.cluster_centers_[cluster_id]
            cluster_indices = np.where(all_labels == cluster_id)[0]
            distances = np.linalg.norm(all_features[cluster_indices] - cluster_center, axis=1)
            sorted_indices = cluster_indices[np.argsort(distances)][:6]  # Take the first 6
            sorted_indices_per_cluster[cluster_id] = sorted_indices

        vis_save_dir = os.path.join(output_dir, "tsne_visualization")
        plot_tsne_and_samples(
            features=all_features,
            labels=all_labels,
            domains=all_domains,
            cluster_centers=kmeans.cluster_centers_,
            all_imgs=all_imgs,
            sorted_indices_per_cluster=sorted_indices_per_cluster,
            save_dir=vis_save_dir
        )


    print("\n==== Clustering Statistics ====")
    for cluster_id in range(k):
        src_count = sum(
            1 for label, dom in zip(all_labels, all_domains)
            if label == cluster_id and dom == "src"
        )
        tgt_count = sum(
            1 for label, dom in zip(all_labels, all_domains)
            if label == cluster_id and dom == "tgt"
        )
        print(f"Cluster {cluster_id}: SRC={src_count}, TGT={tgt_count}")

import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from matplotlib.patches import Ellipse

def plot_tsne_and_samples(
    features,
    labels,
    domains,
    cluster_centers,
    all_imgs,  # Tensor [N, C, F, H, W]
    sorted_indices_per_cluster,
    save_dir
):
    os.makedirs(save_dir, exist_ok=True)

    
    print("Drawing t-SNE...")
    tsne = TSNE(n_components=2, random_state=42, perplexity=30)
    features_2d = tsne.fit_transform(features)

    domain_colors = {"src": "royalblue", "tgt": "seagreen"}
    cluster_colors = ["royalblue", "seagreen"]

    fig, ax = plt.subplots(figsize=(8, 8))

    for domain in ["src", "tgt"]:
        indices = [i for i, d in enumerate(domains) if d == domain]
        ax.scatter(
            features_2d[indices, 0],
            features_2d[indices, 1],
            label=domain.upper(),
            c=domain_colors[domain],
            s=200,
            alpha=0.6
        )

    # Cluster center mapping to 2D
    centers_2d = tsne.fit_transform(np.vstack([features, cluster_centers]))[-len(cluster_centers):]
    for i, (x, y) in enumerate(centers_2d):
        ax.scatter(x, y, marker='*', s=1000, c=cluster_colors[i], edgecolors='k', linewidths = 2)

        # The clustering region is outlined by a gray dashed circle (simplified by using an ellipse).
        ax.add_patch(Ellipse(
            (x, y),
            width=25, height=25,
            angle=0,
            edgecolor='gray',
            linestyle='--',
            linewidth= 5,
            facecolor='none'
        ))

    ax.set_title("t-SNE of clustered samples")
    ax.axis('off')
    ax.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, "tsne_plot.png"))
    plt.close()

    # === Clustering Sample Visualization ===
    print("Save the visualization of typical samples...")
    for cluster_id, indices in sorted_indices_per_cluster.items():
        fig, axs = plt.subplots(3, 1, figsize=(6, 9))
        for i, idx in enumerate(indices):
            img = all_imgs[idx].float().mean(dim=1).squeeze(0).numpy()
            img = np.clip(img, 0, 255).astype(np.uint8)

            # Calculate vertical/horizontal grayscale histograms
            vertical_sum = np.sum(img, axis=1)
            horizontal_sum = np.sum(img, axis=0)
            vertical_norm = vertical_sum / (np.max(vertical_sum) + 1e-5)
            horizontal_norm = horizontal_sum / (np.max(horizontal_sum) + 1e-5)

            color = "royalblue" if domains[idx] == "src" else "seagreen"
            light_color = "#add8e6" if domains[idx] == "src" else "#90ee90"

            # color = "seagreen" if domains[idx] == "src" else "royalblue"
            # light_color = "#90ee90" if domains[idx] == "src" else "#add8e6"

            fig2 = plt.figure(figsize=(4, 4))
            gs = fig2.add_gridspec(2, 2, width_ratios=(4, 1), height_ratios=(1, 4), wspace=0.05, hspace=0.05)
            ax_img = fig2.add_subplot(gs[1, 0])
            ax_top = fig2.add_subplot(gs[0, 0], sharex=ax_img)
            ax_right = fig2.add_subplot(gs[1, 1], sharey=ax_img)

            ax_img.imshow(img, cmap='gray')
            ax_img.axis('off')

            ax_top.bar(range(len(horizontal_norm)), horizontal_norm, color=light_color)
            ax_top.plot(horizontal_norm, color=color, linewidth=1.5)
            ax_top.axis('off')

            ax_right.barh(range(len(vertical_norm)), vertical_norm, color=light_color)
            ax_right.plot(vertical_norm, range(len(vertical_norm)), color=color, linewidth=1.5)
            # ax_right.invert_xaxis()
            ax_right.axis('off')

            fig2.tight_layout()
            fig2.savefig(os.path.join(save_dir, f"cluster{cluster_id}_sample{i}.png"), dpi=200)
            plt.close(fig2)

    print("Saved.")


class Trainer():
    def __init__(self, config):
 
        self.size_crop = {
        "Camus2Echo": [(136,128), (128,128)],
        'Camus2Cardiac_uda': [(272,256), (328, 256)],
        'Echo2Camus':[(256,256), (272,256)],
        'Echo2Cardiac_uda':[(256,256), (328,256)],
        'Cardiac_uda2Camus':[(328, 256),(272,256),],
        'Cardiac_uda2Echo':[(164,128), (128,128)]
        }
 
        self.config = config


        os.environ['CUDA_VISIBLE_DEVICES'] = config['train']['enable_GPU_id']
   

        spatial_size_src, crop_size_src, spatial_size_tgt, crop_size_tgt = self.get_size(config['source']['dataset'], config['target']['dataset'])
        dataset_src = getattr(importlib.import_module("datasets." + config['source']['dataset'].lower()), config['source']['dataset'])
        dataset_tgt = getattr(importlib.import_module("datasets." + config['target']['dataset'].lower()), config['target']['dataset'])
   
        self.train_dataset_src = dataset_src(**config['source']['config'], stage = 'train',  spatial_size = spatial_size_src, crop_size = crop_size_src, )

        self.train_dataset_tgt = dataset_tgt(**config['target']['config'], stage  = 'train',   spatial_size = spatial_size_tgt, crop_size = crop_size_tgt,  )

        self.val_dataset_src  = dataset_src(**config['source']['config'], stage = 'val',  spatial_size = spatial_size_src, crop_size = crop_size_src, )  
        self.val_dataset_tgt  = dataset_tgt(**config['target']['config'], stage = 'val',  spatial_size = spatial_size_tgt, crop_size = crop_size_tgt, ) 

        self.test_dataset_src  = dataset_src(**config['source']['config'], stage = 'test',  spatial_size = spatial_size_src, crop_size = crop_size_src, )  
        self.test_dataset_tgt  = dataset_tgt(**config['target']['config'], stage = 'test',  spatial_size = spatial_size_tgt, crop_size = crop_size_tgt, ) 

        self.train_loader_src = DataLoader(self.train_dataset_src, batch_size=config['train']['batch_size'], shuffle=True, num_workers=config['train']['num_workers'], drop_last=True) 
        self.train_loader_tgt = DataLoader(self.train_dataset_tgt, batch_size=config['train']['batch_size'], shuffle=True, num_workers=config['train']['num_workers'], drop_last=True)        
      

    
    def get_size(self, source_dataset, target_dataset):
        transform_size = self.size_crop["{}2{}".format(source_dataset,target_dataset)]
        return transform_size[0][0],transform_size[0][1], transform_size[1][0], transform_size[1][1], 
        

    def train(self):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        output_dir = "/home/fangxinyan/StaffEcho_jn/cluster_samples/cd2cm"

        joint_cluster_and_save(
            src_loader=self.train_loader_src,
            tgt_loader=self.train_loader_tgt,
            output_dir=output_dir,
            device=device,
            k=2,
            samples_per_cluster=100
        )


if __name__ == "__main__":
    from utils.tools import set_seed 
    from pathlib import Path
    basedir = "/home/fangxinyan/public4jbhi/config/STAffEcho"
    configs = os.listdir(basedir)
    seeds = [42]

    for seed in seeds:
        set_seed(seed)
        for config in configs:
            config_path = os.path.join(basedir, config)
            with open(config_path, "r") as file:
                config = json.load(file) 
            
            config['train']["log_dir"] = str(Path(*Path(config['train']["log_dir"]).parts[:-2], 'seed{}'.format(seed), *Path(config['train']["log_dir"]).parts[-1:]))
            config['train']["save_dir"] = str(Path(*Path(config['train']["save_dir"]).parts[:-2], 'seed{}'.format(seed), *Path(config['train']["save_dir"]).parts[-1:]))
            config['train']["img_dir"] = str(Path(*Path(config['train']["img_dir"]).parts[:-2], 'seed{}'.format(seed),  *Path(config['train']["img_dir"]).parts[-1:]))

            Train_ = Trainer(config)
            Train_.train()