class Realsense:
def __init__(
self,
element_name,
transform_file_path,
calibration_client_path,
depth_shape,
color_shape,
fps,
disparity_shift,
depth_units,
rotation,
retry_delay
):
self._transform_file_path = transform_file_path
self._calibration_client_path = calibration_client_path
self._depth_shape = depth_shape
self._color_shape = color_shape
self._fps = fps
self.disparity_shift = disparity_shift
self.depth_units = depth_units
self._rotation = rotation
self._retry_delay = retry_delay
self._status_is_running = False
self._status_lock = Lock()
self._pipeline = rs.pipeline()
self._rs_pc = rs.pointcloud()
self._transform = TransformStreamContract(x=0, y=0, z=0, qx=0, qy=0, qz=0, qw=1)
self._transform_last_loaded = 0
# Create an align object: rs.align allows us to perform alignment of depth frames to other frames
self._align = rs.align(rs.stream.color)
# Init element
self._element = Element(element_name)
self._element.healthcheck_set(self.is_healthy)
self._element.command_add(
CalculateTransformCommand.COMMAND_NAME,
self.run_transform_estimator,
timeout=2000,
deserialize=CalculateTransformCommand.Request.SERIALIZE
)
# Run command loop
thread = Thread(target=self._element.command_loop, daemon=True)
thread.start()
def is_healthy(self):
"""
Reports whether the realsense is connected and streaming or not
"""
try:
self._status_lock.acquire()
if self._status_is_running:
return Response(err_code=0, err_str="Realsense online")
else:
return Response(err_code=1, err_str="Waiting for realsense")
finally:
self._status_lock.release()
def load_transform_from_file(self, fname):
"""
Opens specified file, reads transform, and returns as list.
Args:
fname (str): CSV file that stores the transform
"""
with open(fname, "r") as f:
transform_list = [float(v) for v in f.readlines()[-1].split(",")]
return TransformStreamContract(
x=transform_list[0],
y=transform_list[1],
z=transform_list[2],
qx=transform_list[3],
qy=transform_list[4],
qz=transform_list[5],
qw=transform_list[6]
)
def run_transform_estimator(self, *args):
"""
Runs the transform estimation procedure, which saves the transform to disk.
"""
process = subprocess.Popen(self._calibration_client_path, stderr=subprocess.PIPE)
out, err = process.communicate()
return Response(
data=CalculateTransformCommand.Response().to_data(),
err_code=process.returncode,
err_str=err.decode(),
serialize=CalculateTransformCommand.Response.SERIALIZE
)
def run_camera_stream(self):
while True:
try:
# Try to establish realsense connection
self._element.log(LogLevel.INFO, "Attempting to connect to Realsense")
# Set disparity shift
device = rs.context().query_devices()[0]
advnc_mode = rs.rs400_advanced_mode(device)
depth_table_control_group = advnc_mode.get_depth_table()
depth_table_control_group.disparityShift = self.disparity_shift
advnc_mode.set_depth_table(depth_table_control_group)
# Attempt to stream accel and gyro data, which requires d435i
# If we can't then we revert to only streaming depth and color
try:
config = rs.config()
config.enable_stream(
rs.stream.depth, self._depth_shape[0], self._depth_shape[1], rs.format.z16, self._fps
)
config.enable_stream(
rs.stream.color, self._color_shape[0], self._color_shape[1], rs.format.bgr8, self._fps
)
config.enable_stream(rs.stream.accel)
config.enable_stream(rs.stream.gyro)
profile = self._pipeline.start(config)
is_d435i = True
except RuntimeError:
config = rs.config()
config.enable_stream(
rs.stream.depth, self._depth_shape[0], self._depth_shape[1], rs.format.z16, self._fps
)
config.enable_stream(
rs.stream.color, self._color_shape[0], self._color_shape[1], rs.format.bgr8, self._fps
)
profile = self._pipeline.start(config)
is_d435i = False
# Set depth units
depth_sensor = profile.get_device().first_depth_sensor()
depth_sensor.set_option(rs.option.depth_units, self.depth_units)
# Publish intrinsics
rs_intrinsics = profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()
intrinsics = IntrinsicsStreamContract(
width=rs_intrinsics.width,
height=rs_intrinsics.height,
ppx=rs_intrinsics.ppx,
ppy=rs_intrinsics.ppy,
fx=rs_intrinsics.fx,
fy=rs_intrinsics.fy
)
self._element.entry_write(
IntrinsicsStreamContract.STREAM_NAME,
intrinsics.to_dict(),
serialize=IntrinsicsStreamContract.SERIALIZE,
maxlen=self._fps
)
try:
self._status_lock.acquire()
self._status_is_running = True
finally:
self._status_lock.release()
self._element.log(LogLevel.INFO, "Realsense connected and streaming")
while True:
start_time = time.time()
frames = self._pipeline.wait_for_frames()
aligned_frames = self._align.process(frames)
depth_frame = aligned_frames.get_depth_frame()
color_frame = aligned_frames.get_color_frame()
# Validate that frames are valid
if not depth_frame or not color_frame:
continue
# Generate realsense pointcloud
self._rs_pc.map_to(color_frame)
points = self._rs_pc.calculate(depth_frame)
# Convert data to numpy arrays
depth_image = np.asanyarray(depth_frame.get_data())
color_image = np.asanyarray(color_frame.get_data())
vertices = np.asanyarray(points.get_vertices())
vertices = vertices.view(np.float32).reshape(vertices.shape + (-1,))
if self._rotation is not None:
depth_image = np.rot90(depth_image, k=self._rotation / 90)
color_image = np.rot90(color_image, k=self._rotation / 90)
# TODO: Apply rotation to pointcloud
_, color_serialized = cv2.imencode(".tif", color_image)
_, depth_serialized = cv2.imencode(".tif", depth_image)
_, pc_serialized = cv2.imencode(".tif", vertices)
if is_d435i:
accel = frames[2].as_motion_frame().get_motion_data()
gyro = frames[3].as_motion_frame().get_motion_data()
accel_data = AccelStreamContract(x=accel.x, y=accel.y, z=accel.z)
gyro_data = GyroStreamContract(x=gyro.x, y=gyro.y, z=gyro.z)
self._element.entry_write(
AccelStreamContract.STREAM_NAME,
accel_data.to_dict(),
serialize=AccelStreamContract.SERIALIZE,
maxlen=self._fps
)
self._element.entry_write(
GyroStreamContract.STREAM_NAME,
gyro_data.to_dict(),
serialize=GyroStreamContract.SERIALIZE,
maxlen=self._fps
)
color_contract = ColorStreamContract(data=color_serialized.tobytes())
depth_contract = DepthStreamContract(data=depth_serialized.tobytes())
pc_contract = PointCloudStreamContract(data=pc_serialized.tobytes())
self._element.entry_write(
ColorStreamContract.STREAM_NAME,
color_contract.to_dict(),
serialize=ColorStreamContract.SERIALIZE,
maxlen=self._fps
)
self._element.entry_write(
DepthStreamContract.STREAM_NAME,
depth_contract.to_dict(),
serialize=DepthStreamContract.SERIALIZE,
maxlen=self._fps
)
self._element.entry_write(
PointCloudStreamContract.STREAM_NAME,
pc_contract.to_dict(),
serialize=PointCloudStreamContract.SERIALIZE,
maxlen=self._fps
)
# Load transform from file if the file exists
# and has been modified since we last checked
if os.path.exists(self._transform_file_path):
transform_last_modified = os.stat(self._transform_file_path).st_mtime
if transform_last_modified > self._transform_last_loaded:
try:
self._transform = self.load_transform_from_file(self._transform_file_path)
self._transform_last_loaded = time.time()
except Exception as e:
self._element.log(LogLevel.ERR, str(e))
self._element.entry_write(
TransformStreamContract.STREAM_NAME,
self._transform.to_dict(),
serialize=TransformStreamContract.SERIALIZE,
maxlen=self._fps
)
time.sleep(max(1 / self._fps - (time.time() - start_time), 0))
except:
self._element.log(LogLevel.INFO, "Camera loop threw exception: %s" % (sys.exc_info()[1]))
finally:
# If camera fails to init or crashes, update status and retry connection
try:
self._status_lock.acquire()
self._status_is_running = False
finally:
self._status_lock.release()
try:
self._pipeline.stop()
except:
pass
time.sleep(self._retry_delay)
class Picamera:
def __init__(self, element_name, width, height, fps, retry_delay):
self._width = width
self._height = height
self._fps = fps
self._retry_delay = retry_delay
self._status_is_running = False
self._status_lock = Lock()
# Init element
self._element = Element(element_name)
self._element.healthcheck_set(self.is_healthy)
#self._element.command_add(command_name, command_func_ptr, timeout, serialize)
# Run command loop
thread = Thread(target=self._element.command_loop, daemon=True)
thread.start()
def is_healthy(self):
# Reports whether the camera is connected or not
try:
self._status_lock.acquire()
if self._status_is_running:
return Response(err_code=0, err_str="Camera is good")
else:
return Response(err_code=1, err_str="Camera is not good")
except:
return Response(err_code=0, err_str="Could not reach thread")
def run_camera_stream(self):
while True:
try:
# try to open up camera
self._element.log(LogLevel.INFO, "Opening PiCamera")
self._camera = PiCamera()
self._color_array = PiRGBArray(self._camera)
# set camera configs
self._camera.resolution = (self._width, self._height)
self._camera.framerate = self._fps
# allow the camera to warm up
time.sleep(.5)
try:
self._status_lock.acquire()
self._status_is_running = True
finally:
self._status_lock.release()
self._element.log(LogLevel.INFO, "Picamera connected and streaming")
while True:
start_time = time.time()
self._camera.capture(self._color_array, format = 'bgr')
color_image = self._color_array.array
#do some rotation here
_, color_serialized = cv2.imencode(".tif", color_image)
color_contract = ColorStreamContract(data=color_serialized.tobytes())
self._element.entry_write(
ColorStreamContract.STREAM_NAME,
color_contract.to_dict(),
serialize=ColorStreamContract.SERIALIZE,
maxlen=self._fps
)
time.sleep(max(1 / self._fps - (time.time() - start_time),0))
self._color_array.truncate(0)
except:
self._element.log(LogLevel.INFO, "Camera threw exception: %s" % (sys.exc_info()[1]))
finally:
try:
self._status_lock.acquire()
self._status_is_running = False
self._camera.close()
finally:
self._status_lock.release()
time.sleep(self._retry_delay)
timeout=50,
deserialize=True)
# Transform takes no inputs, so there's nothing to deserialize
element.command_add("transform", atombot.transform, timeout=50)
# We create a thread and run the command loop which will constantly check for incoming commands from atom
# We use a thread so we don't hang on the command_loop function because we will be performing other tasks
thread = Thread(target=element.command_loop, daemon=True)
thread.start()
# This will block until every element in the list reports it is healthy. Useful if you depend on other elements.
element.wait_for_elements_healthy(['atombot'])
# Create an infinite loop that publishes the position of atombot to a stream as well as a visual of its position
while True:
# We write our position data and the visual of atombot's position to their respective streams
# The maxlen parameter will determine how many entries atom stores
# This data is serialized using msgpack
element.entry_write("pos", {"data": atombot.get_pos()},
maxlen=10,
serialize=True)
element.entry_write("pos_map", {"data": atombot.get_pos_map()},
maxlen=10,
serialize=True)
# We can also choose to write binary data directly without serializing it
element.entry_write("pos_binary", {"data": atombot.get_pos()},
maxlen=10)
# Sleep so that we aren't consuming all of our CPU resources
sleep(0.01)
class SDMaskRCNNEvaluator:
def __init__(self,
mode="both",
input_size=512,
scaling_factor=2,
config_path="sd-maskrcnn/cfg/benchmark.yaml"):
self.element = Element("instance-segmentation")
self.input_size = input_size
self.scaling_factor = scaling_factor
self.config_path = config_path
self.mode = mode
# Streaming of masks is disabled by default to prevent consumption of resources
self.stream_enabled = False
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.gpu_options.visible_device_list = "0"
set_session(tf.Session(config=config))
self.set_mode(b"both")
# Initiate tensorflow graph before running threads
self.get_masks()
self.element.command_add("segment", self.segment, 10000)
self.element.command_add("get_mode", self.get_mode, 100)
self.element.command_add("set_mode", self.set_mode, 10000)
self.element.command_add("stream", self.set_stream, 100)
t = Thread(target=self.element.command_loop, daemon=True)
t.start()
self.publish_segments()
def get_mode(self, _):
"""
Returns the current mode of the algorithm (both or depth).
"""
return Response(self.mode)
def set_mode(self, data):
"""
Sets the mode of the algorithm and loads the corresponding weights.
'both' means that the algorithm is considering grayscale and depth data.
'depth' means that the algorithm only considers depth data.
"""
mode = data.decode().strip().lower()
if mode not in MODES:
return Response(f"Invalid mode {mode}")
self.mode = mode
config = YamlConfig(self.config_path)
inference_config = MaskConfig(config['model']['settings'])
inference_config.GPU_COUNT = 1
inference_config.IMAGES_PER_GPU = 1
model_path = MODEL_PATHS[self.mode]
model_dir, _ = os.path.split(model_path)
self.model = modellib.MaskRCNN(mode=config['model']['mode'],
config=inference_config,
model_dir=model_dir)
self.model.load_weights(model_path, by_name=True)
self.element.log(LogLevel.INFO, f"Loaded weights from {model_path}")
return Response(f"Mode switched to {self.mode}")
def set_stream(self, data):
"""
Sets streaming of segmented masks to true or false.
"""
data = data.decode().strip().lower()
if data == "true":
self.stream_enabled = True
elif data == "false":
self.stream_enabled = False
else:
return Response(f"Expected bool, got {type(data)}.")
return Response(f"Streaming set to {self.stream_enabled}")
def inpaint(self, img, missing_value=0):
"""
Fills the missing values of the depth data.
"""
# cv2 inpainting doesn't handle the border properly
# https://stackoverflow.com/questions/25974033/inpainting-depth-map-still-a-black-image-border
img = cv2.copyMakeBorder(img, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
mask = (img == missing_value).astype(np.uint8)
# Scale to keep as float, but has to be in bounds -1:1 to keep opencv happy.
scale = np.abs(img).max()
img = img.astype(
np.float32) / scale # Has to be float32, 64 not supported.
img = cv2.inpaint(img, mask, 1, cv2.INPAINT_NS)
# Back to original size and value range.
img = img[1:-1, 1:-1]
img = img * scale
return img
def normalize(self, img, max_dist=1000):
"""
Scales the range of the data to be in 8-bit.
Also shifts the values so that maximum is 255.
"""
img = np.clip(img / max_dist, 0, 1) * 255
img = np.clip(img + (255 - img.max()), 0, 255)
return img.astype(np.uint8)
def scale_and_square(self, img, scaling_factor, size):
"""
Scales the image by scaling_factor and creates a border around the image to match size.
Reducing the size of the image tends to improve the output of the model.
"""
img = cv2.resize(img, (int(img.shape[1] / scaling_factor),
int(img.shape[0] / scaling_factor)),
interpolation=cv2.INTER_NEAREST)
v_pad, h_pad = (size - img.shape[0]) // 2, (size - img.shape[1]) // 2
img = cv2.copyMakeBorder(img, v_pad, v_pad, h_pad, h_pad,
cv2.BORDER_REPLICATE)
return img
def unscale(self, results, scaling_factor, size):
"""
Takes the results of the model and transforms them back into the original dimensions of the input image.
"""
masks = results["masks"].astype(np.uint8)
masks = cv2.resize(masks, (int(masks.shape[1] * scaling_factor),
int(masks.shape[0] * scaling_factor)),
interpolation=cv2.INTER_NEAREST)
v_pad, h_pad = (masks.shape[0] - size[0]) // 2, (masks.shape[1] -
size[1]) // 2
masks = masks[v_pad:-v_pad, h_pad:-h_pad]
rois = results["rois"] * scaling_factor
for roi in rois:
roi[0] = min(max(0, roi[0] - v_pad), size[0])
roi[1] = min(max(0, roi[1] - h_pad), size[1])
roi[2] = min(max(0, roi[2] - v_pad), size[0])
roi[3] = min(max(0, roi[3] - h_pad), size[1])
return masks, rois
def publish_segments(self):
"""
Publishes visualization of segmentation masks continuously.
"""
self.colors = []
for i in range(NUM_OF_COLORS):
self.colors.append((np.random.rand(3) * 255).astype(int))
while True:
if not self.stream_enabled:
time.sleep(1 / PUBLISH_RATE)
continue
start_time = time.time()
scores, masks, rois, color_img = self.get_masks()
masked_img = np.zeros(color_img.shape).astype("uint8")
contour_img = np.zeros(color_img.shape).astype("uint8")
if masks is not None and scores.size != 0:
number_of_masks = masks.shape[-1]
# Calculate the areas of masks
mask_areas = []
for i in range(number_of_masks):
width = np.abs(rois[i][0] - rois[i][2])
height = np.abs(rois[i][1] - rois[i][3])
mask_area = width * height
mask_areas.append(mask_area)
np_mask_areas = np.array(mask_areas)
mask_indices = np.argsort(np_mask_areas)
# Add masks in the order of there areas.
for i in mask_indices:
if (scores[i] > SEGMENT_SCORE):
indices = np.where(masks[:, :, i] == 1)
masked_img[indices[0], indices[1], :] = self.colors[i]
# Smoothen masks
masked_img = cv2.medianBlur(masked_img, 15)
# find countours and draw boundaries.
gray_image = cv2.cvtColor(masked_img, cv2.COLOR_BGR2GRAY)
ret, thresh = cv2.threshold(gray_image, 50, 255,
cv2.THRESH_BINARY)
contours, hierarchy = cv2.findContours(thresh,
cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_NONE)
# Draw contours:
for contour in contours:
area = cv2.contourArea(contour)
cv2.drawContours(contour_img, contour, -1, (255, 255, 255),
5)
masked_img = cv2.addWeighted(color_img, 0.6, masked_img, 0.4,
0)
masked_img = cv2.bitwise_or(masked_img, contour_img)
_, color_serialized = cv2.imencode(".tif", masked_img)
self.element.entry_write("color_mask",
{"data": color_serialized.tobytes()},
maxlen=30)
time.sleep(max(0, (1 / PUBLISH_RATE) - (time.time() - start_time)))
def get_masks(self):
"""
Gets the latest data from the realsense, preprocesses it and returns the
segmentation masks, bounding boxes, and scores for each detected object.
"""
color_data = self.element.entry_read_n("realsense", "color", 1)
depth_data = self.element.entry_read_n("realsense", "depth", 1)
try:
color_data = color_data[0]["data"]
depth_data = depth_data[0]["data"]
except IndexError or KeyError:
raise Exception(
"Could not get data. Is the realsense element running?")
depth_img = cv2.imdecode(np.frombuffer(depth_data, dtype=np.uint16),
-1)
original_size = depth_img.shape[:2]
depth_img = self.scale_and_square(depth_img, self.scaling_factor,
self.input_size)
depth_img = self.inpaint(depth_img)
depth_img = self.normalize(depth_img)
if self.mode == "both":
gray_img = cv2.imdecode(np.frombuffer(color_data, dtype=np.uint16),
0)
color_img = cv2.imdecode(
np.frombuffer(color_data, dtype=np.uint16), 1)
gray_img = self.scale_and_square(gray_img, self.scaling_factor,
self.input_size)
input_img = np.zeros((self.input_size, self.input_size, 3))
input_img[..., 0] = gray_img
input_img[..., 1] = depth_img
input_img[..., 2] = depth_img
else:
input_img = np.stack((depth_img, ) * 3, axis=-1)
# Get results and unscale
results = self.model.detect([input_img], verbose=0)[0]
masks, rois = self.unscale(results, self.scaling_factor, original_size)
if masks.ndim < 2 or results["scores"].size == 0:
masks = None
results["scores"] = None
elif masks.ndim == 2:
masks = np.expand_dims(masks, axis=-1)
return results["scores"], masks, rois, color_img
def segment(self, _):
"""
Command for getting the latest segmentation masks and returning the results.
"""
scores, masks, rois, color_img = self.get_masks()
# Encoded masks in TIF format and package everything in dictionary
encoded_masks = []
if masks is not None and scores is not None:
for i in range(masks.shape[-1]):
_, encoded_mask = cv2.imencode(".tif", masks[..., i])
encoded_masks.append(encoded_mask.tobytes())
response_data = {
"rois": rois.tolist(),
"scores": scores.tolist(),
"masks": encoded_masks
}
else:
response_data = {"rois": [], "scores": [], "masks": []}
return Response(response_data, serialize=True)