def fit_object_model(model_cloud, object_cloud): # generate the model kdtree mark = time() object_tree = scipy.spatial.KDTree(object_cloud) logger.info('generated object tree in {:.3f}s'.format(time() - mark)) # align the object cloud with ICP mark = time() diff_xform = icp.match(model_cloud[::100], object_tree, iterations=50, initial_threshold=0.10) model_cloud_aligned = model_cloud.dot(diff_xform[0].T) + diff_xform[1] logger.info('aligned model cloud in {:.3f}s'.format(time() - mark)) return diff_xform, model_cloud_aligned
def shelf_subtract(shelf_cloud, camera_cloud, shelf_model, bounds, downsample=100): # generate the shelf kdtree mark = time() shelf_tree = scipy.spatial.KDTree(shelf_cloud) logger.info('generated shelf tree in {:.3f}s'.format(time() - mark)) # align the camera cloud with ICP mark = time() diff_xform = icp.match(camera_cloud[::downsample], shelf_tree, initial_threshold=0.05) camera_cloud_aligned = camera_cloud.dot(diff_xform[0].T) + diff_xform[1] logger.info('aligned camera cloud in {:.3f}s'.format(time() - mark)) # remove all points near the shelf mark = time() # _, ind = shelf_tree.query(camera_cloud_aligned, distance_upper_bound=0.005) # mask = ind >= len(shelf_tree.data) # segmented_cloud = camera_cloud_aligned[mask] test_point = GeometricPrimitive() test_point.setPoint((0, 0, 0)) test_geom = Geometry3D() test_geom.setGeometricPrimitive(test_point) def collides(point): test_geom.setCurrentTransform(so3.identity(), point) return test_geom.withinDistance(shelf_model.geometry(), 0.01) mask = map(lambda p: not collides(p), camera_cloud_aligned) # segmented_cloud = [ p for (i, p) in enumerate(camera_cloud_aligned) if mask[i] ] camera_cloud_aligned_local = (camera_cloud_aligned - shelf_model.getTransform()[1]).dot(so3.matrix(shelf_model.getTransform()[0])) mask2 = reduce(numpy.bitwise_and, [ (bounds[0][i] < camera_cloud_aligned_local[:, i]) & (bounds[1][i] > camera_cloud_aligned_local[:, i]) for i in range(3) ]) segmented_cloud = camera_cloud_aligned[mask & mask2] logger.info('segmented camera cloud in {:.3f}s: {} -> {}'.format(time() - mark, len(camera_cloud_aligned), len(segmented_cloud))) return diff_xform, camera_cloud_aligned, segmented_cloud, mask & mask2