Skip to content

feat(writer): Allow to write additional data to coco_annotations.json #1209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

feat(writer): Allow to write additional data to coco_annotations.json #1209

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
05ea1c8
Refactor: Improve COCO annotation generation with custom props
probabilisticrobotics Jul 24, 2025
138442e
Merge branch 'main' of github.com:DLR-RM/BlenderProc into extend_coco…
probabilisticrobotics Jul 29, 2025
166001a
Merge remote-tracking branch 'origin/main' into extend_coco_writer
cornerfarmer Aug 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 50 additions & 25 deletions blenderproc/python/writer/CocoWriterUtility.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,18 @@
from blenderproc.python.utility.LabelIdMapping import LabelIdMapping


def write_coco_annotations(output_dir: str, instance_segmaps: List[np.ndarray],
def write_coco_annotations(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: str = "rle", supercategory: str = "coco_annotations",
colors: List[np.ndarray],
color_file_format: str = "PNG",
mask_encoding_format: str = "rle",
supercategory: str = "coco_annotations",
append_to_existing_output: bool = True,
jpg_quality: int = 95, label_mapping: Optional[LabelIdMapping] = None,
file_prefix: str = "", indent: Optional[Union[int, str]] = None):
jpg_quality: int = 95,
label_mapping: Optional[LabelIdMapping] = None,
file_prefix: str = "",
indent: Optional[Union[int, str]] = None):
""" Writes coco annotations in the following steps:
1. Locate the seg images
2. Locate the rgb maps
Expand Down Expand Up @@ -90,15 +95,11 @@ def write_coco_annotations(output_dir: str, instance_segmaps: List[np.ndarray],
else:
raise RuntimeError(f'Unknown color_file_format={color_file_format}. Try "PNG" or "JPEG"')


new_coco_image_paths.append(target_base_path)

coco_output = _CocoWriterUtility.generate_coco_annotations(instance_segmaps,
instance_attribute_maps,
new_coco_image_paths,
supercategory,
mask_encoding_format,
existing_coco_annotations,
coco_output = _CocoWriterUtility.generate_coco_annotations(instance_segmaps, instance_attribute_maps,
new_coco_image_paths, supercategory,
mask_encoding_format, existing_coco_annotations,
label_mapping)

print("Writing coco annotations to " + coco_annotations_path)
Expand Down Expand Up @@ -144,8 +145,12 @@ def rle_to_binary_mask(rle: Dict[str, List[int]]) -> np.ndarray:
class _CocoWriterUtility:

@staticmethod
def generate_coco_annotations(inst_segmaps, inst_attribute_maps, image_paths, supercategory,
mask_encoding_format, existing_coco_annotations=None,
def generate_coco_annotations(inst_segmaps,
inst_attribute_maps,
image_paths,
supercategory,
mask_encoding_format,
existing_coco_annotations=None,
label_mapping: LabelIdMapping = None):
"""Generates coco annotations for images

Expand Down Expand Up @@ -180,8 +185,10 @@ def generate_coco_annotations(inst_segmaps, inst_attribute_maps, image_paths, su

if supercategory in [inst_supercategory, 'coco_annotations']:
if int(inst["category_id"]) not in visited_categories:
cat_dict: Dict[str, Union[str, int]] = {'id': int(inst["category_id"]),
'supercategory': inst_supercategory}
cat_dict: Dict[str, Union[str, int]] = {
'id': int(inst["category_id"]),
'supercategory': inst_supercategory
}
# Determine name of category based on label_mapping, name or category_id
if label_mapping is not None:
cat_dict["name"] = label_mapping.label_from_id(cat_dict['id'])
Expand Down Expand Up @@ -212,8 +219,12 @@ def generate_coco_annotations(inst_segmaps, inst_attribute_maps, image_paths, su
images: List[Dict[str, Union[str, int]]] = []
annotations: List[Dict[str, Union[str, int]]] = []

for inst_segmap, image_path, instance_2_category_map in zip(inst_segmaps, image_paths,
instance_2_category_maps):
for inst_segmap, image_path, instance_2_category_map, inst_attribute_map in zip(
inst_segmaps,
image_paths,
instance_2_category_maps,
inst_attribute_maps,
):

# Add coco info for image
image_id = len(images)
Expand All @@ -225,14 +236,18 @@ def generate_coco_annotations(inst_segmaps, inst_attribute_maps, image_paths, su
instances = np.delete(instances, np.where(instances == 0))
for inst in instances:
if inst in instance_2_category_map:
inst_attributes = next((item for item in inst_attribute_map if item["idx"] == inst), None)
Copy link
Member

@cornerfarmer cornerfarmer Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we maybe exclude idx and name here? They should be part of the coco annotations already anyway right?

# Calc object mask
binary_inst_mask = np.where(inst_segmap == inst, 1, 0)
# Add coco info for object in this image
annotation = _CocoWriterUtility.create_annotation_info(len(annotations) + 1,
image_id,
instance_2_category_map[inst],
binary_inst_mask,
mask_encoding_format)
annotation = _CocoWriterUtility.create_annotation_info(
len(annotations) + 1,
image_id,
instance_2_category_map[inst],
binary_inst_mask,
mask_encoding_format,
custom_props=inst_attributes,
)
if annotation is not None:
annotations.append(annotation)

Expand Down Expand Up @@ -307,8 +322,13 @@ def create_image_info(image_id: int, file_name: str, image_size: Tuple[int, int]
return image_info

@staticmethod
def create_annotation_info(annotation_id: int, image_id: int, category_id: int, binary_mask: np.ndarray,
mask_encoding_format: str, tolerance: int = 2) -> Optional[Dict[str, Union[str, int]]]:
def create_annotation_info(annotation_id: int,
image_id: int,
category_id: int,
binary_mask: np.ndarray,
mask_encoding_format: str,
tolerance: int = 2,
custom_props: Optional[Dict] = None) -> Optional[Dict[str, Union[str, int]]]:
"""Creates info section of coco annotation

:param annotation_id: integer to uniquly identify the annotation
Expand All @@ -317,6 +337,7 @@ def create_annotation_info(annotation_id: int, image_id: int, category_id: int,
:param binary_mask: A binary image mask of the object with the shape [H, W].
:param mask_encoding_format: Encoding format of the mask. Type: string.
:param tolerance: The tolerance for fitting polygons to the objects mask.
:param custom_props: A dictionary of custom properties to add to the annotation.
"""

area = _CocoWriterUtility.calc_binary_mask_area(binary_mask)
Expand Down Expand Up @@ -345,6 +366,10 @@ def create_annotation_info(annotation_id: int, image_id: int, category_id: int,
"width": binary_mask.shape[1],
"height": binary_mask.shape[0],
}

if custom_props:
annotation_info.update(custom_props)

return annotation_info

@staticmethod
Expand Down