infer_determine_thre.py

"""
Assess the checkpoint performance on the validation set to determine the optimal threshold. 
Utilize this identified threshold to evaluate on the test set.
"""

from dataset.dataload import func_getdataloader, func_getdataloader_16
from model.choose_net import func_getnetwork
from utils.data import get_coordinate_list
from utils.data import get_coordinates, get_probabilities, get_fullheatmap_from_fold
from creterion.f1 import compute_metrics_once
import torch
import torch.nn as nn
import numpy as np
import os
import cv2
import time
import scipy
import pandas as pd
from dataset.dataprocess import func_normlize
import argparse
import shutil


__author__ = "Yudong Zhang"


def save_args_to_file(args, path):
    with open(path, "a+") as file:
        for arg, value in vars(args).items():
            if isinstance(value, list):
                value = ", ".join(map(str, value))
            file.write(f"{arg}: {value}\n")
        file.write("--------------------------\n")


def parse_args_():
    parser = argparse.ArgumentParser(description="Train keypoints network")
    # model
    parser.add_argument(
        "--model_mode",
        choices=["deepBlink", "DetNet", "superpoint", "PointDet"],
        default="deepBlink",
    )

    # dataset
    parser.add_argument(
        "--test_datapath",
        type=str,
        default="/data/ldap_shared/synology_shared/zyd/data/20220611_detparticle/testdataset/test_VESICLE/SNR4/",
    )
    parser.add_argument(
        "--val_datapath",
        type=str,
        default="/data/ldap_shared/synology_shared/zyd/data/20220611_detparticle/val_VESICLE/SNR4/",
    )
    parser.add_argument("--datatype", choices=["8bit", "16bit"], default="8bit")

    # optimizer
    parser.add_argument("--gpu_list", nargs="+", default=[1])

    # If resume
    parser.add_argument("--ckpt_path", type=str, required=True)

    # Only for DetNet
    parser.add_argument("--alpha", type=float, default=0.1)

    # Only for PointDet
    parser.add_argument("--cfg", type=str, default="./config/inference_demo_coco.yaml")
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )

    # Log and save
    parser.add_argument("--log_root", type=str, default="./Log/")
    parser.add_argument("--exp_name", type=str, default="VESICEL_SNR4_deepBlink")
    parser.add_argument("--use_visdom", type=bool, default=False)
    parser.add_argument("--port", type=int, default=4006)

    args = parser.parse_args()

    return args


if __name__ == "__main__":
    opt = parse_args_()
    # model
    model_mode = opt.model_mode

    # dataset
    test_datapath = opt.test_datapath
    val_datapath = opt.val_datapath
    datatype = opt.datatype

    # optimizer
    gpu_list = opt.gpu_list

    # if resume
    ckp_path = opt.ckpt_path

    # log and save
    Log_path = opt.log_root

    now = int(round(time.time() * 1000))
    nowname = time.strftime("%Y%m%d_%H_%M_%S", time.localtime(now / 1000))
    expname = nowname + "_" + opt.exp_name + "_eval"
    # Makedirs and Save files
    if not os.path.exists(Log_path + expname):
        os.makedirs(Log_path + expname)

    # save this file
    thisfilepath = os.path.abspath(__file__)
    shutil.copy(thisfilepath, Log_path + expname + "/eval_determinethre_code.py")
    # record log
    logtxt_path = Log_path + expname + "/log.txt"
    logtxt = open(logtxt_path, "a+")
    logtxt.write("\n\n")
    logtxt.write("===============Eval determine thre===============\n")
    logtxt.write("==============={}===============\n".format(expname))
    logtxt.close()
    save_args_to_file(opt, logtxt_path)

    # load data model
    if datatype == "16bit":
        dataloader_ins_val = func_getdataloader_16(
            model_mode,
            val_datapath,
            batch_size=1,
            shuffle=False,
            num_workers=16,
            training=False,
        )
        dataloader_ins_test = func_getdataloader_16(
            model_mode,
            test_datapath,
            batch_size=1,
            shuffle=False,
            num_workers=16,
            training=False,
        )
    else:
        dataloader_ins_val = func_getdataloader(
            model_mode,
            val_datapath,
            batch_size=1,
            shuffle=False,
            num_workers=16,
            training=False,
        )
        dataloader_ins_test = func_getdataloader(
            model_mode,
            test_datapath,
            batch_size=1,
            shuffle=False,
            num_workers=16,
            training=False,
        )

    model_ins = func_getnetwork(model_mode, opt)

    # if use GPU
    device = torch.device(
        "cuda:{}".format(gpu_list[0]) if torch.cuda.is_available() else "cpu"
    )
    model_ins.to(device)
    if torch.cuda.device_count() > 1 and len(gpu_list) > 1:
        model_ins = nn.DataParallel(model_ins, device_ids=gpu_list)

    # load checkpoint
    c_checkpoint = torch.load(ckp_path, map_location="cuda:{}".format(gpu_list[0]))
    model_ins.load_state_dict(c_checkpoint["model_state_dict"])
    print("==> Loaded pretrianed model checkpoint '{}'.".format(ckp_path))

    thre = None
    # start inferance
    print("====>>>Choose Threshold")
    model_ins.eval()
    f1_max = 0
    thre_max = 0.1
    for thre in range(1, 10):
        loss_ = 0
        f1_list = []
        precis_list = []
        recall_list = []
        abs_euclideans_list = []

        since = time.time()
        for data in dataloader_ins_val:

            inp = data[0].to(device)
            if model_mode == "superpoint":
                lab = data[1]
                lab_heatmap = get_fullheatmap_from_fold(lab)[0]
                lab_coords = get_coordinates(lab_heatmap, thre=0.5)

                pred = model_ins(inp)
                pred_heatmap = get_fullheatmap_from_fold(pred)[0].detach().cpu().numpy()
                pred_coords = get_coordinates(pred_heatmap, thre=thre * 0.1)
            elif model_mode == "DetNet":
                lab_coords = data[1][0].numpy()[:, ::-1]
                # lab_coords = get_coordinates(lab, thre=0.5)
                pred = model_ins(inp)[0].permute(1, 2, 0).detach().cpu().numpy()
                pred_coords = get_coordinates(pred, thre=thre * 0.1)
            elif model_mode == "deepBlink":
                if datatype == "16bit":
                    lab = data[1][0]
                    lab_coords, _ = get_coordinate_list(
                        lab, image_size=max(inp.shape), probability=0.5
                    )
                else:
                    lab_coords = data[1][0].numpy()[:, ::-1]
                pred = model_ins(inp)[0].permute(1, 2, 0).detach().cpu().numpy()
                pred_coords, scores = get_coordinate_list(
                    pred, image_size=max(inp.shape), probability=thre * 0.1
                )
            elif model_mode == "PointDet":
                lab_coords = data[1][0].numpy()
                pheatmap, poffset, psegment = model_ins(inp)
                psegment = psegment[0, 0, :, :].detach().cpu().numpy()

                pred_coords = get_coordinates(psegment, thre * 0.1)

            if pred_coords.shape[0] == 0 or lab_coords.shape == 0:
                f1_, precis_, recall_, abs_euclideans = 0, 0, 0, 1e10
            else:
                f1_, precis_, recall_, abs_euclideans = compute_metrics_once(
                    pred=pred_coords, true=lab_coords, mdist=3.0
                )
            f1_list.append(f1_)
            precis_list.append(precis_)
            recall_list.append(recall_)
            abs_euclideans_list.append(abs_euclideans)

        if np.array(f1_list).mean() > f1_max:
            f1_max = np.array(f1_list).mean()
            thre_max = thre * 0.1

        # claculate time
        time_elapsed = time.time() - since
        # record loss time
        message = "threthold:{:.1f} f1:{:.3f} precision{:.3f} recall{:.3f} rmse{:.3f} elapse:{:.0f}m {:.0f}s".format(
            thre * 0.1,
            np.array(f1_list).mean(),
            np.array(precis_list).mean(),
            np.array(recall_list).mean(),
            np.array(abs_euclideans_list).mean(),
            time_elapsed // 60,
            time_elapsed % 60,
        )
        print(message)
        logtxt = open(logtxt_path, "a+")
        logtxt.write(message + "\n")
        logtxt.close()

    print("best f1:{:.3f},with threshold:{:.1f}".format(f1_max, thre_max))

    print("===>Start prediction")

    # make save folder
    inf_dir = Log_path + expname + "/prediction_" + str(thre_max)
    if not os.path.exists(inf_dir):
        os.makedirs(inf_dir)

    loss_ = 0
    f1_list = []
    precis_list = []
    recall_list = []
    abs_euclideans_list = []
    for data in dataloader_ins_test:
        inp = data[0].to(device)
        name = data[2][0]
        inputimage = data[3][0].numpy()
        if model_mode == "superpoint":
            lab = data[1]
            lab_heatmap = get_fullheatmap_from_fold(lab)[0]
            lab_coords = get_coordinates(lab_heatmap, thre=0.5)

            pred = model_ins(inp)
            pred_heatmap = get_fullheatmap_from_fold(pred)[0].detach().cpu().numpy()
            pred_coords = get_coordinates(pred_heatmap, thre=thre_max)
        elif model_mode == "DetNet":
            lab_coords = data[1][0].numpy()[:, ::-1]
            # lab_coords = get_coordinates(lab,thre=0.5)
            pred = model_ins(inp)[0].permute(1, 2, 0).detach().cpu().numpy()
            pred_coords = get_coordinates(pred, thre=thre_max)
        elif model_mode == "deepBlink":
            if datatype == "16bit":
                lab = data[1][0]
                lab_coords, _ = get_coordinate_list(
                    lab, image_size=max(inp.shape), probability=0.5
                )
            else:
                lab_coords = data[1][0].numpy()[:, ::-1]
            pred = model_ins(inp)[0].permute(1, 2, 0).detach().cpu().numpy()
            pred_coords, scores = get_coordinate_list(
                pred, image_size=max(inp.shape), probability=thre_max
            )
        elif model_mode == "PointDet":
            lab_coords = data[1][0].numpy()
            pheatmap, poffset, psegment = model_ins(inp)
            psegment = psegment[0, 0, :, :].detach().cpu().numpy()

            pred_coords = get_coordinates(psegment, thre * 0.1)
        f1_, precis_, recall_, abs_euclideans = compute_metrics_once(
            pred=pred_coords, true=lab_coords, mdist=3.0
        )
        f1_list.append(f1_)
        precis_list.append(precis_)
        recall_list.append(recall_)
        abs_euclideans_list.append(abs_euclideans)

        pred_coords_pd = pd.DataFrame(pred_coords, columns=["pos_y", "pos_x"])
        pred_coords_pd.to_csv(inf_dir + "/" + name + ".csv", index=None)

        inputimage = func_normlize(inputimage, mode="maxmin_norm")
        inputimage = np.clip(np.round(inputimage * 255), 0, 255).astype(np.uint8)

        # inputimage[:,:,0] = 0
        # inputimage[:,:,2] = 0

        # cv2.imwrite(inf_dir+'/'+name+'.png',inputimage)
        if model_mode == "PointDet":
            for y, x in pred_coords:
                cv2.circle(inputimage, (int(y), int(x)), 5, (0, 255, 255), 1)
            for y, x in lab_coords:
                cv2.circle(inputimage, (int(y), int(x)), 1, (0, 0, 255), 1)
        else:
            for x, y in pred_coords:
                cv2.circle(inputimage, (int(y), int(x)), 5, (0, 255, 255), 1)
            for x, y in lab_coords:
                cv2.circle(inputimage, (int(y), int(x)), 1, (0, 0, 255), 1)

        # print('save:'+name)
        cv2.imwrite(inf_dir + "/" + name + "_f1{:.3f}.png".format(f1_), inputimage)

    message = "==>>[TEST]threthold:{:.1f} f1:{:.3f} precision:{:.3f} recall:{:.3f} rmse:{:.3f}".format(
        thre_max,
        np.array(f1_list).mean(),
        np.array(precis_list).mean(),
        np.array(recall_list).mean(),
        np.array(abs_euclideans_list).mean(),
    )
    print(message)
    logtxt_ = open(logtxt_path, "a+")
    logtxt_.write(message)
    logtxt_.close()
    # print('=====>>>>'+nowname)
    # print(test_datapath.split('/')[-3])
    # print(test_datapath.split('/')[-2])
    # print(test_datapath.split('/')[-1])