engine/crog_engine_nomask.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.wandb_utils import fix_metric_names  # keep
from utils.dataset import tokenize  # keep
from utils.misc import (
    AverageMeter,
    ProgressMeter,
    concat_all_gather,
    trainMetricGPU,   # NOTE: no longer used (ins mask metrics)
    get_seg_image     # NOTE: no longer used (ins mask vis)
)
from utils.grasp_eval import (
    detect_grasps,
    calculate_iou,
    calculate_max_iou,
    calculate_jacquard_index,
    visualization,
)


def _all_ranks_finite_loss(loss, device):
    local_finite = torch.tensor(
        [1 if torch.isfinite(loss.detach()).item() else 0],
        device=device,
        dtype=torch.int32,
    )
    if dist.is_available() and dist.is_initialized():
        dist.all_reduce(local_finite, op=dist.ReduceOp.MIN)
    return bool(local_finite.item())


def _format_tensor_stats(name, tensor):
    if not torch.is_tensor(tensor):
        return f"{name}=non_tensor"
    with torch.no_grad():
        data = tensor.detach()
        finite = torch.isfinite(data)
        finite_ratio = finite.float().mean().item()
        if finite.any():
            finite_data = data[finite]
            min_value = finite_data.min().item()
            max_value = finite_data.max().item()
            mean_value = finite_data.float().mean().item()
            return (
                f"{name}: finite={finite_ratio:.4f}, "
                f"min={min_value:.4g}, max={max_value:.4g}, mean={mean_value:.4g}"
            )
        return f"{name}: finite={finite_ratio:.4f}"


def train_with_grasp(train_loader, model, optimizer, scheduler, scaler, epoch, args):
    rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
    batch_time = AverageMeter('Batch', ':2.2f')
    data_time = AverageMeter('Data', ':2.2f')
    lr = AverageMeter('Lr', ':1.6f')
    loss_meter = AverageMeter('Loss', ':2.4f')
    qua_loss_metter = AverageMeter('Loss_qua', ':2.4f')
    sin_loss_metter = AverageMeter('Loss_sin', ':2.4f')
    cos_loss_metter = AverageMeter('Loss_cos', ':2.4f')
    wid_loss_metter = AverageMeter('Loss_wid', ':2.4f')

    # WebDataset / IterableDataset may not support len(train_loader)
    steps_per_epoch = getattr(args, 'steps_per_epoch', None)
    if steps_per_epoch is None:
        try:
            steps_per_epoch = len(train_loader)
        except TypeError:
            raise ValueError(
                "train_loader has no length. Please set `steps_per_epoch` in yaml."
            )

    progress = ProgressMeter(
        steps_per_epoch,
        [
            batch_time, data_time, lr, loss_meter,
            qua_loss_metter, sin_loss_metter, cos_loss_metter, wid_loss_metter,
        ],
        prefix="Training: Epoch=[{}/{}] ".format(epoch, args.epochs)
    )

    model.train()
    time.sleep(2)
    end = time.time()
    debug_steps = max(0, int(getattr(args, "ddp_debug_steps", 0)))
    debug_every = max(0, int(getattr(args, "ddp_debug_every", 0)))
    loader_iter = iter(train_loader)

    for i in range(steps_per_epoch):
        trace_step = i < debug_steps or (debug_every > 0 and ((i + 1) % debug_every == 0))
        if trace_step:
            logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: waiting for next batch")

        try:
            data = next(loader_iter)
        except StopIteration:
            logger.warning(f"[Rank {rank}] dataloader exhausted early at step {i + 1}/{steps_per_epoch}")
            break

        if trace_step:
            batch_size = data["img"].size(0) if "img" in data and hasattr(data["img"], "size") else "unknown"
            sample_ids = data.get("sample_id", [])
            sample_preview = sample_ids[:2] if isinstance(sample_ids, list) else sample_ids
            logger.info(
                f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: dataloader batch received "
                f"(batch_size={batch_size}, sample_preview={sample_preview})"
            )

        try:
            image = data["img"]
            text = data["word_vec"] if not getattr(args, 'use_raw_text', False) else data["sentence"]

            # ====== no instance mask ======
            grasp_qua_mask = data["grasp_masks"]["qua"]
            grasp_sin_mask = data["grasp_masks"]["sin"]
            grasp_cos_mask = data["grasp_masks"]["cos"]
            grasp_wid_mask = data["grasp_masks"]["wid"]

            data_time.update(time.time() - end)

            image = image.cuda(non_blocking=True)
            if isinstance(text, torch.Tensor):
                text = text.cuda(non_blocking=True)

            grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1)
            grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1)
            grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1)
            grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1)

            if trace_step:
                logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: forward start")

            # forward
            with amp.autocast():
                pred, target, loss, loss_dict = model(
                    image, text,
                    grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask
                )

            if not _all_ranks_finite_loss(loss, image.device):
                optimizer.zero_grad(set_to_none=True)
                sample_ids = data.get("sample_id", [])
                sample_preview = sample_ids[:2] if isinstance(sample_ids, list) else sample_ids
                logger.warning(
                    "[Rank {}] non-finite loss at epoch {} step {}; skipping optimizer step. "
                    "loss_dict={} sample_preview={} {} {} {} {}".format(
                        rank,
                        epoch,
                        i + 1,
                        loss_dict,
                        sample_preview,
                        _format_tensor_stats("pred_qua", pred[0]),
                        _format_tensor_stats("pred_sin", pred[1]),
                        _format_tensor_stats("pred_cos", pred[2]),
                        _format_tensor_stats("pred_wid", pred[3]),
                    )
                )
                batch_time.update(time.time() - end)
                end = time.time()
                continue

            if trace_step:
                logger.info(
                    f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: forward done "
                    f"(loss={float(loss.detach().item()):.6f})"
                )
                logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: backward start")

            # backward
            optimizer.zero_grad()
            scaler.scale(loss).backward()
            if trace_step:
                logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: backward done")
            if args.max_norm:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(),
                    args.max_norm,
                    error_if_nonfinite=False,
                )
            scaler.step(optimizer)
            scaler.update()
            if trace_step:
                logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: optimizer step done")

            # scheduler step
            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 hasattr(scheduler, '__class__') and scheduler.__class__.__name__ == 'OneCycleLR':
                scheduler.step()

            if trace_step:
                logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: all_reduce start")

            # distributed reduce
            loss_detached = loss.detach()
            dist.all_reduce(loss_detached)
            loss_detached = loss_detached / dist.get_world_size()

            if trace_step:
                logger.info(f"[Rank {rank}] step {i + 1}/{steps_per_epoch}: all_reduce done")
        except Exception:
            logger.exception(f"[Rank {rank}] step {i + 1}/{steps_per_epoch} failed before collective finished")
            raise

        loss_meter.update(loss_detached.item(), image.size(0))
        qua_loss_metter.update(loss_dict["m_qua"], image.size(0))
        sin_loss_metter.update(loss_dict["m_sin"], image.size(0))
        cos_loss_metter.update(loss_dict["m_cos"], image.size(0))
        wid_loss_metter.update(loss_dict["m_wid"], image.size(0))
        lr.update(scheduler.get_last_lr()[-1] if scheduler is not None else optimizer.param_groups[0]['lr'])
        batch_time.update(time.time() - end)
        end = time.time()

        if (i + 1) % args.print_freq == 0 or (i + 1) == steps_per_epoch:
            progress.display(i + 1)
            if dist.get_rank() in [-1, 0]:
                metrics = {
                    "time/batch": batch_time.val,
                    "time/data": data_time.val,
                    "training/lr": lr.val,
                    "training/loss": loss_meter.val,
                    "training/loss_qua": qua_loss_metter.val,
                    "training/loss_sin": sin_loss_metter.val,
                    "training/loss_cos": cos_loss_metter.val,
                    "training/loss_wid": wid_loss_metter.val,
                }

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

                global_step = (epoch - 1) * steps_per_epoch + 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)

        if torch.is_tensor(mat):
            mat = mat.detach().cpu().numpy()
        mat = np.asarray(mat, dtype=np.float32)

        if torch.is_tensor(w):
            w = int(w.item())
        else:
            w = int(w)

        if torch.is_tensor(h):
            h = int(h.item())
        else:
            h = int(h)

        inv_img = cv2.warpAffine(
            img, mat, (w, h),
            flags=cv2.INTER_CUBIC,
            borderValue=0.
        )
        return inv_img

    model.eval()
    time.sleep(2)

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

    val_steps_per_epoch = getattr(args, 'val_steps_per_epoch', None)
    pbar = tqdm(val_loader, total=val_steps_per_epoch)

    for step, data in enumerate(pbar):
        # stop manually for iterable-style dataloader
        if val_steps_per_epoch is not None and step >= val_steps_per_epoch:
            break

        image = data["img"]
        text = data["word_vec"] if not getattr(args, 'use_raw_text', False) else data["sentence"]

        # ====== no instance mask ======
        grasp_qua_mask = data["grasp_masks"]["qua"]
        grasp_sin_mask = data["grasp_masks"]["sin"]
        grasp_cos_mask = data["grasp_masks"]["cos"]
        grasp_wid_mask = data["grasp_masks"]["wid"]

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

        image = image.cuda(non_blocking=True)
        if isinstance(text, torch.Tensor):
            text = text.cuda(non_blocking=True)

        grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1)

        with amp.autocast():
            pred, target = model(
                image, text,
                grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask
            )

        # predictions (no ins mask)
        grasp_qua_mask_preds = pred[0]
        grasp_sin_mask_preds = pred[1]
        grasp_cos_mask_preds = pred[2]
        grasp_wid_mask_preds = pred[3]

        # targets (no ins mask)
        grasp_qua_mask_targets = target[0]
        grasp_sin_mask_targets = target[1]
        grasp_cos_mask_targets = target[2]
        grasp_wid_mask_targets = target[3]

        # sigmoid where needed (qua, wid)
        grasp_qua_mask_preds = torch.sigmoid(grasp_qua_mask_preds)
        grasp_wid_mask_preds = torch.sigmoid(grasp_wid_mask_preds)

        # resize preds to input image size
        if grasp_qua_mask_preds.shape[-2:] != image.shape[-2:]:
            grasp_qua_mask_preds = F.interpolate(
                grasp_qua_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)

            grasp_sin_mask_preds = F.interpolate(
                grasp_sin_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)

            grasp_cos_mask_preds = F.interpolate(
                grasp_cos_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)

            grasp_wid_mask_preds = F.interpolate(
                grasp_wid_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
        else:
            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)

        # iterate over batch
        for idx in range(grasp_qua_mask_preds.shape[0]):
            inv_mat = inverse_matrix[idx]
            ori_size = ori_sizes[idx]
            h, w = ori_size

            grasp_qua_mask_pred = grasp_qua_mask_preds[idx].squeeze().cpu().numpy()
            grasp_sin_mask_pred = grasp_sin_mask_preds[idx].squeeze().cpu().numpy()
            grasp_cos_mask_pred = grasp_cos_mask_preds[idx].squeeze().cpu().numpy()
            grasp_wid_mask_pred = grasp_wid_mask_preds[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()

            # inverse to original size
            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)

            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)

            for gi in range(len(num_grasps)):
                num_g = num_grasps[gi]
                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

        current_j1 = 100.0 * num_correct_grasps[0] / max(num_total_grasps[0], 1)
        current_j5 = 100.0 * num_correct_grasps[1] / max(num_total_grasps[1], 1)
        pbar.set_description(
            f"Evaluation: Epoch=[{epoch}/{args.epochs}] J@1: {current_j1:.2f} J@5: {current_j5:.2f}"
        )

    J_index = [0, 0]
    for i in range(len(num_grasps)):
        J_index[i] = num_correct_grasps[i] / max(num_total_grasps[i], 1)

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

    return 0.0, {}, J_index



@torch.no_grad()
def validate_without_grasp(val_loader, model, epoch, args):
    model.eval()
    time.sleep(2)

    pbar = tqdm(val_loader)
    for data in pbar:
        image = data["img"]
        text = data["word_vec"]

        grasp_qua_mask = data["grasp_masks"]["qua"]
        grasp_sin_mask = data["grasp_masks"]["sin"]
        grasp_cos_mask = data["grasp_masks"]["cos"]
        grasp_wid_mask = data["grasp_masks"]["wid"]

        image = image.cuda(non_blocking=True)
        text = text.cuda(non_blocking=True)

        grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1)

        _pred, _target = model(image, text, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask)

    head = 'Evaluation: Epoch=[{}/{}]  (no ins mask; validate_without_grasp does not compute metrics)'.format(
        epoch, args.epochs
    )
    logger.info(head)
    return 0.0, {}, [0.0, 0.0]


@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)

        if torch.is_tensor(mat):
            mat = mat.detach().cpu().numpy()
        mat = np.asarray(mat, dtype=np.float32)

        if torch.is_tensor(w):
            w = int(w.item())
        else:
            w = int(w)

        if torch.is_tensor(h):
            h = int(h.item())
        else:
            h = int(h)

        inv_img = cv2.warpAffine(
            img, mat, (w, h),
            flags=cv2.INTER_CUBIC,
            borderValue=0.
        )
        return inv_img

    model.eval()
    time.sleep(2)

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

    iou_list = []  # keep for minimal diff

    tbar = tqdm(test_loader, desc='Inference:', ncols=100)
    for cnt, data in enumerate(tbar):
        image = data["img"]

        # ✅ 对齐 train/val：支持 raw text
        text = data["word_vec"] if not getattr(args, 'use_raw_text', False) else data["sentence"]

        grasp_qua_mask = data["grasp_masks"]["qua"]
        grasp_sin_mask = data["grasp_masks"]["sin"]
        grasp_cos_mask = data["grasp_masks"]["cos"]
        grasp_wid_mask = data["grasp_masks"]["wid"]

        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)

        image = image.cuda(non_blocking=True)
        if isinstance(text, torch.Tensor):
            text = text.cuda(non_blocking=True)

        grasp_qua_mask = grasp_qua_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_sin_mask = grasp_sin_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_cos_mask = grasp_cos_mask.cuda(non_blocking=True).unsqueeze(1)
        grasp_wid_mask = grasp_wid_mask.cuda(non_blocking=True).unsqueeze(1)

        with amp.autocast():
            pred, target = model(
                image, text,
                grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask
            )

        grasp_qua_mask_preds = torch.sigmoid(pred[0])
        grasp_sin_mask_preds = pred[1]
        grasp_cos_mask_preds = pred[2]
        grasp_wid_mask_preds = torch.sigmoid(pred[3])

        # resize preds to input image size
        if grasp_qua_mask_preds.shape[-2:] != image.shape[-2:]:
            grasp_qua_mask_preds = F.interpolate(
                grasp_qua_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)

            grasp_sin_mask_preds = F.interpolate(
                grasp_sin_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)

            grasp_cos_mask_preds = F.interpolate(
                grasp_cos_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)

            grasp_wid_mask_preds = F.interpolate(
                grasp_wid_mask_preds, size=image.shape[-2:],
                mode='bicubic', align_corners=True
            ).squeeze(1)
        else:
            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(grasp_qua_mask_preds.shape[0]):
            inv_mat = inverse_matrix[idx]
            h, w = ori_sizes[idx]

            sent = sentences[idx] if sentences is not None else ""
            img_path = img_paths[idx] if img_paths is not None else None

            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()

            grasp_target = grasp_targets[idx]

            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)

            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)
                    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

                    visualization(
                        img,
                        None,  # ins_mask_pred removed
                        (grasp_qua_mask_pred, grasp_ang_mask_pred, grasp_wid_mask_pred),
                        grasp_preds,
                        sent,
                        save_path=os.path.join("./results", args.exp_name, f"results_{cnt}_{num_g}_grasps.png")
                    )

    J_index = [0.0, 0.0]
    for i in range(len(num_grasps)):
        J_index[i] = num_correct_grasps[i] / max(num_total_grasps[i], 1)

    logger.info("J@1: {:.2f}, J@5: {:.2f}".format(100. * J_index[0], 100. * J_index[1]))
    return 0.0, {}, J_index