Created
January 31, 2018 22:13
-
-
Save dexter800/d6f5b3c0adf99f338a3baad9a5dc3f11 to your computer and use it in GitHub Desktop.
RoboND-Perception-Project
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
# Import modules | |
import numpy as np | |
import sklearn | |
from sklearn.preprocessing import LabelEncoder | |
import pickle | |
from sensor_stick.srv import GetNormals | |
from sensor_stick.features import compute_color_histograms | |
from sensor_stick.features import compute_normal_histograms | |
from visualization_msgs.msg import Marker | |
from sensor_stick.marker_tools import * | |
from sensor_stick.msg import DetectedObjectsArray | |
from sensor_stick.msg import DetectedObject | |
from sensor_stick.pcl_helper import * | |
import rospy | |
import tf | |
from geometry_msgs.msg import Pose | |
from std_msgs.msg import Float64 | |
from std_msgs.msg import Int32 | |
from std_msgs.msg import String | |
from pr2_robot.srv import * | |
from rospy_message_converter import message_converter | |
import yaml | |
# Helper function to get surface normals | |
def get_normals(cloud): | |
get_normals_prox = rospy.ServiceProxy('/feature_extractor/get_normals', GetNormals) | |
return get_normals_prox(cloud).cluster | |
# Helper function to create a yaml friendly dictionary from ROS messages | |
def make_yaml_dict(test_scene_num, arm_name, object_name, pick_pose, place_pose): | |
yaml_dict = {} | |
yaml_dict["test_scene_num"] = test_scene_num.data | |
yaml_dict["arm_name"] = arm_name.data | |
yaml_dict["object_name"] = object_name.data | |
yaml_dict["pick_pose"] = message_converter.convert_ros_message_to_dictionary(pick_pose) | |
yaml_dict["place_pose"] = message_converter.convert_ros_message_to_dictionary(place_pose) | |
return yaml_dict | |
# Helper function to output to yaml file | |
def send_to_yaml(yaml_filename, dict_list): | |
data_dict = {"object_list": dict_list} | |
with open(yaml_filename, 'w') as outfile: | |
yaml.dump(data_dict, outfile, default_flow_style=False) | |
# Callback function for your Point Cloud Subscriber | |
def pcl_callback(pcl_msg): | |
# Exercise-2 TODOs: | |
# Convert ROS msg to PCL data | |
cloud = ros_to_pcl(pcl_msg) | |
# Voxel Grid Downsampling | |
vox = cloud.make_voxel_grid_filter() | |
# Choose a voxel (also known as leaf) size | |
# Note: this (1) is a poor choice of leaf size | |
# Experiment and find the appropriate size! | |
LEAF_SIZE = 0.01 | |
# Set the voxel (or leaf) size | |
vox.set_leaf_size(LEAF_SIZE, LEAF_SIZE, LEAF_SIZE) | |
# Call the filter function to obtain the resultant downsampled point cloud | |
cloud_filtered = vox.filter() | |
# PassThrough Filter | |
# Create a PassThrough filter object. | |
passthrough = cloud_filtered.make_passthrough_filter() | |
# Assign axis and range to the passthrough filter object. | |
filter_axis = 'z' | |
passthrough.set_filter_field_name(filter_axis) | |
axis_min = 0.5 | |
axis_max = 0.8 | |
passthrough.set_filter_limits(axis_min, axis_max) | |
# Finally use the filter function to obtain the resultant point cloud. | |
cloud_filtered = passthrough.filter() | |
passthrough = cloud_filtered.make_passthrough_filter() | |
# Assign axis and range to the passthrough filter object. | |
filter_axis = 'x' | |
passthrough.set_filter_field_name(filter_axis) | |
axis_min = 0.4 | |
axis_max = 2 | |
passthrough.set_filter_limits(axis_min, axis_max) | |
# Finally use the filter function to obtain the resultant point cloud. | |
cloud_filtered = passthrough.filter() | |
# Extract outliers to denoise | |
# Much like the previous filters, we start by creating a filter object: | |
outlier_filter = cloud_filtered.make_statistical_outlier_filter() | |
# Set the number of neighboring points to analyze for any given point | |
outlier_filter.set_mean_k(4) | |
# Set threshold scale factor | |
x = 0.05 | |
# Any point with a mean distance larger than global (mean distance+x*std_dev) will be considered outlier | |
outlier_filter.set_std_dev_mul_thresh(x) | |
# Finally call the filter function for magic | |
cloud_filtered = outlier_filter.filter() | |
cloud_no_noise = cloud_filtered | |
# RANSAC Plane Segmentation | |
# Create the segmentation object | |
seg = cloud_filtered.make_segmenter() | |
# Set the model you wish to fit | |
seg.set_model_type(pcl.SACMODEL_PLANE) | |
seg.set_method_type(pcl.SAC_RANSAC) | |
# Max distance for a point to be considered fitting the model | |
# Experiment with different values for max_distance | |
# for segmenting the table | |
max_distance = 0.01 | |
seg.set_distance_threshold(max_distance) | |
# Call the segment function to obtain set of inlier indices and model coefficients | |
inliers, coefficients = seg.segment() | |
# Extract inliers and outliers | |
cloud_table = cloud_filtered.extract(inliers, negative=False) | |
cloud_objects = cloud_filtered.extract(inliers, negative=True) | |
# Euclidean Clustering | |
white_cloud = XYZRGB_to_XYZ(cloud_objects) | |
tree = white_cloud.make_kdtree() | |
# Create a cluster extraction object | |
ec = white_cloud.make_EuclideanClusterExtraction() | |
# Set tolerances for distance threshold | |
# as well as minimum and maximum cluster size (in points) | |
# NOTE: These are poor choices of clustering parameters | |
# Your task is to experiment and find values that work for segmenting objects. | |
ec.set_ClusterTolerance(0.05) | |
ec.set_MinClusterSize(50) | |
ec.set_MaxClusterSize(200000) | |
# Search the k-d tree for clusters | |
ec.set_SearchMethod(tree) | |
# Extract indices for each of the discovered clusters | |
cluster_indices = ec.Extract() | |
# Create Cluster-Mask Point Cloud to visualize each cluster separately | |
cluster_color = get_color_list(len(cluster_indices)) | |
color_cluster_point_list = [] | |
for j, indices in enumerate(cluster_indices): | |
for i, indice in enumerate(indices): | |
color_cluster_point_list.append([white_cloud[indice][0], | |
white_cloud[indice][1], | |
white_cloud[indice][2], | |
rgb_to_float(cluster_color[j])]) | |
#Create new cloud containing all clusters, each with unique color | |
cluster_cloud = pcl.PointCloud_PointXYZRGB() | |
cluster_cloud.from_list(color_cluster_point_list) | |
# Exercise-3 TODOs: | |
# Classify the clusters! (loop through each detected cluster one at a time) | |
detected_objects_labels = [] | |
detected_objects = [] | |
for index, pts_list in enumerate(cluster_indices): | |
# Grab the points for the cluster from the extracted outliers (cloud_objects) | |
pcl_cluster = cloud_objects.extract(pts_list) | |
# TODO: convert the cluster from pcl to ROS using helper function | |
ros_cluster = pcl_to_ros(pcl_cluster) | |
# Extract histogram features | |
# TODO: complete this step just as is covered in capture_features.py | |
chists = compute_color_histograms(ros_cluster, using_hsv=True) | |
normals = get_normals(ros_cluster) | |
nhists = compute_normal_histograms(normals) | |
feature = np.concatenate((chists, nhists)) | |
# Make the prediction, retrieve the label for the result | |
# and add it to detected_objects_labels list | |
prediction = clf.predict(scaler.transform(feature.reshape(1, -1))) | |
label = encoder.inverse_transform(prediction)[0] | |
detected_objects_labels.append(label) | |
# Publish a label into RViz | |
label_pos = list(white_cloud[pts_list[0]]) | |
label_pos[2] += .4 | |
object_markers_pub.publish(make_label(label,label_pos, index)) | |
# Add the detected object to the list of detected objects. | |
do = DetectedObject() | |
do.label = label | |
do.cloud = ros_cluster | |
detected_objects.append(do) | |
rospy.loginfo('Detected {} objects: {}'.format(len(detected_objects_labels), detected_objects_labels)) | |
# Publish the list of detected objects | |
# This is the output you'll need to complete the upcoming project! | |
detected_objects_pub.publish(detected_objects) | |
# Suggested location for where to invoke your pr2_mover() function within pcl_callback() | |
# Could add some logic to determine whether or not your object detections are robust | |
# before calling pr2_mover() | |
# Convert PCL data to ROS messages | |
ros_cloud_tabel = pcl_to_ros(cloud_table) | |
ros_cloud_objects = pcl_to_ros(cloud_objects) | |
ros_cluster_cloud = pcl_to_ros(cluster_cloud) | |
# Publish ROS messages | |
pcl_table_pub.publish(ros_cloud_tabel) | |
pcl_objects_pub.publish(ros_cloud_objects) | |
pcl_cluster_pub.publish(ros_cluster_cloud) | |
try: | |
pr2_mover(detected_objects) | |
except rospy.ROSInterruptException: | |
pass | |
# Suggested location for where to invoke your pr2_mover() function within pcl_callback() | |
# Could add some logic to determine whether or not your object detections are robust | |
# before calling pr2_mover() | |
# function to load parameters and request PickPlace service | |
def pr2_mover(object_list): | |
# TODO: Initialize variables | |
# TODO: Get/Read parameters | |
# TODO: Parse parameters into individual variables | |
# TODO: Rotate PR2 in place to capture side tables for the collision map | |
# TODO: Loop through the pick list | |
# TODO: Get the PointCloud for a given object and obtain it's centroid | |
# TODO: Create 'place_pose' for the object | |
# TODO: Assign the arm to be used for pick_place | |
# TODO: Create a list of dictionaries (made with make_yaml_dict()) for later output to yaml format | |
# Wait for 'pick_place_routine' service to come up | |
rospy.wait_for_service('pick_place_routine') | |
# try: | |
#pick_place_routine = rospy.ServiceProxy('pick_place_routine', PickPlace) | |
# TODO: Insert your message variables to be sent as a service request | |
#resp = pick_place_routine(TEST_SCENE_NUM, OBJECT_NAME, WHICH_ARM, PICK_POSE, PLACE_POSE) | |
# print ("Response: ",resp.success) | |
# except rospy.ServiceException, e: | |
# print "Service call failed: %s"%e | |
# TODO: Output your request parameters into output yaml file | |
if __name__ == '__main__': | |
# ROS node initialization | |
rospy.init_node('recognition', anonymous=True) | |
# Create Subscriber to receive the published data coming from the | |
#pcl_callback() function that will be processing the point clouds | |
pcl_sub = rospy.Subscriber('/pr2/world/points', pc2.PointCloud2, | |
pcl_callback, queue_size=1) | |
# Create Publishers | |
object_markers_pub = rospy.Publisher('/object_markers', Marker, | |
queue_size=1) | |
detected_objects_pub = rospy.Publisher('/detected_objects', | |
DetectedObjectsArray, | |
queue_size=1) | |
# Isolated object point cloud with the object's original colors | |
# Isolated object point cloud with random colors | |
# Table point cloud without the objects | |
pcl_objects_pub = rospy.Publisher('/pcl_objects', PointCloud2, queue_size=1) | |
pcl_table_pub = rospy.Publisher('/pcl_table', PointCloud2, queue_size=1) | |
pcl_cluster_pub = rospy.Publisher('/pcl_cluster', PointCloud2, queue_size=1) | |
# Load model from disk | |
model = pickle.load(open('model.sav', 'rb')) | |
clf = model['classifier'] | |
encoder = LabelEncoder() | |
encoder.classes_ = model['classes'] | |
scaler = model['scaler'] | |
# Initialize color_list | |
get_color_list.color_list = [] | |
# Spin while node is not shutdown | |
while not rospy.is_shutdown(): | |
rospy.spin() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment