183 lines
6.7 KiB
Python
183 lines
6.7 KiB
Python
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.")
|