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