198 lines
7.3 KiB
Python
198 lines
7.3 KiB
Python
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
from torchvision import models, transforms
|
|
from PIL import Image
|
|
import numpy as np
|
|
import random
|
|
import pandas as pd
|
|
from sklearn.cluster import KMeans
|
|
from sklearn.decomposition import PCA
|
|
from sklearn.metrics import silhouette_score, davies_bouldin_score, calinski_harabasz_score, adjusted_rand_score
|
|
import matplotlib.pyplot as plt
|
|
from scipy.stats import ttest_ind
|
|
import torch.nn.functional as F
|
|
|
|
# 0. REPRODUCIBILITY
|
|
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)
|
|
|
|
# CONFIGURATION
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
IMG_DIR = os.path.join(BASE_DIR, "..", "datos_para_ai")
|
|
MODEL_PATH = os.path.join(BASE_DIR, "..", "models", "msi_encoder_trained.pth")
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
K = 10
|
|
|
|
class MSI_Structural_Encoder(nn.Module):
|
|
def __init__(self):
|
|
super(MSI_Structural_Encoder, self).__init__()
|
|
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):
|
|
return self.features(x).squeeze(-1).squeeze(-1)
|
|
|
|
# 1. DATA LOADING
|
|
def load_images():
|
|
img_names = sorted([f for f in os.listdir(IMG_DIR) if f.endswith('.png')])
|
|
raw_pixels = []
|
|
mzs = []
|
|
|
|
print(f"Loading {len(img_names)} images for validation...")
|
|
for name in img_names:
|
|
img_path = os.path.join(IMG_DIR, name)
|
|
img = Image.open(img_path).convert("L")
|
|
raw_pixels.append(np.array(img).flatten())
|
|
|
|
try:
|
|
mz = float(name.split('_')[1].replace('.png', ''))
|
|
mzs.append(mz)
|
|
except:
|
|
mzs.append(0.0)
|
|
|
|
return np.array(raw_pixels), np.array(mzs), img_names
|
|
|
|
def extract_simclr_features(img_names):
|
|
model = MSI_Structural_Encoder()
|
|
if os.path.exists(MODEL_PATH):
|
|
model.features.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
|
|
model.to(DEVICE)
|
|
model.eval()
|
|
|
|
transform = transforms.Compose([
|
|
transforms.Resize((256, 256)),
|
|
transforms.ToTensor(),
|
|
])
|
|
|
|
features = []
|
|
print(f"Extracting SimCLR features on {DEVICE}...")
|
|
with torch.no_grad():
|
|
for name in img_names:
|
|
img_path = os.path.join(IMG_DIR, name)
|
|
img = Image.open(img_path).convert("L")
|
|
img_t = transform(img).unsqueeze(0).to(DEVICE)
|
|
feat = model(img_t)
|
|
# CRITICAL FIX: Project explicitly onto the unit hypersphere
|
|
feat_norm = F.normalize(feat, p=2, dim=1)
|
|
features.append(feat_norm.cpu().numpy().flatten())
|
|
return np.array(features)
|
|
|
|
# 2. VALIDATION TASKS
|
|
def baseline_comparison(X_raw, X_simclr):
|
|
print("\n--- TASK 1: BASELINE COMPARISON (PCA vs SimCLR) ---")
|
|
pca_raw = PCA(n_components=50, random_state=42)
|
|
X_pca_raw = pca_raw.fit_transform(X_raw)
|
|
|
|
pca_simclr = PCA(n_components=50, random_state=42)
|
|
X_pca_simclr = pca_simclr.fit_transform(X_simclr)
|
|
|
|
km_raw = KMeans(n_clusters=K, random_state=42, n_init=10)
|
|
labels_raw = km_raw.fit_predict(X_pca_raw)
|
|
|
|
km_sim = KMeans(n_clusters=K, random_state=42, n_init=10)
|
|
labels_sim = km_sim.fit_predict(X_pca_simclr)
|
|
|
|
metrics = {
|
|
"Method": ["Raw PCA", "SimCLR Encoder"],
|
|
"Silhouette": [silhouette_score(X_pca_raw, labels_raw), silhouette_score(X_pca_simclr, labels_sim)],
|
|
"Davies-Bouldin": [davies_bouldin_score(X_pca_raw, labels_raw), davies_bouldin_score(X_pca_simclr, labels_sim)],
|
|
"Calinski-Harabasz": [calinski_harabasz_score(X_pca_raw, labels_raw), calinski_harabasz_score(X_pca_simclr, labels_sim)]
|
|
}
|
|
|
|
df_metrics = pd.DataFrame(metrics)
|
|
print(df_metrics.to_string(index=False))
|
|
return df_metrics
|
|
|
|
def stability_analysis(X_features):
|
|
print("\n--- TASK 2: STABILITY ANALYSIS (Multiple Seeds) ---")
|
|
seeds = [10, 20, 30, 42, 50, 60, 70, 80, 90, 100]
|
|
all_labels = []
|
|
metrics_list = []
|
|
|
|
pca = PCA(n_components=50, random_state=42)
|
|
X_pca = pca.fit_transform(X_features)
|
|
|
|
for seed in seeds:
|
|
km = KMeans(n_clusters=K, random_state=seed, n_init=10)
|
|
labels = km.fit_predict(X_pca)
|
|
all_labels.append(labels)
|
|
metrics_list.append(silhouette_score(X_pca, labels))
|
|
|
|
ari_scores = []
|
|
for i in range(len(all_labels)):
|
|
for j in range(i + 1, len(all_labels)):
|
|
ari_scores.append(adjusted_rand_score(all_labels[i], all_labels[j]))
|
|
|
|
print(f"Mean Silhouette Score: {np.mean(metrics_list):.4f} ± {np.std(metrics_list):.4f}")
|
|
print(f"Mean Adjusted Rand Index (ARI) between runs: {np.mean(ari_scores):.4f} (1.0 is perfect stability)")
|
|
|
|
def isotope_validation(X_features, mzs):
|
|
print("\n--- TASK 3: ISOTOPE CO-LOCALIZATION VALIDATION ---")
|
|
X_norm = X_features / np.linalg.norm(X_features, axis=1, keepdims=True)
|
|
|
|
isotope_dist = []
|
|
random_dist = []
|
|
|
|
for i in range(len(mzs)):
|
|
for j in range(i + 1, len(mzs)):
|
|
diff = abs(mzs[i] - mzs[j])
|
|
if 0.95 <= diff <= 1.05:
|
|
dist = np.linalg.norm(X_norm[i] - X_norm[j])
|
|
isotope_dist.append(dist)
|
|
|
|
if len(isotope_dist) > 0:
|
|
for _ in range(len(isotope_dist)):
|
|
idx1 = random.randint(0, len(mzs)-1)
|
|
idx2 = random.randint(0, len(mzs)-1)
|
|
if idx1 != idx2:
|
|
dist = np.linalg.norm(X_norm[idx1] - X_norm[idx2])
|
|
random_dist.append(dist)
|
|
|
|
print(f"Number of isotopic pairs found: {len(isotope_dist)}")
|
|
print(f"Mean distance (Isotopes): {np.mean(isotope_dist):.4f}")
|
|
print(f"Mean distance (Random): {np.mean(random_dist):.4f}")
|
|
|
|
t_stat, p_val = ttest_ind(isotope_dist, random_dist)
|
|
print(f"T-test p-value: {p_val:.4e} (Significant if < 0.05)")
|
|
|
|
plt.figure(figsize=(8, 6))
|
|
plt.boxplot([isotope_dist, random_dist], tick_labels=['Isotopic Pairs', 'Random Pairs'])
|
|
plt.title('Latent Space Distances: Isotopes vs Random')
|
|
plt.ylabel('Euclidean Distance (Normalized Latent Space)')
|
|
plt.savefig(os.path.join(BASE_DIR, "..", "figures", "isotope_validation_plot.png"))
|
|
print(f"Isotope validation plot saved to figures/isotope_validation_plot.png")
|
|
|
|
if __name__ == "__main__":
|
|
if not os.path.exists(IMG_DIR):
|
|
print(f"Error: IMG_DIR {IMG_DIR} not found. Generate ion images first.")
|
|
else:
|
|
X_raw, mzs, img_names = load_images()
|
|
X_simclr = extract_simclr_features(img_names)
|
|
|
|
baseline_comparison(X_raw, X_simclr)
|
|
stability_analysis(X_simclr)
|
|
isotope_validation(X_simclr, mzs)
|