46 from tools
import ImageSubscriber, StatisticsEngine, DrawingTools, Camera, VideoSaver, exit_handler, str2bool
50 rospy.logfatal(
"Could not import the submodules.")
54 ex_type, ex, tb = sys.exc_info()
55 traceback.print_tb(tb)
61 def __init__(self, source, video_save, logdir, cpu, display, tensor_io, mode, gpu_percent=1):
70 tmp_image, tmp_frame = self.
camera.get_image()
74 if isinstance(video_name, bool):
75 video_name = self.
camera.get_from_camera_video_name()
78 video_name = get_from_video_out_video_name(args)
81 tmp_image, video_name=video_name)
86 'Video results of network run will be saved as %s', video_name)
89 rospy.loginfo(
'Video results of network run will not be saved')
95 if self.
mode ==
'segmentation':
96 from segmentation
import MaskPublisher
97 from segmentation
import DeepLabSegmenter
as SegmentationProcessor
98 self.
cnn_processor = SegmentationProcessor(logdir, self.
camera.original_image_size, tensor_io, runCPU=cpu, gpu_percent=gpu_percent)
100 elif self.
mode ==
'detection':
101 from detection
import YoloDetector
as YoloProcessor
102 from detection
import DetectionPublisher
103 self.
cnn_processor = YoloProcessor(logdir, self.
camera.original_image_size, tensor_io, runCPU=cpu, gpu_percent=gpu_percent)
108 atexit.register(exit_handler, self.
camera,
113 if self.
mode ==
'segmentation':
115 self.
camera.original_image_size[0], self.
camera.original_image_size[1])
117 elif self.
mode ==
'detection':
124 while not rospy.is_shutdown():
126 time_iter_start = timeit.default_timer()
128 image_frame, image_header = self.
camera.get_image()
130 if image_frame
is None:
134 if self.
DO_VIDEO_SAVE is not False and count % video_out_grid == 0
and count > 0:
137 time_iter_start_inner = timeit.default_timer()
140 cnn_result, time_pure_tf = self.
cnn_processor.run_model_on_image(
144 if self.
mode ==
'segmentation':
146 elif self.
mode ==
'detection':
150 time_iter_no_drawing = timeit.default_timer() - time_iter_start_inner
154 if self.
mode ==
'segmentation':
155 image_green = self.
drawing_tools.overlay_segmentation(image_frame.copy(), cnn_result)
156 elif self.
mode ==
'detection':
158 image_frame.copy(), cnn_result[
'boxes'], cnn_result[
'scores'], cnn_result[
'classes'], [320,320])
161 time_iter_inner = timeit.default_timer() - time_iter_start_inner
165 frame_green = cv2.cvtColor(image_green, cv2.COLOR_RGB2BGR)
167 cv2.imshow(
'Path prediction', cv2.resize(
168 frame_green, dsize=self.
camera.image_size_to_show))
173 if self.
camera.camera_type
is "cv2vid":
178 cv2.cvtColor(image_frame, cv2.COLOR_RGB2BGR))
181 time_iter = timeit.default_timer() - time_iter_start
183 str_to_print =
'%d: Iteration spent %.4g inner %.4g' % (
184 count, time_iter, time_iter_inner)
185 if time_pure_tf
is not None:
186 str_to_print +=
'\t(in tf was %.4g)' % time_pure_tf
190 rospy.logdebug(str_to_print)
193 time_iter_no_drawing)
202 time_iter - time_iter_inner)
204 if time_pure_tf
is not None:
212 if cv2.waitKey(1) & 0xFF == ord(
'q'):
215 rospy.logfatal(
'CNN stopped!')
216 rospy.signal_shutdown(
'Segmentation closed. Shutting down...')
def __init__(self, source, video_save, logdir, cpu, display, tensor_io, mode, gpu_percent=1)