class AtomCLI:
def __init__(self):
self.element = Element(f"atom-cli_{uname().nodename}_{uuid4().hex}")
self.indent = 2
self.style = Style.from_dict({
"logo_color": "#6039C8",
})
self.session = PromptSession(style=self.style)
self.print_atom_os_logo()
self.serialization = "msgpack"
self.cmd_map = {
"help": self.cmd_help,
"list": self.cmd_list,
"records": self.cmd_records,
"command": self.cmd_command,
"read": self.cmd_read,
"exit": self.cmd_exit,
"serialization": self.cmd_serialization,
}
self.usage = {
"cmd_help": cleandoc("""
Displays available commands and shows usage for commands.
Usage:
help [<command>]"""),
"cmd_list": cleandoc("""
Displays available elements, streams, or commands.
Can filter streams and commands based on element.
Usage:
list elements
list streams [<element>]
list commands [<element>]"""),
"cmd_records": cleandoc("""
Displays log records or command and response records.
Can filter records from the last N seconds or from certain elements.
Usage:
records log [<last_N_seconds>] [<element>...]
records cmdres [<last_N_seconds>] <element>..."""),
"cmd_command": cleandoc("""
Sends a command to an element and displays the response.
Usage:
command <element> <element_command> [<data>]"""),
"cmd_read": cleandoc("""
Displays the entries of an element's stream.
Can provide a rate to print the entries for ease of reading.
Usage:
read <element> <stream> [<rate_hz>]"""),
"cmd_exit": cleandoc("""
Exits the atom-cli tool.
Can also use the shortcut CTRL+D.
Usage:
exit"""),
"cmd_serialization": cleandoc("""
Sets serialization/deserialization setting to either use msgpack,
Apache arrow, or no (de)serialization. Defaults to msgpack serialization.
This setting is overriden by deserialization keys received in stream.
Usage:
serialization (msgpack | arrow | none)"""),
}
def run(self):
"""
The main loop of the CLI.
Reads the user input, verifies the command exists and calls the command.
"""
while True:
try:
inp = self.session.prompt(
"\n> ", auto_suggest=AutoSuggestFromHistory()).split(" ")
if not inp:
continue
command, args = inp[0], inp[1:]
if command not in self.cmd_map.keys():
print("Invalid command. Type 'help' for valid commands.")
else:
self.cmd_map[command](*args)
# Handle CTRL+C so user can break loops without exiting
except KeyboardInterrupt:
pass
# Exit on CTRL+D
except EOFError:
self.cmd_exit()
except Exception as e:
print(str(type(e)) + " " + str(e))
def print_atom_os_logo(self):
f = Figlet(font="slant")
logo = f.renderText("ATOM OS")
print(HTML(f"<logo_color>{logo}</logo_color>"), style=self.style)
def format_record(self, record):
"""
Takes a record out of Redis, decodes the keys and values (if possible)
and returns a formatted json string sorted by keys.
"""
formatted_record = {}
for k, v in record.items():
if type(k) is bytes:
k = k.decode()
if not self.serialization:
try:
v = v.decode()
except:
v = str(v)
formatted_record[k] = v
sorted_record = {k: v for k, v in sorted(
formatted_record.items(), key=lambda x: x[0])}
try:
ret = json.dumps(sorted_record, indent=self.indent)
except TypeError as te:
ret = sorted_record
finally:
return ret
def cmd_help(self, *args):
usage = self.usage["cmd_help"]
if len(args) > 1:
print(usage)
print("\nToo many arguments to 'help'.")
return
if args:
# Prints the usage of the command
if args[0] in self.cmd_map.keys():
print(self.usage[f"cmd_{args[0]}"])
else:
print(f"Command {args[0]} does not exist.")
else:
print("Try 'help <command>' for usage on a command")
print("Available commands:")
for command in self.cmd_map.keys():
print(f" {command}")
def cmd_list(self, *args):
usage = self.usage["cmd_list"]
mode_map = {
"elements": self.element.get_all_elements,
"streams": self.element.get_all_streams,
"commands": self.element.get_all_commands
}
if not args:
print(usage)
print("\n'list' must have an argument.")
return
mode = args[0]
if mode not in mode_map.keys():
print(usage)
print("\nInvalid argument to 'list'.")
return
if len(args) > 1 and mode == "elements":
print(usage)
print(f"\nInvalid number of arguments for command 'list elements'.")
return
if len(args) > 2:
print(usage)
print("\n'list' takes at most 2 arguments.")
return
items = mode_map[mode](*args[1:])
if not items:
print(f"No {mode} exist.")
return
for item in items:
print(item)
def cmd_records(self, *args):
usage = self.usage["cmd_records"]
if not args:
print(usage)
print("\n'records' must have an argument.")
return
mode = args[0]
# Check for start time
if len(args) > 1 and args[1].isdigit():
ms = int(args[1]) * 1000
start_time = str(int(self.element._get_redis_timestamp()) - ms)
elements = set(args[2:])
# If no start time, go from the very beginning
else:
start_time = "0"
elements = set(args[1:])
if mode == "log":
records = self.mode_log(start_time, elements)
elif mode == "cmdres":
if not elements:
print(usage)
print(
"\nMust provide elements from which to get command response streams from.")
return
records = self.mode_cmdres(start_time, elements)
else:
print(usage)
print("\nInvalid argument to 'records'.")
return
if not records:
print("No records.")
return
for record in records:
print(self.format_record(record))
def mode_log(self, start_time, elements):
"""
Reads the logs from Atom's log stream.
Args:
start_time (str): The time from which to start reading logs.
elements (list): The elements on which to filter the logs for.
"""
records = []
all_records = self.element.entry_read_since(
None, "log", start_time, serialization=None)
for record in all_records:
if not elements or record["element"].decode() in elements:
record = {key: (value if isinstance(value, str) else value.decode(
)) for key, value in record.items()} # Decode strings only which are required to
records.append(record)
return records
def mode_cmdres(self, start_time, elements):
"""
Reads the command and response records from the provided elements.
Args:
start_time (str): The time from which to start reading logs.
elements (list): The elements to get the command and response records from.
"""
streams, records = [], []
for element in elements:
streams.append(self.element._make_response_id(element))
streams.append(self.element._make_command_id(element))
for stream in streams:
cur_records = self.element.entry_read_since(
None, stream, start_time, serialization=None)
for record in cur_records:
for key, value in record.items():
try:
if not isinstance(value, str):
value = value.decode()
except:
try:
value = ser.deserialize(value, method=self.serialization)
except:
pass
finally:
record[key] = value
record["type"], record["element"] = stream.split(":")
records.append(record)
return sorted(records, key=lambda x: (x["id"], x["type"]))
def cmd_command(self, *args):
usage = self.usage["cmd_command"]
if len(args) < 2:
print(usage)
print("\nToo few arguments.")
return
element_name = args[0]
command_name = args[1]
if len(args) >= 3:
data = str(" ".join(args[2:]))
if self.serialization:
try:
data = json.loads(data)
except:
print("Received improperly formatted data!")
return
else:
data = ""
resp = self.element.command_send(element_name,
command_name,
data,
serialize=(self.serialization is not None),
deserialize=(self.serialization is not None),
serialization=self.serialization) # shouldn't be used if it's None
print(self.format_record(resp))
def cmd_read(self, *args):
usage = self.usage["cmd_read"]
if len(args) < 2:
print(usage)
print("\nToo few arguments.")
return
if len(args) > 3:
print(usage)
print("\nToo many arguments.")
return
if len(args) == 3:
try:
rate = float(args[2])
if rate < 0:
raise ValueError()
except ValueError:
print("rate must be an float greater than 0.")
return
else:
rate = None
element_name, stream_name = args[:2]
last_timestamp = None
while True:
start_time = time.time()
entries = self.element.entry_read_n(element_name,
stream_name,
1,
deserialize=(self.serialization is not None),
serialization=self.serialization) # shouldn't be used if it's None
if not entries:
print(f"No data from {element_name} {stream_name}.")
return
entry = entries[0]
timestamp = entry["id"]
# Only print the entry if it is different from the previous one
if timestamp != last_timestamp:
last_timestamp = timestamp
print(self.format_record(entry))
if rate:
time.sleep(max(1 / rate - (time.time() - start_time), 0))
def cmd_serialization(self, *args):
usage = self.usage["cmd_serialization"]
if (len(args) != 1):
print(usage)
print(f"\nPass one argument: {ser.Serializations.print_values()}.")
return
# Otherwise try to get the new setting
if ser.is_valid_serialization(args[0].lower()):
self.serialization = args[0].lower() if args[0].lower() != "none" else None
else:
print(f"\nArgument must be one of {ser.Serializations.print_values()}.")
print("Current serialization status is {}".format(self.serialization))
def cmd_exit(*args):
print("Exiting.")
sys.exit()
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)