Upload 37 files

.gitattributes

@@ -33,3 +33,13 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+contact_graspnet_train_and_test/model.ckpt-143887.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_2048_bs3_rad2_32/model.ckpt-144144.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_001/model.ckpt-144144.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_001/model.ckpt-45045.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_001/model.ckpt-54054.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_001/model.ckpt-72072.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_0025/model.ckpt-144144.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_0025/model.ckpt-45045.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_0025/model.ckpt-54054.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+scene_test_2048_bs3_hor_sigma_0025/model.ckpt-72072.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text

contact_graspnet_train_and_test/checkpoint.txt

@@ -0,0 +1,6 @@
+model_checkpoint_path: "model.ckpt-143887"
+all_model_checkpoint_paths: "model.ckpt-107915"
+all_model_checkpoint_paths: "model.ckpt-116908"
+all_model_checkpoint_paths: "model.ckpt-125901"
+all_model_checkpoint_paths: "model.ckpt-134894"
+all_model_checkpoint_paths: "model.ckpt-143887"

contact_graspnet_train_and_test/config.yaml

@@ -0,0 +1,192 @@
+DATA:
+  classes: null
+  data_path: /volume/pekdat/datasets/public/acronym/original
+  depth_augm:
+    clip: 0.005
+    gaussian_kernel: 0
+    sigma: 0.001
+  gripper_width: 0.08
+  input_normals: false
+  intrinsics: kinect_azure
+  labels:
+    bin_weights:
+    - 0.16652107
+    - 0.21488856
+    - 0.37031708
+    - 0.55618503
+    - 0.75124664
+    - 0.93943357
+    - 1.07824539
+    - 1.19423112
+    - 1.55731375
+    - 3.17161779
+    filter_z: true
+    k: 1
+    max_radius: 0.005
+    num_neg_contacts: 0
+    num_pos_contacts: 8000
+    offset_bins:
+    - 0
+    - 0.00794435329
+    - 0.0158887021
+    - 0.0238330509
+    - 0.0317773996
+    - 0.0397217484
+    - 0.0476660972
+    - 0.055610446
+    - 0.0635547948
+    - 0.0714991435
+    - 0.08
+    to_gpu: false
+    z_val: -0.1
+  ndataset_points: 20000
+  num_point: 2048
+  num_test_scenes: 1000
+  pc_augm:
+    clip: 0.005
+    occlusion_dropout_rate: 0.0
+    occlusion_nclusters: 0
+    sigma: 0.0
+  raw_num_points: 20000
+  scene_contacts_path: scene_contacts_train_test
+  train_and_test: false
+  train_on_scenes: true
+  use_farthest_point: false
+  use_uniform_quaternions: false
+  view_sphere:
+    distance_range:
+    - 0.9
+    - 1.3
+    elevation:
+    - 30
+    - 150
+LOSS:
+  min_geom_loss_divisor: 1.0
+  offset_loss_type: sigmoid_cross_entropy
+  too_small_offset_pred_bin_factor: 0
+  topk_confidence: 512
+MODEL:
+  asymmetric_model: true
+  bin_offsets: true
+  dir_vec_length_offset: false
+  grasp_conf_head:
+    conv1d: 1
+    dropout_keep: 0.5
+  grasp_dir_head:
+    conv1d: 3
+    dropout_keep: 0.7
+  joint_head:
+    conv1d: 4
+    dropout_keep: 0.7
+  joint_heads: false
+  model: contact_graspnet
+  pointnet_fp_modules:
+  - mlp:
+    - 256
+    - 256
+  - mlp:
+    - 256
+    - 128
+  - mlp:
+    - 128
+    - 128
+    - 128
+  pointnet_sa_module:
+    group_all: true
+    mlp:
+    - 256
+    - 512
+    - 1024
+  pointnet_sa_modules_msg:
+  - mlp_list:
+    - - 32
+      - 32
+      - 64
+    - - 64
+      - 64
+      - 128
+    - - 64
+      - 96
+      - 128
+    npoint: 2048
+    nsample_list:
+    - 32
+    - 64
+    - 128
+    radius_list:
+    - 0.02
+    - 0.04
+    - 0.08
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 512
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.04
+    - 0.08
+    - 0.16
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 128
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.08
+    - 0.16
+    - 0.32
+  pred_contact_approach: false
+  pred_contact_base: false
+  pred_contact_offset: true
+  pred_contact_success: true
+  pred_grasps_adds: true
+  pred_grasps_adds_gt2pred: false
+OPTIMIZER:
+  adds_gt2pred_loss_weight: 1
+  adds_loss_weight: 10
+  approach_cosine_loss_weight: 1
+  batch_size: 3
+  bn_decay_clip: 0.99
+  bn_decay_decay_rate: 0.5
+  bn_decay_decay_step: 200000
+  bn_init_decay: 0.5
+  decay_rate: 0.7
+  decay_step: 200000
+  dir_cosine_loss_weight: 1
+  learning_rate: 0.001
+  max_epoch: 16
+  momentum: 0.9
+  offset_loss_weight: 1
+  optimizer: adam
+  score_ce_loss_weight: 1
+TEST:
+  allow_zero_margin: 0
+  bin_vals: max
+  center_to_tip: 0.0
+  extra_opening: 0.005
+  filter_thres: 0.0001
+  first_thres: 0.23
+  max_farthest_points: 150
+  num_samples: 200
+  second_thres: 0.19
+  with_replacement: false

contact_graspnet_train_and_test/contact_graspnet.py

@@ -0,0 +1,438 @@
+import os
+import sys
+import numpy as np
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+TF2 = True
+
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR))
+
+sys.path.append(os.path.join(BASE_DIR))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2', 'utils'))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2'))
+
+import tf_util
+from pointnet_util import pointnet_sa_module, pointnet_fp_module, pointnet_sa_module_msg
+from tf_sampling import farthest_point_sample, gather_point
+from tf_grouping import query_ball_point, group_point, knn_point
+
+import mesh_utils
+
+def placeholder_inputs(batch_size, num_input_points=20000, input_normals=False):
+    """
+    Creates placeholders for input pointclouds, scene indices, camera poses and training/eval mode 
+
+    Arguments:
+        batch_size {int} -- batch size
+        num_input_points {int} -- number of input points to the network (default: 20000)
+
+    Keyword Arguments:
+        input_normals {bool} -- whether to use normals as input (default: {False})
+
+    Returns:
+        dict[str:tf.placeholder] -- dict of placeholders
+    """
+
+    pl_dict = {}
+    dim = 6 if input_normals else 3
+    pl_dict['pointclouds_pl'] = tf.placeholder(tf.float32, shape=(batch_size, num_input_points, dim))
+    pl_dict['scene_idx_pl'] = tf.placeholder(tf.int32, ())
+    pl_dict['cam_poses_pl'] = tf.placeholder(tf.float32, shape=(batch_size, 4, 4))
+    pl_dict['is_training_pl'] = tf.placeholder(tf.bool, shape=())
+
+    return pl_dict
+
+def get_bin_vals(global_config):
+    """
+    Creates bin values for grasping widths according to bounds defined in config
+
+    Arguments:
+        global_config {dict} -- config
+
+    Returns:
+        tf.constant -- bin value tensor 
+    """
+    bins_bounds = np.array(global_config['DATA']['labels']['offset_bins'])
+    if global_config['TEST']['bin_vals'] == 'max':
+        bin_vals = (bins_bounds[1:] + bins_bounds[:-1])/2
+        bin_vals[-1] = bins_bounds[-1]
+    elif global_config['TEST']['bin_vals'] == 'mean':
+        bin_vals = bins_bounds[1:]
+    else:
+        raise NotImplementedError
+
+    if not global_config['TEST']['allow_zero_margin']:
+        bin_vals = np.minimum(bin_vals, global_config['DATA']['gripper_width']-global_config['TEST']['extra_opening'])
+        
+    tf_bin_vals = tf.constant(bin_vals, tf.float32)
+    return tf_bin_vals
+
+def get_model(point_cloud, is_training, global_config, bn_decay=None):
+    """
+    Contact-GraspNet model consisting of a PointNet++ backbone and multiple output heads
+
+    Arguments:
+        point_cloud {tf.placeholder} -- batch of point clouds
+        is_training {bool} -- train or eval mode
+        global_config {dict} -- config
+
+    Keyword Arguments:
+        bn_decay {tf.variable} -- batch norm decay (default: {None})
+
+    Returns:
+        [dict] -- endpoints of the network
+    """
+
+    model_config = global_config['MODEL']
+    data_config = global_config['DATA']
+
+    radius_list_0 = model_config['pointnet_sa_modules_msg'][0]['radius_list']
+    radius_list_1 = model_config['pointnet_sa_modules_msg'][1]['radius_list']
+    radius_list_2 = model_config['pointnet_sa_modules_msg'][2]['radius_list']
+    
+    nsample_list_0 = model_config['pointnet_sa_modules_msg'][0]['nsample_list']
+    nsample_list_1 = model_config['pointnet_sa_modules_msg'][1]['nsample_list']
+    nsample_list_2 = model_config['pointnet_sa_modules_msg'][2]['nsample_list']
+    
+    mlp_list_0 = model_config['pointnet_sa_modules_msg'][0]['mlp_list']
+    mlp_list_1 = model_config['pointnet_sa_modules_msg'][1]['mlp_list']
+    mlp_list_2 = model_config['pointnet_sa_modules_msg'][2]['mlp_list']
+    
+    npoint_0 = model_config['pointnet_sa_modules_msg'][0]['npoint']
+    npoint_1 = model_config['pointnet_sa_modules_msg'][1]['npoint']
+    npoint_2 = model_config['pointnet_sa_modules_msg'][2]['npoint']
+    
+    fp_mlp_0 = model_config['pointnet_fp_modules'][0]['mlp']
+    fp_mlp_1 = model_config['pointnet_fp_modules'][1]['mlp']
+    fp_mlp_2 = model_config['pointnet_fp_modules'][2]['mlp']
+
+    input_normals = data_config['input_normals']
+    offset_bins = data_config['labels']['offset_bins']
+    joint_heads = model_config['joint_heads']
+
+    # expensive, rather use random only
+    if 'raw_num_points' in data_config and data_config['raw_num_points'] != data_config['ndataset_points']:
+        point_cloud = gather_point(point_cloud, farthest_point_sample(data_config['ndataset_points'], point_cloud))
+
+    end_points = {}
+    l0_xyz = tf.slice(point_cloud, [0,0,0], [-1,-1,3])
+    l0_points = tf.slice(point_cloud, [0,0,3], [-1,-1,3]) if input_normals else None 
+
+    # Set abstraction layers
+    l1_xyz, l1_points = pointnet_sa_module_msg(l0_xyz, l0_points, npoint_0, radius_list_0, nsample_list_0, mlp_list_0, is_training, bn_decay, scope='layer1')
+    l2_xyz, l2_points = pointnet_sa_module_msg(l1_xyz, l1_points, npoint_1, radius_list_1, nsample_list_1,mlp_list_1, is_training, bn_decay, scope='layer2')
+
+    if 'asymmetric_model' in model_config and model_config['asymmetric_model']:
+        l3_xyz, l3_points = pointnet_sa_module_msg(l2_xyz, l2_points, npoint_2, radius_list_2, nsample_list_2,mlp_list_2, is_training, bn_decay, scope='layer3')
+        l4_xyz, l4_points, _ = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer4')
+
+        # Feature Propagation layers
+        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+        l0_points = l1_points
+        pred_points = l1_xyz
+    else:
+        l3_xyz, l3_points, _ = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer3')
+
+        # Feature Propagation layers
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+        l0_points = tf.concat([l0_xyz, l0_points],axis=-1) if input_normals else l0_xyz 
+        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+        pred_points = l0_xyz
+
+    if joint_heads:
+        head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
+        head = tf_util.dropout(head, keep_prob=0.7, is_training=is_training, scope='dp1')
+        head = tf_util.conv1d(head, 4, 1, padding='VALID', activation_fn=None, scope='fc2')
+        grasp_dir_head = tf.slice(head, [0,0,0], [-1,-1,3])
+        grasp_dir_head = tf.math.l2_normalize(grasp_dir_head, axis=2)
+        binary_seg_head = tf.slice(head, [0,0,3], [-1,-1,1])
+    else:
+        # Head for grasp direction
+        grasp_dir_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
+        grasp_dir_head = tf_util.dropout(grasp_dir_head, keep_prob=0.7, is_training=is_training, scope='dp1')
+        grasp_dir_head = tf_util.conv1d(grasp_dir_head, 3, 1, padding='VALID', activation_fn=None, scope='fc3')
+        grasp_dir_head_normed = tf.math.l2_normalize(grasp_dir_head, axis=2)
+
+        # Head for grasp approach
+        approach_dir_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_app', bn_decay=bn_decay)
+        approach_dir_head = tf_util.dropout(approach_dir_head, keep_prob=0.7, is_training=is_training, scope='dp1_app')
+        approach_dir_head = tf_util.conv1d(approach_dir_head, 3, 1, padding='VALID', activation_fn=None, scope='fc3_app')
+        approach_dir_head_orthog = tf.math.l2_normalize(approach_dir_head - tf.reduce_sum(tf.multiply(grasp_dir_head_normed, approach_dir_head), axis=2, keepdims=True)*grasp_dir_head_normed, axis=2)
+        
+        # Head for grasp width
+        if model_config['dir_vec_length_offset']:
+            grasp_offset_head = tf.norm(grasp_dir_head, axis=2, keepdims=True)
+        elif model_config['bin_offsets']:
+            grasp_offset_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_off', bn_decay=bn_decay)
+            grasp_offset_head = tf_util.conv1d(grasp_offset_head, len(offset_bins)-1, 1, padding='VALID', activation_fn=None, scope='fc2_off')
+        else:
+            grasp_offset_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_off', bn_decay=bn_decay)
+            grasp_offset_head = tf_util.dropout(grasp_offset_head, keep_prob=0.7, is_training=is_training, scope='dp1_off')
+            grasp_offset_head = tf_util.conv1d(grasp_offset_head, 1, 1, padding='VALID', activation_fn=None, scope='fc2_off')
+
+        # Head for contact points
+        binary_seg_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_seg', bn_decay=bn_decay)
+        binary_seg_head = tf_util.dropout(binary_seg_head, keep_prob=0.5, is_training=is_training, scope='dp1_seg')
+        binary_seg_head = tf_util.conv1d(binary_seg_head, 1, 1, padding='VALID', activation_fn=None, scope='fc2_seg')
+
+    end_points['grasp_dir_head'] = grasp_dir_head_normed
+    end_points['binary_seg_head'] = binary_seg_head
+    end_points['binary_seg_pred'] = tf.math.sigmoid(binary_seg_head)
+    end_points['grasp_offset_head'] = grasp_offset_head
+    end_points['grasp_offset_pred'] = tf.math.sigmoid(grasp_offset_head) if model_config['bin_offsets'] else grasp_offset_head
+    end_points['approach_dir_head'] = approach_dir_head_orthog
+    end_points['pred_points'] = pred_points
+
+    return end_points
+
+def build_6d_grasp(approach_dirs, base_dirs, contact_pts, thickness, use_tf=False, gripper_depth = 0.1034):
+    """
+    Build 6-DoF grasps + width from point-wise network predictions
+
+    Arguments:
+        approach_dirs {np.ndarray/tf.tensor} -- Nx3 approach direction vectors
+        base_dirs {np.ndarray/tf.tensor} -- Nx3 base direction vectors
+        contact_pts {np.ndarray/tf.tensor} -- Nx3 contact points
+        thickness {np.ndarray/tf.tensor} -- Nx1 grasp width
+
+    Keyword Arguments:
+        use_tf {bool} -- whether inputs and outputs are tf tensors (default: {False})
+        gripper_depth {float} -- distance from gripper coordinate frame to gripper baseline in m (default: {0.1034})
+
+    Returns:
+        np.ndarray -- Nx4x4 grasp poses in camera coordinates
+    """
+    if use_tf:
+        grasps_R = tf.stack([base_dirs, tf.linalg.cross(approach_dirs,base_dirs),approach_dirs], axis=3)        
+        grasps_t = contact_pts + tf.expand_dims(thickness,2)/2 * base_dirs - gripper_depth * approach_dirs
+        ones = tf.ones((contact_pts.shape[0], contact_pts.shape[1], 1, 1), dtype=tf.float32)
+        zeros = tf.zeros((contact_pts.shape[0], contact_pts.shape[1], 1, 3), dtype=tf.float32)
+        homog_vec = tf.concat([zeros, ones], axis=3)
+        grasps = tf.concat([tf.concat([grasps_R,  tf.expand_dims(grasps_t, 3)], axis=3), homog_vec], axis=2)
+    else:
+        grasps = []
+        for i in range(len(contact_pts)):
+            grasp = np.eye(4)
+
+            grasp[:3,0] = base_dirs[i] / np.linalg.norm(base_dirs[i])
+            grasp[:3,2] = approach_dirs[i] / np.linalg.norm(approach_dirs[i])
+            grasp_y = np.cross( grasp[:3,2],grasp[:3,0])
+            grasp[:3,1] = grasp_y / np.linalg.norm(grasp_y)
+            # base_gripper xyz = contact + thickness / 2 * baseline_dir - gripper_d * approach_dir 
+            grasp[:3,3] = contact_pts[i] + thickness[i] / 2 * grasp[:3,0] - gripper_depth * grasp[:3,2]
+            # grasp[0,3] = finger_width
+            grasps.append(grasp)
+        grasps = np.array(grasps)
+
+    return grasps
+
+def get_losses(pointclouds_pl, end_points, dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam, global_config):
+    """
+    Computes loss terms from pointclouds, network predictions and labels 
+
+    Arguments:
+        pointclouds_pl {tf.placeholder} -- bxNx3 input point clouds
+        end_points {dict[str:tf.variable]} -- endpoints of the network containing predictions
+        dir_labels_pc_cam {tf.variable} -- base direction labels in camera coordinates (bxNx3)
+        offset_labels_pc {tf.variable} -- grasp width labels (bxNx1) 
+        grasp_success_labels_pc {tf.variable} -- contact success labels (bxNx1) 
+        approach_labels_pc_cam {tf.variable} -- approach direction labels in camera coordinates (bxNx3)
+        global_config {dict} -- config dict 
+        
+    Returns:
+        [dir_cosine_loss, bin_ce_loss, offset_loss, approach_cosine_loss, adds_loss, 
+        adds_loss_gt2pred, gt_control_points, pred_control_points, pos_grasps_in_view] -- All losses (not all are used for training)
+    """
+
+    grasp_dir_head = end_points['grasp_dir_head']
+    grasp_offset_head = end_points['grasp_offset_head']
+    approach_dir_head = end_points['approach_dir_head']
+
+    bin_weights = global_config['DATA']['labels']['bin_weights']
+    tf_bin_weights = tf.constant(bin_weights)
+    
+    min_geom_loss_divisor = tf.constant(float(global_config['LOSS']['min_geom_loss_divisor'])) if 'min_geom_loss_divisor' in global_config['LOSS'] else tf.constant(1.)
+    pos_grasps_in_view = tf.math.maximum(tf.reduce_sum(grasp_success_labels_pc, axis=1), min_geom_loss_divisor)   
+
+    ### ADS Gripper PC Loss
+    if global_config['MODEL']['bin_offsets']:
+        thickness_pred = tf.gather_nd(get_bin_vals(global_config), tf.expand_dims(tf.argmax(grasp_offset_head, axis=2), axis=2))
+        thickness_gt = tf.gather_nd(get_bin_vals(global_config), tf.expand_dims(tf.argmax(offset_labels_pc, axis=2), axis=2))
+    else:
+        thickness_pred = grasp_offset_head[:,:,0]
+        thickness_gt = offset_labels_pc[:,:,0]
+    pred_grasps = build_6d_grasp(approach_dir_head, grasp_dir_head, pointclouds_pl, thickness_pred, use_tf=True) # b x num_point x 4 x 4
+    gt_grasps_proj = build_6d_grasp(approach_labels_pc_cam, dir_labels_pc_cam, pointclouds_pl, thickness_gt, use_tf=True) # b x num_point x 4 x 4
+    pos_gt_grasps_proj = tf.where(tf.broadcast_to(tf.expand_dims(tf.expand_dims(tf.cast(grasp_success_labels_pc, tf.bool),2),3), gt_grasps_proj.shape), gt_grasps_proj, tf.ones_like(gt_grasps_proj)*100000)
+    # pos_gt_grasps_proj = tf.reshape(pos_gt_grasps_proj, (global_config['OPTIMIZER']['batch_size'], -1, 4, 4)) 
+
+    gripper = mesh_utils.create_gripper('panda')
+    gripper_control_points = gripper.get_control_point_tensor(global_config['OPTIMIZER']['batch_size']) # b x 5 x 3
+    sym_gripper_control_points = gripper.get_control_point_tensor(global_config['OPTIMIZER']['batch_size'], symmetric=True)
+
+    gripper_control_points_homog =  tf.concat([gripper_control_points, tf.ones((global_config['OPTIMIZER']['batch_size'], gripper_control_points.shape[1], 1))], axis=2)  # b x 5 x 4
+    sym_gripper_control_points_homog =  tf.concat([sym_gripper_control_points, tf.ones((global_config['OPTIMIZER']['batch_size'], gripper_control_points.shape[1], 1))], axis=2)  # b x 5 x 4
+    
+    # only use per point pred grasps but not per point gt grasps
+    control_points = tf.keras.backend.repeat_elements(tf.expand_dims(gripper_control_points_homog,1), gt_grasps_proj.shape[1], axis=1)  # b x num_point x 5 x 4
+    sym_control_points = tf.keras.backend.repeat_elements(tf.expand_dims(sym_gripper_control_points_homog,1), gt_grasps_proj.shape[1], axis=1)  # b x num_point x 5 x 4
+    pred_control_points = tf.matmul(control_points, tf.transpose(pred_grasps, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_point x 5 x 3
+
+    ### Pred Grasp to GT Grasp ADD-S Loss
+    gt_control_points = tf.matmul(control_points, tf.transpose(pos_gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+    sym_gt_control_points = tf.matmul(sym_control_points, tf.transpose(pos_gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+
+    squared_add = tf.reduce_sum((tf.expand_dims(pred_control_points,2)-tf.expand_dims(gt_control_points,1))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+    sym_squared_add = tf.reduce_sum((tf.expand_dims(pred_control_points,2)-tf.expand_dims(sym_gt_control_points,1))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+
+    # symmetric ADD-S
+    neg_squared_adds = -tf.concat([squared_add,sym_squared_add], axis=2) # b x num_point x 2num_pos_grasp_point
+    neg_squared_adds_k = tf.math.top_k(neg_squared_adds, k=1, sorted=False)[0] # b x num_point
+    # If any pos grasp exists
+    min_adds = tf.minimum(tf.reduce_sum(grasp_success_labels_pc, axis=1, keepdims=True), tf.ones_like(neg_squared_adds_k[:,:,0])) * tf.sqrt(-neg_squared_adds_k[:,:,0])#tf.minimum(tf.sqrt(-neg_squared_adds_k), tf.ones_like(neg_squared_adds_k)) # b x num_point
+    adds_loss = tf.reduce_mean(end_points['binary_seg_pred'][:,:,0] * min_adds)
+
+    ### GT Grasp to pred Grasp ADD-S Loss
+    gt_control_points = tf.matmul(control_points, tf.transpose(gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+    sym_gt_control_points = tf.matmul(sym_control_points, tf.transpose(gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+
+    neg_squared_adds = -tf.reduce_sum((tf.expand_dims(pred_control_points,1)-tf.expand_dims(gt_control_points,2))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+    neg_squared_adds_sym = -tf.reduce_sum((tf.expand_dims(pred_control_points,1)-tf.expand_dims(sym_gt_control_points,2))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+
+    neg_squared_adds_k_gt2pred, pred_grasp_idcs = tf.math.top_k(neg_squared_adds, k=1, sorted=False) # b x num_pos_grasp_point
+    neg_squared_adds_k_sym_gt2pred, pred_grasp_sym_idcs = tf.math.top_k(neg_squared_adds_sym, k=1, sorted=False) # b x num_pos_grasp_point
+    pred_grasp_idcs_joined = tf.where(neg_squared_adds_k_gt2pred<neg_squared_adds_k_sym_gt2pred, pred_grasp_sym_idcs, pred_grasp_idcs)
+    min_adds_gt2pred = tf.minimum(-neg_squared_adds_k_gt2pred, -neg_squared_adds_k_sym_gt2pred) # b x num_pos_grasp_point x 1
+    # min_adds_gt2pred = tf.math.exp(-min_adds_gt2pred)
+    masked_min_adds_gt2pred = tf.multiply(min_adds_gt2pred[:,:,0], grasp_success_labels_pc)
+    
+    batch_idcs = tf.meshgrid(tf.range(pred_grasp_idcs_joined.shape[1]), tf.range(pred_grasp_idcs_joined.shape[0]))
+    gather_idcs = tf.stack((batch_idcs[1],pred_grasp_idcs_joined[:,:,0]), axis=2)
+    nearest_pred_grasp_confidence = tf.gather_nd(end_points['binary_seg_pred'][:,:,0], gather_idcs)
+    adds_loss_gt2pred = tf.reduce_mean(tf.reduce_sum(nearest_pred_grasp_confidence*masked_min_adds_gt2pred, axis=1) / pos_grasps_in_view)
+ 
+    ### Grasp baseline Loss
+    cosine_distance = tf.constant(1.)-tf.reduce_sum(tf.multiply(dir_labels_pc_cam, grasp_dir_head),axis=2)
+    # only pass loss where we have labeled contacts near pc points 
+    masked_cosine_loss = tf.multiply(cosine_distance, grasp_success_labels_pc)
+    dir_cosine_loss = tf.reduce_mean(tf.reduce_sum(masked_cosine_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Approach Loss
+    approach_labels_orthog = tf.math.l2_normalize(approach_labels_pc_cam - tf.reduce_sum(tf.multiply(grasp_dir_head, approach_labels_pc_cam), axis=2, keepdims=True)*grasp_dir_head, axis=2)
+    cosine_distance_approach = tf.constant(1.)-tf.reduce_sum(tf.multiply(approach_labels_orthog, approach_dir_head), axis=2)
+    masked_approach_loss = tf.multiply(cosine_distance_approach, grasp_success_labels_pc)
+    approach_cosine_loss = tf.reduce_mean(tf.reduce_sum(masked_approach_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Offset/Thickness Loss
+    if global_config['MODEL']['bin_offsets']:
+        if global_config['LOSS']['offset_loss_type'] == 'softmax_cross_entropy':
+            offset_loss = tf.losses.softmax_cross_entropy(offset_labels_pc, grasp_offset_head)
+        else:
+            offset_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=offset_labels_pc, logits=grasp_offset_head)
+            
+            if 'too_small_offset_pred_bin_factor' in global_config['LOSS'] and global_config['LOSS']['too_small_offset_pred_bin_factor']:
+                too_small_offset_pred_bin_factor = tf.constant(global_config['LOSS']['too_small_offset_pred_bin_factor'], tf.float32)
+                collision_weight = tf.math.cumsum(offset_labels_pc, axis=2, reverse=True)*too_small_offset_pred_bin_factor + tf.constant(1.)
+                offset_loss = tf.multiply(collision_weight, offset_loss)
+
+            offset_loss = tf.reduce_mean(tf.multiply(tf.reshape(tf_bin_weights,(1,1,-1)), offset_loss),axis=2)
+    else:
+        offset_loss = (grasp_offset_head[:,:,0] - offset_labels_pc[:,:,0])**2
+    masked_offset_loss = tf.multiply(offset_loss, grasp_success_labels_pc)        
+    offset_loss = tf.reduce_mean(tf.reduce_sum(masked_offset_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Confidence Loss
+    bin_ce_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.expand_dims(grasp_success_labels_pc,axis=2), logits=end_points['binary_seg_head'])
+    if 'topk_confidence' in global_config['LOSS'] and global_config['LOSS']['topk_confidence']:
+        bin_ce_loss,_ = tf.math.top_k(tf.squeeze(bin_ce_loss), k=global_config['LOSS']['topk_confidence'])
+    bin_ce_loss = tf.reduce_mean(bin_ce_loss)
+
+    return dir_cosine_loss, bin_ce_loss, offset_loss, approach_cosine_loss, adds_loss, adds_loss_gt2pred
+
+def multi_bin_labels(cont_labels, bin_boundaries):
+    """
+    Computes binned grasp width labels from continous labels and bin boundaries
+
+    Arguments:
+        cont_labels {tf.Variable} -- continouos labels
+        bin_boundaries {list} -- bin boundary values
+
+    Returns:
+        tf.Variable -- one/multi hot bin labels
+    """
+    bins = []
+    for b in range(len(bin_boundaries)-1):
+        bins.append(tf.math.logical_and(tf.greater_equal(cont_labels, bin_boundaries[b]), tf.less(cont_labels,bin_boundaries[b+1])))
+    multi_hot_labels = tf.concat(bins, axis=2)
+    multi_hot_labels = tf.cast(multi_hot_labels, tf.float32)
+
+    return multi_hot_labels
+
+
+def compute_labels(pos_contact_pts_mesh, pos_contact_dirs_mesh, pos_contact_approaches_mesh, pos_finger_diffs, pc_cam_pl, camera_pose_pl, global_config):
+    """
+    Project grasp labels defined on meshes onto rendered point cloud from a camera pose via nearest neighbor contacts within a maximum radius. 
+    All points without nearby successful grasp contacts are considered negativ contact points.
+
+    Arguments:
+        pos_contact_pts_mesh {tf.constant} -- positive contact points on the mesh scene (Mx3)
+        pos_contact_dirs_mesh {tf.constant} -- respective contact base directions in the mesh scene (Mx3)
+        pos_contact_approaches_mesh {tf.constant} -- respective contact approach directions in the mesh scene (Mx3)
+        pos_finger_diffs {tf.constant} -- respective grasp widths in the mesh scene (Mx1)
+        pc_cam_pl {tf.placeholder} -- bxNx3 rendered point clouds
+        camera_pose_pl {tf.placeholder} -- bx4x4 camera poses
+        global_config {dict} -- global config
+
+    Returns:
+        [dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam] -- Per-point contact success labels and per-contact pose labels in rendered point cloud
+    """
+    label_config = global_config['DATA']['labels']
+    model_config = global_config['MODEL']
+
+    nsample = label_config['k']
+    radius = label_config['max_radius']
+    filter_z = label_config['filter_z']
+    z_val = label_config['z_val']
+
+    xyz_cam = pc_cam_pl[:,:,:3]
+    pad_homog = tf.ones((xyz_cam.shape[0],xyz_cam.shape[1], 1)) 
+    pc_mesh = tf.matmul(tf.concat([xyz_cam, pad_homog], 2), tf.transpose(tf.linalg.inv(camera_pose_pl),perm=[0, 2, 1]))[:,:,:3]
+
+    contact_point_offsets_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_finger_diffs,0), pc_mesh.shape[0], axis=0)
+
+    pad_homog2 = tf.ones((pc_mesh.shape[0], pos_contact_dirs_mesh.shape[0], 1)) 
+    contact_point_dirs_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0)
+    contact_point_dirs_batch_cam = tf.matmul(contact_point_dirs_batch, tf.transpose(camera_pose_pl[:,:3,:3], perm=[0, 2, 1]))[:,:,:3]
+
+    pos_contact_approaches_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_approaches_mesh,0), pc_mesh.shape[0], axis=0)
+    pos_contact_approaches_batch_cam = tf.matmul(pos_contact_approaches_batch, tf.transpose(camera_pose_pl[:,:3,:3], perm=[0, 2, 1]))[:,:,:3]
+    
+    contact_point_batch_mesh = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_pts_mesh,0), pc_mesh.shape[0], axis=0)
+    contact_point_batch_cam = tf.matmul(tf.concat([contact_point_batch_mesh, pad_homog2], 2), tf.transpose(camera_pose_pl, perm=[0, 2, 1]))[:,:,:3]
+
+    if filter_z:
+        dir_filter_passed = tf.keras.backend.repeat_elements(tf.math.greater(contact_point_dirs_batch_cam[:,:,2:3], tf.constant([z_val])), 3, axis=2)
+        contact_point_batch_mesh = tf.where(dir_filter_passed, contact_point_batch_mesh, tf.ones_like(contact_point_batch_mesh)*100000)
+
+    squared_dists_all = tf.reduce_sum((tf.expand_dims(contact_point_batch_cam,1)-tf.expand_dims(xyz_cam,2))**2,axis=3)
+    neg_squared_dists_k, close_contact_pt_idcs = tf.math.top_k(-squared_dists_all, k=nsample, sorted=False)
+    squared_dists_k = -neg_squared_dists_k
+
+    # Nearest neighbor mapping
+    grasp_success_labels_pc = tf.cast(tf.less(tf.reduce_mean(squared_dists_k, axis=2), radius*radius), tf.float32) # (batch_size, num_point)
+
+    grouped_dirs_pc_cam = group_point(contact_point_dirs_batch_cam, close_contact_pt_idcs)
+    grouped_approaches_pc_cam = group_point(pos_contact_approaches_batch_cam, close_contact_pt_idcs)
+    grouped_offsets = group_point(tf.expand_dims(contact_point_offsets_batch,2), close_contact_pt_idcs)
+
+    dir_labels_pc_cam = tf.math.l2_normalize(tf.reduce_mean(grouped_dirs_pc_cam, axis=2),axis=2) # (batch_size, num_point, 3)
+    approach_labels_pc_cam = tf.math.l2_normalize(tf.reduce_mean(grouped_approaches_pc_cam, axis=2),axis=2) # (batch_size, num_point, 3)
+    offset_labels_pc = tf.reduce_mean(grouped_offsets, axis=2)
+        
+    return dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam
+

contact_graspnet_train_and_test/model.ckpt-143887.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b92cf08a978dee3d2316a2dae685c26eb506a5ce5ee5a14c0a3840c1ceeba0e
+size 27961560

contact_graspnet_train_and_test/train.py

@@ -0,0 +1,224 @@
+from genericpath import exists
+import os
+import sys
+import argparse
+from datetime import datetime
+import numpy as np
+import time
+from tqdm import tqdm
+
+CONTACT_DIR = os.path.dirname(os.path.abspath(__file__))
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR))
+sys.path.append(os.path.join(BASE_DIR))
+sys.path.append(os.path.join(ROOT_DIR))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2',  'models'))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2',  'utils'))
+
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+TF2 = True
+physical_devices = tf.config.experimental.list_physical_devices('GPU')
+print(physical_devices)
+tf.config.experimental.set_memory_growth(physical_devices[0], True)
+
+import config_utils
+from data import PointCloudReader, load_scene_contacts, center_pc_convert_cam
+from summaries import build_summary_ops, build_file_writers
+from tf_train_ops import load_labels_and_losses, build_train_op
+from contact_grasp_estimator import GraspEstimator
+
+def train(global_config, log_dir):
+    """
+    Trains Contact-GraspNet
+
+    Arguments:
+        global_config {dict} -- config dict
+        log_dir {str} -- Checkpoint directory
+    """
+
+    contact_infos = load_scene_contacts(global_config['DATA']['data_path'],
+                                        scene_contacts_path=global_config['DATA']['scene_contacts_path'])
+    
+    num_train_samples = len(contact_infos)-global_config['DATA']['num_test_scenes']
+    num_test_samples = global_config['DATA']['num_test_scenes']
+        
+    print('using %s meshes' % (num_train_samples + num_test_samples))
+    if 'train_and_test' in global_config['DATA'] and global_config['DATA']['train_and_test']:
+        num_train_samples = num_train_samples + num_test_samples
+        num_test_samples = 0
+        print('using train and test data')
+
+    pcreader = PointCloudReader(
+        root_folder=global_config['DATA']['data_path'],
+        batch_size=global_config['OPTIMIZER']['batch_size'],
+        estimate_normals=global_config['DATA']['input_normals'],
+        raw_num_points=global_config['DATA']['raw_num_points'],
+        use_uniform_quaternions = global_config['DATA']['use_uniform_quaternions'],
+        scene_obj_scales = [c['obj_scales'] for c in contact_infos],
+        scene_obj_paths = [c['obj_paths'] for c in contact_infos],
+        scene_obj_transforms = [c['obj_transforms'] for c in contact_infos],
+        num_train_samples = num_train_samples,
+        num_test_samples = num_test_samples,
+        use_farthest_point = global_config['DATA']['use_farthest_point'],
+        intrinsics=global_config['DATA']['intrinsics'],
+        elevation=global_config['DATA']['view_sphere']['elevation'],
+        distance_range=global_config['DATA']['view_sphere']['distance_range'],
+        pc_augm_config=global_config['DATA']['pc_augm'],
+        depth_augm_config=global_config['DATA']['depth_augm']
+    )
+
+    with tf.Graph().as_default():
+        
+        # Build the model
+        grasp_estimator = GraspEstimator(global_config)
+        ops = grasp_estimator.build_network()
+        
+        # contact_tensors = load_contact_grasps(contact_infos, global_config['DATA'])
+        
+        loss_ops = load_labels_and_losses(grasp_estimator, contact_infos, global_config)
+
+        ops.update(loss_ops)
+        ops['train_op'] = build_train_op(ops['loss'], ops['step'], global_config)
+
+        # Add ops to save and restore all the variables.
+        saver = tf.train.Saver(save_relative_paths=True, keep_checkpoint_every_n_hours=4)
+
+        # Create a session
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        config.allow_soft_placement = True
+        sess = tf.Session(config=config)
+
+        # Log summaries
+        summary_ops = build_summary_ops(ops, sess, global_config)
+
+        # Init/Load weights
+        grasp_estimator.load_weights(sess, saver, log_dir, mode='train')
+
+        # sess = tf_debug.LocalCLIDebugWrapperSession(sess)
+        # sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
+        file_writers = build_file_writers(sess, log_dir)
+
+    ## define: epoch = arbitrary number of views of every training scene
+    cur_epoch = sess.run(ops['step']) // num_train_samples
+    for epoch in range(cur_epoch, global_config['OPTIMIZER']['max_epoch']):
+        log_string('**** EPOCH %03d ****' % (epoch))
+        
+        sess.run(ops['iterator'].initializer)
+        epoch_time = time.time()
+        step = train_one_epoch(sess, ops, summary_ops, file_writers, pcreader)
+        log_string('trained epoch {} in: {}'.format(epoch, time.time()-epoch_time))
+
+        # Save the variables to disk.
+        save_path = saver.save(sess, os.path.join(log_dir, "model.ckpt"), global_step=step, write_meta_graph=False)
+        log_string("Model saved in file: %s" % save_path)
+
+        if num_test_samples > 0:
+            eval_time = time.time()
+            eval_validation_scenes(sess, ops, summary_ops, file_writers, pcreader)
+            log_string('evaluation time: {}'.format(time.time()-eval_time))
+
+def train_one_epoch(sess, ops, summary_ops, file_writers, pcreader):
+    """ ops: dict mapping from string to tf ops """
+    
+    log_string(str(datetime.now()))
+    loss_log = np.zeros((10,7))
+    get_time = time.time()
+    
+    for batch_idx in range(pcreader._num_train_samples):
+
+        batch_data, cam_poses, scene_idx = pcreader.get_scene_batch(scene_idx=batch_idx)
+        
+        # OpenCV OpenGL conversion
+        cam_poses, batch_data = center_pc_convert_cam(cam_poses, batch_data)
+        
+        feed_dict = {ops['pointclouds_pl']: batch_data, ops['cam_poses_pl']: cam_poses,
+                     ops['scene_idx_pl']: scene_idx, ops['is_training_pl']: True}
+
+        step, summary, _, loss_val, dir_loss, bin_ce_loss, \
+        offset_loss, approach_loss, adds_loss, adds_gt2pred_loss, scene_idx = sess.run([ops['step'], summary_ops['merged'], ops['train_op'], ops['loss'], ops['dir_loss'], 
+                                                                                        ops['bin_ce_loss'], ops['offset_loss'], ops['approach_loss'], ops['adds_loss'], 
+                                                                                        ops['adds_gt2pred_loss'], ops['scene_idx']], feed_dict=feed_dict)
+        assert scene_idx[0] == scene_idx     
+        
+        loss_log[batch_idx%10,:] = loss_val, dir_loss, bin_ce_loss, offset_loss, approach_loss, adds_loss, adds_gt2pred_loss
+        
+        if (batch_idx+1)%10 == 0:
+            file_writers['train_writer'].add_summary(summary, step)
+            f = tuple(np.mean(loss_log, axis=0)) + ((time.time() - get_time) / 10., )
+            log_string('total loss: %f \t dir loss: %f \t ce loss: %f \t off loss: %f \t app loss: %f adds loss: %f \t adds_gt2pred loss: %f \t batch time: %f' % f)
+            get_time = time.time()
+            
+    return step
+
+def eval_validation_scenes(sess, ops, summary_ops, file_writers, pcreader, max_eval_objects=500):
+    """ ops: dict mapping from string to tf ops """
+    is_training = False
+    log_string(str(datetime.now()))
+    loss_log = np.zeros((min(pcreader._num_test_samples, max_eval_objects),7))
+
+    # resets accumulation of pr and auc data
+    sess.run(summary_ops['pr_reset_op'])
+    
+    for batch_idx in np.arange(min(pcreader._num_test_samples, max_eval_objects)):
+
+        batch_data, cam_poses, scene_idx = pcreader.get_scene_batch(scene_idx=pcreader._num_train_samples + batch_idx)
+
+        # OpenCV OpenGL conversion
+        cam_poses, batch_data = center_pc_convert_cam(cam_poses, batch_data)
+
+        feed_dict = {ops['pointclouds_pl']: batch_data, ops['cam_poses_pl']: cam_poses,
+                     ops['scene_idx_pl']: scene_idx, ops['is_training_pl']: False}
+
+        scene_idx, step, loss_val, dir_loss, bin_ce_loss, offset_loss, approach_loss, adds_loss, adds_gt2pred_loss, pr_summary,_,_,_ = sess.run([ops['scene_idx'], ops['step'], ops['loss'], ops['dir_loss'], ops['bin_ce_loss'],
+                                                                                                        ops['offset_loss'], ops['approach_loss'], ops['adds_loss'], ops['adds_gt2pred_loss'],
+                                                                                                        summary_ops['merged_eval'], summary_ops['pr_update_op'], 
+                                                                                                        summary_ops['auc_update_op']] + [summary_ops['acc_update_ops']], feed_dict=feed_dict)
+        assert scene_idx[0] == (pcreader._num_train_samples + batch_idx)
+        
+        loss_log[batch_idx,:] = loss_val, dir_loss, bin_ce_loss, offset_loss, approach_loss, adds_loss, adds_gt2pred_loss
+
+    file_writers['test_writer'].add_summary(pr_summary, step)
+    f = tuple(np.mean(loss_log, axis=0))
+    log_string('mean val loss: %f \t mean val dir loss: %f \t mean val ce loss: %f \t mean off loss: %f \t mean app loss: %f \t mean adds loss: %f \t mean adds_gt2pred loss: %f' % f)
+
+    return step
+
+if __name__ == "__main__":
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt_dir', default='checkpoints/contact_graspnet', help='Checkpoint dir')
+    parser.add_argument('--data_path', type=str, default=None, help='Grasp data root dir')
+    parser.add_argument('--max_epoch', type=int, default=None, help='Epochs to run')
+    parser.add_argument('--batch_size', type=int, default=None, help='Batch Size during training')
+    parser.add_argument('--arg_configs', nargs="*", type=str, default=[], help='overwrite config parameters')
+    FLAGS = parser.parse_args()
+
+    ckpt_dir = FLAGS.ckpt_dir
+    if not os.path.exists(ckpt_dir): 
+        if not os.path.exists(os.path.dirname(ckpt_dir)):
+            ckpt_dir = os.path.join(BASE_DIR, ckpt_dir)
+        os.makedirs(ckpt_dir, exist_ok=True) 
+        
+    os.system('cp {} {}'.format(os.path.join(CONTACT_DIR, 'contact_graspnet.py'), ckpt_dir)) # bkp of model def
+    os.system('cp {} {}'.format(os.path.join(CONTACT_DIR, 'train.py'), ckpt_dir)) # bkp of train procedure
+
+    LOG_FOUT = open(os.path.join(ckpt_dir, 'log_train.txt'), 'w')
+    LOG_FOUT.write(str(FLAGS)+'\n')
+    def log_string(out_str):
+        LOG_FOUT.write(out_str+'\n')
+        LOG_FOUT.flush()
+        print(out_str)
+
+    global_config = config_utils.load_config(ckpt_dir, batch_size=FLAGS.batch_size, max_epoch=FLAGS.max_epoch, 
+                                          data_path= FLAGS.data_path, arg_configs=FLAGS.arg_configs, save=True)
+    
+    log_string(str(global_config))
+    log_string('pid: %s'%(str(os.getpid())))
+
+    train(global_config, ckpt_dir)
+
+    LOG_FOUT.close()

scene_2048_bs3_rad2_32/checkpoint

@@ -0,0 +1,6 @@
+model_checkpoint_path: "model.ckpt-144144"
+all_model_checkpoint_paths: "model.ckpt-108108"
+all_model_checkpoint_paths: "model.ckpt-117117"
+all_model_checkpoint_paths: "model.ckpt-126126"
+all_model_checkpoint_paths: "model.ckpt-135135"
+all_model_checkpoint_paths: "model.ckpt-144144"

scene_2048_bs3_rad2_32/config.yaml

@@ -0,0 +1,186 @@
+DATA:
+  gripper_width: 0.08
+  input_normals: false
+  intrinsics: realsense
+  labels:
+    bin_weights:
+    - 0.16652107
+    - 0.21488856
+    - 0.37031708
+    - 0.55618503
+    - 0.75124664
+    - 0.93943357
+    - 1.07824539
+    - 1.19423112
+    - 1.55731375
+    - 3.17161779
+    contact_gather: knn
+    filter_z: true
+    k: 1
+    max_radius: 0.005
+    min_unique_pos_contacts: 1
+    num_neg_contacts: 0
+    num_pos_contacts: 4000
+    offset_bins:
+    - 0
+    - 0.00794435329
+    - 0.0158887021
+    - 0.0238330509
+    - 0.0317773996
+    - 0.0397217484
+    - 0.0476660972
+    - 0.055610446
+    - 0.0635547948
+    - 0.0714991435
+    - 0.08
+    to_gpu: false
+    z_val: -0.1
+  ndataset_points: 20000
+  num_point: 2048
+  num_test_scenes: 1000
+  raw_num_points: 20000
+  train_and_test: false
+  train_on_scenes: true
+  use_farthest_point: false
+  use_uniform_quaternions: false
+LOSS:
+  max_geom_loss_divisor: 100.0
+  min_geom_loss_divisor: 1.0
+  offset_loss_type: sigmoid_cross_entropy
+  too_small_offset_pred_bin_factor: 0
+  topk_confidence: 512
+MODEL:
+  asymmetric_model: true
+  bin_offsets: true
+  contact_distance_offset: true
+  dir_vec_length_offset: false
+  grasp_conf_head:
+    conv1d: 1
+    dropout_keep: 0.5
+  grasp_dir_head:
+    conv1d: 3
+    dropout_keep: 0.7
+  joint_head:
+    conv1d: 4
+    dropout_keep: 0.7
+  joint_heads: false
+  larger_model: false
+  model: contact_graspnet
+  pointnet_fp_modules:
+  - mlp:
+    - 256
+    - 256
+  - mlp:
+    - 256
+    - 128
+  - mlp:
+    - 128
+    - 128
+    - 128
+  pointnet_sa_module:
+    group_all: true
+    mlp:
+    - 256
+    - 512
+    - 1024
+  pointnet_sa_modules_msg:
+  - mlp_list:
+    - - 32
+      - 32
+      - 64
+    - - 64
+      - 64
+      - 128
+    - - 64
+      - 96
+      - 128
+    npoint: 2048
+    nsample_list:
+    - 32
+    - 64
+    - 128
+    radius_list:
+    - 0.02
+    - 0.04
+    - 0.08
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 512
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.04
+    - 0.08
+    - 0.16
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 128
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.08
+    - 0.16
+    - 0.32
+  pred_contact_approach: false
+  pred_contact_base: false
+  pred_contact_offset: true
+  pred_contact_success: true
+  pred_grasps_adds: true
+  pred_grasps_adds_gt2pred: false
+OPTIMIZER:
+  adds_gt2pred_loss_weight: 1
+  adds_loss_weight: 10
+  approach_cosine_loss_weight: 1
+  batch_size: 1
+  bn_decay_clip: 0.99
+  bn_decay_decay_rate: 0.5
+  bn_decay_decay_step: 200000
+  bn_init_decay: 0.5
+  decay_rate: 0.7
+  decay_step: 200000
+  dir_cosine_loss_weight: 1
+  learning_rate: 0.001
+  max_epoch: 16
+  momentum: 0.9
+  offset_loss_weight: 1
+  optimizer: adam
+  score_ce_loss_weight: 1
+TEST:
+  center_to_tip: 0.0
+  allow_zero_margin: 0
+  bin_vals: max
+  extra_opening: 0.005
+  first_thres: 0.23
+  second_thres: 0.19
+  max_farthest_points: 150
+  num_samples: 200
+  save: false
+  scale_fac:
+  - 1.25
+  - 1.0
+  - 0.75
+  - 0.5
+  scales: false
+  with_replacement: false
+  filter_thres: 0.0001
+  

scene_2048_bs3_rad2_32/model.ckpt-144144.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9811b79d977ad2f1cf18a78755befc8421189652f1179c73febbaea05ceaf9b1
+size 27860696

scene_test_2048_bs3_hor_sigma_001/checkpoint

@@ -0,0 +1,6 @@
+model_checkpoint_path: "model.ckpt-54054"
+all_model_checkpoint_paths: "model.ckpt-108108"
+all_model_checkpoint_paths: "model.ckpt-117117"
+all_model_checkpoint_paths: "model.ckpt-126126"
+all_model_checkpoint_paths: "model.ckpt-135135"
+all_model_checkpoint_paths: "model.ckpt-144144"

scene_test_2048_bs3_hor_sigma_001/config.yaml

@@ -0,0 +1,187 @@
+DATA:
+  gripper_width: 0.08
+  input_normals: false
+  use_uniform_quaternions: False
+  train_on_scenes: True
+  labels:
+    to_gpu: False
+    bin_weights:
+    - 0.16652107
+    - 0.21488856
+    - 0.37031708
+    - 0.55618503
+    - 0.75124664
+    - 0.93943357
+    - 1.07824539
+    - 1.19423112
+    - 1.55731375
+    - 3.17161779
+    contact_gather: knn
+    filter_z: true
+    k: 1
+    max_radius: 0.005
+    min_unique_pos_contacts: 1
+    num_neg_contacts: 0
+    num_pos_contacts: 10000
+    offset_bins:
+    - 0
+    - 0.00794435329
+    - 0.0158887021
+    - 0.0238330509
+    - 0.0317773996
+    - 0.0397217484
+    - 0.0476660972
+    - 0.055610446
+    - 0.0635547948
+    - 0.0714991435
+    - 0.08
+    z_val: -0.1
+  raw_num_points: 20000
+  ndataset_points: 20000
+  num_point: 2048
+  sigma: 0.001
+  clip: 0.005
+  use_farthest_point: false
+  train_and_test: false
+  num_test_scenes: 1000
+  intrinsics: 'realsense'
+LOSS:
+  min_geom_loss_divisor: 1.0
+  max_geom_loss_divisor: 100.0
+  offset_loss_type: sigmoid_cross_entropy
+  too_small_offset_pred_bin_factor: 0
+  topk_confidence: 512
+MODEL:
+  bin_offsets: true
+  contact_distance_offset: true
+  dir_vec_length_offset: false
+  grasp_conf_head:
+    conv1d: 1
+    dropout_keep: 0.5
+  grasp_dir_head:
+    conv1d: 3
+    dropout_keep: 0.7
+  joint_head:
+    conv1d: 4
+    dropout_keep: 0.7
+  joint_heads: false
+  larger_model: false
+  asymmetric_model: true
+  model: contact_graspnet
+  pointnet_fp_modules:
+  - mlp:
+    - 256
+    - 256
+  - mlp:
+    - 256
+    - 128
+  - mlp:
+    - 128
+    - 128
+    - 128
+  pointnet_sa_module:
+    group_all: true
+    mlp:
+    - 256
+    - 512
+    - 1024
+  pointnet_sa_modules_msg:
+  - mlp_list:
+    - - 32
+      - 32
+      - 64
+    - - 64
+      - 64
+      - 128
+    - - 64
+      - 96
+      - 128
+    npoint: 2048
+    nsample_list:
+    - 32
+    - 64
+    - 128
+    radius_list:
+    - 0.02
+    - 0.04
+    - 0.08
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 512
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.04
+    - 0.08
+    - 0.16
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 128
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.08
+    - 0.16
+    - 0.32
+  pred_contact_approach: true
+  pred_contact_base: true
+  pred_contact_offset: true
+  pred_contact_success: true
+  pred_grasps_adds: false
+  pred_grasps_adds_gt2pred: false
+OPTIMIZER:
+  adds_gt2pred_loss_weight: 1
+  adds_loss_weight: 10
+  approach_cosine_loss_weight: 1
+  batch_size: 14
+  bn_decay_clip: 0.99
+  bn_decay_decay_rate: 0.5
+  bn_decay_decay_step: 200000
+  bn_init_decay: 0.5
+  decay_rate: 0.7
+  decay_step: 200000
+  dir_cosine_loss_weight: 1
+  learning_rate: 0.001
+  max_epoch: 16
+  momentum: 0.9
+  offset_loss_weight: 1
+  optimizer: adam
+  score_ce_loss_weight: 1
+TEST:
+  center_to_tip: 0.0
+  allow_zero_margin: 0
+  bin_vals: max
+  extra_opening: 0.005
+  first_thres: 0.23
+  second_thres: 0.18
+  max_farthest_points: 150
+  num_samples: 200
+  save: false
+  scale_fac:
+  - 1.25
+  - 1.0
+  - 0.75
+  - 0.5
+  scales: false
+  with_replacement: false
+  filter_thres: 0.0001

scene_test_2048_bs3_hor_sigma_001/model.ckpt-144144.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffbd3968ed8526e5b27982eb8e6533bb697e72a5efd56c2f7c2f321f7dc5c9a6
+size 27860696

scene_test_2048_bs3_hor_sigma_001/model.ckpt-45045.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60575929774dd116b539a10680a123c4117f666e4750c404f41b00a4723f72a5
+size 27860696

scene_test_2048_bs3_hor_sigma_001/model.ckpt-54054.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:992e15e60e353796542f44acb33f062cbb7d9a313f42a13b420cf4d0c0275de2
+size 27860696

scene_test_2048_bs3_hor_sigma_001/model.ckpt-72072.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c760637a20f3cd271aba8c97b91c14acd6648dca854705fe4c817af0a2fdc52d
+size 27860696

scene_test_2048_bs3_hor_sigma_001/pointnet2_grasp_direct.py

@@ -0,0 +1,437 @@
+import os
+import sys
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(os.path.join(BASE_DIR))
+
+ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR))
+sys.path.append(os.path.join(BASE_DIR, 'utils'))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2', 'utils'))
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2'))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2'))
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_eager_execution()
+    TF2 = True
+except:
+    import tensorflow as tf
+    TF2 = False
+    
+import numpy as np
+import tf_util
+from pointnet_util import pointnet_sa_module, pointnet_fp_module, pointnet_sa_module_msg
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2', 'tf_ops/sampling'))
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2', 'tf_ops/grouping'))
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2', 'tf_ops/3d_interpolation'))
+from tf_sampling import farthest_point_sample, gather_point
+from tf_grouping import query_ball_point, group_point, knn_point
+from tf_interpolate import three_nn, three_interpolate
+import tf_utils
+
+def placeholder_inputs(global_config):
+    batch_size = global_config['OPTIMIZER']['batch_size']
+    if 'raw_num_points' in global_config['DATA']:
+        num_point = global_config['DATA']['raw_num_points']
+    else:
+        num_point = global_config['DATA']['num_point']
+    
+    input_normals = global_config['DATA']['input_normals']
+
+    pl_dict = {}
+    dim = 6 if input_normals else 3
+    pl_dict['pointclouds_pl'] = tf.placeholder(tf.float32, shape=(batch_size, num_point, dim))
+    pl_dict['obj_idx_pl'] = tf.placeholder(tf.int32, ())
+    pl_dict['cam_poses_pl'] = tf.placeholder(tf.float32, shape=(batch_size, 4, 4))
+    pl_dict['is_training_pl'] = tf.placeholder(tf.bool, shape=())
+
+    return pl_dict #, labels_pl
+
+def get_bin_vals(global_config):
+    bins_bounds = np.array(global_config['DATA']['labels']['offset_bins'])
+    if global_config['TEST']['bin_vals'] == 'max':
+        bin_vals = (bins_bounds[1:] + bins_bounds[:-1])/2
+        bin_vals[-1] = bins_bounds[-1]
+    elif global_config['TEST']['bin_vals'] == 'mean':
+        bin_vals = bins_bounds[1:]
+    else:
+        raise NotImplementedError
+
+    if not global_config['TEST']['allow_zero_margin']:
+        bin_vals = np.minimum(bin_vals, global_config['DATA']['gripper_width']-global_config['TEST']['extra_opening'])
+        
+    tf_bin_vals = tf.constant(bin_vals, tf.float32)
+    return tf_bin_vals
+
+def get_model(point_cloud, is_training, global_config, bn_decay=None):
+    """ Grasp direction PointNet++ """
+
+    model_config = global_config['MODEL']
+    data_config = global_config['DATA']
+
+    radius_list_0 = model_config['pointnet_sa_modules_msg'][0]['radius_list']
+    radius_list_1 = model_config['pointnet_sa_modules_msg'][1]['radius_list']
+    nsample_list_0 = model_config['pointnet_sa_modules_msg'][0]['nsample_list']
+    nsample_list_1 = model_config['pointnet_sa_modules_msg'][1]['nsample_list']
+    mlp_list_0 = model_config['pointnet_sa_modules_msg'][0]['mlp_list']
+    mlp_list_1 = model_config['pointnet_sa_modules_msg'][1]['mlp_list']
+    npoint_0 = model_config['pointnet_sa_modules_msg'][0]['npoint']
+    npoint_1 = model_config['pointnet_sa_modules_msg'][1]['npoint']
+    fp_mlp_0 = model_config['pointnet_fp_modules'][0]['mlp']
+    fp_mlp_1 = model_config['pointnet_fp_modules'][1]['mlp']
+    fp_mlp_2 = model_config['pointnet_fp_modules'][2]['mlp']
+
+    #larger model
+    radius_list_2 = model_config['pointnet_sa_modules_msg'][2]['radius_list']
+    nsample_list_2 = model_config['pointnet_sa_modules_msg'][2]['nsample_list']
+    npoint_2 = model_config['pointnet_sa_modules_msg'][2]['npoint']
+    mlp_list_2 = model_config['pointnet_sa_modules_msg'][2]['mlp_list']
+
+    input_normals = data_config['input_normals']
+    gripper_width = data_config['gripper_width']
+    offset_bins = data_config['labels']['offset_bins']
+    joint_heads = model_config['joint_heads']
+
+    # expensive, rather use random only
+    if 'raw_num_points' in data_config and data_config['raw_num_points'] != data_config['ndataset_points']:
+        point_cloud = gather_point(point_cloud, farthest_point_sample(data_config['ndataset_points'], point_cloud))
+
+    batch_size = point_cloud.get_shape()[0]
+    num_point = point_cloud.get_shape()[1]
+    end_points = {}
+    l0_xyz = tf.slice(point_cloud, [0,0,0], [-1,-1,3])
+    l0_points = tf.slice(point_cloud, [0,0,3], [-1,-1,3]) if input_normals else None 
+
+    l_xyz = [l0_xyz]
+    l_points = [l0_points]
+    pred_points = l0_xyz
+
+    # Set abstraction layers
+    l1_xyz, l1_points = pointnet_sa_module_msg(l0_xyz, l0_points, npoint_0, radius_list_0, nsample_list_0, mlp_list_0, is_training, bn_decay, scope='layer1')
+    l2_xyz, l2_points = pointnet_sa_module_msg(l1_xyz, l1_points, npoint_1, radius_list_1, nsample_list_1,mlp_list_1, is_training, bn_decay, scope='layer2')
+
+    # large model
+    if 'larger_model' in model_config and model_config['larger_model']:
+        fp_mlp_3 = model_config['pointnet_fp_modules'][3]['mlp']
+
+        l3_xyz, l3_points = pointnet_sa_module_msg(l2_xyz, l2_points, npoint_2, radius_list_2, nsample_list_2,mlp_list_2, is_training, bn_decay, scope='layer3')
+        l4_xyz, l4_points, _ = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer4')
+
+        # Feature Propagation layers
+        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+        l0_points = tf.concat([l0_xyz, l0_points],axis=-1) if input_normals else l0_xyz 
+        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, fp_mlp_3, is_training, bn_decay, scope='fa_layer4')
+    elif 'asymmetric_model' in model_config and model_config['asymmetric_model']:
+        l3_xyz, l3_points = pointnet_sa_module_msg(l2_xyz, l2_points, npoint_2, radius_list_2, nsample_list_2,mlp_list_2, is_training, bn_decay, scope='layer3')
+        l4_xyz, l4_points, _ = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer4')
+
+        # Feature Propagation layers
+        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+        l0_points = l1_points
+        pred_points = l1_xyz
+    else:
+        l3_xyz, l3_points, _ = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer3')
+
+        # Feature Propagation layers
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+
+        l0_points = tf.concat([l0_xyz, l0_points],axis=-1) if input_normals else l0_xyz 
+        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+    if joint_heads:
+        head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
+        # end_points['feats'] = head 
+        head = tf_util.dropout(head, keep_prob=0.7, is_training=is_training, scope='dp1')
+        head = tf_util.conv1d(head, 4, 1, padding='VALID', activation_fn=None, scope='fc2')
+        grasp_dir_head = tf.slice(head, [0,0,0], [-1,-1,3])
+        grasp_dir_head = tf.math.l2_normalize(grasp_dir_head, axis=2)
+        binary_seg_head = tf.slice(head, [0,0,3], [-1,-1,1])
+    else:
+        # FC layers for grasp direction
+        grasp_dir_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
+        # end_points['feats'] = grasp_dir_head 
+        grasp_dir_head = tf_util.dropout(grasp_dir_head, keep_prob=0.7, is_training=is_training, scope='dp1')
+        grasp_dir_head = tf_util.conv1d(grasp_dir_head, 3, 1, padding='VALID', activation_fn=None, scope='fc3')
+        grasp_dir_head_normed = tf.math.l2_normalize(grasp_dir_head, axis=2)
+
+
+        # FC layers for grasp direction
+        approach_dir_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_app', bn_decay=bn_decay)
+        # end_points['feats'] = approach_dir_head 
+        approach_dir_head = tf_util.dropout(approach_dir_head, keep_prob=0.7, is_training=is_training, scope='dp1_app')
+        approach_dir_head = tf_util.conv1d(approach_dir_head, 3, 1, padding='VALID', activation_fn=None, scope='fc3_app')
+        approach_dir_head_orthog = tf.math.l2_normalize(approach_dir_head - tf.reduce_sum(tf.multiply(grasp_dir_head_normed, approach_dir_head), axis=2, keepdims=True)*grasp_dir_head_normed, axis=2)
+        
+        if model_config['dir_vec_length_offset']:
+            grasp_offset_head = tf.norm(grasp_dir_head, axis=2, keepdims=True)
+        elif model_config['bin_offsets']:
+            grasp_offset_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_off', bn_decay=bn_decay)
+            # end_points['feats'] = grasp_offset_head 
+            # grasp_offset_head = tf_util.dropout(grasp_offset_head, keep_prob=0.7, is_training=is_training, scope='dp1_off')
+            grasp_offset_head = tf_util.conv1d(grasp_offset_head, len(offset_bins)-1, 1, padding='VALID', activation_fn=None, scope='fc2_off')
+        else:
+            grasp_offset_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_off', bn_decay=bn_decay)
+            # end_points['feats'] = grasp_offset_head 
+            grasp_offset_head = tf_util.dropout(grasp_offset_head, keep_prob=0.7, is_training=is_training, scope='dp1_off')
+            grasp_offset_head = tf_util.conv1d(grasp_offset_head, 1, 1, padding='VALID', activation_fn=None, scope='fc2_off')
+
+        binary_seg_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_seg', bn_decay=bn_decay)
+        binary_seg_head = tf_util.dropout(binary_seg_head, keep_prob=0.5, is_training=is_training, scope='dp1_seg')
+        binary_seg_head = tf_util.conv1d(binary_seg_head, 1, 1, padding='VALID', activation_fn=None, scope='fc2_seg')
+
+    end_points['grasp_dir_head'] = grasp_dir_head_normed
+    end_points['binary_seg_head'] = binary_seg_head
+    end_points['binary_seg_pred'] = tf.math.sigmoid(binary_seg_head)
+    end_points['grasp_offset_head'] = grasp_offset_head
+    end_points['grasp_offset_pred'] = tf.math.sigmoid(grasp_offset_head) if model_config['bin_offsets'] else grasp_offset_head
+    end_points['approach_dir_head'] = approach_dir_head_orthog
+    end_points['pred_points'] = pred_points
+
+    return end_points
+
+def build_6d_grasp(approach_dirs, base_dirs, contact_pts, thickness, use_tf=False, gripper_depth = 0.1034):
+    
+    if use_tf:
+        grasps_R = tf.stack([base_dirs, tf.linalg.cross(approach_dirs,base_dirs),approach_dirs], axis=3)        
+        grasps_t = contact_pts + tf.expand_dims(thickness,2)/2 * base_dirs - gripper_depth * approach_dirs
+        ones = tf.ones((contact_pts.shape[0], contact_pts.shape[1], 1, 4), dtype=tf.float32)
+        grasps = tf.concat([tf.concat([grasps_R,  tf.expand_dims(grasps_t, 3)], axis=3), ones], axis=2)
+    else:
+        grasps = []
+        for i in range(len(contact_pts)):
+            grasp = np.eye(4)
+
+            grasp[:3,0] = base_dirs[i] / np.linalg.norm(base_dirs[i])
+            grasp[:3,2] = approach_dirs[i] / np.linalg.norm(approach_dirs[i])
+            grasp_y = np.cross( grasp[:3,2],grasp[:3,0])
+            grasp[:3,1] = grasp_y / np.linalg.norm(grasp_y)
+            # base_gripper xyz = contact + thickness / 2 * baseline_dir - gripper_d * approach_dir 
+            grasp[:3,3] = contact_pts[i] + thickness[i] / 2 * grasp[:3,0] - gripper_depth * grasp[:3,2]
+            # grasp[0,3] = finger_width
+            grasps.append(grasp)
+        grasps = np.array(grasps)
+
+    return grasps
+
+def multi_bin_labels(cont_labels, bin_boundaries):
+
+    # multi_hot_labels = tf.zeros_like((cont_labels.shape[0], cont_labels.shape[1], len(bin_boundaries)-1), tf.bool)
+    bins = []
+    for b in range(len(bin_boundaries)-1):
+        bins.append(tf.math.logical_and(tf.greater_equal(cont_labels, bin_boundaries[b]), tf.less(cont_labels,bin_boundaries[b+1])))
+        # cont_labels = tf.where(condition, b*tf.ones_like(condition, tf.float32), cont_labels)
+    multi_hot_labels = tf.concat(bins, axis=2)
+    multi_hot_labels = tf.cast(multi_hot_labels, tf.float32)
+
+    return multi_hot_labels
+
+
+def get_losses(pointclouds_pl, end_points, dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam, global_config):
+    """ pred: (batch_size, num_point, 3),
+        label: (batch_size, num_point, 3), 
+        grasp_success_labels_pc: (batch_size, num_point)
+        """
+
+    grasp_dir_head = end_points['grasp_dir_head']
+    grasp_offset_head = end_points['grasp_offset_head']
+    approach_dir_head = end_points['approach_dir_head']
+
+    bin_weights = global_config['DATA']['labels']['bin_weights']
+    tf_bin_weights = tf.constant(bin_weights)
+    
+    min_geom_loss_divisor = tf.constant(float(global_config['LOSS']['min_geom_loss_divisor'])) if 'min_geom_loss_divisor' in global_config['LOSS'] else tf.constant(1.)
+    pos_grasps_in_view = tf.math.maximum(tf.reduce_sum(grasp_success_labels_pc, axis=1), min_geom_loss_divisor)   
+
+    ### ADS Gripper PC Loss
+    if global_config['MODEL']['bin_offsets']:
+        thickness_pred = tf.gather_nd(get_bin_vals(global_config), tf.expand_dims(tf.argmax(grasp_offset_head, axis=2), axis=2))
+        thickness_gt = tf.gather_nd(get_bin_vals(global_config), tf.expand_dims(tf.argmax(offset_labels_pc, axis=2), axis=2))
+    else:
+        thickness_pred = grasp_offset_head[:,:,0]
+        thickness_gt = offset_labels_pc[:,:,0]
+    pred_grasps = build_6d_grasp(approach_dir_head, grasp_dir_head, pointclouds_pl, thickness_pred, use_tf=True) # b x num_point x 4 x 4
+    gt_grasps_proj = build_6d_grasp(approach_labels_pc_cam, dir_labels_pc_cam, pointclouds_pl, thickness_gt, use_tf=True) # b x num_point x 4 x 4
+    pos_gt_grasps_proj = tf.where(tf.broadcast_to(tf.expand_dims(tf.expand_dims(tf.cast(grasp_success_labels_pc, tf.bool),2),3), gt_grasps_proj.shape), gt_grasps_proj, tf.ones_like(gt_grasps_proj)*100000)
+    # pos_gt_grasps_proj = tf.reshape(pos_gt_grasps_proj, (global_config['OPTIMIZER']['batch_size'], -1, 4, 4)) 
+
+    gripper_control_points = tf_utils.get_control_point_tensor(global_config['OPTIMIZER']['batch_size'])[:,1:,:] # b x 5 x 3
+    sym_gripper_control_points = tf_utils.get_control_point_tensor(global_config['OPTIMIZER']['batch_size'], symmetric=True)[:,1:,:]
+
+    gripper_control_points_homog =  tf.concat([gripper_control_points, tf.ones((global_config['OPTIMIZER']['batch_size'], gripper_control_points.shape[1], 1))], axis=2)  # b x 5 x 4
+    sym_gripper_control_points_homog =  tf.concat([sym_gripper_control_points, tf.ones((global_config['OPTIMIZER']['batch_size'], gripper_control_points.shape[1], 1))], axis=2)  # b x 5 x 4
+    
+    # only use per point pred grasps but not per point gt grasps
+    control_points = tf.keras.backend.repeat_elements(tf.expand_dims(gripper_control_points_homog,1), gt_grasps_proj.shape[1], axis=1)  # b x num_point x 5 x 4
+    sym_control_points = tf.keras.backend.repeat_elements(tf.expand_dims(sym_gripper_control_points_homog,1), gt_grasps_proj.shape[1], axis=1)  # b x num_point x 5 x 4
+    pred_control_points = tf.matmul(control_points, tf.transpose(pred_grasps, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_point x 5 x 3
+
+    ### Pred Grasp to GT Grasp ADD-S Loss
+    gt_control_points = tf.matmul(control_points, tf.transpose(pos_gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+    sym_gt_control_points = tf.matmul(sym_control_points, tf.transpose(pos_gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+
+    squared_add = tf.reduce_sum((tf.expand_dims(pred_control_points,2)-tf.expand_dims(gt_control_points,1))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+    sym_squared_add = tf.reduce_sum((tf.expand_dims(pred_control_points,2)-tf.expand_dims(sym_gt_control_points,1))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+
+    # symmetric ADD-S
+    neg_squared_adds = -tf.concat([squared_add,sym_squared_add], axis=2) # b x num_point x 2num_pos_grasp_point
+    neg_squared_adds_k = tf.math.top_k(neg_squared_adds, k=1, sorted=False)[0] # b x num_point
+    # If any pos grasp exists
+    min_adds = tf.minimum(tf.reduce_sum(grasp_success_labels_pc, axis=1, keepdims=True), tf.ones_like(neg_squared_adds_k[:,:,0])) * tf.sqrt(-neg_squared_adds_k[:,:,0])#tf.minimum(tf.sqrt(-neg_squared_adds_k), tf.ones_like(neg_squared_adds_k)) # b x num_point
+    adds_loss = tf.reduce_mean(end_points['binary_seg_pred'][:,:,0] * min_adds)
+
+    ### GT Grasp to pred Grasp ADD-S Loss
+    gt_control_points = tf.matmul(control_points, tf.transpose(gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+    sym_gt_control_points = tf.matmul(sym_control_points, tf.transpose(gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+
+    neg_squared_adds = -tf.reduce_sum((tf.expand_dims(pred_control_points,1)-tf.expand_dims(gt_control_points,2))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+    neg_squared_adds_sym = -tf.reduce_sum((tf.expand_dims(pred_control_points,1)-tf.expand_dims(sym_gt_control_points,2))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+
+    neg_squared_adds_k_gt2pred, pred_grasp_idcs = tf.math.top_k(neg_squared_adds, k=1, sorted=False) # b x num_pos_grasp_point
+    neg_squared_adds_k_sym_gt2pred, pred_grasp_sym_idcs = tf.math.top_k(neg_squared_adds_sym, k=1, sorted=False) # b x num_pos_grasp_point
+    pred_grasp_idcs_joined = tf.where(neg_squared_adds_k_gt2pred<neg_squared_adds_k_sym_gt2pred, pred_grasp_sym_idcs, pred_grasp_idcs)
+    min_adds_gt2pred = tf.minimum(-neg_squared_adds_k_gt2pred, -neg_squared_adds_k_sym_gt2pred) # b x num_pos_grasp_point x 1
+    # min_adds_gt2pred = tf.math.exp(-min_adds_gt2pred)
+    masked_min_adds_gt2pred = tf.multiply(min_adds_gt2pred[:,:,0], grasp_success_labels_pc)
+    
+
+    batch_idcs = tf.meshgrid(tf.range(pred_grasp_idcs_joined.shape[1]), tf.range(pred_grasp_idcs_joined.shape[0]))
+    gather_idcs = tf.stack((batch_idcs[1],pred_grasp_idcs_joined[:,:,0]), axis=2)
+    nearest_pred_grasp_confidence = tf.gather_nd(end_points['binary_seg_pred'][:,:,0], gather_idcs)
+    adds_loss_gt2pred = tf.reduce_mean(tf.reduce_sum(nearest_pred_grasp_confidence*masked_min_adds_gt2pred, axis=1) / pos_grasps_in_view)
+ 
+    ### Grasp baseline Loss
+    cosine_distance = tf.constant(1.)-tf.reduce_sum(tf.multiply(dir_labels_pc_cam, grasp_dir_head),axis=2)
+    # only pass loss where we have labeled contacts near pc points 
+    masked_cosine_loss = tf.multiply(cosine_distance, grasp_success_labels_pc)
+    dir_cosine_loss = tf.reduce_mean(tf.reduce_sum(masked_cosine_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Approach Loss
+    approach_labels_orthog = tf.math.l2_normalize(approach_labels_pc_cam - tf.reduce_sum(tf.multiply(grasp_dir_head, approach_labels_pc_cam), axis=2, keepdims=True)*grasp_dir_head, axis=2)
+    cosine_distance_approach = tf.constant(1.)-tf.reduce_sum(tf.multiply(approach_labels_orthog, approach_dir_head), axis=2)
+    masked_approach_loss = tf.multiply(cosine_distance_approach, grasp_success_labels_pc)
+    approach_cosine_loss = tf.reduce_mean(tf.reduce_sum(masked_approach_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Offset/Thickness Loss
+    if global_config['MODEL']['bin_offsets']:
+        if global_config['LOSS']['offset_loss_type'] == 'softmax_cross_entropy':
+            offset_loss = tf.losses.softmax_cross_entropy(offset_labels_pc, grasp_offset_head)
+        else:
+            offset_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=offset_labels_pc, logits=grasp_offset_head)
+            
+            if 'too_small_offset_pred_bin_factor' in global_config['LOSS'] and global_config['LOSS']['too_small_offset_pred_bin_factor']:
+                too_small_offset_pred_bin_factor = tf.constant(global_config['LOSS']['too_small_offset_pred_bin_factor'], tf.float32)
+                collision_weight = tf.math.cumsum(offset_labels_pc, axis=2, reverse=True)*too_small_offset_pred_bin_factor + tf.constant(1.)
+                offset_loss = tf.multiply(collision_weight, offset_loss)
+
+            offset_loss = tf.reduce_mean(tf.multiply(tf.reshape(tf_bin_weights,(1,1,-1)), offset_loss),axis=2)
+    else:
+        offset_loss = (grasp_offset_head[:,:,0] - offset_labels_pc[:,:,0])**2
+    masked_offset_loss = tf.multiply(offset_loss, grasp_success_labels_pc)        
+    offset_loss = tf.reduce_mean(tf.reduce_sum(masked_offset_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Confidence Loss
+    bin_ce_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.expand_dims(grasp_success_labels_pc,axis=2), logits=end_points['binary_seg_head'])
+    if 'topk_confidence' in global_config['LOSS'] and global_config['LOSS']['topk_confidence']:
+        bin_ce_loss,_ = tf.math.top_k(tf.squeeze(bin_ce_loss), k=global_config['LOSS']['topk_confidence'])
+    bin_ce_loss = tf.reduce_mean(bin_ce_loss)
+
+    return dir_cosine_loss, bin_ce_loss, offset_loss, approach_cosine_loss, adds_loss, adds_loss_gt2pred, gt_control_points, pred_control_points, pos_grasps_in_view
+
+def compute_labels(pos_contact_pts_mesh, pos_contact_dirs_mesh, pos_contact_offsets, pos_contact_approaches_mesh, pos_finger_diffs, neg_contact_pts_mesh, 
+                  neg_contact_dirs_mesh, neg_contact_offsets, neg_contact_approaches_mesh, neg_finger_diffs, pc_cam_pl, camera_pose_pl, global_config):
+    """
+    for every point find nearest contact points
+    group and average labels
+    """
+    label_config = global_config['DATA']['labels']
+    model_config = global_config['MODEL']
+
+    contact_distance_offset = model_config['contact_distance_offset']
+    bin_offsets = model_config['bin_offsets']
+
+    mode = label_config['contact_gather']
+    nsample = label_config['k']
+    radius = label_config['max_radius']
+    filter_z = label_config['filter_z']
+    z_val = label_config['z_val']
+
+    xyz_cam = pc_cam_pl[:,:,:3]
+    pad_homog = tf.ones((xyz_cam.shape[0],xyz_cam.shape[1], 1)) 
+    pc_mesh = tf.matmul(tf.concat([xyz_cam, pad_homog], 2), tf.transpose(tf.linalg.inv(camera_pose_pl),perm=[0, 2, 1]))[:,:,:3]
+
+    if contact_distance_offset:
+        # finger_diff = tf.norm(pos_contact_pts_mesh[1::2,:]-pos_contact_pts_mesh[::2,:] + 0.0000001, axis=1, keepdims=True)
+        # finger_diff = tf.keras.backend.repeat_elements(finger_diff, 2, axis=0)
+        contact_point_offsets_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_finger_diffs,0), pc_mesh.shape[0], axis=0)
+    else:
+        contact_point_offsets_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_offsets,0), pc_mesh.shape[0], axis=0)
+
+
+    pad_homog2 = tf.ones((pc_mesh.shape[0], pos_contact_dirs_mesh.shape[0], 1)) 
+    contact_point_dirs_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0)
+    contact_point_dirs_batch_cam = tf.matmul(contact_point_dirs_batch, tf.transpose(camera_pose_pl[:,:3,:3], perm=[0, 2, 1]))[:,:,:3]
+
+    pos_contact_approaches_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_approaches_mesh,0), pc_mesh.shape[0], axis=0)
+    pos_contact_approaches_batch_cam = tf.matmul(pos_contact_approaches_batch, tf.transpose(camera_pose_pl[:,:3,:3], perm=[0, 2, 1]))[:,:,:3]
+    
+    contact_point_batch_mesh = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_pts_mesh,0), pc_mesh.shape[0], axis=0)
+    contact_point_batch_cam = tf.matmul(tf.concat([contact_point_batch_mesh, pad_homog2], 2), tf.transpose(camera_pose_pl, perm=[0, 2, 1]))[:,:,:3]
+
+
+    # orthogonal = tf.reduce_sum(pos_contact_approaches_batch_cam * contact_point_dirs_batch_cam, axis=2)
+    if filter_z:
+        dir_filter_passed = tf.keras.backend.repeat_elements(tf.math.greater(contact_point_dirs_batch_cam[:,:,2:3], tf.constant([z_val])), 3, axis=2)
+        contact_point_batch_mesh = tf.where(dir_filter_passed, contact_point_batch_mesh, tf.ones_like(contact_point_batch_mesh)*100000)
+        # contact_point_dirs_batch_cam = tf.boolean_mask(contact_point_dirs_batch_cam, dir_filter_passed)
+        # print(contact_point_dirs_batch_cam.shape)
+        # contact_point_batch_mesh = tf.boolean_mask(contact_point_batch_mesh, dir_filter_passed)
+    # else:
+        # contact_point_dirs_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0)
+
+    # TODO: Divide into positive and negative grasps; for now only insert positive contacts (otherwise mean direction does not make sense)
+    if mode == 'mean':
+        close_contact_pt_idcs,_ = query_ball_point(radius, nsample, contact_point_batch_mesh, pc_mesh) # (batch_size, npoint, nsample)
+
+        grasp_success_labels_pc = tf.cast(tf.reduce_sum(close_contact_pt_idcs, axis=2)>0, tf.float32) # (batch_size, num_point)
+
+        # contact_point to contact_dir mapping
+        grouped_dirs_mesh = group_point(contact_point_dirs_batch, close_contact_pt_idcs) # (batch_size, num_point, nsample, 3)
+        
+        # grouped_dirs_mesh = tf.gather_nd(tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0), close_contact_pt_idcs)
+    elif mode == 'knn':
+        # dists, close_contact_pt_idcs = knn_point(nsample, contact_point_batch_mesh, pc_mesh) # (batch_size, num_point, nsample)
+
+        # squared_dists_all = tf.reduce_sum((tf.expand_dims(contact_point_batch_mesh,1)-tf.expand_dims(pc_mesh,2))**2,axis=3)
+        squared_dists_all = tf.reduce_sum((tf.expand_dims(contact_point_batch_cam,1)-tf.expand_dims(xyz_cam,2))**2,axis=3)
+        neg_squared_dists_k, close_contact_pt_idcs = tf.math.top_k(-squared_dists_all, k=nsample, sorted=False)
+        squared_dists_k = -neg_squared_dists_k
+
+        # Distance check, careful it depends on point density
+        grasp_success_labels_pc = tf.cast(tf.less(tf.reduce_mean(squared_dists_k, axis=2), radius*radius), tf.float32) # (batch_size, num_point)
+
+        grouped_dirs_pc_cam = group_point(contact_point_dirs_batch_cam, close_contact_pt_idcs)
+        grouped_approaches_pc_cam = group_point(pos_contact_approaches_batch_cam, close_contact_pt_idcs)
+        grouped_offsets = group_point(tf.expand_dims(contact_point_offsets_batch,2), close_contact_pt_idcs)
+        # not sure for nsample > 1
+        dir_labels_pc_cam = tf.math.l2_normalize(tf.reduce_mean(grouped_dirs_pc_cam, axis=2),axis=2) # (batch_size, num_point, 3)
+        approach_labels_pc_cam = tf.math.l2_normalize(tf.reduce_mean(grouped_approaches_pc_cam, axis=2),axis=2) # (batch_size, num_point, 3)
+        offset_labels_pc = tf.reduce_mean(grouped_offsets, axis=2)
+        
+    else:
+        raise NotImplementedError
+        # if filter_z:
+        #     pad_homog3 = tf.ones((dir_labels_mesh.shape[0], dir_labels_mesh.shape[1], 1)) 
+        #     dir_labels_pc_cam = tf.matmul(tf.concat([dir_labels_mesh,pad_homog3], 2), tf.transpose(camera_pose_pl, perm=[0, 2, 1]))[:,:,:3] # (batch_size, num_point, 3)
+
+    return dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam
+
+if __name__=='__main__':
+    with tf.Graph().as_default():
+        inputs = tf.zeros((32,2048,6))
+        net, _ = get_model(inputs, tf.constant(True))
+        print(net)

scene_test_2048_bs3_hor_sigma_001/train_grasp_direct.py

@@ -0,0 +1,292 @@
+import os
+import sys
+import argparse
+import math
+from datetime import datetime
+import numpy as np
+import socket
+import importlib
+import time
+from tqdm import tqdm
+import glob
+import json
+import yaml
+
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR))
+sys.path.append(os.path.join(BASE_DIR))
+sys.path.append(os.path.join(ROOT_DIR))
+
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_eager_execution()
+    TF2 = True
+    physical_devices = tf.config.experimental.list_physical_devices('GPU')
+    print(physical_devices)
+    tf.config.experimental.set_memory_growth(physical_devices[0], True)
+except:
+    import tensorflow as tf
+    TF2 = False
+
+POINT_DIR_NGC = os.path.join(ROOT_DIR, 'pointnet2')
+if os.path.exists(POINT_DIR_NGC):
+    sys.path.append(os.path.join(POINT_DIR_NGC, 'models'))
+    sys.path.append(os.path.join(POINT_DIR_NGC, 'utils'))
+else:
+    sys.path.append(os.path.join(BASE_DIR, 'pointnet2',  'models'))
+    sys.path.append(os.path.join(BASE_DIR, 'pointnet2',  'utils'))
+    
+import provider
+import sample
+import utilities
+from data import PointCloudReader, load_mesh_path_and_scale, preprocess_pc_for_inference, load_filtered_contact_data, load_scene_contacts, load_obj_scales_cats, inverse_transform, center_pc_convert_cam
+from summaries import top_grasp_acc_summaries, build_summary_ops, build_file_writers
+from tf_train_ops import load_labels_and_losses, build_train_op
+from surface_grasp_estimator import GraspEstimator
+
+
+def train(global_config, LOG_DIR):
+
+    if 'train_on_scenes' in global_config['DATA'] and global_config['DATA']['train_on_scenes']:
+        mesh_scales, mesh_cats = load_obj_scales_cats(global_config['DATA']['data_path'])
+        contact_infos, scene_obj_paths, scene_obj_transforms = load_scene_contacts(global_config['DATA']['data_path'])
+        num_train_samples = len(contact_infos)-global_config['DATA']['num_test_scenes']
+        num_test_samples = global_config['DATA']['num_test_scenes']
+    else:
+        scene_obj_paths, scene_obj_transforms = None, None
+        train_contact_paths, test_contact_paths, contact_infos, mesh_scales = load_filtered_contact_data(global_config['DATA']['data_path'], min_pos_contacts=1, classes=global_config['DATA']['classes'])
+        num_train_samples = len(train_contact_paths)
+        num_test_samples = len(test_contact_paths)
+        
+    print('using %s meshes' % (num_train_samples + num_test_samples))
+    
+    if 'train_and_test' in global_config['DATA'] and global_config['DATA']['train_and_test']:
+        num_train_samples = num_train_samples + num_test_samples
+        num_test_samples = 0
+        print('using train and test data')
+
+    pcreader = PointCloudReader(
+        root_folder=global_config['DATA']['data_path'],
+        batch_size=global_config['OPTIMIZER']['batch_size'],
+        num_grasp_clusters=None,
+        estimate_normals=global_config['DATA']['input_normals'],
+        npoints=global_config['DATA']['num_point'],
+        raw_num_points=global_config['DATA']['raw_num_points'],
+        use_uniform_quaternions = global_config['DATA']['use_uniform_quaternions'],
+        run_in_another_process = False,
+        mesh_scales = mesh_scales,
+        scene_obj_paths = scene_obj_paths,
+        scene_obj_transforms = scene_obj_transforms,
+        num_train_samples = num_train_samples,
+        num_test_samples = num_test_samples,
+        use_farthest_point = global_config['DATA']['use_farthest_point'],
+        intrinsics=global_config['DATA']['intrinsics']
+    )
+
+    with tf.Graph().as_default():
+        
+        # Build the model
+        grasp_estimator = GraspEstimator(global_config)
+
+        ops = grasp_estimator.build_network()
+        
+        # contact_tensors = load_contact_grasps(contact_infos, global_config['DATA'])
+        
+        loss_ops = load_labels_and_losses(grasp_estimator, contact_infos, global_config)
+
+        ops.update(loss_ops)
+        ops['train_op'] = build_train_op(ops['loss'], ops['step'], global_config)
+
+        # Add ops to save and restore all the variables.
+        saver = tf.train.Saver(save_relative_paths=True, keep_checkpoint_every_n_hours=4)
+
+        # Create a session
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        config.allow_soft_placement = True
+        # config.log_device_placement = False
+        sess = tf.Session(config=config)
+
+
+        summary_ops = build_summary_ops(ops, sess, global_config)
+
+        grasp_estimator.load_weights(sess, saver, LOG_DIR, mode='train')
+        file_writers = build_file_writers(sess, LOG_DIR)
+
+    batches_per_epoch = num_train_samples #// global_config['OPTIMIZER']['batch_size'] + 1
+    cur_epoch = sess.run(ops['step']) // (batches_per_epoch * global_config['OPTIMIZER']['batch_size'])
+    for epoch in range(cur_epoch, global_config['OPTIMIZER']['max_epoch']):
+        log_string('**** EPOCH %03d ****' % (epoch))
+        
+        sess.run(ops['iterator'].initializer)
+        epoch_time = time.time()
+        step = train_one_epoch(sess, ops, summary_ops, file_writers, pcreader)
+        print('trained %s batches in: ' % batches_per_epoch, time.time()-epoch_time)
+
+        # Save the variables to disk.
+        if (epoch+1) % 1 == 0:
+            save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt"), global_step=step, write_meta_graph=False)
+            log_string("Model saved in file: %s" % save_path)
+
+        if epoch % 1 == 0 and num_test_samples > 0:
+            eval_time = time.time()
+            eval_test_objects(sess, ops, summary_ops, file_writers, pcreader)
+            print('evaluation time: ', time.time()-eval_time)
+
+def train_one_epoch(sess, ops, summary_ops, file_writers, pcreader):
+    """ ops: dict mapping from string to tf ops """
+    is_training = True
+    
+    log_string(str(datetime.now()))
+
+    loss_sum, loss_sum_dir, loss_sum_ce, loss_sum_off, loss_sum_app, loss_sum_adds, loss_sum_adds_gt2pred, time_sum = 8 * [0]
+
+    # batches_per_epoch = pcreader._num_train_samples // pcreader._batch_size
+    ## define one epoch = all objects/scenes seen
+    batches_per_epoch = pcreader._num_train_samples
+
+    # run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, report_tensor_allocations_upon_oom = True)
+
+    for batch_idx in range(batches_per_epoch):
+        get_time = time.time()
+
+        batch_data, cam_poses, obj_idx = pcreader.get_batch(batch_idx)
+        print(time.time()- get_time)
+        if 'train_on_scenes' in global_config['DATA'] and global_config['DATA']['train_on_scenes']:
+            # OpenCV OpenGL conversion
+            cam_poses, batch_data = center_pc_convert_cam(cam_poses, batch_data)
+        print(time.time() - get_time)
+        # Augment batched point clouds by rotation and jittering
+        # aug_data = provider.random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25)
+        if 'sigma' in global_config['DATA'] and global_config['DATA']['sigma'] > 0:
+            batch_data[:,:,0:3] = provider.jitter_point_cloud(batch_data[:,:,0:3], 
+                                                            sigma=global_config['DATA']['sigma'], 
+                                                            clip=global_config['DATA']['clip']*2)
+        
+        feed_dict = {ops['pointclouds_pl']: batch_data,
+                     ops['cam_poses_pl']: cam_poses,
+                     ops['obj_idx_pl']: obj_idx,
+                    #  ops['labels_pl']: batch_label,
+                     ops['is_training_pl']: is_training}
+
+        step, summary, _, loss_val, dir_loss, bin_ce_loss, \
+        offset_loss, approach_loss, adds_loss, adds_gt2pred_loss,pos_grasps_in_view, scene_idx = sess.run([ops['step'], summary_ops['merged'], ops['train_op'], 
+                                                                            ops['loss'], ops['dir_loss'], ops['bin_ce_loss'], 
+                                                                            ops['offset_loss'], ops['approach_loss'], ops['adds_loss'], 
+                                                                            ops['adds_gt2pred_loss'], ops['pos_grasps_in_view'], ops['scene_idx']], feed_dict=feed_dict)
+        print(time.time()- get_time)
+        print(pos_grasps_in_view)
+        print(scene_idx, obj_idx)
+        assert scene_idx[0] == obj_idx
+
+        loss_sum += loss_val
+        loss_sum_dir += dir_loss
+        loss_sum_ce += bin_ce_loss
+        loss_sum_off += offset_loss
+        loss_sum_app += approach_loss
+        loss_sum_adds += adds_loss
+        loss_sum_adds_gt2pred += adds_gt2pred_loss
+        time_sum += time.time() - get_time
+        
+        if (batch_idx+1)%10 == 0:
+            file_writers['train_writer'].add_summary(summary, step)
+            log_string('total loss: %f \t dir loss: %f \t ce loss: %f \t off loss: %f \t app loss: %f adds loss: %f \t adds_gt2pred loss: %f \t batch time: %f' % (loss_sum/10,loss_sum_dir/10,loss_sum_ce/10, loss_sum_off/10, loss_sum_app/10, loss_sum_adds/10, loss_sum_adds_gt2pred/10, time_sum/10))
+            # log_string('accuracy: %f' % (total_correct / float(total_seen)))
+            loss_sum, loss_sum_dir, loss_sum_ce, loss_sum_off, loss_sum_app, loss_sum_adds, loss_sum_adds_gt2pred, time_sum = 8 * [0]
+            
+    return step
+
+def eval_test_objects(sess, ops, summary_ops, file_writers, pcreader, max_eval_objects=500):
+    """ ops: dict mapping from string to tf ops """
+    is_training = False
+
+    log_string(str(datetime.now()))
+    losses = []
+    losses_dir = []
+    losses_ce = []
+    losses_off = []
+    losses_app = []
+    losses_add = []
+    losses_add_gt2pred = []
+
+    # resets accumulation of pr and auc data
+    sess.run(summary_ops['pr_reset_op'])
+    
+    for batch_idx in np.arange(min(pcreader._num_test_samples, max_eval_objects)):
+
+        batch_data, cam_poses, obj_idx = pcreader.get_batch(obj_idx=pcreader._num_train_samples + batch_idx)
+
+        if 'train_on_scenes' in global_config['DATA'] and global_config['DATA']['train_on_scenes']:
+            # OpenCV OpenGL conversion
+            cam_poses, batch_data = center_pc_convert_cam(cam_poses, batch_data)
+        # Augment batched point clouds by rotation and jittering
+        # aug_data = provider.random_scale_point_cloud(batch_data)
+        # batch_data[:,:,0:3] = provider.jitter_point_cloud(batch_data[:,:,0:3])
+        feed_dict = {ops['pointclouds_pl']: batch_data,
+                     ops['cam_poses_pl']: cam_poses,
+                     ops['obj_idx_pl']: obj_idx,
+                     ops['is_training_pl']: is_training}
+
+        scene_idx, step, loss_val, dir_loss, bin_ce_loss, offset_loss, approach_loss, adds_loss, adds_gt2pred_loss, pr_summary,_,_,_ = sess.run([ops['scene_idx'], ops['step'], ops['loss'], ops['dir_loss'], ops['bin_ce_loss'],
+                                                                                                        ops['offset_loss'], ops['approach_loss'], ops['adds_loss'], ops['adds_gt2pred_loss'],
+                                                                                                        summary_ops['merged_eval'], summary_ops['pr_update_op'], 
+                                                                                                        summary_ops['auc_update_op']] + [summary_ops['acc_update_ops']], feed_dict=feed_dict)
+        assert scene_idx[0] == (pcreader._num_train_samples + batch_idx)
+        losses.append(loss_val)
+        losses_dir.append(dir_loss)
+        losses_ce.append(bin_ce_loss)
+        losses_off.append(offset_loss)
+        losses_app.append(approach_loss)
+        losses_add.append(adds_loss)
+        losses_add_gt2pred.append(adds_gt2pred_loss)
+
+    loss_mean = np.mean(losses)
+    losses_dir_mean = np.mean(losses_dir)
+    loss_ce_mean = np.mean(losses_ce)
+    loss_off_mean = np.mean(losses_off)
+    loss_app_mean = np.mean(losses_app)
+    loss_add_mean = np.mean(losses_add)
+    loss_add_gt2pred_mean = np.mean(losses_add_gt2pred)
+
+    file_writers['test_writer'].add_summary(pr_summary, step)
+
+    log_string('mean val loss: %f \t mean val dir loss: %f \t mean val ce loss: %f \t mean off loss: %f \t mean app loss: %f \t mean adds loss: %f \t mean adds_gt2pred loss:  %f' % (loss_mean, losses_dir_mean, loss_ce_mean, loss_off_mean, loss_app_mean, loss_add_mean, loss_add_gt2pred_mean))
+
+    return step
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
+    parser.add_argument('--log_dir', default='/result', help='Log dir [default: log]')
+    parser.add_argument('--data_path', type=str, default=None, help='internal grasp root dir')
+    parser.add_argument('--max_epoch', type=int, default=None, help='Epoch to run [default: 201]')
+    parser.add_argument('--batch_size', type=int, default=None, help='Batch Size during training [default: 32]')
+    parser.add_argument('--classes', nargs="*", type=str, default=None, help='train or test classes')
+    parser.add_argument('--arg_configs', nargs="*", type=str, default=[], help='overwrite config parameters')
+    FLAGS = parser.parse_args()
+
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu)
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+    if not os.path.exists(FLAGS.log_dir): 
+        os.makedirs(FLAGS.log_dir)
+
+    os.system('cp pointnet2_grasp_direct.py %s' % (FLAGS.log_dir)) # bkp of model def
+    os.system('cp train_grasp_direct.py %s' % (FLAGS.log_dir)) # bkp of train procedure
+
+    LOG_FOUT = open(os.path.join(FLAGS.log_dir, 'log_train.txt'), 'w')
+    LOG_FOUT.write(str(FLAGS)+'\n')
+    def log_string(out_str):
+        LOG_FOUT.write(out_str+'\n')
+        LOG_FOUT.flush()
+        print(out_str)
+
+    global_config = utilities.load_config(FLAGS.log_dir, batch_size=FLAGS.batch_size, max_epoch=FLAGS.max_epoch, data_path= FLAGS.data_path, classes=FLAGS.classes, arg_configs=FLAGS.arg_configs)
+    
+    log_string(str(global_config))
+    log_string('pid: %s'%(str(os.getpid())))
+
+    train(global_config, FLAGS.log_dir)
+
+    LOG_FOUT.close()

scene_test_2048_bs3_hor_sigma_0025/checkpoint

@@ -0,0 +1,6 @@
+model_checkpoint_path: "model.ckpt-45045"
+all_model_checkpoint_paths: "model.ckpt-108108"
+all_model_checkpoint_paths: "model.ckpt-117117"
+all_model_checkpoint_paths: "model.ckpt-126126"
+all_model_checkpoint_paths: "model.ckpt-135135"
+all_model_checkpoint_paths: "model.ckpt-144144"

scene_test_2048_bs3_hor_sigma_0025/config.yaml

@@ -0,0 +1,188 @@
+DATA:
+  gripper_width: 0.08
+  input_normals: false
+  use_uniform_quaternions: False
+  train_on_scenes: True
+  labels:
+    to_gpu: False
+    bin_weights:
+    - 0.16652107
+    - 0.21488856
+    - 0.37031708
+    - 0.55618503
+    - 0.75124664
+    - 0.93943357
+    - 1.07824539
+    - 1.19423112
+    - 1.55731375
+    - 3.17161779
+    contact_gather: knn
+    filter_z: true
+    k: 1
+    max_radius: 0.005
+    min_unique_pos_contacts: 1
+    num_neg_contacts: 0
+    num_pos_contacts: 10000
+    offset_bins:
+    - 0
+    - 0.00794435329
+    - 0.0158887021
+    - 0.0238330509
+    - 0.0317773996
+    - 0.0397217484
+    - 0.0476660972
+    - 0.055610446
+    - 0.0635547948
+    - 0.0714991435
+    - 0.08
+    z_val: -0.1
+  raw_num_points: 20000
+  ndataset_points: 20000
+  num_point: 2048
+  sigma: 0.0025
+  clip: 0.005
+  use_farthest_point: false
+  train_and_test: false
+  num_test_scenes: 1000
+  intrinsics: 'realsense'
+LOSS:
+  min_geom_loss_divisor: 1.0
+  max_geom_loss_divisor: 100.0
+  offset_loss_type: sigmoid_cross_entropy
+  too_small_offset_pred_bin_factor: 0
+  topk_confidence: 512
+MODEL:
+  bin_offsets: true
+  contact_distance_offset: true
+  dir_vec_length_offset: false
+  grasp_conf_head:
+    conv1d: 1
+    dropout_keep: 0.5
+  grasp_dir_head:
+    conv1d: 3
+    dropout_keep: 0.7
+  joint_head:
+    conv1d: 4
+    dropout_keep: 0.7
+  joint_heads: false
+  larger_model: false
+  asymmetric_model: true
+  model: contact_graspnet
+  pointnet_fp_modules:
+  - mlp:
+    - 256
+    - 256
+  - mlp:
+    - 256
+    - 128
+  - mlp:
+    - 128
+    - 128
+    - 128
+  pointnet_sa_module:
+    group_all: true
+    mlp:
+    - 256
+    - 512
+    - 1024
+  pointnet_sa_modules_msg:
+  - mlp_list:
+    - - 32
+      - 32
+      - 64
+    - - 64
+      - 64
+      - 128
+    - - 64
+      - 96
+      - 128
+    npoint: 2048
+    nsample_list:
+    - 32
+    - 64
+    - 128
+    radius_list:
+    - 0.02
+    - 0.04
+    - 0.08
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 512
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.04
+    - 0.08
+    - 0.16
+  - mlp_list:
+    - - 64
+      - 64
+      - 128
+    - - 128
+      - 128
+      - 256
+    - - 128
+      - 128
+      - 256
+    npoint: 128
+    nsample_list:
+    - 64
+    - 64
+    - 128
+    radius_list:
+    - 0.08
+    - 0.16
+    - 0.32
+  pred_contact_approach: true
+  pred_contact_base: true
+  pred_contact_offset: true
+  pred_contact_success: true
+  pred_grasps_adds: false
+  pred_grasps_adds_gt2pred: false
+OPTIMIZER:
+  adds_gt2pred_loss_weight: 1
+  adds_loss_weight: 10
+  approach_cosine_loss_weight: 1
+  batch_size: 14
+  bn_decay_clip: 0.99
+  bn_decay_decay_rate: 0.5
+  bn_decay_decay_step: 200000
+  bn_init_decay: 0.5
+  decay_rate: 0.7
+  decay_step: 200000
+  dir_cosine_loss_weight: 1
+  learning_rate: 0.001
+  max_epoch: 16
+  momentum: 0.9
+  offset_loss_weight: 1
+  optimizer: adam
+  score_ce_loss_weight: 1
+TEST:
+  center_to_tip: 0.0
+  allow_zero_margin: 0
+  bin_vals: max
+  extra_opening: 0.005
+  first_thres: 0.23
+  second_thres: 0.18
+  max_farthest_points: 150
+  num_samples: 200
+  save: false
+  scale_fac:
+  - 1.25
+  - 1.0
+  - 0.75
+  - 0.5
+  scales: false
+  with_replacement: false
+  filter_thres: 0.0001
+  

scene_test_2048_bs3_hor_sigma_0025/model.ckpt-144144.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60e005e6497eddeb241dae9c665ca9a43031d3c7e9816038b3bcde4d9619a40f
+size 27860696

scene_test_2048_bs3_hor_sigma_0025/model.ckpt-45045.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d11992fc4ba7b17891e03788669efcd4795f62f53111137e1c2099c2d8e0dbda
+size 27860696

scene_test_2048_bs3_hor_sigma_0025/model.ckpt-54054.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4954c0868a1fde7f6d9cf5da594356538bd17691f34697573701132c9dd0095
+size 27860696

scene_test_2048_bs3_hor_sigma_0025/model.ckpt-72072.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4eb1a8476d86815d6ac2aeeb9e8f0ae3dc665024b51fa46bbc6dfad006398226
+size 27860696

scene_test_2048_bs3_hor_sigma_0025/pointnet2_grasp_direct.py

@@ -0,0 +1,437 @@
+import os
+import sys
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(os.path.join(BASE_DIR))
+
+ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR))
+sys.path.append(os.path.join(BASE_DIR, 'utils'))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2', 'utils'))
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2'))
+sys.path.append(os.path.join(BASE_DIR, 'pointnet2'))
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_eager_execution()
+    TF2 = True
+except:
+    import tensorflow as tf
+    TF2 = False
+    
+import numpy as np
+import tf_util
+from pointnet_util import pointnet_sa_module, pointnet_fp_module, pointnet_sa_module_msg
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2', 'tf_ops/sampling'))
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2', 'tf_ops/grouping'))
+sys.path.append(os.path.join(ROOT_DIR, 'pointnet2', 'tf_ops/3d_interpolation'))
+from tf_sampling import farthest_point_sample, gather_point
+from tf_grouping import query_ball_point, group_point, knn_point
+from tf_interpolate import three_nn, three_interpolate
+import tf_utils
+
+def placeholder_inputs(global_config):
+    batch_size = global_config['OPTIMIZER']['batch_size']
+    if 'raw_num_points' in global_config['DATA']:
+        num_point = global_config['DATA']['raw_num_points']
+    else:
+        num_point = global_config['DATA']['num_point']
+    
+    input_normals = global_config['DATA']['input_normals']
+
+    pl_dict = {}
+    dim = 6 if input_normals else 3
+    pl_dict['pointclouds_pl'] = tf.placeholder(tf.float32, shape=(batch_size, num_point, dim))
+    pl_dict['obj_idx_pl'] = tf.placeholder(tf.int32, ())
+    pl_dict['cam_poses_pl'] = tf.placeholder(tf.float32, shape=(batch_size, 4, 4))
+    pl_dict['is_training_pl'] = tf.placeholder(tf.bool, shape=())
+
+    return pl_dict #, labels_pl
+
+def get_bin_vals(global_config):
+    bins_bounds = np.array(global_config['DATA']['labels']['offset_bins'])
+    if global_config['TEST']['bin_vals'] == 'max':
+        bin_vals = (bins_bounds[1:] + bins_bounds[:-1])/2
+        bin_vals[-1] = bins_bounds[-1]
+    elif global_config['TEST']['bin_vals'] == 'mean':
+        bin_vals = bins_bounds[1:]
+    else:
+        raise NotImplementedError
+
+    if not global_config['TEST']['allow_zero_margin']:
+        bin_vals = np.minimum(bin_vals, global_config['DATA']['gripper_width']-global_config['TEST']['extra_opening'])
+        
+    tf_bin_vals = tf.constant(bin_vals, tf.float32)
+    return tf_bin_vals
+
+def get_model(point_cloud, is_training, global_config, bn_decay=None):
+    """ Grasp direction PointNet++ """
+
+    model_config = global_config['MODEL']
+    data_config = global_config['DATA']
+
+    radius_list_0 = model_config['pointnet_sa_modules_msg'][0]['radius_list']
+    radius_list_1 = model_config['pointnet_sa_modules_msg'][1]['radius_list']
+    nsample_list_0 = model_config['pointnet_sa_modules_msg'][0]['nsample_list']
+    nsample_list_1 = model_config['pointnet_sa_modules_msg'][1]['nsample_list']
+    mlp_list_0 = model_config['pointnet_sa_modules_msg'][0]['mlp_list']
+    mlp_list_1 = model_config['pointnet_sa_modules_msg'][1]['mlp_list']
+    npoint_0 = model_config['pointnet_sa_modules_msg'][0]['npoint']
+    npoint_1 = model_config['pointnet_sa_modules_msg'][1]['npoint']
+    fp_mlp_0 = model_config['pointnet_fp_modules'][0]['mlp']
+    fp_mlp_1 = model_config['pointnet_fp_modules'][1]['mlp']
+    fp_mlp_2 = model_config['pointnet_fp_modules'][2]['mlp']
+
+    #larger model
+    radius_list_2 = model_config['pointnet_sa_modules_msg'][2]['radius_list']
+    nsample_list_2 = model_config['pointnet_sa_modules_msg'][2]['nsample_list']
+    npoint_2 = model_config['pointnet_sa_modules_msg'][2]['npoint']
+    mlp_list_2 = model_config['pointnet_sa_modules_msg'][2]['mlp_list']
+
+    input_normals = data_config['input_normals']
+    gripper_width = data_config['gripper_width']
+    offset_bins = data_config['labels']['offset_bins']
+    joint_heads = model_config['joint_heads']
+
+    # expensive, rather use random only
+    if 'raw_num_points' in data_config and data_config['raw_num_points'] != data_config['ndataset_points']:
+        point_cloud = gather_point(point_cloud, farthest_point_sample(data_config['ndataset_points'], point_cloud))
+
+    batch_size = point_cloud.get_shape()[0]
+    num_point = point_cloud.get_shape()[1]
+    end_points = {}
+    l0_xyz = tf.slice(point_cloud, [0,0,0], [-1,-1,3])
+    l0_points = tf.slice(point_cloud, [0,0,3], [-1,-1,3]) if input_normals else None 
+
+    l_xyz = [l0_xyz]
+    l_points = [l0_points]
+    pred_points = l0_xyz
+
+    # Set abstraction layers
+    l1_xyz, l1_points = pointnet_sa_module_msg(l0_xyz, l0_points, npoint_0, radius_list_0, nsample_list_0, mlp_list_0, is_training, bn_decay, scope='layer1')
+    l2_xyz, l2_points = pointnet_sa_module_msg(l1_xyz, l1_points, npoint_1, radius_list_1, nsample_list_1,mlp_list_1, is_training, bn_decay, scope='layer2')
+
+    # large model
+    if 'larger_model' in model_config and model_config['larger_model']:
+        fp_mlp_3 = model_config['pointnet_fp_modules'][3]['mlp']
+
+        l3_xyz, l3_points = pointnet_sa_module_msg(l2_xyz, l2_points, npoint_2, radius_list_2, nsample_list_2,mlp_list_2, is_training, bn_decay, scope='layer3')
+        l4_xyz, l4_points, _ = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer4')
+
+        # Feature Propagation layers
+        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+        l0_points = tf.concat([l0_xyz, l0_points],axis=-1) if input_normals else l0_xyz 
+        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, fp_mlp_3, is_training, bn_decay, scope='fa_layer4')
+    elif 'asymmetric_model' in model_config and model_config['asymmetric_model']:
+        l3_xyz, l3_points = pointnet_sa_module_msg(l2_xyz, l2_points, npoint_2, radius_list_2, nsample_list_2,mlp_list_2, is_training, bn_decay, scope='layer3')
+        l4_xyz, l4_points, _ = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer4')
+
+        # Feature Propagation layers
+        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+        l0_points = l1_points
+        pred_points = l1_xyz
+    else:
+        l3_xyz, l3_points, _ = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=model_config['pointnet_sa_module']['mlp'], mlp2=None, group_all=model_config['pointnet_sa_module']['group_all'], is_training=is_training, bn_decay=bn_decay, scope='layer3')
+
+        # Feature Propagation layers
+        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, fp_mlp_0, is_training, bn_decay, scope='fa_layer1')
+        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, fp_mlp_1, is_training, bn_decay, scope='fa_layer2')
+
+        l0_points = tf.concat([l0_xyz, l0_points],axis=-1) if input_normals else l0_xyz 
+        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, fp_mlp_2, is_training, bn_decay, scope='fa_layer3')
+
+    if joint_heads:
+        head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
+        # end_points['feats'] = head 
+        head = tf_util.dropout(head, keep_prob=0.7, is_training=is_training, scope='dp1')
+        head = tf_util.conv1d(head, 4, 1, padding='VALID', activation_fn=None, scope='fc2')
+        grasp_dir_head = tf.slice(head, [0,0,0], [-1,-1,3])
+        grasp_dir_head = tf.math.l2_normalize(grasp_dir_head, axis=2)
+        binary_seg_head = tf.slice(head, [0,0,3], [-1,-1,1])
+    else:
+        # FC layers for grasp direction
+        grasp_dir_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
+        # end_points['feats'] = grasp_dir_head 
+        grasp_dir_head = tf_util.dropout(grasp_dir_head, keep_prob=0.7, is_training=is_training, scope='dp1')
+        grasp_dir_head = tf_util.conv1d(grasp_dir_head, 3, 1, padding='VALID', activation_fn=None, scope='fc3')
+        grasp_dir_head_normed = tf.math.l2_normalize(grasp_dir_head, axis=2)
+
+
+        # FC layers for grasp direction
+        approach_dir_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_app', bn_decay=bn_decay)
+        # end_points['feats'] = approach_dir_head 
+        approach_dir_head = tf_util.dropout(approach_dir_head, keep_prob=0.7, is_training=is_training, scope='dp1_app')
+        approach_dir_head = tf_util.conv1d(approach_dir_head, 3, 1, padding='VALID', activation_fn=None, scope='fc3_app')
+        approach_dir_head_orthog = tf.math.l2_normalize(approach_dir_head - tf.reduce_sum(tf.multiply(grasp_dir_head_normed, approach_dir_head), axis=2, keepdims=True)*grasp_dir_head_normed, axis=2)
+        
+        if model_config['dir_vec_length_offset']:
+            grasp_offset_head = tf.norm(grasp_dir_head, axis=2, keepdims=True)
+        elif model_config['bin_offsets']:
+            grasp_offset_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_off', bn_decay=bn_decay)
+            # end_points['feats'] = grasp_offset_head 
+            # grasp_offset_head = tf_util.dropout(grasp_offset_head, keep_prob=0.7, is_training=is_training, scope='dp1_off')
+            grasp_offset_head = tf_util.conv1d(grasp_offset_head, len(offset_bins)-1, 1, padding='VALID', activation_fn=None, scope='fc2_off')
+        else:
+            grasp_offset_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_off', bn_decay=bn_decay)
+            # end_points['feats'] = grasp_offset_head 
+            grasp_offset_head = tf_util.dropout(grasp_offset_head, keep_prob=0.7, is_training=is_training, scope='dp1_off')
+            grasp_offset_head = tf_util.conv1d(grasp_offset_head, 1, 1, padding='VALID', activation_fn=None, scope='fc2_off')
+
+        binary_seg_head = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1_seg', bn_decay=bn_decay)
+        binary_seg_head = tf_util.dropout(binary_seg_head, keep_prob=0.5, is_training=is_training, scope='dp1_seg')
+        binary_seg_head = tf_util.conv1d(binary_seg_head, 1, 1, padding='VALID', activation_fn=None, scope='fc2_seg')
+
+    end_points['grasp_dir_head'] = grasp_dir_head_normed
+    end_points['binary_seg_head'] = binary_seg_head
+    end_points['binary_seg_pred'] = tf.math.sigmoid(binary_seg_head)
+    end_points['grasp_offset_head'] = grasp_offset_head
+    end_points['grasp_offset_pred'] = tf.math.sigmoid(grasp_offset_head) if model_config['bin_offsets'] else grasp_offset_head
+    end_points['approach_dir_head'] = approach_dir_head_orthog
+    end_points['pred_points'] = pred_points
+
+    return end_points
+
+def build_6d_grasp(approach_dirs, base_dirs, contact_pts, thickness, use_tf=False, gripper_depth = 0.1034):
+    
+    if use_tf:
+        grasps_R = tf.stack([base_dirs, tf.linalg.cross(approach_dirs,base_dirs),approach_dirs], axis=3)        
+        grasps_t = contact_pts + tf.expand_dims(thickness,2)/2 * base_dirs - gripper_depth * approach_dirs
+        ones = tf.ones((contact_pts.shape[0], contact_pts.shape[1], 1, 4), dtype=tf.float32)
+        grasps = tf.concat([tf.concat([grasps_R,  tf.expand_dims(grasps_t, 3)], axis=3), ones], axis=2)
+    else:
+        grasps = []
+        for i in range(len(contact_pts)):
+            grasp = np.eye(4)
+
+            grasp[:3,0] = base_dirs[i] / np.linalg.norm(base_dirs[i])
+            grasp[:3,2] = approach_dirs[i] / np.linalg.norm(approach_dirs[i])
+            grasp_y = np.cross( grasp[:3,2],grasp[:3,0])
+            grasp[:3,1] = grasp_y / np.linalg.norm(grasp_y)
+            # base_gripper xyz = contact + thickness / 2 * baseline_dir - gripper_d * approach_dir 
+            grasp[:3,3] = contact_pts[i] + thickness[i] / 2 * grasp[:3,0] - gripper_depth * grasp[:3,2]
+            # grasp[0,3] = finger_width
+            grasps.append(grasp)
+        grasps = np.array(grasps)
+
+    return grasps
+
+def multi_bin_labels(cont_labels, bin_boundaries):
+
+    # multi_hot_labels = tf.zeros_like((cont_labels.shape[0], cont_labels.shape[1], len(bin_boundaries)-1), tf.bool)
+    bins = []
+    for b in range(len(bin_boundaries)-1):
+        bins.append(tf.math.logical_and(tf.greater_equal(cont_labels, bin_boundaries[b]), tf.less(cont_labels,bin_boundaries[b+1])))
+        # cont_labels = tf.where(condition, b*tf.ones_like(condition, tf.float32), cont_labels)
+    multi_hot_labels = tf.concat(bins, axis=2)
+    multi_hot_labels = tf.cast(multi_hot_labels, tf.float32)
+
+    return multi_hot_labels
+
+
+def get_losses(pointclouds_pl, end_points, dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam, global_config):
+    """ pred: (batch_size, num_point, 3),
+        label: (batch_size, num_point, 3), 
+        grasp_success_labels_pc: (batch_size, num_point)
+        """
+
+    grasp_dir_head = end_points['grasp_dir_head']
+    grasp_offset_head = end_points['grasp_offset_head']
+    approach_dir_head = end_points['approach_dir_head']
+
+    bin_weights = global_config['DATA']['labels']['bin_weights']
+    tf_bin_weights = tf.constant(bin_weights)
+    
+    min_geom_loss_divisor = tf.constant(float(global_config['LOSS']['min_geom_loss_divisor'])) if 'min_geom_loss_divisor' in global_config['LOSS'] else tf.constant(1.)
+    pos_grasps_in_view = tf.math.maximum(tf.reduce_sum(grasp_success_labels_pc, axis=1), min_geom_loss_divisor)   
+
+    ### ADS Gripper PC Loss
+    if global_config['MODEL']['bin_offsets']:
+        thickness_pred = tf.gather_nd(get_bin_vals(global_config), tf.expand_dims(tf.argmax(grasp_offset_head, axis=2), axis=2))
+        thickness_gt = tf.gather_nd(get_bin_vals(global_config), tf.expand_dims(tf.argmax(offset_labels_pc, axis=2), axis=2))
+    else:
+        thickness_pred = grasp_offset_head[:,:,0]
+        thickness_gt = offset_labels_pc[:,:,0]
+    pred_grasps = build_6d_grasp(approach_dir_head, grasp_dir_head, pointclouds_pl, thickness_pred, use_tf=True) # b x num_point x 4 x 4
+    gt_grasps_proj = build_6d_grasp(approach_labels_pc_cam, dir_labels_pc_cam, pointclouds_pl, thickness_gt, use_tf=True) # b x num_point x 4 x 4
+    pos_gt_grasps_proj = tf.where(tf.broadcast_to(tf.expand_dims(tf.expand_dims(tf.cast(grasp_success_labels_pc, tf.bool),2),3), gt_grasps_proj.shape), gt_grasps_proj, tf.ones_like(gt_grasps_proj)*100000)
+    # pos_gt_grasps_proj = tf.reshape(pos_gt_grasps_proj, (global_config['OPTIMIZER']['batch_size'], -1, 4, 4)) 
+
+    gripper_control_points = tf_utils.get_control_point_tensor(global_config['OPTIMIZER']['batch_size'])[:,1:,:] # b x 5 x 3
+    sym_gripper_control_points = tf_utils.get_control_point_tensor(global_config['OPTIMIZER']['batch_size'], symmetric=True)[:,1:,:]
+
+    gripper_control_points_homog =  tf.concat([gripper_control_points, tf.ones((global_config['OPTIMIZER']['batch_size'], gripper_control_points.shape[1], 1))], axis=2)  # b x 5 x 4
+    sym_gripper_control_points_homog =  tf.concat([sym_gripper_control_points, tf.ones((global_config['OPTIMIZER']['batch_size'], gripper_control_points.shape[1], 1))], axis=2)  # b x 5 x 4
+    
+    # only use per point pred grasps but not per point gt grasps
+    control_points = tf.keras.backend.repeat_elements(tf.expand_dims(gripper_control_points_homog,1), gt_grasps_proj.shape[1], axis=1)  # b x num_point x 5 x 4
+    sym_control_points = tf.keras.backend.repeat_elements(tf.expand_dims(sym_gripper_control_points_homog,1), gt_grasps_proj.shape[1], axis=1)  # b x num_point x 5 x 4
+    pred_control_points = tf.matmul(control_points, tf.transpose(pred_grasps, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_point x 5 x 3
+
+    ### Pred Grasp to GT Grasp ADD-S Loss
+    gt_control_points = tf.matmul(control_points, tf.transpose(pos_gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+    sym_gt_control_points = tf.matmul(sym_control_points, tf.transpose(pos_gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+
+    squared_add = tf.reduce_sum((tf.expand_dims(pred_control_points,2)-tf.expand_dims(gt_control_points,1))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+    sym_squared_add = tf.reduce_sum((tf.expand_dims(pred_control_points,2)-tf.expand_dims(sym_gt_control_points,1))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+
+    # symmetric ADD-S
+    neg_squared_adds = -tf.concat([squared_add,sym_squared_add], axis=2) # b x num_point x 2num_pos_grasp_point
+    neg_squared_adds_k = tf.math.top_k(neg_squared_adds, k=1, sorted=False)[0] # b x num_point
+    # If any pos grasp exists
+    min_adds = tf.minimum(tf.reduce_sum(grasp_success_labels_pc, axis=1, keepdims=True), tf.ones_like(neg_squared_adds_k[:,:,0])) * tf.sqrt(-neg_squared_adds_k[:,:,0])#tf.minimum(tf.sqrt(-neg_squared_adds_k), tf.ones_like(neg_squared_adds_k)) # b x num_point
+    adds_loss = tf.reduce_mean(end_points['binary_seg_pred'][:,:,0] * min_adds)
+
+    ### GT Grasp to pred Grasp ADD-S Loss
+    gt_control_points = tf.matmul(control_points, tf.transpose(gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+    sym_gt_control_points = tf.matmul(sym_control_points, tf.transpose(gt_grasps_proj, perm=[0, 1, 3, 2]))[:,:,:,:3] #  b x num_pos_grasp_point x 5 x 3
+
+    neg_squared_adds = -tf.reduce_sum((tf.expand_dims(pred_control_points,1)-tf.expand_dims(gt_control_points,2))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+    neg_squared_adds_sym = -tf.reduce_sum((tf.expand_dims(pred_control_points,1)-tf.expand_dims(sym_gt_control_points,2))**2, axis=(3,4)) # b x num_point x num_pos_grasp_point x ( 5 x 3)
+
+    neg_squared_adds_k_gt2pred, pred_grasp_idcs = tf.math.top_k(neg_squared_adds, k=1, sorted=False) # b x num_pos_grasp_point
+    neg_squared_adds_k_sym_gt2pred, pred_grasp_sym_idcs = tf.math.top_k(neg_squared_adds_sym, k=1, sorted=False) # b x num_pos_grasp_point
+    pred_grasp_idcs_joined = tf.where(neg_squared_adds_k_gt2pred<neg_squared_adds_k_sym_gt2pred, pred_grasp_sym_idcs, pred_grasp_idcs)
+    min_adds_gt2pred = tf.minimum(-neg_squared_adds_k_gt2pred, -neg_squared_adds_k_sym_gt2pred) # b x num_pos_grasp_point x 1
+    # min_adds_gt2pred = tf.math.exp(-min_adds_gt2pred)
+    masked_min_adds_gt2pred = tf.multiply(min_adds_gt2pred[:,:,0], grasp_success_labels_pc)
+    
+
+    batch_idcs = tf.meshgrid(tf.range(pred_grasp_idcs_joined.shape[1]), tf.range(pred_grasp_idcs_joined.shape[0]))
+    gather_idcs = tf.stack((batch_idcs[1],pred_grasp_idcs_joined[:,:,0]), axis=2)
+    nearest_pred_grasp_confidence = tf.gather_nd(end_points['binary_seg_pred'][:,:,0], gather_idcs)
+    adds_loss_gt2pred = tf.reduce_mean(tf.reduce_sum(nearest_pred_grasp_confidence*masked_min_adds_gt2pred, axis=1) / pos_grasps_in_view)
+ 
+    ### Grasp baseline Loss
+    cosine_distance = tf.constant(1.)-tf.reduce_sum(tf.multiply(dir_labels_pc_cam, grasp_dir_head),axis=2)
+    # only pass loss where we have labeled contacts near pc points 
+    masked_cosine_loss = tf.multiply(cosine_distance, grasp_success_labels_pc)
+    dir_cosine_loss = tf.reduce_mean(tf.reduce_sum(masked_cosine_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Approach Loss
+    approach_labels_orthog = tf.math.l2_normalize(approach_labels_pc_cam - tf.reduce_sum(tf.multiply(grasp_dir_head, approach_labels_pc_cam), axis=2, keepdims=True)*grasp_dir_head, axis=2)
+    cosine_distance_approach = tf.constant(1.)-tf.reduce_sum(tf.multiply(approach_labels_orthog, approach_dir_head), axis=2)
+    masked_approach_loss = tf.multiply(cosine_distance_approach, grasp_success_labels_pc)
+    approach_cosine_loss = tf.reduce_mean(tf.reduce_sum(masked_approach_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Offset/Thickness Loss
+    if global_config['MODEL']['bin_offsets']:
+        if global_config['LOSS']['offset_loss_type'] == 'softmax_cross_entropy':
+            offset_loss = tf.losses.softmax_cross_entropy(offset_labels_pc, grasp_offset_head)
+        else:
+            offset_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=offset_labels_pc, logits=grasp_offset_head)
+            
+            if 'too_small_offset_pred_bin_factor' in global_config['LOSS'] and global_config['LOSS']['too_small_offset_pred_bin_factor']:
+                too_small_offset_pred_bin_factor = tf.constant(global_config['LOSS']['too_small_offset_pred_bin_factor'], tf.float32)
+                collision_weight = tf.math.cumsum(offset_labels_pc, axis=2, reverse=True)*too_small_offset_pred_bin_factor + tf.constant(1.)
+                offset_loss = tf.multiply(collision_weight, offset_loss)
+
+            offset_loss = tf.reduce_mean(tf.multiply(tf.reshape(tf_bin_weights,(1,1,-1)), offset_loss),axis=2)
+    else:
+        offset_loss = (grasp_offset_head[:,:,0] - offset_labels_pc[:,:,0])**2
+    masked_offset_loss = tf.multiply(offset_loss, grasp_success_labels_pc)        
+    offset_loss = tf.reduce_mean(tf.reduce_sum(masked_offset_loss, axis=1) / pos_grasps_in_view)
+
+    ### Grasp Confidence Loss
+    bin_ce_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.expand_dims(grasp_success_labels_pc,axis=2), logits=end_points['binary_seg_head'])
+    if 'topk_confidence' in global_config['LOSS'] and global_config['LOSS']['topk_confidence']:
+        bin_ce_loss,_ = tf.math.top_k(tf.squeeze(bin_ce_loss), k=global_config['LOSS']['topk_confidence'])
+    bin_ce_loss = tf.reduce_mean(bin_ce_loss)
+
+    return dir_cosine_loss, bin_ce_loss, offset_loss, approach_cosine_loss, adds_loss, adds_loss_gt2pred, gt_control_points, pred_control_points, pos_grasps_in_view
+
+def compute_labels(pos_contact_pts_mesh, pos_contact_dirs_mesh, pos_contact_offsets, pos_contact_approaches_mesh, pos_finger_diffs, neg_contact_pts_mesh, 
+                  neg_contact_dirs_mesh, neg_contact_offsets, neg_contact_approaches_mesh, neg_finger_diffs, pc_cam_pl, camera_pose_pl, global_config):
+    """
+    for every point find nearest contact points
+    group and average labels
+    """
+    label_config = global_config['DATA']['labels']
+    model_config = global_config['MODEL']
+
+    contact_distance_offset = model_config['contact_distance_offset']
+    bin_offsets = model_config['bin_offsets']
+
+    mode = label_config['contact_gather']
+    nsample = label_config['k']
+    radius = label_config['max_radius']
+    filter_z = label_config['filter_z']
+    z_val = label_config['z_val']
+
+    xyz_cam = pc_cam_pl[:,:,:3]
+    pad_homog = tf.ones((xyz_cam.shape[0],xyz_cam.shape[1], 1)) 
+    pc_mesh = tf.matmul(tf.concat([xyz_cam, pad_homog], 2), tf.transpose(tf.linalg.inv(camera_pose_pl),perm=[0, 2, 1]))[:,:,:3]
+
+    if contact_distance_offset:
+        # finger_diff = tf.norm(pos_contact_pts_mesh[1::2,:]-pos_contact_pts_mesh[::2,:] + 0.0000001, axis=1, keepdims=True)
+        # finger_diff = tf.keras.backend.repeat_elements(finger_diff, 2, axis=0)
+        contact_point_offsets_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_finger_diffs,0), pc_mesh.shape[0], axis=0)
+    else:
+        contact_point_offsets_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_offsets,0), pc_mesh.shape[0], axis=0)
+
+
+    pad_homog2 = tf.ones((pc_mesh.shape[0], pos_contact_dirs_mesh.shape[0], 1)) 
+    contact_point_dirs_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0)
+    contact_point_dirs_batch_cam = tf.matmul(contact_point_dirs_batch, tf.transpose(camera_pose_pl[:,:3,:3], perm=[0, 2, 1]))[:,:,:3]
+
+    pos_contact_approaches_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_approaches_mesh,0), pc_mesh.shape[0], axis=0)
+    pos_contact_approaches_batch_cam = tf.matmul(pos_contact_approaches_batch, tf.transpose(camera_pose_pl[:,:3,:3], perm=[0, 2, 1]))[:,:,:3]
+    
+    contact_point_batch_mesh = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_pts_mesh,0), pc_mesh.shape[0], axis=0)
+    contact_point_batch_cam = tf.matmul(tf.concat([contact_point_batch_mesh, pad_homog2], 2), tf.transpose(camera_pose_pl, perm=[0, 2, 1]))[:,:,:3]
+
+
+    # orthogonal = tf.reduce_sum(pos_contact_approaches_batch_cam * contact_point_dirs_batch_cam, axis=2)
+    if filter_z:
+        dir_filter_passed = tf.keras.backend.repeat_elements(tf.math.greater(contact_point_dirs_batch_cam[:,:,2:3], tf.constant([z_val])), 3, axis=2)
+        contact_point_batch_mesh = tf.where(dir_filter_passed, contact_point_batch_mesh, tf.ones_like(contact_point_batch_mesh)*100000)
+        # contact_point_dirs_batch_cam = tf.boolean_mask(contact_point_dirs_batch_cam, dir_filter_passed)
+        # print(contact_point_dirs_batch_cam.shape)
+        # contact_point_batch_mesh = tf.boolean_mask(contact_point_batch_mesh, dir_filter_passed)
+    # else:
+        # contact_point_dirs_batch = tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0)
+
+    # TODO: Divide into positive and negative grasps; for now only insert positive contacts (otherwise mean direction does not make sense)
+    if mode == 'mean':
+        close_contact_pt_idcs,_ = query_ball_point(radius, nsample, contact_point_batch_mesh, pc_mesh) # (batch_size, npoint, nsample)
+
+        grasp_success_labels_pc = tf.cast(tf.reduce_sum(close_contact_pt_idcs, axis=2)>0, tf.float32) # (batch_size, num_point)
+
+        # contact_point to contact_dir mapping
+        grouped_dirs_mesh = group_point(contact_point_dirs_batch, close_contact_pt_idcs) # (batch_size, num_point, nsample, 3)
+        
+        # grouped_dirs_mesh = tf.gather_nd(tf.keras.backend.repeat_elements(tf.expand_dims(pos_contact_dirs_mesh,0), pc_mesh.shape[0], axis=0), close_contact_pt_idcs)
+    elif mode == 'knn':
+        # dists, close_contact_pt_idcs = knn_point(nsample, contact_point_batch_mesh, pc_mesh) # (batch_size, num_point, nsample)
+
+        # squared_dists_all = tf.reduce_sum((tf.expand_dims(contact_point_batch_mesh,1)-tf.expand_dims(pc_mesh,2))**2,axis=3)
+        squared_dists_all = tf.reduce_sum((tf.expand_dims(contact_point_batch_cam,1)-tf.expand_dims(xyz_cam,2))**2,axis=3)
+        neg_squared_dists_k, close_contact_pt_idcs = tf.math.top_k(-squared_dists_all, k=nsample, sorted=False)
+        squared_dists_k = -neg_squared_dists_k
+
+        # Distance check, careful it depends on point density
+        grasp_success_labels_pc = tf.cast(tf.less(tf.reduce_mean(squared_dists_k, axis=2), radius*radius), tf.float32) # (batch_size, num_point)
+
+        grouped_dirs_pc_cam = group_point(contact_point_dirs_batch_cam, close_contact_pt_idcs)
+        grouped_approaches_pc_cam = group_point(pos_contact_approaches_batch_cam, close_contact_pt_idcs)
+        grouped_offsets = group_point(tf.expand_dims(contact_point_offsets_batch,2), close_contact_pt_idcs)
+        # not sure for nsample > 1
+        dir_labels_pc_cam = tf.math.l2_normalize(tf.reduce_mean(grouped_dirs_pc_cam, axis=2),axis=2) # (batch_size, num_point, 3)
+        approach_labels_pc_cam = tf.math.l2_normalize(tf.reduce_mean(grouped_approaches_pc_cam, axis=2),axis=2) # (batch_size, num_point, 3)
+        offset_labels_pc = tf.reduce_mean(grouped_offsets, axis=2)
+        
+    else:
+        raise NotImplementedError
+        # if filter_z:
+        #     pad_homog3 = tf.ones((dir_labels_mesh.shape[0], dir_labels_mesh.shape[1], 1)) 
+        #     dir_labels_pc_cam = tf.matmul(tf.concat([dir_labels_mesh,pad_homog3], 2), tf.transpose(camera_pose_pl, perm=[0, 2, 1]))[:,:,:3] # (batch_size, num_point, 3)
+
+    return dir_labels_pc_cam, offset_labels_pc, grasp_success_labels_pc, approach_labels_pc_cam
+
+if __name__=='__main__':
+    with tf.Graph().as_default():
+        inputs = tf.zeros((32,2048,6))
+        net, _ = get_model(inputs, tf.constant(True))
+        print(net)

scene_test_2048_bs3_hor_sigma_0025/train_grasp_direct.py

@@ -0,0 +1,292 @@
+import os
+import sys
+import argparse
+import math
+from datetime import datetime
+import numpy as np
+import socket
+import importlib
+import time
+from tqdm import tqdm
+import glob
+import json
+import yaml
+
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR))
+sys.path.append(os.path.join(BASE_DIR))
+sys.path.append(os.path.join(ROOT_DIR))
+
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_eager_execution()
+    TF2 = True
+    physical_devices = tf.config.experimental.list_physical_devices('GPU')
+    print(physical_devices)
+    tf.config.experimental.set_memory_growth(physical_devices[0], True)
+except:
+    import tensorflow as tf
+    TF2 = False
+
+POINT_DIR_NGC = os.path.join(ROOT_DIR, 'pointnet2')
+if os.path.exists(POINT_DIR_NGC):
+    sys.path.append(os.path.join(POINT_DIR_NGC, 'models'))
+    sys.path.append(os.path.join(POINT_DIR_NGC, 'utils'))
+else:
+    sys.path.append(os.path.join(BASE_DIR, 'pointnet2',  'models'))
+    sys.path.append(os.path.join(BASE_DIR, 'pointnet2',  'utils'))
+    
+import provider
+import sample
+import utilities
+from data import PointCloudReader, load_mesh_path_and_scale, preprocess_pc_for_inference, load_filtered_contact_data, load_scene_contacts, load_obj_scales_cats, inverse_transform, center_pc_convert_cam
+from summaries import top_grasp_acc_summaries, build_summary_ops, build_file_writers
+from tf_train_ops import load_labels_and_losses, build_train_op
+from surface_grasp_estimator import GraspEstimator
+
+
+def train(global_config, LOG_DIR):
+
+    if 'train_on_scenes' in global_config['DATA'] and global_config['DATA']['train_on_scenes']:
+        mesh_scales, mesh_cats = load_obj_scales_cats(global_config['DATA']['data_path'])
+        contact_infos, scene_obj_paths, scene_obj_transforms = load_scene_contacts(global_config['DATA']['data_path'])
+        num_train_samples = len(contact_infos)-global_config['DATA']['num_test_scenes']
+        num_test_samples = global_config['DATA']['num_test_scenes']
+    else:
+        scene_obj_paths, scene_obj_transforms = None, None
+        train_contact_paths, test_contact_paths, contact_infos, mesh_scales = load_filtered_contact_data(global_config['DATA']['data_path'], min_pos_contacts=1, classes=global_config['DATA']['classes'])
+        num_train_samples = len(train_contact_paths)
+        num_test_samples = len(test_contact_paths)
+        
+    print('using %s meshes' % (num_train_samples + num_test_samples))
+    
+    if 'train_and_test' in global_config['DATA'] and global_config['DATA']['train_and_test']:
+        num_train_samples = num_train_samples + num_test_samples
+        num_test_samples = 0
+        print('using train and test data')
+
+    pcreader = PointCloudReader(
+        root_folder=global_config['DATA']['data_path'],
+        batch_size=global_config['OPTIMIZER']['batch_size'],
+        num_grasp_clusters=None,
+        estimate_normals=global_config['DATA']['input_normals'],
+        npoints=global_config['DATA']['num_point'],
+        raw_num_points=global_config['DATA']['raw_num_points'],
+        use_uniform_quaternions = global_config['DATA']['use_uniform_quaternions'],
+        run_in_another_process = False,
+        mesh_scales = mesh_scales,
+        scene_obj_paths = scene_obj_paths,
+        scene_obj_transforms = scene_obj_transforms,
+        num_train_samples = num_train_samples,
+        num_test_samples = num_test_samples,
+        use_farthest_point = global_config['DATA']['use_farthest_point'],
+        intrinsics=global_config['DATA']['intrinsics']
+    )
+
+    with tf.Graph().as_default():
+        
+        # Build the model
+        grasp_estimator = GraspEstimator(global_config)
+
+        ops = grasp_estimator.build_network()
+        
+        # contact_tensors = load_contact_grasps(contact_infos, global_config['DATA'])
+        
+        loss_ops = load_labels_and_losses(grasp_estimator, contact_infos, global_config)
+
+        ops.update(loss_ops)
+        ops['train_op'] = build_train_op(ops['loss'], ops['step'], global_config)
+
+        # Add ops to save and restore all the variables.
+        saver = tf.train.Saver(save_relative_paths=True, keep_checkpoint_every_n_hours=4)
+
+        # Create a session
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        config.allow_soft_placement = True
+        # config.log_device_placement = False
+        sess = tf.Session(config=config)
+
+
+        summary_ops = build_summary_ops(ops, sess, global_config)
+
+        grasp_estimator.load_weights(sess, saver, LOG_DIR, mode='train')
+        file_writers = build_file_writers(sess, LOG_DIR)
+
+    batches_per_epoch = num_train_samples #// global_config['OPTIMIZER']['batch_size'] + 1
+    cur_epoch = sess.run(ops['step']) // (batches_per_epoch * global_config['OPTIMIZER']['batch_size'])
+    for epoch in range(cur_epoch, global_config['OPTIMIZER']['max_epoch']):
+        log_string('**** EPOCH %03d ****' % (epoch))
+        
+        sess.run(ops['iterator'].initializer)
+        epoch_time = time.time()
+        step = train_one_epoch(sess, ops, summary_ops, file_writers, pcreader)
+        print('trained %s batches in: ' % batches_per_epoch, time.time()-epoch_time)
+
+        # Save the variables to disk.
+        if (epoch+1) % 1 == 0:
+            save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt"), global_step=step, write_meta_graph=False)
+            log_string("Model saved in file: %s" % save_path)
+
+        if epoch % 1 == 0 and num_test_samples > 0:
+            eval_time = time.time()
+            eval_test_objects(sess, ops, summary_ops, file_writers, pcreader)
+            print('evaluation time: ', time.time()-eval_time)
+
+def train_one_epoch(sess, ops, summary_ops, file_writers, pcreader):
+    """ ops: dict mapping from string to tf ops """
+    is_training = True
+    
+    log_string(str(datetime.now()))
+
+    loss_sum, loss_sum_dir, loss_sum_ce, loss_sum_off, loss_sum_app, loss_sum_adds, loss_sum_adds_gt2pred, time_sum = 8 * [0]
+
+    # batches_per_epoch = pcreader._num_train_samples // pcreader._batch_size
+    ## define one epoch = all objects/scenes seen
+    batches_per_epoch = pcreader._num_train_samples
+
+    # run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, report_tensor_allocations_upon_oom = True)
+
+    for batch_idx in range(batches_per_epoch):
+        get_time = time.time()
+
+        batch_data, cam_poses, obj_idx = pcreader.get_batch(batch_idx)
+        print(time.time()- get_time)
+        if 'train_on_scenes' in global_config['DATA'] and global_config['DATA']['train_on_scenes']:
+            # OpenCV OpenGL conversion
+            cam_poses, batch_data = center_pc_convert_cam(cam_poses, batch_data)
+        print(time.time() - get_time)
+        # Augment batched point clouds by rotation and jittering
+        # aug_data = provider.random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25)
+        if 'sigma' in global_config['DATA'] and global_config['DATA']['sigma'] > 0:
+            batch_data[:,:,0:3] = provider.jitter_point_cloud(batch_data[:,:,0:3], 
+                                                            sigma=global_config['DATA']['sigma'], 
+                                                            clip=global_config['DATA']['clip']*2)
+        
+        feed_dict = {ops['pointclouds_pl']: batch_data,
+                     ops['cam_poses_pl']: cam_poses,
+                     ops['obj_idx_pl']: obj_idx,
+                    #  ops['labels_pl']: batch_label,
+                     ops['is_training_pl']: is_training}
+
+        step, summary, _, loss_val, dir_loss, bin_ce_loss, \
+        offset_loss, approach_loss, adds_loss, adds_gt2pred_loss,pos_grasps_in_view, scene_idx = sess.run([ops['step'], summary_ops['merged'], ops['train_op'], 
+                                                                            ops['loss'], ops['dir_loss'], ops['bin_ce_loss'], 
+                                                                            ops['offset_loss'], ops['approach_loss'], ops['adds_loss'], 
+                                                                            ops['adds_gt2pred_loss'], ops['pos_grasps_in_view'], ops['scene_idx']], feed_dict=feed_dict)
+        print(time.time()- get_time)
+        print(pos_grasps_in_view)
+        print(scene_idx, obj_idx)
+        assert scene_idx[0] == obj_idx
+
+        loss_sum += loss_val
+        loss_sum_dir += dir_loss
+        loss_sum_ce += bin_ce_loss
+        loss_sum_off += offset_loss
+        loss_sum_app += approach_loss
+        loss_sum_adds += adds_loss
+        loss_sum_adds_gt2pred += adds_gt2pred_loss
+        time_sum += time.time() - get_time
+        
+        if (batch_idx+1)%10 == 0:
+            file_writers['train_writer'].add_summary(summary, step)
+            log_string('total loss: %f \t dir loss: %f \t ce loss: %f \t off loss: %f \t app loss: %f adds loss: %f \t adds_gt2pred loss: %f \t batch time: %f' % (loss_sum/10,loss_sum_dir/10,loss_sum_ce/10, loss_sum_off/10, loss_sum_app/10, loss_sum_adds/10, loss_sum_adds_gt2pred/10, time_sum/10))
+            # log_string('accuracy: %f' % (total_correct / float(total_seen)))
+            loss_sum, loss_sum_dir, loss_sum_ce, loss_sum_off, loss_sum_app, loss_sum_adds, loss_sum_adds_gt2pred, time_sum = 8 * [0]
+            
+    return step
+
+def eval_test_objects(sess, ops, summary_ops, file_writers, pcreader, max_eval_objects=500):
+    """ ops: dict mapping from string to tf ops """
+    is_training = False
+
+    log_string(str(datetime.now()))
+    losses = []
+    losses_dir = []
+    losses_ce = []
+    losses_off = []
+    losses_app = []
+    losses_add = []
+    losses_add_gt2pred = []
+
+    # resets accumulation of pr and auc data
+    sess.run(summary_ops['pr_reset_op'])
+    
+    for batch_idx in np.arange(min(pcreader._num_test_samples, max_eval_objects)):
+
+        batch_data, cam_poses, obj_idx = pcreader.get_batch(obj_idx=pcreader._num_train_samples + batch_idx)
+
+        if 'train_on_scenes' in global_config['DATA'] and global_config['DATA']['train_on_scenes']:
+            # OpenCV OpenGL conversion
+            cam_poses, batch_data = center_pc_convert_cam(cam_poses, batch_data)
+        # Augment batched point clouds by rotation and jittering
+        # aug_data = provider.random_scale_point_cloud(batch_data)
+        # batch_data[:,:,0:3] = provider.jitter_point_cloud(batch_data[:,:,0:3])
+        feed_dict = {ops['pointclouds_pl']: batch_data,
+                     ops['cam_poses_pl']: cam_poses,
+                     ops['obj_idx_pl']: obj_idx,
+                     ops['is_training_pl']: is_training}
+
+        scene_idx, step, loss_val, dir_loss, bin_ce_loss, offset_loss, approach_loss, adds_loss, adds_gt2pred_loss, pr_summary,_,_,_ = sess.run([ops['scene_idx'], ops['step'], ops['loss'], ops['dir_loss'], ops['bin_ce_loss'],
+                                                                                                        ops['offset_loss'], ops['approach_loss'], ops['adds_loss'], ops['adds_gt2pred_loss'],
+                                                                                                        summary_ops['merged_eval'], summary_ops['pr_update_op'], 
+                                                                                                        summary_ops['auc_update_op']] + [summary_ops['acc_update_ops']], feed_dict=feed_dict)
+        assert scene_idx[0] == (pcreader._num_train_samples + batch_idx)
+        losses.append(loss_val)
+        losses_dir.append(dir_loss)
+        losses_ce.append(bin_ce_loss)
+        losses_off.append(offset_loss)
+        losses_app.append(approach_loss)
+        losses_add.append(adds_loss)
+        losses_add_gt2pred.append(adds_gt2pred_loss)
+
+    loss_mean = np.mean(losses)
+    losses_dir_mean = np.mean(losses_dir)
+    loss_ce_mean = np.mean(losses_ce)
+    loss_off_mean = np.mean(losses_off)
+    loss_app_mean = np.mean(losses_app)
+    loss_add_mean = np.mean(losses_add)
+    loss_add_gt2pred_mean = np.mean(losses_add_gt2pred)
+
+    file_writers['test_writer'].add_summary(pr_summary, step)
+
+    log_string('mean val loss: %f \t mean val dir loss: %f \t mean val ce loss: %f \t mean off loss: %f \t mean app loss: %f \t mean adds loss: %f \t mean adds_gt2pred loss:  %f' % (loss_mean, losses_dir_mean, loss_ce_mean, loss_off_mean, loss_app_mean, loss_add_mean, loss_add_gt2pred_mean))
+
+    return step
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
+    parser.add_argument('--log_dir', default='/result', help='Log dir [default: log]')
+    parser.add_argument('--data_path', type=str, default=None, help='internal grasp root dir')
+    parser.add_argument('--max_epoch', type=int, default=None, help='Epoch to run [default: 201]')
+    parser.add_argument('--batch_size', type=int, default=None, help='Batch Size during training [default: 32]')
+    parser.add_argument('--classes', nargs="*", type=str, default=None, help='train or test classes')
+    parser.add_argument('--arg_configs', nargs="*", type=str, default=[], help='overwrite config parameters')
+    FLAGS = parser.parse_args()
+
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu)
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+    if not os.path.exists(FLAGS.log_dir): 
+        os.makedirs(FLAGS.log_dir)
+
+    os.system('cp pointnet2_grasp_direct.py %s' % (FLAGS.log_dir)) # bkp of model def
+    os.system('cp train_grasp_direct.py %s' % (FLAGS.log_dir)) # bkp of train procedure
+
+    LOG_FOUT = open(os.path.join(FLAGS.log_dir, 'log_train.txt'), 'w')
+    LOG_FOUT.write(str(FLAGS)+'\n')
+    def log_string(out_str):
+        LOG_FOUT.write(out_str+'\n')
+        LOG_FOUT.flush()
+        print(out_str)
+
+    global_config = utilities.load_config(FLAGS.log_dir, batch_size=FLAGS.batch_size, max_epoch=FLAGS.max_epoch, data_path= FLAGS.data_path, classes=FLAGS.classes, arg_configs=FLAGS.arg_configs)
+    
+    log_string(str(global_config))
+    log_string('pid: %s'%(str(os.getpid())))
+
+    train(global_config, FLAGS.log_dir)
+
+    LOG_FOUT.close()