-
Notifications
You must be signed in to change notification settings - Fork 0
/
detect.py
62 lines (48 loc) · 2.2 KB
/
detect.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
from sensor_stick.pcl_helper import *
from sensor_stick.marker_tools import *
from pr2_robot.srv import *
from sensor_stick.features import compute_color_histograms
from sensor_stick.features import compute_normal_histograms
from sensor_stick.msg import DetectedObject
from sensor_stick.msg import DetectedObjectsArray
from visualization_msgs.msg import Marker
import rospy
import numpy as np
from sensor_stick.srv import GetNormals
# 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
def detect_objects(cluster_indices, cloud_objects, white_cloud, clf,
object_markers_pub, scaler, encoder):
# Create some empty lists
detected_objects_labels = []
detected_objects = []
for index, pts_list in enumerate(cluster_indices):
# Grab points for the cluster from the extracted outliers(cloud_objects)
pcl_cluster = cloud_objects.extract(pts_list)
# Convert pcl_cluster to ros_msg data type
ros_cluster = pcl_to_ros(pcl_cluster)
# Extract histogram features
chists = compute_color_histograms(ros_cluster, using_hsv=True)
normals = get_normals(ros_cluster)
nhists = compute_normal_histograms(normals)
feature = np.concatenate((chists, nhists))
# labeled_features.append([feature, model_name])
# Make prediction, retrieve the label for the result
# then 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 label to 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))
return detected_objects