class MM2DPerception:
# Set apollo to True if the car selected inside the simulator is the Apollo prefab. Set to false if it is from Autoware.
# Set publish to True if you are using the Autoware prefab and want to visualize detected bounding boxes directly inside the simulator.
def __init__(self, apollo=True, publish=False):
self._cv_bridge = CvBridge()
rospy.init_node('mm_2d_perception_node')
self.yolo = YOLO()
self.seq = 0
self.publish = publish
self.apollo = apollo
if self.apollo:
self.sub_2d = rospy.Subscriber(
'/apollo/sensor/camera/traffic/image_short/compressed',
CompressedImage,
self.callback,
queue_size=1)
else:
self.sub_2d = rospy.Subscriber(
'/simulator/camera_node/image/compressed',
CompressedImage,
self.callback,
queue_size=1)
if self.publish:
self.pub_2d = rospy.Publisher('/detection/vision_objects',
DetectedObjectArray,
queue_size=5)
try:
rospy.spin()
except KeyboardInterrupt:
print("Shutting down")
cv2.destroyAllWindows()
def callback(self, image_msg):
image_np = self._cv_bridge.compressed_imgmsg_to_cv2(image_msg, "rgb8")
if not self.publish:
image_pil = Image.fromarray(image_np)
self.infer_image(image_pil)
else:
self.get_objects_and_publish(image_np)
def infer_image(self, image):
self.seq += 1
r_image = self.yolo.detect_image(image)
cv_image = cv2.cvtColor(np.array(r_image), cv2.COLOR_RGB2BGR)
cv2.imshow("2D Perception", cv_image)
cv2.waitKey(1)
# Uncomment the next two lines if you want to save the images as they are inferred
# r_image.show()
# r_image.save('PATH_TO_FOLDER/{}.jpg'.format(self.seq))
def get_objects_and_publish(self, image):
self.seq += 1
image_pil = Image.fromarray(image)
objects = self.yolo.get_detections(image_pil)
det_obj_arr = DetectedObjectArray()
for obj in objects:
det_obj = DetectedObject()
det_obj.label = obj['label']
det_obj.score = obj['score']
det_obj.x = int(obj['x'])
det_obj.y = int(obj['y'])
det_obj.width = obj['width']
det_obj.height = obj['height']
twist = Twist()
twist.linear.x = 0.0
twist.angular.z = 0.0
det_obj.velocity = twist
det_obj_arr.objects.append(det_obj)
self.pub_2d.publish(det_obj_arr)
class TestSetInference:
labels_mapping = {
"vehicle" : "0",
"pedestrian" : "1"
}
def __init__(self,
test_set_folder,
detections_output_file,
test_set_images=None,
weights=None,
anchors=None,
labels=None,
score=0.1,
iou=0.1):
self.test_set_folder = test_set_folder
self.detections_output_file = detections_output_file
self.test_set_images = test_set_images
self.yolo = YOLO(model_path=weights, anchors_path=anchors, classes_path=labels, score=score, iou=iou)
self.seq = 0
def detect_single_image(self, image, display=False, image_name=None, save_image_folder=None):
self.seq += 1
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
image_pil = Image.fromarray(image)
objects = self.yolo.get_detections(image_pil)
if display:
r_image = self.yolo.detect_image(image_pil)
if save_image_folder and image_name:
r_image.save(os.path.join(save_image_folder, image_name))
# cv_image = cv2.cvtColor(np.array(r_image), cv2.COLOR_RGB2BGR)
# cv2.imshow("Detected Objects", cv_image)
# cv2.waitKey(1)
return_list = ""
for obj in objects:
xmin = obj['x']
ymin = obj['y']
xmax = obj['x'] + obj['width']
ymax = obj['y'] + obj['height']
score = obj['score']
label = TestSetInference.labels_mapping.get(obj['label'])
obj_str = str(xmin) + ',' + str(ymin) + ',' + str(xmax) + ',' + str(ymax) + ',' + str(score) + ',' + label + ' '
return_list += obj_str
return return_list
def run_inference(self, save_folder=None, ext='.jpg'):
# Get a list of images
files = []
if self.test_set_images is not None:
with open(self.test_set_images, 'r') as f_in:
for count, line in enumerate(f_in):
image_name = line.split(' ')[0]
files.append(os.path.join(self.test_set_folder, image_name + ext))
else:
types = ['*.jpg', '*.png']
for t in types:
files.extend(glob.glob(self.test_set_folder + t))
files.sort()
print("Files found: " + str(len(files)))
print(files[0])
# Get detections for each image
for file in files:
fn = os.path.basename(file)
image = cv2.imread(file)
# Output detections to the save_folder
dets = self.detect_single_image(image, display=True, image_name=fn, save_image_folder=save_folder)
# Append to CSV file
with open(self.detections_output_file, 'a') as f_out:
wr = csv.writer(f_out, delimiter=';')
wr.writerow([fn, dets])
print("Processed " + str(self.seq) + " images")
