engine/ssg_engine.py

import os
import time

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 tqdm import tqdm
import matplotlib.pyplot as plt

from utils.dataset import tokenize
from utils.misc import (AverageMeter, ProgressMeter, concat_all_gather, trainMetricGPU, get_seg_image)
from utils.grasp_eval import (ssg_post_processing, detect_grasps, calculate_iou, calculate_max_iou, calculate_jacquard_index, visualization)


def train(train_loader, model, optimizer, scheduler, scaler, epoch, args):
    batch_time = AverageMeter('Batch', ':2.2f')
    data_time = AverageMeter('Data', ':2.2f')
    lr = AverageMeter('Lr', ':1.6f')
    cls_loss_meter = AverageMeter('Loss_cls', ':2.4f')
    box_loss_meter = AverageMeter('Loss_box', ':2.4f')
    ins_loss_meter = AverageMeter('Loss_ins', ':2.4f')
    sem_loss_meter = AverageMeter('Loss_sem', ':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')
    iou_meter = AverageMeter('IoU', ':2.2f')
    pr_meter = AverageMeter('Prec@50', ':2.2f')

    progress = ProgressMeter(
        len(train_loader),
        [
            batch_time, data_time, lr, 
            cls_loss_meter, box_loss_meter, ins_loss_meter, sem_loss_meter,
            qua_loss_metter, sin_loss_metter, cos_loss_metter, wid_loss_metter, 
            iou_meter, pr_meter
        ],
        prefix="Training: Epoch=[{}/{}] ".format(epoch, args.epochs))
    
    model.train()
    time.sleep(2)
    end = time.time()

    for i, data_dict in enumerate(train_loader):
        data_dict["rgb"] = data_dict["rgb"].cuda(non_blocking=True)
        data_dict["depth"] = data_dict["depth"].cuda(non_blocking=True)
        data_dict["labels"] = [x.cuda(non_blocking=True) for x in data_dict["labels"]]
        data_dict["bboxes"] = [x.cuda(non_blocking=True) for x in data_dict["bboxes"]]
        data_dict["sem_mask"] = [x.cuda(non_blocking=True) for x in data_dict["sem_mask"]]
        data_dict["ins_masks"] = [x.cuda(non_blocking=True) for x in data_dict["ins_masks"]]
        data_dict["grasp_masks"]["qua"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["qua"]]
        data_dict["grasp_masks"]["sin"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["sin"]]
        data_dict["grasp_masks"]["cos"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["cos"]]
        data_dict["grasp_masks"]["wid"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["wid"]]
        
        output_dict, loss_dict = model(data_dict)
        loss = 0.0
        for k, v in loss_dict.items():
            loss += v

        optimizer.zero_grad()
        loss.backward()
        if args.max_norm:
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
        optimizer.step()

        iou, pr5 = trainMetricGPU(output_dict["inter_mask_p"], output_dict["inter_mask_gt"], 0.35, 0.5, sigmoid=False)
        dist.all_reduce(loss.detach())
        dist.all_reduce(iou)
        dist.all_reduce(pr5)
        loss = loss / dist.get_world_size()
        iou = iou / dist.get_world_size()
        pr5 = pr5 / dist.get_world_size()

        cls_loss_meter.update(loss_dict["loss_cls"].item(), data_dict["rgb"].size(0))
        box_loss_meter.update(loss_dict["loss_box"].item(), data_dict["rgb"].size(0))
        ins_loss_meter.update(loss_dict["loss_ins"].item(), data_dict["rgb"].size(0))
        sem_loss_meter.update(loss_dict["loss_sem"].item(), data_dict["rgb"].size(0))
        qua_loss_metter.update(loss_dict["loss_qua"].item(), data_dict["rgb"].size(0))
        sin_loss_metter.update(loss_dict["loss_sin"].item(), data_dict["rgb"].size(0))
        cos_loss_metter.update(loss_dict["loss_cos"].item(), data_dict["rgb"].size(0))
        wid_loss_metter.update(loss_dict["loss_wid"].item(), data_dict["rgb"].size(0))
        iou_meter.update(iou, data_dict["rgb"].size(0)) #@TODO add iou & pr evaluation
        pr_meter.update(pr5, data_dict["rgb"].size(0))
        lr.update(scheduler.get_last_lr()[-1])
        batch_time.update(time.time() - end)
        end = time.time()

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

            if dist.get_rank() in [-1, 0]:
                wandb.log(
                    {
                        "time/batch": batch_time.val,
                        "time/data": data_time.val,
                        "training/lr": lr.val,
                        "training/loss_cls": cls_loss_meter.val,
                        "training/loss_box": box_loss_meter.val,
                        "training/loss_ins": ins_loss_meter.val,
                        "training/loss_sem": sem_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,
                        "training/iou": iou_meter.val,
                        "training/prec@50": pr_meter.val,
                    },
                    step=epoch * len(train_loader) + (i + 1))



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

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

    for i, data_dict in enumerate(val_loader):

        data_dict["rgb"] = data_dict["rgb"].cuda(non_blocking=True)
        data_dict["depth"] = data_dict["depth"].cuda(non_blocking=True)
        data_dict["labels"] = [x.cuda(non_blocking=True) for x in data_dict["labels"]]
        data_dict["bboxes"] = [x.cuda(non_blocking=True) for x in data_dict["bboxes"]]
        data_dict["sem_mask"] = [x.cuda(non_blocking=True) for x in data_dict["sem_mask"]]
        data_dict["ins_masks"] = [x.cuda(non_blocking=True) for x in data_dict["ins_masks"]]
        data_dict["grasp_masks"]["qua"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["qua"]]
        data_dict["grasp_masks"]["sin"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["sin"]]
        data_dict["grasp_masks"]["cos"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["cos"]]
        data_dict["grasp_masks"]["wid"] = [x.cuda(non_blocking=True) for x in data_dict["grasp_masks"]["wid"]]
        
        output_dict = model(data_dict)
        post_dict = ssg_post_processing(args, output_dict, data_dict)

        for idx, num_g in enumerate(num_grasps):
            for obj_grasp_rects_gt in data_dict["grasp_rects"][0]:
                flag_top1 = False
                for obj_grasp_rects_p in post_dict["grasps_top1"]:
                    j_index_top1 = calculate_jacquard_index(obj_grasp_rects_p, obj_grasp_rects_gt)
                    if j_index_top1 > 0:
                        flag_top1 = True
                        break
                if flag_top1:
                    num_correct_grasps[0] += 1
                num_total_grasps[0] += 1
                
                flag_top5 = False
                for obj_grasp_rects_p in post_dict["grasps_top5"]:
                    j_index_top5 = calculate_jacquard_index(obj_grasp_rects_p, obj_grasp_rects_gt)
                    if j_index_top5 > 0:
                        flag_top5 =  True
                        break
                if flag_top5:
                    num_correct_grasps[1] += 1
                num_total_grasps[1] += 1
        # Use first 100 samples for validation to save time
        if i > 100:
            break
    
    J_index = [0, 0]
    for i in range(len(num_grasps)):
        J_index[i] = num_correct_grasps[i]/num_total_grasps[i]
    
    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 J_index



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

    tgt_dir = os.path.join(args.output_folder, "vis", f"epoch-{epoch}")
    os.makedirs(tgt_dir, exist_ok=True)

    vis_idx = np.random.randint(len(dataset), size=1)[0]
    data = dataset[vis_idx]
    data_dict = dataset.collate_fn([data])

    data_dict["rgb"] = data_dict["rgb"].cuda()
    data_dict["depth"] = data_dict["depth"].cuda()
    
    output_dict = model(data_dict)
    post_dict = ssg_post_processing(args, output_dict, data_dict)

    cls_pred = post_dict["cls"]
    bboxes_pred = post_dict["bboxes"]
    ins_masks_pred = post_dict["ins_masks"]
    qua_masks_pred, ang_masks_pred, wid_masks_pred = post_dict["grasp_masks"]
    grasps_top1 = post_dict["grasps_top1"]
    grasps_top5 = post_dict["grasps_top5"]

    img_path = os.path.join(dataset.root_dir, data_dict["scene_id"][0], "rgb", data_dict["img_f"][0])
    img = cv2.imread(img_path, cv2.COLOR_BGR2RGB)
    for i in range(bboxes_pred.shape[0]):
        tmp = img.copy()
        cls_p = cls_pred[i]
        bbox = bboxes_pred[i]
        ins_mask = ins_masks_pred[i].astype(int)
        qua_mask_p = qua_masks_pred[i]
        ang_mask_p = ang_masks_pred[i]
        wid_mask_p = wid_masks_pred[i]
        grasp_top5_p = grasps_top5[i]

        fig = plt.figure(figsize=(20, 10))

        ax = fig.add_subplot(1, 4, 1)
        ax.imshow(tmp[:, :, ::-1])
        ax.set_title('RGB')
        ax.axis('off')

        mask = np.asarray([0, 0, 255]).reshape(1, 1, 3).repeat(480, axis=0).repeat(640, axis=1)
        ins_mask = np.expand_dims(ins_mask, axis=-1).repeat(3, axis=-1)
        mask *= ins_mask
        img_with_mask = (tmp * 0.4 + mask * 0.6).astype(int)
        
        ax = fig.add_subplot(1, 4, 2)
        ax.imshow(img_with_mask[:, :, ::-1])
        ax.set_title('ins mask')
        ax.axis('off')

        # Put bounding box
        tmp = cv2.rectangle(tmp, (int(bbox[0]),int(bbox[1])), (int(bbox[2]),int(bbox[3])), (0,255,0), 2)
        tmp = cv2.putText(tmp, dataset.idx_to_class[cls_p], (int(bbox[0]),int(bbox[1])-20), cv2.FONT_HERSHEY_PLAIN, 1, (0,0,0), 2, cv2.LINE_AA)

        ax = fig.add_subplot(1, 4, 3)
        ax.imshow(tmp[:, :, ::-1])
        ax.set_title('bounding box')
        ax.axis('off')

        # Put grasp rects
        grasps = dataset.grasp_transforms.inverse(grasp_top5_p)
        for entry in grasps:
            ptA, ptB, ptC, ptD = [list(map(int, pt.tolist())) for pt in entry]
            tmp = cv2.line(tmp, ptA, ptB, (0,0,0xff), 2)
            tmp = cv2.line(tmp, ptD, ptC, (0,0,0xff), 2)
            tmp = cv2.line(tmp, ptB, ptC, (0xff,0,0), 2)
            tmp = cv2.line(tmp, ptA, ptD, (0xff,0,0), 2)

        ax = fig.add_subplot(1, 4, 4)
        ax.imshow(tmp[:, :, ::-1])
        ax.set_title('grasps')
        ax.axis('off')

        plt.savefig(os.path.join(tgt_dir, f"sample_{i}.png"))
        plt.cla()
        plt.clf()
        plt.close()