knn classifier about mediar HOT 1 CLOSED

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from train_tools.models import MEDIARFormer
from train_tools.data_utils.datasetter import get_dataloaders_labeled

# Load data and model
dataloaders = get_dataloaders_labeled("/home/gihun/MEDIAR", "./train_tools/data_utils/mapping_labeled.json", "./train_tools/data_utils/mapping_tuning.json")
model = MEDIARFormer().to("cuda:7").eval()
model.load_state_dict(torch.load("./weights/pretrained/phase1.pth", map_location="cpu"))

# Extract and process embeddings
embeddings_all = []
for batch in dataloaders["train"]:
    with torch.no_grad():
        embeddings = model.encoder(batch["img"].to("cuda:7"))[-1].mean(dim=[2, 3]).cpu()
        embeddings_all.append(embeddings)
embeddings_all = torch.cat(embeddings_all, dim=0)

# Cluster embeddings
kmeans = KMeans(n_clusters=40, random_state=0).fit(embeddings_all.numpy())
labels, counts = np.unique(kmeans.labels_, return_counts=True)
sorted_indices = np.argsort(counts)[::-1]

# Plot histogram
plt.figure(figsize=(10, 6))
plt.bar(range(len(counts)), counts[sorted_indices], tick_label=labels[sorted_indices])
plt.xlabel('Cluster Label')
plt.ylabel('Count')
plt.title('KMeans Labels Histogram (Desc. Count)')
plt.show()```