import os import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset from torchvision import models, transforms from PIL import Image import numpy as np import random from sklearn.decomposition import PCA from sklearn.manifold import TSNE from sklearn.cluster import KMeans import matplotlib.pyplot as plt import pandas as pd import torch.nn.functional as F import plotly.express as px def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False set_seed(42) BASE_DIR = os.path.dirname(os.path.abspath(__file__)) IMG_DIR = os.path.join(BASE_DIR, "..", "datos_para_ai") ENCODER_PATH = os.path.join(BASE_DIR, "..", "models", "msi_encoder_trained.pth") DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") TARGET_K = 10 # Enforce your target 10 metabolic domains class MSI_Structural_Encoder(nn.Module): def __init__(self): super(MSI_Structural_Encoder, self).__init__() # Matches the high-fidelity GroupNorm architecture self.features = nn.Sequential( nn.Conv2d(1, 32, kernel_size=5, stride=2, padding=2, bias=False), nn.GroupNorm(4, 32), nn.ReLU(), nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False), nn.GroupNorm(8, 64), nn.ReLU(), nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False), nn.GroupNorm(16, 128), nn.ReLU(), nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=False), nn.GroupNorm(32, 256), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)) ) def forward(self, x): h = self.features(x).squeeze(-1).squeeze(-1) return h def load_trained_encoder(path): if not os.path.exists(path): raise FileNotFoundError(f"Missing model checkpoint: {path}") model = MSI_Structural_Encoder() # Loading only the feature extractor weights model.features.load_state_dict(torch.load(path, map_location=DEVICE)) model.to(DEVICE) model.eval() return model class MSIDatasetInference(Dataset): def __init__(self, img_dir): self.img_dir = img_dir self.img_names = sorted([f for f in os.listdir(img_dir) if f.endswith('.png')]) self.transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor() ]) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_path = os.path.join(self.img_dir, self.img_names[idx]) image = Image.open(img_path).convert("L") return self.transform(image), self.img_names[idx] if __name__ == "__main__": print("Loading specialized custom MSI Structural Encoder...") encoder = load_trained_encoder(ENCODER_PATH) dataset = MSIDatasetInference(IMG_DIR) loader = DataLoader(dataset, batch_size=1, shuffle=False) features = [] ion_names = [] print("Extracting and L2-normalizing features from custom manifold...") with torch.no_grad(): for img, name in loader: img = img.to(DEVICE) feat = encoder(img) # Re-enforce unit hypersphere mapping feat = F.normalize(feat, p=2, dim=1) features.append(feat.cpu().numpy().flatten()) ion_names.append(name[0]) features = np.array(features) # Dimensionality Reduction Optimization pca_km = PCA(n_components=min(50, len(features)), random_state=42) features_pca_all = pca_km.fit_transform(features) # FIX: Slice out the first principal component (index 0) # if it continues to hold the global brightness/intensity bias features_pca_km = features_pca_all[:, 1:] # CRITICAL FIX: Re-normalize strictly for K-Means to ensure objective Cosine geometry features_pca_km_norm = features_pca_km / np.linalg.norm(features_pca_km, axis=1, keepdims=True) # ENFORCE EXACTLY 10 METABOLIC DOMAINS print(f"Executing K-Means clustering strictly constrained to K={TARGET_K}...") clusterer = KMeans(n_clusters=TARGET_K, init='k-means++', n_init=20, random_state=42) labels = clusterer.fit_predict(features_pca_km_norm) print("Projecting t-SNE map...") tsne = TSNE( n_components=2, perplexity=15, early_exaggeration=24.0, # Boosted significantly to force tighter, more definitive visual clusters learning_rate='auto', init='pca', random_state=42 ) # Feed the unnormalized PCA vectors to t-SNE to allow intensity variance to separate the visual blobs embeddings_2d = tsne.fit_transform(features_pca_km) # Plotting plt.figure(figsize=(10, 8)) scatter = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=labels, cmap='tab10', alpha=0.9, edgecolors='none', s=45) cbar = plt.colorbar(scatter, ticks=range(TARGET_K)) cbar.set_label('Metabolic Domain ID (K=10)') plt.title("Objectively Segmented Leaf Manifold (Custom Encoder)") plt.xlabel("t-SNE Dimension 1") plt.ylabel("t-SNE Dimension 2") plt.grid(True, linestyle='--', alpha=0.3) plt.savefig(os.path.join(BASE_DIR, "..", "msi_tsne_corregido.png"), dpi=300, bbox_inches='tight') plt.close() # INTERACTIVE PLOTLY EXPORT print("Exporting Interactive Plotly t-SNE...") mz_values = [float(n.replace("ion_", "").replace(".png", "")) for n in ion_names] df_plot = pd.DataFrame({ 't-SNE Dimension 1': embeddings_2d[:, 0], 't-SNE Dimension 2': embeddings_2d[:, 1], 'Cluster': [str(lbl) for lbl in labels], 'm/z': mz_values, 'Filename': ion_names }) fig = px.scatter( df_plot, x='t-SNE Dimension 1', y='t-SNE Dimension 2', color='Cluster', hover_data=['m/z', 'Filename'], title="Interactive Segmented Leaf Manifold (Custom Encoder)", category_orders={"Cluster": [str(i) for i in range(TARGET_K)]} ) plotly_out = os.path.join(BASE_DIR, "..", "figures", "msi_tsne_interactive.html") os.makedirs(os.path.dirname(plotly_out), exist_ok=True) fig.write_html(plotly_out) print(f"Interactive t-SNE saved to: {plotly_out}") # Save outputs print("Exporting synchronized cluster assignments...") df_results = pd.DataFrame({'ion_image': ion_names, 'cluster': labels}) df_results['mz'] = df_results['ion_image'].str.extract(r'(\d+\.\d+)').astype(float) df_results.to_csv(os.path.join(BASE_DIR, "..", "msi_clusters_results.csv"), index=False) print("Process complete.")