110 lines
3.6 KiB
Python
110 lines
3.6 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
|
|
from sklearn.manifold import TSNE
|
|
import matplotlib.pyplot as plt
|
|
import pandas as pd
|
|
|
|
# 1. CONFIGURACIÓN
|
|
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")
|
|
N_CLUSTERS = 10
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# 2. CARGAR EL MODELO ENTRENADO
|
|
def load_trained_encoder(path):
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"No se encontró el modelo en {path}")
|
|
|
|
model = models.efficientnet_b0()
|
|
model.classifier = nn.Identity()
|
|
|
|
state_dict = torch.load(path, map_location=DEVICE)
|
|
model.load_state_dict(state_dict)
|
|
model.to(DEVICE)
|
|
model.eval()
|
|
return model
|
|
|
|
# 3. DATASET PARA INFERENCIA
|
|
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((224, 224)),
|
|
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("RGB")
|
|
return self.transform(image), self.img_names[idx]
|
|
|
|
# --- PROCESO ---
|
|
if __name__ == "__main__":
|
|
print("Extrayendo características con el encoder...")
|
|
try:
|
|
encoder = load_trained_encoder(ENCODER_PATH)
|
|
except Exception as e:
|
|
print(f"Error al cargar el modelo: {e}")
|
|
exit(1)
|
|
|
|
dataset = MSIDatasetInference(IMG_DIR)
|
|
if len(dataset) == 0:
|
|
print(f"No se encontraron imágenes en {IMG_DIR}")
|
|
exit(1)
|
|
|
|
loader = DataLoader(dataset, batch_size=1, shuffle=False)
|
|
|
|
features = []
|
|
ion_names = []
|
|
|
|
with torch.no_grad():
|
|
for img, name in loader:
|
|
img = img.to(DEVICE)
|
|
feat = encoder(img)
|
|
# Normalización L2
|
|
feat = feat / feat.norm(p=2, dim=1, keepdim=True)
|
|
features.append(feat.cpu().numpy().flatten())
|
|
ion_names.append(name[0])
|
|
|
|
features = np.array(features)
|
|
|
|
# 4. CLUSTERING K-MEANS
|
|
print(f"Realizando KMeans con {N_CLUSTERS} clusters...")
|
|
km = KMeans(n_clusters=N_CLUSTERS, n_init=20, random_state=42)
|
|
labels = km.fit_predict(features)
|
|
|
|
# 5. VISUALIZACIÓN t-SNE CORREGIDA
|
|
print("Generando visualización t-SNE con escala corregida...")
|
|
tsne = TSNE(n_components=2, perplexity=30, random_state=42)
|
|
embeddings_2d = tsne.fit_transform(features)
|
|
|
|
plt.figure(figsize=(10, 8))
|
|
# Aplicando tu corrección específica:
|
|
scatter = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=labels, cmap='tab10', alpha=0.8, edgecolors='w', s=50)
|
|
|
|
# Configuración de la barra de colores con los ticks exactos para 10 clusters
|
|
cbar = plt.colorbar(scatter, ticks=range(N_CLUSTERS))
|
|
cbar.set_label('Cluster ID')
|
|
|
|
plt.title(f"t-SNE Visualization of MSI Ions (K={N_CLUSTERS})")
|
|
plt.xlabel("t-SNE dimension 1")
|
|
plt.ylabel("t-SNE dimension 2")
|
|
plt.grid(True, linestyle='--', alpha=0.5)
|
|
|
|
output_file = "msi_tsne_corregido.png"
|
|
plt.savefig(output_file, dpi=300, bbox_inches='tight')
|
|
# plt.show()
|
|
|
|
print(f"Visualización guardada exitosamente en '{output_file}'.")
|