def sound_element_thread(sound_queue):
"""
Registers all of the sound playback commands for this element
and handles when they're called
"""
elem = Element("sound")
# Register our callback to play sounds
elem_class = SoundElement(sound_queue)
elem.command_add("play_sound", elem_class.command_cb)
elem.command_loop()
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)
if __name__ == "__main__":
print("Launching...")
# Create our element and call it "atombot"
element = Element("atombot")
# Instantiate our AtomBot class
atombot = AtomBot()
# We add a healthcheck to our atombot element.
# This is optional. If you don't do this, atombot is assumed healthy as soon as its command_loop executes
element.healthcheck_set(atombot.is_healthy)
# This registers the relevant AtomBot methods as a command in the atom system
# We set the timeout so the caller will know how long to wait for the command to execute
element.command_add("move_left",
atombot.move_left,
timeout=50,
deserialize=True)
element.command_add("move_right",
atombot.move_right,
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'])
run_element = Element("run_demo")
res = run_element.command_send("robot_api",
"home", {
"v": VEL,
"a": ACCEL
},
serialize=True)
if (res['err_code'] != 0):
return Response(err_code=1,
err_str="Failed to move to home position",
serialize=True)
# turn off control
set_control(run_element, False)
return Response("Success", serialize=True)
if __name__ == '__main__':
'''
Mainloop, wait for the command to run the grab and then do
the grab
'''
element = Element("demo")
element.command_add("run", run_demo, timeout=60000)
element.command_add("driving_position", driving_position, timeout=10000)
element.command_add("home", home, timeout=10000)
element.command_loop()
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)
buff = "{},".format(x_val)
# And add each item from the iterable
try:
for v in val:
buff += "{},".format(v)
except:
buff += "{}".format(val)
# Finish off with a newline and write to file
buff += "\n"
files[key].write(buff)
# And note the success
return Response("Success", serialize=True)
if __name__ == '__main__':
elem = Element("record")
elem.command_add("start", start_recording, timeout=1000, deserialize=True)
elem.command_add("stop", stop_recording, timeout=1000, deserialize=True)
elem.command_add("wait", wait_recording, timeout=60000, deserialize=True)
elem.command_add("list", list_recordings, timeout=1000)
elem.command_add("get", get_recording, timeout=1000, deserialize=True)
elem.command_add("plot", plot_recording, timeout=1000000, deserialize=True)
elem.command_add("csv", csv_recording, timeout=10000, deserialize=True)
# Want to launch the plot thread s.t. our plot API can return quickly
elem.command_loop()