def _write_frames(self):
""" Writes images, GT annotations and camera info.
"""
# Paths to the already existing chunk folders (such folders may exist
# when appending to an existing dataset).
chunk_dirs = sorted(glob.glob(os.path.join(self.chunks_dir, '*')))
chunk_dirs = [d for d in chunk_dirs if os.path.isdir(d)]
# Get ID's of the last already existing chunk and frame.
curr_chunk_id = 0
curr_frame_id = 0
if len(chunk_dirs):
last_chunk_dir = sorted(chunk_dirs)[-1]
last_chunk_gt_fpath = os.path.join(last_chunk_dir, 'scene_gt.json')
chunk_gt = load_json(last_chunk_gt_fpath, keys_to_int=True)
# Last chunk and frame ID's.
last_chunk_id = int(os.path.basename(last_chunk_dir))
last_frame_id = int(sorted(chunk_gt.keys())[-1])
# Current chunk and frame ID's.
curr_chunk_id = last_chunk_id
curr_frame_id = last_frame_id + 1
if curr_frame_id % self.frames_per_chunk == 0:
curr_chunk_id += 1
curr_frame_id = 0
# Initialize structures for the GT annotations and camera info.
chunk_gt = {}
chunk_camera = {}
if curr_frame_id != 0:
# Load GT and camera info of the chunk we are appending to.
chunk_gt = load_json(
self.chunk_gt_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
chunk_camera = load_json(
self.chunk_camera_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
# Go through all frames.
num_new_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start
for frame_id in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
# Activate frame.
bpy.context.scene.frame_set(frame_id)
# Reset data structures and prepare folders for a new chunk.
if curr_frame_id == 0:
chunk_gt = {}
chunk_camera = {}
os.makedirs(
os.path.dirname(
self.rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=0,
im_type='PNG')))
os.makedirs(
os.path.dirname(
self.depth_tpath.format(chunk_id=curr_chunk_id,
im_id=0)))
# Get GT annotations and camera info for the current frame.
chunk_gt[curr_frame_id] = self._get_frame_gt()
chunk_camera[curr_frame_id] = self._get_frame_camera()
# Copy the resulting RGB image.
rgb_output = Utility.find_registered_output_by_key("colors")
if rgb_output is None:
raise Exception("RGB image has not been rendered.")
image_type = '.png' if rgb_output['path'].endswith(
'png') else '.jpg'
rgb_fpath = self.rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type=image_type)
shutil.copyfile(rgb_output['path'] % frame_id, rgb_fpath)
# Load the resulting dist image.
dist_output = Utility.find_registered_output_by_key("distance")
if dist_output is None:
raise Exception("Distance image has not been rendered.")
depth, _, _ = self._load_and_postprocess(
dist_output['path'] % frame_id, "distance")
# Scale the depth to retain a higher precision (the depth is saved
# as a 16-bit PNG image with range 0-65535).
depth_mm = 1000.0 * depth # [m] -> [mm]
depth_mm_scaled = depth_mm / float(self.depth_scale)
# Save the scaled depth image.
depth_fpath = self.depth_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id)
save_depth(depth_fpath, depth_mm_scaled)
# Save the chunk info if we are at the end of a chunk or at the last new frame.
if ((curr_frame_id + 1) % self.frames_per_chunk == 0) or\
(frame_id == num_new_frames - 1):
# Save GT annotations.
save_json(self.chunk_gt_tpath.format(chunk_id=curr_chunk_id),
chunk_gt)
# Save camera info.
save_json(
self.chunk_camera_tpath.format(chunk_id=curr_chunk_id),
chunk_camera)
# Update ID's.
curr_chunk_id += 1
curr_frame_id = 0
else:
curr_frame_id += 1
def run(self):
""" Does the stereo global matching in the following steps:
1. Collect camera object and its state,
2. For each frame, load left and right images and call the `sgm()` methode.
3. Write the results to a numpy file.
"""
if self._avoid_output:
print("Avoid output is on, no output produced!")
return
if GlobalStorage.is_in_storage("renderer_distance_end"):
self.depth_max = GlobalStorage.get("renderer_distance_end")
else:
raise RuntimeError(
"A distance rendering has to be executed before this module is executed, "
"else the `renderer_distance_end` is not set!")
self.rgb_output_path = Utility.find_registered_output_by_key(
self.rgb_output_key)["path"]
# Collect camera and camera object
cam_ob = bpy.context.scene.camera
cam = cam_ob.data
self.width = bpy.context.scene.render.resolution_x
self.height = bpy.context.scene.render.resolution_y
print('Resolution: {}, {}'.format(self.width, self.height))
self.baseline = cam.stereo.interocular_distance
if not self.baseline:
raise Exception(
"Stereo parameters are not set. Make sure to enable RGB stereo rendering before this module."
)
if self.config.get_bool("infer_focal_length_from_fov", False):
fov = cam.angle_x if cam.angle_x else cam.angle
if not fov:
raise Exception("Could not obtain field of view angle")
self.focal_length = float(
(1.0 / tan(fov / 2.0)) * (float(self.width) / 2.0))
else:
self.focal_length = self.config.get_float("focal_length", 0.0)
if self.focal_length == 0.0:
raise Exception(
"Focal length set to 0. This is either intentional or because no value was set by the user. Either way, this needs to be corrected by setting a value > 0 or enabling 'infer_focal_length_from_fov'."
)
for frame in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
path_split = self.rgb_output_path.split(".")
path_l = "{}_L.{}".format(path_split[0], path_split[1])
path_r = "{}_R.{}".format(path_split[0], path_split[1])
imgL = load_image(path_l % frame)
imgR = load_image(path_r % frame)
depth, disparity = self.sgm(imgL, imgR)
np.save(
os.path.join(self.output_dir, "stereo-depth_%04d") % frame,
depth)
if self.config.get_bool("output_disparity", False):
np.save(
os.path.join(self.output_dir, "disparity_%04d") % frame,
disparity)
Utility.register_output(self._determine_output_dir(), "stereo-depth_",
"stereo-depth", ".npy", "1.0.0")
if self.config.get_bool("output_disparity", False):
Utility.register_output(self._determine_output_dir(), "disparity_",
"disparity", ".npy", "1.0.0")
def write(output_dir: str,
mask_encoding_format="rle",
supercategory="coco_annotations",
append_to_existing_output=False,
segmap_output_key="segmap",
segcolormap_output_key="segcolormap",
rgb_output_key="colors"):
""" Writes coco annotations in the following steps:
1. Locate the seg images
2. Locate the rgb maps
3. Locate the seg mappings
4. Read color mappings
5. For each frame write the coco annotation
:param output_dir: Output directory to write the coco annotations
:param mask_encoding_format: Encoding format of the binary masks. Default: 'rle'. Available: 'rle', 'polygon'.
:param supercategory: name of the dataset/supercategory to filter for, e.g. a specific BOP dataset set by 'bop_dataset_name' or
any loaded object with specified 'cp_supercategory'
:param append_to_existing_output: If true and if there is already a coco_annotations.json file in the output directory, the new coco
annotations will be appended to the existing file. Also the rgb images will be named such that there are
no collisions. Default: False.
:param segmap_output_key: The output key with which the segmentation images were registered. Should be the same as the output_key
of the SegMapRenderer module. Default: segmap.
:param segcolormap_output_key: The output key with which the csv file for object name/class correspondences was registered. Should be
the same as the colormap_output_key of the SegMapRenderer module. Default: segcolormap.
:param rgb_output_key: The output key with which the rgb images were registered. Should be the same as the output_key of the
RgbRenderer module. Default: colors.
"""
# Create output directory
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Find path pattern of segmentation images
segmentation_map_output = Utility.find_registered_output_by_key(
segmap_output_key)
if segmentation_map_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the SegMapRenderer module "
"before?".format(segmap_output_key))
# Find path pattern of rgb images
rgb_output = Utility.find_registered_output_by_key(rgb_output_key)
if rgb_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the RgbRenderer module "
"before?".format(rgb_output_key))
# collect all segmaps
segmentation_map_paths = []
# Find path of name class mapping csv file
segcolormap_output = Utility.find_registered_output_by_key(
segcolormap_output_key)
if segcolormap_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the SegMapRenderer module "
"with 'map_by' set to 'instance' before?".format(
segcolormap_output_key))
# read colormappings, which include object name/class to integer mapping
inst_attribute_maps = []
with open(segcolormap_output["path"], 'r') as csvfile:
reader = csv.DictReader(csvfile)
for mapping in reader:
inst_attribute_maps.append(mapping)
coco_annotations_path = os.path.join(output_dir,
"coco_annotations.json")
# Calculate image numbering offset, if append_to_existing_output is activated and coco data exists
if append_to_existing_output and os.path.exists(coco_annotations_path):
with open(coco_annotations_path, 'r') as fp:
existing_coco_annotations = json.load(fp)
image_offset = max(
[image["id"]
for image in existing_coco_annotations["images"]]) + 1
else:
image_offset = 0
existing_coco_annotations = None
# collect all RGB paths
new_coco_image_paths = []
# for each rendered frame
for frame in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
segmentation_map_paths.append(segmentation_map_output["path"] %
frame)
source_path = rgb_output["path"] % frame
target_path = os.path.join(
output_dir,
os.path.basename(rgb_output["path"] % (frame + image_offset)))
shutil.copyfile(source_path, target_path)
new_coco_image_paths.append(os.path.basename(target_path))
coco_output = CocoWriterUtility.generate_coco_annotations(
segmentation_map_paths, new_coco_image_paths, inst_attribute_maps,
supercategory, mask_encoding_format, existing_coco_annotations)
print("Writing coco annotations to " + coco_annotations_path)
with open(coco_annotations_path, 'w') as fp:
json.dump(coco_output, fp)
def run(self):
""" Writes coco annotations in the following steps:
1. Locat the seg images
2. Locat the rgb maps
3. Locat the seg maps
4. Read color mappings
5. For each frame write the coco annotation
"""
if self._avoid_rendering:
print("Avoid rendering is on, no output produced!")
return
# Find path pattern of segmentation images
segmentation_map_output = Utility.find_registered_output_by_key(
self.segmap_output_key)
if segmentation_map_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the SegMapRenderer module "
"before?".format(self.segmap_output_key))
# Find path pattern of rgb images
rgb_output = Utility.find_registered_output_by_key(self.rgb_output_key)
if rgb_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the RgbRenderer module "
"before?".format(self.rgb_output_key))
# collect all segmaps
segmentation_map_paths = []
# Find path of name class mapping csv file
segcolormap_output = Utility.find_registered_output_by_key(
self.segcolormap_output_key)
if segcolormap_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the SegMapRenderer module "
"with 'map_by' set to 'instance' before?".format(
self.segcolormap_output_key))
# read colormappings, which include object name/class to integer mapping
inst_attribute_maps = []
with open(segcolormap_output["path"], 'r') as csvfile:
reader = csv.DictReader(csvfile)
for mapping in reader:
inst_attribute_maps.append(mapping)
coco_annotations_path = os.path.join(self._coco_data_dir,
"coco_annotations.json")
# Calculate image numbering offset, if append_to_existing_output is activated and coco data exists
if self.config.get_bool(
"append_to_existing_output",
False) and os.path.exists(coco_annotations_path):
with open(coco_annotations_path, 'r') as fp:
existing_coco_annotations = json.load(fp)
image_offset = max(
[image["id"]
for image in existing_coco_annotations["images"]]) + 1
else:
image_offset = 0
existing_coco_annotations = None
# collect all RGB paths
new_coco_image_paths = []
# for each rendered frame
for frame in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
segmentation_map_paths.append(segmentation_map_output["path"] %
frame)
source_path = rgb_output["path"] % frame
target_path = os.path.join(
self._coco_data_dir,
os.path.basename(rgb_output["path"] % (frame + image_offset)))
shutil.copyfile(source_path, target_path)
new_coco_image_paths.append(os.path.basename(target_path))
coco_output = CocoUtility.generate_coco_annotations(
segmentation_map_paths, new_coco_image_paths, inst_attribute_maps,
self._supercategory, self.mask_encoding_format,
existing_coco_annotations)
print("Writing coco annotations to " + coco_annotations_path)
with open(coco_annotations_path, 'w') as fp:
json.dump(coco_output, fp)
def _write_frames(self):
""" Writes images, GT annotations and camera info.
"""
with Locker():
# Paths to the already existing chunk folders (such folders may exist
# when appending to an existing dataset).
chunk_dirs = sorted(glob.glob(os.path.join(self.dataset_dir, '*')))
chunk_dirs = [d for d in chunk_dirs if os.path.isdir(d)]
# Get ID's of the last already existing chunk and frame.
curr_chunk_id = len(chunk_dirs)
curr_frame_id = 0
#Make subfolders to chunks
os.makedirs(os.path.dirname(
self.rgb_tpath.format(chunk_id=curr_chunk_id, im_id=0, im_type='PNG'))) # only takes the directory
os.makedirs(os.path.dirname(
self.depth_tpath.format(chunk_id=curr_chunk_id, im_id=0)))
os.makedirs(os.path.dirname(
self.camera_poses_tpath.format(chunk_id=curr_chunk_id, im_id=0))) # only takes the directory
os.makedirs(os.path.dirname(
self.object_poses_tpath.format(chunk_id=curr_chunk_id, im_id=0)))
# Initialize structures for the camera info.
chunk_camera = {}
chunk_depth_fpath = []
# Go through all frames.
num_new_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start
end_frame = bpy.context.scene.frame_end if not self._avoid_rendering else bpy.context.scene.frame_start
####DEBUG ###
# image = []
######
for frame_id in range(bpy.context.scene.frame_start, end_frame):
# Activate frame.
bpy.context.scene.frame_set(frame_id)
### DO RGB ###
# Copy the resulting RGB image.
rgb_output = Utility.find_registered_output_by_key("colors")
if rgb_output is None:
raise Exception("RGB image has not been rendered.")
image_type = '.png' if rgb_output['path'].endswith('png') else '.jpg'
rgb_fpath = self.rgb_tpath.format(chunk_id=curr_chunk_id, im_id=curr_frame_id, im_type=image_type)
shutil.copyfile(rgb_output['path'] % frame_id, rgb_fpath)
### DO Object States ###
object_state_output = Utility.find_registered_output_by_key("object_states")
if object_state_output is None:
raise Exception("No Object State")
ob_state = np.load(object_state_output["path"]%frame_id)
np.save(self.object_poses_tpath.format(chunk_id=curr_chunk_id, im_id=curr_frame_id), json.loads(np.string_(ob_state))[0]['matrix_world'])
#### DO DEPTH ###
# Load the resulting dist image.
dist_output = Utility.find_registered_output_by_key("distance")
if dist_output is None:
raise Exception("Distance image has not been rendered.")
depth, _, _ = self._load_and_postprocess(dist_output['path'] % frame_id, "distance")
# Scale the depth to retain a higher precision (the depth is saved
# as a 16-bit PNG image with range 0-65535).
depth_mm = 1000.0 * depth # [m] -> [mm]
depth_mm_scaled = depth_mm / float(self.depth_scale)
# Save the scaled depth image.
depth_fpath = self.depth_tpath.format(chunk_id=curr_chunk_id, im_id=curr_frame_id)
chunk_depth_fpath.append(depth_fpath)
save_depth(depth_fpath, depth_mm_scaled)
########
##### DO Camera Intrinsics ####
# Get GT annotations and camera info for the current frame.
chunk_camera[curr_frame_id] = self._get_intrinsics_camera()
# Save the chunk info if we are at the end of a chunk or at the last new frame.
if (frame_id == num_new_frames - 1):
# Save camera info.
save_json(self.chunk_camera_tpath.format(chunk_id=curr_chunk_id), chunk_camera)
else:
curr_frame_id += 1
# DO camera_extrinsics ####
cam_posemat_in_world = self.get_extrinsics_camera()
np.save(self.camera_poses_tpath.format(chunk_id=curr_chunk_id, im_id=curr_frame_id),
cam_posemat_in_world)
def _write_frames(chunks_dir,
dataset_objects,
depth_scale: float = 1.0,
frames_per_chunk: int = 1000,
m2mm: bool = True,
ignore_dist_thres: float = 100.,
save_world2cam: bool = True):
""" Writes images, GT annotations and camera info.
"""
# Format of the depth images.
depth_ext = '.png'
rgb_tpath = os.path.join(chunks_dir, '{chunk_id:06d}', 'rgb',
'{im_id:06d}' + '{im_type}')
depth_tpath = os.path.join(chunks_dir, '{chunk_id:06d}', 'depth',
'{im_id:06d}' + depth_ext)
chunk_camera_tpath = os.path.join(chunks_dir, '{chunk_id:06d}',
'scene_camera.json')
chunk_gt_tpath = os.path.join(chunks_dir, '{chunk_id:06d}',
'scene_gt.json')
# Paths to the already existing chunk folders (such folders may exist
# when appending to an existing dataset).
chunk_dirs = sorted(glob.glob(os.path.join(chunks_dir, '*')))
chunk_dirs = [d for d in chunk_dirs if os.path.isdir(d)]
# Get ID's of the last already existing chunk and frame.
curr_chunk_id = 0
curr_frame_id = 0
if len(chunk_dirs):
last_chunk_dir = sorted(chunk_dirs)[-1]
last_chunk_gt_fpath = os.path.join(last_chunk_dir, 'scene_gt.json')
chunk_gt = BopWriterUtility._load_json(last_chunk_gt_fpath,
keys_to_int=True)
# Last chunk and frame ID's.
last_chunk_id = int(os.path.basename(last_chunk_dir))
last_frame_id = int(sorted(chunk_gt.keys())[-1])
# Current chunk and frame ID's.
curr_chunk_id = last_chunk_id
curr_frame_id = last_frame_id + 1
if curr_frame_id % frames_per_chunk == 0:
curr_chunk_id += 1
curr_frame_id = 0
# Initialize structures for the GT annotations and camera info.
chunk_gt = {}
chunk_camera = {}
if curr_frame_id != 0:
# Load GT and camera info of the chunk we are appending to.
chunk_gt = BopWriterUtility._load_json(
chunk_gt_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
chunk_camera = BopWriterUtility._load_json(
chunk_camera_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
# Go through all frames.
num_new_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start
for frame_id in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
# Activate frame.
bpy.context.scene.frame_set(frame_id)
# Reset data structures and prepare folders for a new chunk.
if curr_frame_id == 0:
chunk_gt = {}
chunk_camera = {}
os.makedirs(
os.path.dirname(
rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=0,
im_type='PNG')))
os.makedirs(
os.path.dirname(
depth_tpath.format(chunk_id=curr_chunk_id, im_id=0)))
# Get GT annotations and camera info for the current frame.
# Output translation gt in m or mm
unit_scaling = 1000. if m2mm else 1.
chunk_gt[curr_frame_id] = BopWriterUtility._get_frame_gt(
dataset_objects, unit_scaling, ignore_dist_thres)
chunk_camera[curr_frame_id] = BopWriterUtility._get_frame_camera(
save_world2cam, depth_scale, unit_scaling)
# Copy the resulting RGB image.
rgb_output = Utility.find_registered_output_by_key("colors")
if rgb_output is None:
raise Exception("RGB image has not been rendered.")
image_type = '.png' if rgb_output['path'].endswith(
'png') else '.jpg'
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type=image_type)
shutil.copyfile(rgb_output['path'] % frame_id, rgb_fpath)
# Load the resulting dist image.
dist_output = Utility.find_registered_output_by_key("distance")
if dist_output is None:
raise Exception("Distance image has not been rendered.")
distance = WriterUtility.load_output_file(
Utility.resolve_path(dist_output['path'] % frame_id))
depth = PostProcessingUtility.dist2depth(distance)
# Scale the depth to retain a higher precision (the depth is saved
# as a 16-bit PNG image with range 0-65535).
depth_mm = 1000.0 * depth # [m] -> [mm]
depth_mm_scaled = depth_mm / float(depth_scale)
# Save the scaled depth image.
depth_fpath = depth_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id)
BopWriterUtility._save_depth(depth_fpath, depth_mm_scaled)
# Save the chunk info if we are at the end of a chunk or at the last new frame.
if ((curr_frame_id + 1) % frames_per_chunk == 0) or\
(frame_id == num_new_frames - 1):
# Save GT annotations.
BopWriterUtility._save_json(
chunk_gt_tpath.format(chunk_id=curr_chunk_id), chunk_gt)
# Save camera info.
BopWriterUtility._save_json(
chunk_camera_tpath.format(chunk_id=curr_chunk_id),
chunk_camera)
# Update ID's.
curr_chunk_id += 1
curr_frame_id = 0
else:
curr_frame_id += 1
def write(output_dir: str,
instance_segmaps: List[np.ndarray] = [],
instance_attribute_maps: List[dict] = [],
colors: List[np.ndarray] = [],
color_file_format: str = "PNG",
mask_encoding_format="rle",
supercategory="coco_annotations",
append_to_existing_output: bool = True,
segmap_output_key="segmap",
segcolormap_output_key="segcolormap",
rgb_output_key="colors",
jpg_quality: int = 95,
label_mapping: LabelIdMapping = None):
""" Writes coco annotations in the following steps:
1. Locate the seg images
2. Locate the rgb maps
3. Locate the seg mappings
4. Read color mappings
5. For each frame write the coco annotation
:param output_dir: Output directory to write the coco annotations
:param instance_segmaps: List of instance segmentation maps
:param segcolormaps: per-frame mappings with idx, class and optionally supercategory/bop_dataset_name
:param colors: List of color images
:param color_file_format: Format to save color images in
:param mask_encoding_format: Encoding format of the binary masks. Default: 'rle'. Available: 'rle', 'polygon'.
:param supercategory: name of the dataset/supercategory to filter for, e.g. a specific BOP dataset set by 'bop_dataset_name' or
any loaded object with specified 'cp_supercategory'
:param append_to_existing_output: If true and if there is already a coco_annotations.json file in the output directory, the new coco
annotations will be appended to the existing file. Also the rgb images will be named such that there are
no collisions.
:param segmap_output_key: The output key with which the segmentation images were registered. Should be the same as the output_key
of the SegMapRenderer module. Default: segmap.
:param segcolormap_output_key: The output key with which the csv file for object name/class correspondences was registered. Should be
the same as the colormap_output_key of the SegMapRenderer module. Default: segcolormap.
:param rgb_output_key: The output key with which the rgb images were registered. Should be the same as the output_key of the
RgbRenderer module. Default: colors.
:param label_mapping: The label mapping which should be used to label the categories based on their ids.
If None, is given then the `name` field in the csv files is used or - if not existing - the category id itself is used.
"""
# Create output directory
os.makedirs(os.path.join(output_dir, 'coco_data'), exist_ok=True)
if not instance_segmaps:
# Find path pattern of segmentation images
segmentation_map_output = Utility.find_registered_output_by_key(
segmap_output_key)
if segmentation_map_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the SegMapRenderer module "
"before?".format(segmap_output_key))
if not colors:
# Find path pattern of rgb images
rgb_output = Utility.find_registered_output_by_key(rgb_output_key)
if rgb_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the RgbRenderer module "
"before?".format(rgb_output_key))
if not instance_attribute_maps:
# Find path of name class mapping csv file
segcolormap_output = Utility.find_registered_output_by_key(
segcolormap_output_key)
if segcolormap_output is None:
raise Exception(
"There is no output registered with key {}. Are you sure you ran the SegMapRenderer module "
"with 'map_by' set to 'instance' before?".format(
segcolormap_output_key))
coco_annotations_path = os.path.join(
output_dir, "coco_data/coco_annotations.json")
# Calculate image numbering offset, if append_to_existing_output is activated and coco data exists
if append_to_existing_output and os.path.exists(coco_annotations_path):
with open(coco_annotations_path, 'r') as fp:
existing_coco_annotations = json.load(fp)
image_offset = max(
[image["id"]
for image in existing_coco_annotations["images"]]) + 1
else:
image_offset = 0
existing_coco_annotations = None
# collect all RGB paths
new_coco_image_paths = []
# collect all mappings from csv (backwards compat)
segcolormaps = []
# collect all instance segmaps (backwards compat)
inst_segmaps = []
# for each rendered frame
for frame in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
if not instance_attribute_maps:
# read colormappings, which include object name/class to integer mapping
segcolormap = []
with open(segcolormap_output["path"] % frame, 'r') as csvfile:
reader = csv.DictReader(csvfile)
for mapping in reader:
segcolormap.append(mapping)
segcolormaps.append(segcolormap)
if not instance_segmaps:
# Load segmaps (backwards compat)
segmap = np.load(segmentation_map_output["path"] % frame)
inst_channel = int(segcolormap[0]['channel_instance'])
inst_segmaps.append(segmap[:, :, inst_channel])
if colors:
color_rgb = colors[frame]
color_bgr = color_rgb[..., ::-1].copy()
if color_file_format == 'PNG':
target_base_path = 'coco_data/rgb_{:04d}.png'.format(
frame + image_offset)
target_path = os.path.join(output_dir, target_base_path)
cv2.imwrite(target_path, color_bgr)
elif color_file_format == 'JPEG':
target_base_path = 'coco_data/rgb_{:04d}.jpg'.format(
frame + image_offset)
target_path = os.path.join(output_dir, target_base_path)
cv2.imwrite(target_path, color_bgr,
[int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
else:
raise ('Unknown color_file_format={}. Try "PNG" or "JPEG"'.
format(color_file_format))
else:
source_path = rgb_output["path"] % frame
target_base_path = os.path.join(
'coco_data',
os.path.basename(rgb_output["path"] %
(frame + image_offset)))
target_path = os.path.join(output_dir, target_base_path)
shutil.copyfile(source_path, target_path)
new_coco_image_paths.append(target_base_path)
instance_attibute_maps = segcolormaps if segcolormaps else instance_attribute_maps
instance_segmaps = inst_segmaps if inst_segmaps else instance_segmaps
coco_output = CocoWriterUtility.generate_coco_annotations(
instance_segmaps, instance_attibute_maps, new_coco_image_paths,
supercategory, mask_encoding_format, existing_coco_annotations,
label_mapping)
print("Writing coco annotations to " + coco_annotations_path)
with open(coco_annotations_path, 'w') as fp:
json.dump(coco_output, fp)
def _write_frames(chunks_dir: str,
dataset_objects: list,
depths: List[np.ndarray] = [],
colors: List[np.ndarray] = [],
color_file_format: str = "PNG",
depth_scale: float = 1.0,
frames_per_chunk: int = 1000,
m2mm: bool = True,
ignore_dist_thres: float = 100.,
save_world2cam: bool = True,
jpg_quality: int = 95):
"""Write each frame's ground truth into chunk directory in BOP format
:param chunks_dir: Path to the output directory of the current chunk.
:param dataset_objects: Save annotations for these objects.
:param depths: List of depth images in m to save
:param colors: List of color images to save
:param color_file_format: File type to save color images. Available: "PNG", "JPEG"
:param jpg_quality: If color_file_format is "JPEG", save with the given quality.
:param depth_scale: Multiply the uint16 output depth image with this factor to get depth in mm. Used to trade-off between depth accuracy
and maximum depth value. Default corresponds to 65.54m maximum depth and 1mm accuracy.
:param ignore_dist_thres: Distance between camera and object after which object is ignored. Mostly due to failed physics.
:param m2mm: Original bop annotations and models are in mm. If true, we convert the gt annotations to mm here. This
is needed if BopLoader option mm2m is used.
:param frames_per_chunk: Number of frames saved in each chunk (called scene in BOP)
"""
# Format of the depth images.
depth_ext = '.png'
rgb_tpath = os.path.join(chunks_dir, '{chunk_id:06d}', 'rgb',
'{im_id:06d}' + '{im_type}')
depth_tpath = os.path.join(chunks_dir, '{chunk_id:06d}', 'depth',
'{im_id:06d}' + depth_ext)
chunk_camera_tpath = os.path.join(chunks_dir, '{chunk_id:06d}',
'scene_camera.json')
chunk_gt_tpath = os.path.join(chunks_dir, '{chunk_id:06d}',
'scene_gt.json')
# Paths to the already existing chunk folders (such folders may exist
# when appending to an existing dataset).
chunk_dirs = sorted(glob.glob(os.path.join(chunks_dir, '*')))
chunk_dirs = [d for d in chunk_dirs if os.path.isdir(d)]
# Get ID's of the last already existing chunk and frame.
curr_chunk_id = 0
curr_frame_id = 0
if len(chunk_dirs):
last_chunk_dir = sorted(chunk_dirs)[-1]
last_chunk_gt_fpath = os.path.join(last_chunk_dir, 'scene_gt.json')
chunk_gt = BopWriterUtility._load_json(last_chunk_gt_fpath,
keys_to_int=True)
# Last chunk and frame ID's.
last_chunk_id = int(os.path.basename(last_chunk_dir))
last_frame_id = int(sorted(chunk_gt.keys())[-1])
# Current chunk and frame ID's.
curr_chunk_id = last_chunk_id
curr_frame_id = last_frame_id + 1
if curr_frame_id % frames_per_chunk == 0:
curr_chunk_id += 1
curr_frame_id = 0
# Initialize structures for the GT annotations and camera info.
chunk_gt = {}
chunk_camera = {}
if curr_frame_id != 0:
# Load GT and camera info of the chunk we are appending to.
chunk_gt = BopWriterUtility._load_json(
chunk_gt_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
chunk_camera = BopWriterUtility._load_json(
chunk_camera_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
# Go through all frames.
num_new_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start
if len(depths) != len(colors) != num_new_frames:
raise Exception(
"The amount of images stored in the depths/colors does not correspond to the amount"
"of images specified by frame_start to frame_end.")
for frame_id in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
# Activate frame.
bpy.context.scene.frame_set(frame_id)
# Reset data structures and prepare folders for a new chunk.
if curr_frame_id == 0:
chunk_gt = {}
chunk_camera = {}
os.makedirs(
os.path.dirname(
rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=0,
im_type='PNG')))
os.makedirs(
os.path.dirname(
depth_tpath.format(chunk_id=curr_chunk_id, im_id=0)))
# Get GT annotations and camera info for the current frame.
# Output translation gt in m or mm
unit_scaling = 1000. if m2mm else 1.
chunk_gt[curr_frame_id] = BopWriterUtility._get_frame_gt(
dataset_objects, unit_scaling, ignore_dist_thres)
chunk_camera[curr_frame_id] = BopWriterUtility._get_frame_camera(
save_world2cam, depth_scale, unit_scaling)
if colors:
color_rgb = colors[frame_id]
color_bgr = color_rgb[..., ::-1].copy()
if color_file_format == 'PNG':
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type='.png')
cv2.imwrite(rgb_fpath, color_bgr)
elif color_file_format == 'JPEG':
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type='.jpg')
cv2.imwrite(rgb_fpath, color_bgr,
[int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
else:
rgb_output = Utility.find_registered_output_by_key("colors")
if rgb_output is None:
raise Exception("RGB image has not been rendered.")
color_ext = '.png' if rgb_output['path'].endswith(
'png') else '.jpg'
# Copy the resulting RGB image.
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type=color_ext)
shutil.copyfile(rgb_output['path'] % frame_id, rgb_fpath)
if depths:
depth = depths[frame_id]
else:
# Load the resulting dist image.
dist_output = Utility.find_registered_output_by_key("distance")
if dist_output is None:
raise Exception("Distance image has not been rendered.")
distance = WriterUtility.load_output_file(Utility.resolve_path(
dist_output['path'] % frame_id),
remove=False)
depth = PostProcessingUtility.dist2depth(distance)
# Scale the depth to retain a higher precision (the depth is saved
# as a 16-bit PNG image with range 0-65535).
depth_mm = 1000.0 * depth # [m] -> [mm]
depth_mm_scaled = depth_mm / float(depth_scale)
# Save the scaled depth image.
depth_fpath = depth_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id)
BopWriterUtility._save_depth(depth_fpath, depth_mm_scaled)
# Save the chunk info if we are at the end of a chunk or at the last new frame.
if ((curr_frame_id + 1) % frames_per_chunk == 0) or\
(frame_id == num_new_frames - 1):
# Save GT annotations.
BopWriterUtility._save_json(
chunk_gt_tpath.format(chunk_id=curr_chunk_id), chunk_gt)
# Save camera info.
BopWriterUtility._save_json(
chunk_camera_tpath.format(chunk_id=curr_chunk_id),
chunk_camera)
# Update ID's.
curr_chunk_id += 1
curr_frame_id = 0
else:
curr_frame_id += 1
def _write_frames(self):
""" Writes images, GT annotations and camera info.
"""
with Locker():
# Paths to the already existing chunk folders (such folders may exist
# when appending to an existing dataset).
chunk_dirs = sorted(glob.glob(os.path.join(self.dataset_dir, '*')))
chunk_dirs = [d for d in chunk_dirs if os.path.isdir(d)]
# Get ID's of the last already existing chunk and frame.
curr_chunk_id = len(chunk_dirs)
curr_frame_id = 0
#Make subfolders to chunks
os.makedirs(
os.path.dirname(
self.rgb_tpath.format(
chunk_id=curr_chunk_id, im_id=0,
im_type='PNG'))) # only takes the directory
os.makedirs(
os.path.dirname(
self.depth_tpath.format(chunk_id=curr_chunk_id, im_id=0)))
os.makedirs(
os.path.dirname(
self.segmentations_tpath.format(chunk_id=curr_chunk_id,
im_id=0)))
os.makedirs(
os.path.dirname(
self.correspondances_tpath.format(chunk_id=curr_chunk_id,
im_id_from=0,
im_id_to=0)))
os.makedirs(
os.path.dirname(
self.semantic_map_tpath.format(chunk_id=curr_chunk_id,
im_id_from=0,
im_id_to=0)))
# Initialize structures for the camera info.
chunk_camera = {}
chunk_depth_fpath = []
# Go through all frames.
num_new_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start
end_frame = bpy.context.scene.frame_end if not self._avoid_rendering else bpy.context.scene.frame_start
####DEBUG ###
# image = []
######
for frame_id in range(bpy.context.scene.frame_start, end_frame):
# Activate frame.
bpy.context.scene.frame_set(frame_id)
background_objects = []
foreground_objects = []
for obj in bpy.context.scene.objects:
if ("manip_object" in obj):
if (not obj['manip_object']):
background_objects.append(obj.name)
if (obj['manip_object']):
foreground_objects.append(obj.name)
save_json(
self.foreground_background_objects_tpath.format(
chunk_id=curr_chunk_id), {
"foreground_objects": foreground_objects,
"background_objects": background_objects
})
### DO RGB ###
# Copy the resulting RGB image.
rgb_output = Utility.find_registered_output_by_key("colors")
if rgb_output is None:
raise Exception("RGB image has not been rendered.")
image_type = '.png' if rgb_output['path'].endswith(
'png') else '.jpg'
rgb_fpath = self.rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type=image_type)
shutil.copyfile(rgb_output['path'] % frame_id, rgb_fpath)
####~DEBUG #####
# im =cv2.imread(rgb_fpath,-1)
# image.append(im)
########
#####
#### DO DEPTH ###
# Load the resulting dist image.
dist_output = Utility.find_registered_output_by_key("distance")
if dist_output is None:
raise Exception("Distance image has not been rendered.")
depth, _, _ = self._load_and_postprocess(
dist_output['path'] % frame_id, "distance")
# Scale the depth to retain a higher precision (the depth is saved
# as a 16-bit PNG image with range 0-65535).
depth_mm = 1000.0 * depth # [m] -> [mm]
depth_mm_scaled = depth_mm / float(self.depth_scale)
# Save the scaled depth image.
depth_fpath = self.depth_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id)
chunk_depth_fpath.append(depth_fpath)
save_depth(depth_fpath, depth_mm_scaled)
########
#### DO SEG MAP ####
seg_output = Utility.find_registered_output_by_key("segmap")
segmap = self._load_file(
Utility.resolve_path(seg_output['path'] % frame_id))
segmap_fpath = self.segmentations_tpath.format(
chunk_id=curr_chunk_id, im_id=curr_frame_id)
save_segmap(segmap_fpath, segmap)
seg_labels = Utility.find_registered_output_by_key("segcolormap")
seg_labels_fpath = self.segmentations_labels_tpath.format(
chunk_id=curr_chunk_id)
shutil.copyfile(seg_labels['path'], seg_labels_fpath)
##### DO Camera posemat ####
# Get GT annotations and camera info for the current frame.
chunk_camera[curr_frame_id] = self._get_frame_camera()
# Save the chunk info if we are at the end of a chunk or at the last new frame.
if (frame_id == num_new_frames - 1):
# Save camera info.
save_json(
self.chunk_camera_tpath.format(chunk_id=curr_chunk_id),
chunk_camera)
else:
curr_frame_id += 1
#### DEBUG ####
# average_pixel_value = np.mean([np.mean(np.asarray(i).astype(float)) for i in image])
# average_pixel_std = np.std([np.mean(np.asarray(i).astype(float)) for i in image])
# average_in_pixel_std = np.mean([np.std(np.asarray(i).astype(float)) for i in image])
# print('AVERAGE_PIXEL VAL {}'.format(average_pixel_value))
# print('AVERAGE_PIXEL std between images {}'.format(average_pixel_std))
# print('AVERAGE_PIXEL std within images {}'.format(average_in_pixel_std))
#########
#### DO CORRESPONDANCES
for i in range(len(chunk_depth_fpath)):
for j in range(i + 1, len(chunk_depth_fpath)):
# j = i+1
# if(j == len(chunk_depth_fpath)):
# continue
self.get_and_save_correspondances(i, j, chunk_depth_fpath,
chunk_camera, curr_chunk_id)
