vis_output.py

import argparse
import random

import numpy as np
import open3d as o3d
import torch
import torch.nn.functional as F
from matplotlib import pyplot as plt
from PIL import Image

from dataset.evaluation import anchor_output_process, collision_detect, detect_2d_grasp, detect_6d_grasp_multi
from dataset.pc_dataset_tools import data_process, feature_fusion
from dataset.utils import PointCloudHelper

parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint-path', default='resources/checkpoint')
parser.add_argument('--random-seed', type=int, default=1)

# image input
parser.add_argument('--rgb-path', default='resources/demo_rgb.png')
parser.add_argument('--depth-path', default='resources/demo_depth.png')
parser.add_argument('--input-h', type=int, default=160) # target height of input image
parser.add_argument('--input-w', type=int, default=320) # target width of input image

# 2d grasping
parser.add_argument('--sigma', type=int, default=10)
parser.add_argument('--ratio', type=int, default=8) # grasp-attributes prediction downsample ratio, must be 2^N
parser.add_argument('--anchor-k', type=int, default=6) # in-plane rotation anchor number
parser.add_argument('--anchor-w', type=float, default=50.0) # grasp width anchor size
parser.add_argument('--anchor-z', type=float, default=20.0) # grasp depth anchor size
parser.add_argument('--grid-size', type=int, default=8) # grid size for grid-based center sampling

# 6d grasping
parser.add_argument('--heatmap-thres', type=float, default=0.01) # heatmap threshold
parser.add_argument('--local-k', type=int, default=10) # grasp detection number in each local region (localnet)
parser.add_argument('--depth-thres', type=float, default=0.02) # depth threshold for collision detection
parser.add_argument('--max-points', type=int, default=25600) # downsampled max number of points in pc
parser.add_argument('--anchor-num', type=int, default=7) # spatial rotation anchor number
parser.add_argument('--center-num', type=int, default=64) # sampled local center/region number (how many grasps are predicted)
parser.add_argument('--group-num', type=int, default=512) # local region pc number

args = parser.parse_args()

# --------------------------------------------------------------------------- #

if __name__ == '__main__':
    # set torch and gpu setting
    np.set_printoptions(precision=4, suppress=True)
    torch.set_printoptions(precision=4, sci_mode=False)
    gpu = torch.cuda.is_available()
    if gpu:
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = False

    # random seed
    random.seed(args.random_seed)
    np.random.seed(args.random_seed)
    torch.manual_seed(args.random_seed)

    # load data
    pred_2d = [i for i in range(5)]
    with open('output/pred_0.dat', 'rb') as file:
        pred_2d[0] = torch.from_numpy(np.load(file))
    with open('output/pred_1.dat', 'rb') as file:
        pred_2d[1] = torch.from_numpy(np.load(file))
    with open('output/pred_2.dat', 'rb') as file:
        pred_2d[2] = torch.from_numpy(np.load(file))
    with open('output/pred_3.dat', 'rb') as file:
        pred_2d[3] = torch.from_numpy(np.load(file))
    with open('output/pred_4.dat', 'rb') as file:
        pred_2d[4] = torch.from_numpy(np.load(file))
    pred_2d = tuple(pred_2d)
    with open('output/perpoint.dat', 'rb') as file:
        perpoint_features = torch.from_numpy(np.load(file))
    with open('output/output_1.dat', 'rb') as file:
        pred = torch.from_numpy(np.load(file))
    with open('output/output_2.dat', 'rb') as file:
        offset = torch.from_numpy(np.load(file))
    if gpu:
        pred_2d = tuple([i.cuda() for i in pred_2d])
        perpoint_features = perpoint_features.cuda()
        pred = pred.cuda()
        offset = offset.cuda()

    # load anchors
    checkpoint = torch.load(args.checkpoint_path, weights_only=True, map_location=torch.device('cpu'))
    basic_ranges = torch.linspace(-1, 1, args.anchor_num + 1)
    if gpu:
        basic_ranges = basic_ranges.cuda()
    basic_anchors = (basic_ranges[1:] + basic_ranges[:-1]) / 2
    anchors = {'gamma': basic_anchors, 'beta': basic_anchors}
    anchors['gamma'] = checkpoint['gamma']
    anchors['beta'] = checkpoint['beta']
    print('-> loaded checkpoint %s ' % (args.checkpoint_path))

    # read image
    ori_depth = np.array(Image.open(args.depth_path))
    ori_rgb = np.array(Image.open(args.rgb_path)) / 255.0
    ori_depth = np.clip(ori_depth, 0, 1000)
    ori_rgb = torch.from_numpy(ori_rgb).permute(2, 1, 0)[None]
    ori_rgb = ori_rgb.to(device='cuda' if gpu else 'cpu', dtype=torch.float32)
    ori_depth = torch.from_numpy(ori_depth).T[None]
    ori_depth = ori_depth.to(device='cuda' if gpu else 'cpu', dtype=torch.float32)

    # create pc for 3d visualization
    pc_helper = PointCloudHelper(args.max_points, (args.input_w, args.input_h))
    view_points, _, _ = pc_helper.to_scene_points(ori_rgb, ori_depth)
    xyzs = pc_helper.to_xyz_maps(ori_depth)

    # generate 2d input (downscale image and normalize depth)
    rgb = F.interpolate(ori_rgb, (args.input_w, args.input_h))
    depth = F.interpolate(ori_depth[None], (args.input_w, args.input_h))[0]
    depth = depth / 1000.0
    depth = torch.clip((depth - depth.mean()), -1, 1)
    x = torch.concat([depth[None], rgb], 1)
    x = x.to(device='cuda' if gpu else 'cpu', dtype=torch.float32)

    # post process 2d grasp output
    loc_map, cls_mask, theta_offset, height_offset, width_offset = anchor_output_process(*pred_2d, sigma=args.sigma)

    # detect 2d grasps
    rect_gg = detect_2d_grasp(
        loc_map,
        cls_mask,
        theta_offset,
        height_offset,
        width_offset,
        ratio=args.ratio,
        anchor_k=args.anchor_k,
        anchor_w=args.anchor_w,
        anchor_z=args.anchor_z,
        mask_thre=args.heatmap_thres,
        center_num=args.center_num,
        grid_size=args.grid_size,
        grasp_nms=args.grid_size)
    print("# 2d grasps found:", rect_gg.size)

    # show heatmap
    rgb_t = x[0, 1:].cpu().numpy().squeeze().transpose(2, 1, 0)
    depth_t = ori_depth.cpu().numpy().squeeze().T / 1000.0
    resized_rgb = Image.fromarray((rgb_t * 255.0).astype(np.uint8))
    resized_rgb = np.array(resized_rgb.resize((args.input_w, args.input_h))) / 255.0
    rect_rgb = rect_gg.plot_rect_grasp_group(resized_rgb, 0).clip(0, 1)
    plt.subplot(221)
    plt.imshow(rgb_t) # original image
    plt.subplot(222)
    plt.imshow(depth_t) # depth image
    plt.subplot(223)
    plt.imshow(loc_map.squeeze().T, cmap='jet') # heatmap
    plt.subplot(224)
    plt.imshow(rect_rgb) # grasps
    plt.tight_layout()
    plt.show()

    # -------------------------------- #
    # wait for user to close window... #
    # -------------------------------- #

    # fuse heatmap and 2d grasp attributes with point cloud
    points_all = feature_fusion(view_points[..., :3], perpoint_features, xyzs)
    rect_ggs = [rect_gg]
    pc_group, valid_local_centers = data_process(
        points_all,
        ori_depth,
        rect_ggs,
        args.center_num,
        args.group_num,
        (args.input_w, args.input_h),
        min_points=32,
        is_training=False)
    rect_gg = rect_ggs[0] # overwrite rect_gg to update valid mask
    points_all = points_all.squeeze()

    # detect 6d grasp from 2d output and 6d output
    _, pred_rect_gg = detect_6d_grasp_multi(rect_gg,
                                            pred,
                                            offset,
                                            valid_local_centers,
                                            (args.input_w, args.input_h),
                                            anchors,
                                            k=args.local_k)

    # collision detect
    pred_grasp_from_rect = pred_rect_gg.to_6d_grasp_group(depth=args.depth_thres)
    pred_gg, _ = collision_detect(points_all, pred_grasp_from_rect, mode='graspnet')
    pred_gg = pred_gg.nms() # remove redundant grasps
    print("# 6d grasps found:", pred_gg.size)

    # show grasps in 3d
    grasp_geo = pred_gg.to_open3d_geometry_list()
    points = view_points[..., :3].cpu().numpy().squeeze()
    colors = view_points[..., 3:6].cpu().numpy().squeeze()
    vispc = o3d.geometry.PointCloud()
    vispc.points = o3d.utility.Vector3dVector(points)
    vispc.colors = o3d.utility.Vector3dVector(colors)
    o3d.visualization.draw_geometries([vispc] + grasp_geo)