MSI_Julia_CNN/scripts_python/verify_coherence.py

85 lines
3.1 KiB
Python

import os
import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
from PIL import Image
from torchvision import transforms
# Copy the exact encoder structure from your scripts
import torch.nn as nn
class MSI_Structural_Encoder(nn.Module):
def __init__(self):
super().__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)
DEVICE = torch.device("cpu")
IMG_DIR = "datos_para_ai"
MODEL_PATH = "models/msi_encoder_trained.pth"
if not os.path.exists(MODEL_PATH):
print("Model not found. Please ensure you are in the correct directory.")
exit()
model = MSI_Structural_Encoder()
model.features.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.eval()
transform = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])
def get_vector(mz_target):
# Find the file closest to mz_target
files = [f for f in os.listdir(IMG_DIR) if f.endswith('.png')]
mzs = [float(f.replace("ion_", "").replace(".png", "")) for f in files]
idx = np.argmin(np.abs(np.array(mzs) - mz_target))
closest_mz = mzs[idx]
if abs(closest_mz - mz_target) > 0.1:
return None, closest_mz
img_path = os.path.join(IMG_DIR, files[idx])
img = Image.open(img_path).convert("L")
tensor = transform(img).unsqueeze(0)
with torch.no_grad():
feat = model(tensor)
feat_norm = F.normalize(feat, p=2, dim=1)
return feat_norm.numpy().flatten(), closest_mz
print("--- INTERNAL VALIDATION: ISOTOPIC & STRUCTURAL COHERENCE ---")
# Test 1: Isotopic Coherence (Atropine and its +1 Isotope)
v1, mz1 = get_vector(290.02)
v2, mz2 = get_vector(291.00)
if v1 is not None and v2 is not None:
sim = np.dot(v1, v2)
print(f"Isotopic Match -> m/z {mz1:.2f} (Atropine) vs m/z {mz2:.2f} (+1 Isotope): Cosine Similarity = {sim:.4f}")
# Test 2: Structural Coherence (Atropine and Scopolamine - both vascular)
v3, mz3 = get_vector(304.00)
if v1 is not None and v3 is not None:
sim = np.dot(v1, v3)
print(f"Biological Match -> m/z {mz1:.2f} (Atropine) vs m/z {mz3:.2f} (Scopolamine): Cosine Similarity = {sim:.4f}")
# Test 3: Unrelated Coherence (Atropine vs random background or non-vascular)
v4, mz4 = get_vector(142.09) # Tropine (paravascular, should be lower than scopolamine)
if v1 is not None and v4 is not None:
sim = np.dot(v1, v4)
print(f"Divergent Match -> m/z {mz1:.2f} (Atropine) vs m/z {mz4:.2f} (Tropine): Cosine Similarity = {sim:.4f}")