74 lines
2.7 KiB
Python
74 lines
2.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
|
|
from sklearn.cluster import KMeans
|
|
import torch.nn.functional as F
|
|
|
|
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")
|
|
|
|
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)
|
|
|
|
def load_encoder_diagnostic(path):
|
|
model = MSI_Structural_Encoder()
|
|
if os.path.exists(path):
|
|
state_dict = torch.load(path, map_location=DEVICE)
|
|
model.features.load_state_dict(state_dict)
|
|
model.to(DEVICE)
|
|
model.eval()
|
|
return model
|
|
|
|
if __name__ == "__main__":
|
|
encoder = load_encoder_diagnostic(ENCODER_PATH)
|
|
|
|
# Load targets directly to test raw alignment
|
|
img_155 = os.path.join(IMG_DIR, "ion_155.75.png")
|
|
img_156 = os.path.join(IMG_DIR, "ion_156.07.png")
|
|
|
|
transform = transforms.Compose([
|
|
transforms.Resize((256, 256)),
|
|
transforms.ToTensor()
|
|
])
|
|
|
|
if os.path.exists(img_155) and os.path.exists(img_156):
|
|
# Convert to Grayscale "L"
|
|
t_155 = transform(Image.open(img_155).convert("L")).unsqueeze(0).to(DEVICE)
|
|
t_156 = transform(Image.open(img_156).convert("L")).unsqueeze(0).to(DEVICE)
|
|
|
|
with torch.no_grad():
|
|
f_155 = encoder(t_155)
|
|
f_156 = encoder(t_156)
|
|
|
|
# Check directional similarity
|
|
sim = F.cosine_similarity(f_155, f_156)
|
|
print(f"\n--- Feature Diagnostics ---")
|
|
print(f"Cosine Similarity between Vein Ions (155.75 vs 156.07): {sim.item():.4f}")
|
|
print("If this value is below 0.85, the loaded weight file is treating them as different shapes.")
|
|
|