def __getitem__(self, index):
dataset = self._dataset
frame = dataset.get_frame(index)
examples = dataset.get_example(index)
vizs = []
for example in examples:
viz = imgviz.tile(
[
example["rgb"],
imgviz.depth2rgb(example["pcd"][:, :, 0]),
imgviz.depth2rgb(example["pcd"][:, :, 1]),
imgviz.depth2rgb(example["pcd"][:, :, 2]),
],
border=(255, 255, 255),
)
vizs.append(viz)
viz = imgviz.tile(vizs)
del vizs
viz = imgviz.tile([frame["rgb"], viz], shape=(1, 2))
viz = imgviz.resize(viz, width=1000)
return viz
def _callback(self, cam_msg, depth_msg):
bridge = cv_bridge.CvBridge()
depth = bridge.imgmsg_to_cv2(depth_msg)
if depth.dtype == np.uint16:
depth = depth.astype(np.float32) / 1000
depth[depth == 0] = np.nan
assert depth.dtype == np.float32
if self._pub_normal.get_num_connections() > 0:
K = np.array(cam_msg.K).reshape(3, 3)
points = morefusion.geometry.pointcloud_from_depth(depth,
fx=K[0, 0],
fy=K[1, 1],
cx=K[0, 2],
cy=K[1, 2])
normal = morefusion.geometry.estimate_pointcloud_normals(points)
normal = np.uint8((normal + 1) / 2 * 255)
out_msg = bridge.cv2_to_imgmsg(normal, "rgb8")
out_msg.header = cam_msg.header
self._pub_normal.publish(out_msg)
if self._pub_jet.get_num_connections() > 0:
mean = np.nanmean(depth)
std = np.nanstd(depth)
min_value = np.maximum(0.3, mean - 3 * std)
max_value = np.minimum(2.0, mean + 3 * std)
jet = imgviz.depth2rgb(depth,
min_value=min_value,
max_value=max_value)
out_msg = bridge.cv2_to_imgmsg(jet, "rgb8")
out_msg.header = cam_msg.header
self._pub_jet.publish(out_msg)
def test_depth2rgb():
data = imgviz.data.arc2017()
depthviz = imgviz.depth2rgb(data["depth"])
assert depthviz.dtype == np.uint8
H, W = data["depth"].shape[:2]
assert depthviz.shape == (H, W, 3)
def __getitem__(self, index):
example = self._dataset[index]
instance_id = self._dataset._ids[index]
image_id = osp.dirname(instance_id)
index_parent = self._dataset_parent._ids.index(image_id)
frame = self._dataset_parent.get_frame(index_parent)
viz = imgviz.tile(
[
example["rgb"],
imgviz.depth2rgb(example["pcd"][:, :, 0]),
imgviz.depth2rgb(example["pcd"][:, :, 1]),
imgviz.depth2rgb(example["pcd"][:, :, 2]),
],
border=(255, 255, 255),
)
viz = imgviz.tile([frame["rgb"], viz], (1, 2))
return viz
def __getitem__(self, i):
example = self._dataset[i]
rgb = example["color"]
depth_viz = imgviz.depth2rgb(example["depth"])
label_viz = imgviz.label2rgb(
example["result"]["labels"],
label_names=morefusion.datasets.ycb_video.class_names,
)
viz = imgviz.tile(
[rgb, depth_viz, label_viz],
shape=(1, 3),
border=(255, 255, 255),
)
viz = imgviz.resize(viz, width=1000)
return viz
def depth2rgb():
data = imgviz.data.arc2017()
depthviz = imgviz.depth2rgb(data["depth"], min_value=0.3, max_value=1)
# -------------------------------------------------------------------------
plt.figure(dpi=200)
plt.subplot(121)
plt.title("rgb")
plt.imshow(data["rgb"])
plt.axis("off")
plt.subplot(122)
plt.title("depth (colorized)")
plt.imshow(depthviz)
plt.axis("off")
img = imgviz.io.pyplot_to_numpy()
plt.close()
return img
import matplotlib.pyplot as plt
here = osp.dirname(osp.abspath(__file__)) # NOQA
# -----------------------------------------------------------------------------
# GETTING_STARTED {{
import imgviz
# sample data of rgb, depth, class label and instance masks
data = imgviz.data.arc2017()
# colorize depth image with JET colormap
depth = data["depth"]
depthviz = imgviz.depth2rgb(depth, min_value=0.3, max_value=1)
# colorize label image
class_label = data["class_label"]
labelviz = imgviz.label2rgb(class_label, label_names=data["class_names"])
# instance bboxes
rgb = data["rgb"]
bboxes = data["bboxes"].astype(int)
labels = data["labels"]
captions = [data["class_names"][l] for l in labels]
bboxviz = imgviz.instances2rgb(image=rgb, bboxes=bboxes, labels=labels, captions=captions)
# instance masks
masks = data["masks"] == 1
maskviz = imgviz.instances2rgb(image=rgb, masks=masks, labels=labels, captions=captions)
morefusion.extra.pybullet.init_world()
pybullet.resetDebugVisualizerCamera(
cameraDistance=0.5,
cameraYaw=45,
cameraPitch=-45,
cameraTargetPosition=(0, 0, 0),
)
morefusion.extra.pybullet.add_model(
visual_file=cad_file,
position=(0, 0, 0.3),
)
for _ in range(1000):
pybullet.stepSimulation()
rgb, depth, segm = morefusion.extra.pybullet.get_debug_visualizer_image()
morefusion.extra.pybullet.del_world()
viz = imgviz.tile(
[rgb, imgviz.depth2rgb(depth),
imgviz.label2rgb(segm)],
shape=(1, 3),
border=(255, 255, 255),
)
viz = imgviz.resize(viz, width=1500)
imgviz.io.pyglet_imshow(viz)
imgviz.io.pyglet_run()
def get_scene(dataset):
camera = trimesh.scene.Camera(fov=(30, 22.5))
index = 0
frame = dataset.get_frame(index)
examples = dataset.get_example(index)
scenes = {
"rgb": None,
}
camera_transform = morefusion.extra.trimesh.to_opengl_transform()
vizs = [frame["rgb"]]
for i, example in enumerate(examples):
viz = imgviz.tile(
[
example["rgb"],
imgviz.depth2rgb(example["pcd"][:, :, 0]),
imgviz.depth2rgb(example["pcd"][:, :, 1]),
imgviz.depth2rgb(example["pcd"][:, :, 2]),
],
border=(255, 255, 255),
)
viz = imgviz.draw.text_in_rectangle(
viz,
"lt",
f"visibility: {example['visibility']:.0%}",
size=30,
background=(0, 255, 0),
color=(0, 0, 0),
)
vizs.append(viz)
geom = trimesh.voxel.VoxelGrid(
example["grid_target"],
ttf.scale_and_translate(example["pitch"], example["origin"]),
).as_boxes(colors=(1.0, 0, 0, 0.5))
scenes[f"occupied_{i:04d}"] = trimesh.Scene(
geom, camera=camera, camera_transform=camera_transform)
geom = trimesh.voxel.VoxelGrid(
example["grid_nontarget"],
ttf.scale_and_translate(example["pitch"], example["origin"]),
).as_boxes(colors=(0, 1.0, 0, 0.5))
scenes[f"occupied_{i:04d}"].add_geometry(geom)
geom = trimesh.voxel.VoxelGrid(
example["grid_empty"],
ttf.scale_and_translate(example["pitch"], example["origin"]),
).as_boxes(colors=(0.5, 0.5, 0.5, 0.5))
scenes[f"empty_{i:04d}"] = trimesh.Scene(
geom, camera=camera, camera_transform=camera_transform)
scenes[f"full_occupied_{i:04d}"] = trimesh.Scene(
camera=camera, camera_transform=camera_transform)
if (example["grid_target_full"] > 0).any():
geom = trimesh.voxel.VoxelGrid(
example["grid_target_full"],
ttf.scale_and_translate(example["pitch"], example["origin"]),
).as_boxes(colors=(1.0, 0, 0, 0.5))
scenes[f"full_occupied_{i:04d}"].add_geometry(geom)
if (example["grid_nontarget_full"] > 0).any():
colors = imgviz.label2rgb(
example["grid_nontarget_full"].reshape(1, -1) +
1).reshape(example["grid_nontarget_full"].shape + (3, ))
geom = trimesh.voxel.VoxelGrid(
example["grid_nontarget_full"],
ttf.scale_and_translate(example["pitch"], example["origin"]),
).as_boxes(colors=colors)
scenes[f"full_occupied_{i:04d}"].add_geometry(geom)
dim = example["grid_target"].shape[0]
extents = np.array([dim, dim, dim]) * example["pitch"]
geom = trimesh.path.creation.box_outline(extents)
geom.apply_translation(example["origin"] +
(dim / 2 - 0.5) * example["pitch"])
scenes[f"occupied_{i:04d}"].add_geometry(geom)
scenes[f"empty_{i:04d}"].add_geometry(geom)
viz = imgviz.tile(vizs)
scenes["rgb"] = viz
return scenes
dim = 16
pitch = max(geom.extents) / dim * 1.1
origin = (-pitch * dim / 2, ) * 3
sdf = models.get_cad(class_id=2).nearest.signed_distance(points)
print(f"[{name}] pitch: {pitch}")
print(f"[{name}] dim: {dim}")
grid, _, _ = morefusion.functions.pseudo_occupancy_voxelization(
points=cuda.to_gpu(points),
sdf=cuda.to_gpu(sdf),
pitch=pitch,
origin=origin,
dims=(dim, ) * 3,
threshold=2,
)
grid = cuda.to_cpu(grid.array)
colors = imgviz.depth2rgb(grid.reshape(1, -1), min_value=0, max_value=1)
colors = colors.reshape(dim, dim, dim, 3)
colors = np.concatenate((colors, np.full((dim, dim, dim, 1), 127)), axis=3)
voxel = trimesh.voxel.VoxelGrid(grid,
ttf.scale_and_translate(pitch, origin))
geom = voxel.as_boxes()
I, J, K = zip(*np.argwhere(grid))
geom.visual.face_colors = colors[I, J, K].repeat(12, axis=0)
scene.add_geometry(geom)
scenes[name] = scene
morefusion.extra.trimesh.display_scenes(scenes, tile=(1, 2))
def get_images():
data = imgviz.data.arc2017()
yield data["rgb"]
yield imgviz.depth2rgb(data["depth"], min_value=0.3, max_value=1)
yield imgviz.label2rgb(data["class_label"])
def __getitem__(self, index):
index = self.indices[index]
result_file = args.result / f"{index:04d}.mat" # NOQA
print(result_file)
result = scipy.io.loadmat(result_file,
chars_as_strings=True,
squeeze_me=True)
frame_id = "/".join(result["frame_id"].split("/")[1:])
frame = morefusion.datasets.YCBVideoDataset.get_frame(frame_id)
rgb = frame["color"]
depth = frame["depth"]
bboxes = result["bboxes"]
K = frame["meta"]["intrinsic_matrix"]
labels = result["labels"].astype(np.int32)
masks = result["masks"].astype(bool)
keep = np.isin(labels, frame["meta"]["cls_indexes"])
bboxes = bboxes[keep]
labels = labels[keep]
masks = masks[keep]
captions = [
morefusion.datasets.ycb_video.class_names[l] for l in labels
]
detections_viz = imgviz.instances2rgb(
rgb,
labels=labels,
bboxes=bboxes,
masks=masks,
captions=captions,
font_size=15,
)
camera = trimesh.scene.Camera(resolution=(640, 480),
focal=(K[0, 0], K[1, 1]))
pybullet.connect(pybullet.DIRECT)
for class_id, pose in zip(labels, result["poses"]):
cad_file = morefusion.datasets.YCBVideoModels().get_cad_file(
class_id=class_id)
morefusion.extra.pybullet.add_model(
cad_file,
position=pose[4:],
orientation=pose[:4][[1, 2, 3, 0]],
)
(
rgb_rend,
depth_rend,
segm_rend,
) = morefusion.extra.pybullet.render_camera(np.eye(4),
fovy=camera.fov[1],
height=480,
width=640)
pybullet.disconnect()
min_value = 0.3
max_value = 2 * np.nanmedian(depth) - min_value
depth = imgviz.depth2rgb(depth,
min_value=min_value,
max_value=max_value)
depth_rend = imgviz.depth2rgb(depth_rend,
min_value=min_value,
max_value=max_value)
viz = imgviz.tile(
[rgb, depth, detections_viz, rgb_rend, depth_rend],
(2, 3),
border=(255, ) * 3,
)
viz = imgviz.resize(viz, width=1500)
return viz
def get_images():
data = imgviz.data.arc2017()
yield data['rgb']
yield imgviz.depth2rgb(data['depth'], min_value=0.3, max_value=1)
yield imgviz.label2rgb(data['class_label'])
def _process_frame(self, frame):
meta = frame["meta"]
color = frame["color"]
depth = frame["depth"]
depth_viz = imgviz.depth2rgb(depth, min_value=0, max_value=2)
label = frame["label"]
label_viz = imgviz.label2rgb(label)
labels = meta["cls_indexes"].astype(np.int32)
# NOTE: cls_mask is the same as ins_mask in YCB_Video_Dataset
masks = np.asarray([label == cls_id for cls_id in labels])
bboxes = morefusion.geometry.masks_to_bboxes(masks)
keep = ~(bboxes == 0).all(axis=1)
labels = labels[keep]
bboxes = bboxes[keep]
masks = masks[keep]
gray = imgviz.gray2rgb(imgviz.rgb2gray(color))
ins_viz = imgviz.instances2rgb(gray,
labels=labels,
bboxes=bboxes,
masks=masks)
vertmap = meta["vertmap"]
vertmap[label == 0] = np.nan
vert_viz_x = imgviz.depth2rgb(vertmap[:, :, 0])
vert_viz_y = imgviz.depth2rgb(vertmap[:, :, 1])
vert_viz_z = imgviz.depth2rgb(vertmap[:, :, 2])
roi_viz_color = []
roi_viz_depth = []
roi_viz_label = []
for bbox, mask in zip(bboxes, masks):
y1, x1, y2, x2 = bbox.round().astype(int)
mask_roi = mask[y1:y2, x1:x2]
color_roi = color[y1:y2, x1:x2].copy()
color_roi[~mask_roi] = 0
depth_roi = depth_viz[y1:y2, x1:x2].copy()
depth_roi[~mask_roi] = 0
label_roi = label_viz[y1:y2, x1:x2].copy()
label_roi[~mask_roi] = 0
roi_viz_color.append(color_roi)
roi_viz_depth.append(depth_roi)
roi_viz_label.append(label_roi)
roi_viz_color = imgviz.tile(roi_viz_color, border=(255, 255, 255))
roi_viz_depth = imgviz.tile(roi_viz_depth, border=(255, 255, 255))
roi_viz_label = imgviz.tile(roi_viz_label, border=(255, 255, 255))
viz = imgviz.tile(
[
color,
depth_viz,
label_viz,
ins_viz,
vert_viz_x,
vert_viz_y,
vert_viz_z,
np.zeros_like(color),
roi_viz_color,
roi_viz_depth,
roi_viz_label,
np.zeros_like(roi_viz_color),
],
shape=(3, 4),
border=(255, 255, 255),
)
viz = imgviz.centerize(viz, (1000, 1000))
return viz
import os.path as osp
import matplotlib.pyplot as plt
import imgviz
here = osp.dirname(osp.abspath(__file__))
if __name__ == '__main__':
data = imgviz.data.arc2017()
depthviz = imgviz.depth2rgb(data['depth'], min_value=0.3, max_value=1)
# -------------------------------------------------------------------------
fig = plt.figure(dpi=200)
plt.subplot(121)
plt.title('rgb')
plt.imshow(data['rgb'])
plt.axis('off')
plt.subplot(122)
plt.title('depth (colorized)')
plt.imshow(depthviz)
plt.axis('off')
out_file = osp.join(here, '.readme/depth2rgb.jpg')
plt.savefig(
def _callback(self,
cam_msg,
rgb_msg,
depth_msg,
ins_msg,
cls_msg,
noentry_msg=None):
bridge = cv_bridge.CvBridge()
rgb = bridge.imgmsg_to_cv2(rgb_msg, desired_encoding="rgb8")
depth = bridge.imgmsg_to_cv2(depth_msg)
if depth.dtype == np.uint16:
depth = depth.astype(np.float32) / 1000
depth[depth == 0] = np.nan
assert depth.dtype == np.float32
K = np.array(cam_msg.K).reshape(3, 3)
pcd = morefusion.geometry.pointcloud_from_depth(
depth, K[0, 0], K[1, 1], K[0, 2], K[1, 2])
ins = bridge.imgmsg_to_cv2(ins_msg)
grids_noentry = {}
if noentry_msg:
for grid in noentry_msg.grids:
instance_id = grid.instance_id
dims = (grid.dims.x, grid.dims.y, grid.dims.z)
indices = np.array(grid.indices)
k = indices % grid.dims.z
j = indices // grid.dims.z % grid.dims.y
i = indices // grid.dims.z // grid.dims.y
grid_nontarget_empty = np.zeros(dims, dtype=bool)
grid_nontarget_empty[i, j, k] = True
origin = np.array(
[grid.origin.x, grid.origin.y, grid.origin.z],
dtype=np.float32,
)
grids_noentry[instance_id] = dict(
origin=origin,
pitch=grid.pitch,
matrix=grid_nontarget_empty,
)
instance_ids = []
class_ids = []
for cls in cls_msg.classes:
instance_ids.append(cls.instance_id)
class_ids.append(cls.class_id)
instance_ids = np.array(instance_ids)
class_ids = np.array(class_ids)
examples = []
keep = []
nanmask = np.isnan(pcd).any(axis=2)
for i, (ins_id, cls_id) in enumerate(zip(instance_ids, class_ids)):
if self._with_occupancy and ins_id not in grids_noentry:
# it is inactive in some reason (e.g., on the edge)
continue
mask = ins == ins_id
if (~nanmask & mask).sum() < 50:
continue
bbox = morefusion.geometry.masks_to_bboxes([mask])[0]
y1, x1, y2, x2 = bbox.round().astype(int)
rgb_ins = rgb[y1:y2, x1:x2].copy()
rgb_ins[~mask[y1:y2, x1:x2]] = 0
rgb_ins = imgviz.centerize(rgb_ins, (256, 256), cval=0)
pcd_ins = pcd[y1:y2, x1:x2].copy()
pcd_ins[~mask[y1:y2, x1:x2]] = np.nan
pcd_ins = imgviz.centerize(pcd_ins, (256, 256),
cval=np.nan,
interpolation="nearest")
example = dict(
class_id=cls_id,
rgb=rgb_ins,
pcd=pcd_ins,
)
if grids_noentry:
example["origin"] = grids_noentry[ins_id]["origin"]
example["pitch"] = grids_noentry[ins_id]["pitch"]
example["grid_nontarget_empty"] = grids_noentry[ins_id][
"matrix"]
examples.append(example)
keep.append(i)
if not examples:
return
inputs = chainer.dataset.concat_examples(examples, device=0)
instance_ids = instance_ids[keep]
del class_ids
if self._pub_debug_rgbd.get_num_connections() > 0:
debug_rgbd = [
imgviz.tile(
[e["rgb"], imgviz.depth2rgb(e["pcd"][:, :, 2])], (1, 2))
for e in examples
]
debug_rgbd = imgviz.tile(debug_rgbd, border=(255, 255, 255))
debug_rgbd_msg = bridge.cv2_to_imgmsg(debug_rgbd, encoding="rgb8")
debug_rgbd_msg.header = rgb_msg.header
self._pub_debug_rgbd.publish(debug_rgbd_msg)
with chainer.no_backprop_mode(), chainer.using_config("train", False):
quaternion, translation, confidence = self._model.predict(**inputs)
indices = confidence.array.argmax(axis=1)
B = quaternion.shape[0]
confidence = confidence[np.arange(B), indices]
quaternion = quaternion[np.arange(B), indices]
translation = translation[np.arange(B), indices]
confidence = chainer.cuda.to_cpu(confidence.array)
quaternion = chainer.cuda.to_cpu(quaternion.array)
translation = chainer.cuda.to_cpu(translation.array)
"""
transforms = morefusion.functions.transformation_matrix(
quaternion, translation
).array
for i in range(B):
pcd_cad = self._models.get_pcd(examples[i]['class_id'])
pcd_depth = examples[i]['pcd']
pcd_depth = pcd_depth[~np.isnan(pcd_depth).any(axis=2)]
icp = morefusion.contrib.ICPRegistration(
pcd_depth=pcd_depth,
pcd_cad=pcd_cad,
transform_init=transforms[i],
)
transform = icp.register()
quaternion[i] = ttf.quaternion_from_matrix(transform)
translation[i] = ttf.translation_from_matrix(transform)
del transforms
"""
poses = ObjectPoseArray()
poses.header = rgb_msg.header
for i, (ins_id, example) in enumerate(zip(instance_ids, examples)):
"""
cls_id = example['class_id']
class_name = morefusion.datasets.ycb_video.class_names[cls_id]
morefusion.ros.loginfo_green(
f'instance_id={ins_id}, class_id={cls_id}, '
f'class_name={class_name}, confidence={confidence[i].item()}'
)
"""
if confidence[i].item() < 0.9:
continue
pose = ObjectPose()
pose.pose.position.x = translation[i][0]
pose.pose.position.y = translation[i][1]
pose.pose.position.z = translation[i][2]
pose.pose.orientation.w = quaternion[i][0]
pose.pose.orientation.x = quaternion[i][1]
pose.pose.orientation.y = quaternion[i][2]
pose.pose.orientation.z = quaternion[i][3]
pose.instance_id = ins_id
pose.class_id = examples[i]["class_id"]
poses.poses.append(pose)
self._pub_poses.publish(poses)