# ==========================================================================================================
# All the functions in this file are made to draw a correlation heat map.
# The project can run successfully without this file unless you need to draw a correlation heat map as well.
# ==========================================================================================================
import numpy as np
import torch
from torch.cuda.amp import autocast

from models.repvgg import repvgg_model_convert


def record_logits(eval_loader, model):
    if hasattr(model, "repvgg_flag") and model.repvgg_flag:  # repvgg
        deploy_model = repvgg_model_convert(model)
    else:
        deploy_model = model
    deploy_model.eval()
    outputs = []
    for i, (input, target) in enumerate(eval_loader):
        input = input.cuda()
        with autocast():
            output = torch.sigmoid(deploy_model(input).detach())
        outputs.append(output.cpu().detach().numpy())
    outputs = np.concatenate(outputs)
    return outputs
