engine/crog_engine_withdepth.py

import os
import time
from tqdm import tqdm

import cv2
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.distributed as dist
import torch.nn.functional as F
import wandb
from loguru import logger

from utils.misc import AverageMeter, ProgressMeter, concat_all_gather, trainMetricGPU
from utils.grasp_eval import detect_grasps, calculate_jacquard_index, visualization


def _to_cuda(batch_tensor):
    return batch_tensor.cuda(non_blocking=True)


def _get_text_input(data, args):
    if getattr(args, "use_raw_text", True):
        return data["sentence"]
    return data["word_vec"]


def _move_batch_to_device(data, args):
    image = _to_cuda(data["img"])
    depth = _to_cuda(data["depth"])

    text = _get_text_input(data, args)
    if isinstance(text, torch.Tensor):
        text = _to_cuda(text)

    ins_mask = _to_cuda(data["mask"]).unsqueeze(1)
    grasp_qua_mask = _to_cuda(data["grasp_masks"]["qua"]).unsqueeze(1)
    grasp_sin_mask = _to_cuda(data["grasp_masks"]["sin"]).unsqueeze(1)
    grasp_cos_mask = _to_cuda(data["grasp_masks"]["cos"]).unsqueeze(1)
    grasp_wid_mask = _to_cuda(data["grasp_masks"]["wid"]).unsqueeze(1)

    return {
        "image": image,
        "depth": depth,
        "text": text,
        "ins_mask": ins_mask,
        "grasp_qua_mask": grasp_qua_mask,
        "grasp_sin_mask": grasp_sin_mask,
        "grasp_cos_mask": grasp_cos_mask,
        "grasp_wid_mask": grasp_wid_mask,
    }


def train_with_grasp(train_loader, model, optimizer, scheduler, scaler, epoch, args):
    batch_time = AverageMeter('Batch', ':2.2f')
    data_time = AverageMeter('Data', ':2.2f')
    lr_meter = AverageMeter('Lr', ':1.6f')
    loss_meter = AverageMeter('Loss', ':2.4f')
    qua_loss_meter = AverageMeter('Loss_qua', ':2.4f')
    sin_loss_meter = AverageMeter('Loss_sin', ':2.4f')
    cos_loss_meter = AverageMeter('Loss_cos', ':2.4f')
    wid_loss_meter = AverageMeter('Loss_wid', ':2.4f')
    iou_meter = AverageMeter('IoU', ':2.2f')
    pr_meter = AverageMeter('Prec@50', ':2.2f')

    progress = ProgressMeter(
        len(train_loader),
        [
            batch_time, data_time, lr_meter, loss_meter,
            qua_loss_meter, sin_loss_meter, cos_loss_meter, wid_loss_meter,
            iou_meter, pr_meter
        ],
        prefix="Training: Epoch=[{}/{}] ".format(epoch, args.epochs)
    )

    model.train()
    time.sleep(2)
    end = time.time()

    for i, data in enumerate(train_loader):
        data_time.update(time.time() - end)

        batch = _move_batch_to_device(data, args)

        with amp.autocast():
            pred, target, loss, loss_dict = model(
                batch["image"],
                batch["text"],
                mask=batch["ins_mask"],
                grasp_qua_mask=batch["grasp_qua_mask"],
                grasp_sin_mask=batch["grasp_sin_mask"],
                grasp_cos_mask=batch["grasp_cos_mask"],
                grasp_wid_mask=batch["grasp_wid_mask"],
                dp_image=batch["depth"],
            )

        ins_mask_pred = pred[0]
        ins_mask_target = target[0]

        optimizer.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()

        if getattr(args, "max_norm", None):
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)

        scaler.step(optimizer)
        scaler.update()

        if scheduler is not None:
            if hasattr(scheduler, '_scheduler_type'):
                if scheduler._scheduler_type == 'OneCycleLR':
                    scheduler.step()
                elif hasattr(scheduler, '_warmup_steps') and scheduler._step_count < scheduler._warmup_steps:
                    scheduler.step()
            elif scheduler.__class__.__name__ == 'OneCycleLR':
                scheduler.step()

        iou, pr5 = trainMetricGPU(ins_mask_pred, ins_mask_target, 0.35, 0.5)

        if dist.is_available() and dist.is_initialized():
            loss_reduced = loss.detach().clone()
            iou_reduced = iou.clone()
            pr5_reduced = pr5.clone()
            dist.all_reduce(loss_reduced)
            dist.all_reduce(iou_reduced)
            dist.all_reduce(pr5_reduced)
            world_size = dist.get_world_size()
            loss_val = loss_reduced / world_size
            iou_val = iou_reduced / world_size
            pr5_val = pr5_reduced / world_size
        else:
            loss_val = loss.detach()
            iou_val = iou
            pr5_val = pr5

        bs = batch["image"].size(0)
        loss_meter.update(loss_val.item(), bs)
        qua_loss_meter.update(loss_dict["m_qua"], bs)
        sin_loss_meter.update(loss_dict["m_sin"], bs)
        cos_loss_meter.update(loss_dict["m_cos"], bs)
        wid_loss_meter.update(loss_dict["m_wid"], bs)
        iou_meter.update(iou_val.item(), bs)
        pr_meter.update(pr5_val.item(), bs)

        if scheduler is not None:
            lr_meter.update(scheduler.get_last_lr()[-1])
        else:
            lr_meter.update(optimizer.param_groups[-1]["lr"])

        batch_time.update(time.time() - end)
        end = time.time()

        if (i + 1) % args.print_freq == 0:
            progress.display(i + 1)

            should_log = True
            if dist.is_available() and dist.is_initialized():
                should_log = dist.get_rank() == 0

            if should_log:
                metrics = {
                    "time/batch": batch_time.val,
                    "time/data": data_time.val,
                    "training/lr": lr_meter.val,
                    "training/loss": loss_meter.val,
                    "training/loss_qua": qua_loss_meter.val,
                    "training/loss_sin": sin_loss_meter.val,
                    "training/loss_cos": cos_loss_meter.val,
                    "training/loss_wid": wid_loss_meter.val,
                    "training/iou": iou_meter.val,
                    "training/prec@50": pr_meter.val,
                }

                for idx, param_group in enumerate(optimizer.param_groups):
                    metrics[f"lr/param_group_{idx}"] = param_group["lr"]

                global_step = (epoch - 1) * len(train_loader) + i
                wandb.log(metrics, step=global_step)


@torch.no_grad()
def validate_with_grasp(val_loader, model, epoch, args):
    def inverse(img, mat, w, h):
        img = np.asarray(img)
        if img.dtype != np.float32:
            img = img.astype(np.float32)
        return cv2.warpAffine(
            img, mat, (w, h),
            flags=cv2.INTER_CUBIC,
            borderValue=0.
        )

    iou_list = []
    model.eval()
    time.sleep(2)

    num_grasps = [1, 5]
    num_correct_grasps = [0, 0]
    num_total_grasps = [0, 0]

    pbar = tqdm(val_loader)
    for data in pbar:
        batch = _move_batch_to_device(data, args)

        inverse_matrix = data["inverse"]
        ori_sizes = data["ori_size"]
        grasp_targets = data["grasps"]

        with amp.autocast():
            pred, target = model(
                batch["image"],
                batch["text"],
                mask=batch["ins_mask"],
                grasp_qua_mask=batch["grasp_qua_mask"],
                grasp_sin_mask=batch["grasp_sin_mask"],
                grasp_cos_mask=batch["grasp_cos_mask"],
                grasp_wid_mask=batch["grasp_wid_mask"],
                dp_image=batch["depth"],
            )

        ins_mask_preds = pred[0]
        grasp_qua_mask_preds = pred[1]
        grasp_sin_mask_preds = pred[2]
        grasp_cos_mask_preds = pred[3]
        grasp_wid_mask_preds = pred[4]

        ins_mask_targets = target[0]
        grasp_qua_mask_targets = target[1]
        grasp_sin_mask_targets = target[2]
        grasp_cos_mask_targets = target[3]
        grasp_wid_mask_targets = target[4]

        ins_mask_preds = torch.sigmoid(ins_mask_preds)
        grasp_qua_mask_preds = torch.sigmoid(grasp_qua_mask_preds)
        grasp_wid_mask_preds = torch.sigmoid(grasp_wid_mask_preds)

        if ins_mask_preds.shape[-2:] != batch["image"].shape[-2:]:
            ins_mask_preds = F.interpolate(
                ins_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_qua_mask_preds = F.interpolate(
                grasp_qua_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_sin_mask_preds = F.interpolate(
                grasp_sin_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_cos_mask_preds = F.interpolate(
                grasp_cos_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_wid_mask_preds = F.interpolate(
                grasp_wid_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
        else:
            ins_mask_preds = ins_mask_preds.squeeze(1)
            grasp_qua_mask_preds = grasp_qua_mask_preds.squeeze(1)
            grasp_sin_mask_preds = grasp_sin_mask_preds.squeeze(1)
            grasp_cos_mask_preds = grasp_cos_mask_preds.squeeze(1)
            grasp_wid_mask_preds = grasp_wid_mask_preds.squeeze(1)

        for idx in range(ins_mask_preds.shape[0]):
            inv_mat = inverse_matrix[idx]
            ori_size = ori_sizes[idx]
            h, w = ori_size

            ins_mask_pred = ins_mask_preds[idx].cpu().numpy()
            grasp_qua_mask_pred = grasp_qua_mask_preds[idx].cpu().numpy()
            grasp_sin_mask_pred = grasp_sin_mask_preds[idx].cpu().numpy()
            grasp_cos_mask_pred = grasp_cos_mask_preds[idx].cpu().numpy()
            grasp_wid_mask_pred = grasp_wid_mask_preds[idx].cpu().numpy()

            ins_mask_target = ins_mask_targets[idx].squeeze().cpu().numpy()
            grasp_target = grasp_targets[idx]
            grasp_qua_mask_target = grasp_qua_mask_targets[idx].squeeze().cpu().numpy()
            grasp_sin_mask_target = grasp_sin_mask_targets[idx].squeeze().cpu().numpy()
            grasp_cos_mask_target = grasp_cos_mask_targets[idx].squeeze().cpu().numpy()
            grasp_wid_mask_target = grasp_wid_mask_targets[idx].squeeze().cpu().numpy()

            ins_mask_pred = inverse(ins_mask_pred, inv_mat, w, h) > 0.35
            grasp_qua_mask_pred = inverse(grasp_qua_mask_pred, inv_mat, w, h)
            grasp_sin_mask_pred = inverse(grasp_sin_mask_pred, inv_mat, w, h)
            grasp_cos_mask_pred = inverse(grasp_cos_mask_pred, inv_mat, w, h)
            grasp_wid_mask_pred = inverse(grasp_wid_mask_pred, inv_mat, w, h)

            ins_mask_target = inverse(ins_mask_target, inv_mat, w, h)
            grasp_qua_mask_target = inverse(grasp_qua_mask_target, inv_mat, w, h)
            grasp_sin_mask_target = inverse(grasp_sin_mask_target, inv_mat, w, h)
            grasp_cos_mask_target = inverse(grasp_cos_mask_target, inv_mat, w, h)
            grasp_wid_mask_target = inverse(grasp_wid_mask_target, inv_mat, w, h)

            inter = np.logical_and(ins_mask_pred, ins_mask_target)
            union = np.logical_or(ins_mask_pred, ins_mask_target)
            iou = np.sum(inter) / (np.sum(union) + 1e-6)
            iou_list.append(iou)

            for gi, num_g in enumerate(num_grasps):
                grasp_preds, _ = detect_grasps(
                    grasp_qua_mask_pred, grasp_sin_mask_pred,
                    grasp_cos_mask_pred, grasp_wid_mask_pred, num_g
                )
                j_index = calculate_jacquard_index(grasp_preds, grasp_target)
                num_correct_grasps[gi] += j_index
                num_total_grasps[gi] += 1

    grasp_stats = torch.tensor(
        [
            value
            for pair in zip(num_correct_grasps, num_total_grasps)
            for value in pair
        ],
        device=batch["image"].device,
        dtype=torch.float64,
    )
    if dist.is_available() and dist.is_initialized():
        dist.all_reduce(grasp_stats, op=dist.ReduceOp.SUM)
    J_index = [
        (grasp_stats[2 * i] / grasp_stats[2 * i + 1].clamp_min(1.0)).item()
        for i in range(len(num_grasps))
    ]

    iou_list = np.stack(iou_list)
    iou_list = torch.from_numpy(iou_list).to(batch["image"].device)
    if dist.is_available() and dist.is_initialized():
        iou_list = concat_all_gather(iou_list)

    prec_list = []
    for thres in torch.arange(0.5, 1.0, 0.1):
        prec_list.append((iou_list > thres).float().mean())

    iou = iou_list.mean()
    prec = {}
    temp = '  '
    for i, thres in enumerate(range(5, 10)):
        key = f'Pr@{thres * 10}'
        value = prec_list[i].item()
        prec[key] = value
        temp += f"{key}: {100. * value:.2f}  "

    head = 'Evaluation: Epoch=[{}/{}]  IoU={:.2f}  J_index@1: {:.2f}  J_index@5: {:.2f}'.format(
        epoch, args.epochs, 100. * iou.item(), 100. * J_index[0], 100. * J_index[1]
    )
    logger.info(head + temp)

    return iou.item(), prec, J_index


@torch.no_grad()
def validate_without_grasp(val_loader, model, epoch, args):
    def inverse(img, mat, w, h):
        img = np.asarray(img)
        if img.dtype != np.float32:
            img = img.astype(np.float32)
        return cv2.warpAffine(
            img, mat, (w, h),
            flags=cv2.INTER_CUBIC,
            borderValue=0.
        )

    iou_list = []
    model.eval()
    time.sleep(2)

    pbar = tqdm(val_loader)
    for data in pbar:
        batch = _move_batch_to_device(data, args)
        inverse_matrix = data["inverse"]
        ori_sizes = data["ori_size"]

        with amp.autocast():
            pred, target = model(
                batch["image"],
                batch["text"],
                mask=batch["ins_mask"],
                grasp_qua_mask=batch["grasp_qua_mask"],
                grasp_sin_mask=batch["grasp_sin_mask"],
                grasp_cos_mask=batch["grasp_cos_mask"],
                grasp_wid_mask=batch["grasp_wid_mask"],
                dp_image=batch["depth"],
            )

        ins_mask_preds = torch.sigmoid(pred[0])
        ins_mask_targets = target[0]

        if ins_mask_preds.shape[-2:] != batch["image"].shape[-2:]:
            ins_mask_preds = F.interpolate(
                ins_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
        else:
            ins_mask_preds = ins_mask_preds.squeeze(1)

        for idx in range(ins_mask_preds.shape[0]):
            inv_mat = inverse_matrix[idx]
            ori_size = ori_sizes[idx]
            h, w = ori_size

            ins_mask_pred = ins_mask_preds[idx].cpu().numpy()
            ins_mask_target = ins_mask_targets[idx].squeeze().cpu().numpy()

            ins_mask_pred = inverse(ins_mask_pred, inv_mat, w, h) > 0.35
            ins_mask_target = inverse(ins_mask_target, inv_mat, w, h)

            inter = np.logical_and(ins_mask_pred, ins_mask_target)
            union = np.logical_or(ins_mask_pred, ins_mask_target)
            iou = np.sum(inter) / (np.sum(union) + 1e-6)
            iou_list.append(iou)

    J_index = [0, 0]

    iou_list = np.stack(iou_list)
    iou_list = torch.from_numpy(iou_list).to(batch["image"].device)
    if dist.is_available() and dist.is_initialized():
        iou_list = concat_all_gather(iou_list)

    prec_list = []
    for thres in torch.arange(0.5, 1.0, 0.1):
        prec_list.append((iou_list > thres).float().mean())

    iou = iou_list.mean()
    prec = {}
    temp = '  '
    for i, thres in enumerate(range(5, 10)):
        key = f'Pr@{thres * 10}'
        value = prec_list[i].item()
        prec[key] = value
        temp += f"{key}: {100. * value:.2f}  "

    head = 'Evaluation: Epoch=[{}/{}]  IoU={:.2f}  J_index@1: {:.2f}  J_index@5: {:.2f}'.format(
        epoch, args.epochs, 100. * iou.item(), 100. * J_index[0], 100. * J_index[1]
    )
    logger.info(head + temp)
    return iou.item(), prec, J_index


@torch.no_grad()
def inference_with_grasp(test_loader, model, args):
    def inverse(img, mat, w, h):
        img = np.asarray(img)
        if img.dtype != np.float32:
            img = img.astype(np.float32)
        return cv2.warpAffine(
            img, mat, (w, h),
            flags=cv2.INTER_CUBIC,
            borderValue=0.
        )

    iou_list = []
    model.eval()
    time.sleep(2)

    num_grasps = [1, 5]
    num_correct_grasps = [0, 0]
    num_total_grasps = [0, 0]

    pbar = tqdm(test_loader)
    for cnt, data in enumerate(pbar):
        batch = _move_batch_to_device(data, args)

        inverse_matrix = data["inverse"]
        ori_sizes = data["ori_size"]
        grasp_targets = data["grasps"]
        sentences = data.get("sentence", None)
        img_paths = data.get("img_path", None)

        with amp.autocast():
            pred, target = model(
                batch["image"],
                batch["text"],
                mask=batch["ins_mask"],
                grasp_qua_mask=batch["grasp_qua_mask"],
                grasp_sin_mask=batch["grasp_sin_mask"],
                grasp_cos_mask=batch["grasp_cos_mask"],
                grasp_wid_mask=batch["grasp_wid_mask"],
                dp_image=batch["depth"],
            )

        ins_mask_preds = pred[0]
        grasp_qua_mask_preds = pred[1]
        grasp_sin_mask_preds = pred[2]
        grasp_cos_mask_preds = pred[3]
        grasp_wid_mask_preds = pred[4]

        ins_mask_targets = target[0]
        grasp_qua_mask_targets = target[1]
        grasp_sin_mask_targets = target[2]
        grasp_cos_mask_targets = target[3]
        grasp_wid_mask_targets = target[4]

        ins_mask_preds = torch.sigmoid(ins_mask_preds)
        grasp_qua_mask_preds = torch.sigmoid(grasp_qua_mask_preds)
        grasp_wid_mask_preds = torch.sigmoid(grasp_wid_mask_preds)

        if ins_mask_preds.shape[-2:] != batch["image"].shape[-2:]:
            ins_mask_preds = F.interpolate(
                ins_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_qua_mask_preds = F.interpolate(
                grasp_qua_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_sin_mask_preds = F.interpolate(
                grasp_sin_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_cos_mask_preds = F.interpolate(
                grasp_cos_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
            grasp_wid_mask_preds = F.interpolate(
                grasp_wid_mask_preds, size=batch["image"].shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
        else:
            ins_mask_preds = ins_mask_preds.squeeze(1)
            grasp_qua_mask_preds = grasp_qua_mask_preds.squeeze(1)
            grasp_sin_mask_preds = grasp_sin_mask_preds.squeeze(1)
            grasp_cos_mask_preds = grasp_cos_mask_preds.squeeze(1)
            grasp_wid_mask_preds = grasp_wid_mask_preds.squeeze(1)

        for idx in range(ins_mask_preds.shape[0]):
            inv_mat = inverse_matrix[idx]
            ori_size = ori_sizes[idx]
            h, w = ori_size
            sent = sentences[idx] if sentences is not None else ""
            img_path = img_paths[idx] if img_paths is not None else None

            ins_mask_pred = ins_mask_preds[idx].cpu().numpy()
            grasp_qua_mask_pred = grasp_qua_mask_preds[idx].cpu().numpy()
            grasp_sin_mask_pred = grasp_sin_mask_preds[idx].cpu().numpy()
            grasp_cos_mask_pred = grasp_cos_mask_preds[idx].cpu().numpy()
            grasp_wid_mask_pred = grasp_wid_mask_preds[idx].cpu().numpy()

            ins_mask_target = ins_mask_targets[idx].squeeze().cpu().numpy()
            grasp_target = grasp_targets[idx]
            grasp_qua_mask_target = grasp_qua_mask_targets[idx].squeeze().cpu().numpy()
            grasp_sin_mask_target = grasp_sin_mask_targets[idx].squeeze().cpu().numpy()
            grasp_cos_mask_target = grasp_cos_mask_targets[idx].squeeze().cpu().numpy()
            grasp_wid_mask_target = grasp_wid_mask_targets[idx].squeeze().cpu().numpy()

            ins_mask_pred = inverse(ins_mask_pred, inv_mat, w, h) > 0.35
            grasp_qua_mask_pred = inverse(grasp_qua_mask_pred, inv_mat, w, h)
            grasp_sin_mask_pred = inverse(grasp_sin_mask_pred, inv_mat, w, h)
            grasp_cos_mask_pred = inverse(grasp_cos_mask_pred, inv_mat, w, h)
            grasp_wid_mask_pred = inverse(grasp_wid_mask_pred, inv_mat, w, h)

            ins_mask_target = inverse(ins_mask_target, inv_mat, w, h)
            grasp_qua_mask_target = inverse(grasp_qua_mask_target, inv_mat, w, h)
            grasp_sin_mask_target = inverse(grasp_sin_mask_target, inv_mat, w, h)
            grasp_cos_mask_target = inverse(grasp_cos_mask_target, inv_mat, w, h)
            grasp_wid_mask_target = inverse(grasp_wid_mask_target, inv_mat, w, h)

            inter = np.logical_and(ins_mask_pred, ins_mask_target)
            union = np.logical_or(ins_mask_pred, ins_mask_target)
            iou = np.sum(inter) / (np.sum(union) + 1e-6)
            iou_list.append(iou)

            for gi, num_g in enumerate(num_grasps):
                grasp_preds, grasp_ang_mask_pred = detect_grasps(
                    grasp_qua_mask_pred, grasp_sin_mask_pred,
                    grasp_cos_mask_pred, grasp_wid_mask_pred, num_g
                )
                j_index = calculate_jacquard_index(grasp_preds, grasp_target)
                num_correct_grasps[gi] += j_index
                num_total_grasps[gi] += 1

                if getattr(args, "visualize", False) and img_path is not None:
                    img_bgr = cv2.imread(img_path)
                    if img_bgr is None:
                        logger.warning(f"Failed to read image for visualization: {img_path}")
                        continue

                    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
                    vis_dir = getattr(
                        args,
                        "vis_dir",
                        os.path.join("./results", args.exp_name)
                    )
                    os.makedirs(vis_dir, exist_ok=True)
                    visualization(
                        img,
                        ins_mask_pred,
                        (grasp_qua_mask_pred, grasp_ang_mask_pred, grasp_wid_mask_pred),
                        grasp_preds,
                        sent,
                        save_path=os.path.join(
                            vis_dir,
                            f"results_{cnt}_{idx}_{num_g}_grasps.png"
                        )
                    )

    J_index = [num_correct_grasps[i] / num_total_grasps[i] for i in range(len(num_grasps))]

    iou_list = np.stack(iou_list)
    iou_list = torch.from_numpy(iou_list).to(batch["image"].device)

    prec_list = []
    for thres in torch.arange(0.5, 1.0, 0.1):
        prec_list.append((iou_list > thres).float().mean())

    iou = iou_list.mean()
    prec = {}
    temp = '  '
    for i, thres in enumerate(range(5, 10)):
        key = f'Pr@{thres * 10}'
        value = prec_list[i].item()
        prec[key] = value
        temp += f"{key}: {100. * value:.2f}  "

    head = 'Evaluation: Epoch=[{}/{}]  IoU={:.2f}  J_index@1: {:.2f}  J_index@5: {:.2f}'.format(
        0, getattr(args, 'epochs', 0), 100. * iou.item(), 100. * J_index[0], 100. * J_index[1]
    )
    logger.info(head + temp)

    return iou.item(), prec, J_index