From 4d35458e7c00ba36e48a210babd30c3475ce9f62 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Sun, 11 Aug 2024 22:53:20 +0500 Subject: [PATCH] Update `convert_segment_masks_to_yolo_seg` to support custom datasets (#15176) --- docs/en/usage/simple-utilities.md | 2 +- ultralytics/data/converter.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/en/usage/simple-utilities.md b/docs/en/usage/simple-utilities.md index 694ecf1df..2f10dc448 100644 --- a/docs/en/usage/simple-utilities.md +++ b/docs/en/usage/simple-utilities.md @@ -63,7 +63,7 @@ The converted masks will be saved in the specified output directory. ```python from ultralytics.data.converter import convert_segment_masks_to_yolo_seg -# For COCO dataset we have 80 classes +# The classes here is the total classes in the dataset, for COCO dataset we have 80 classes convert_segment_masks_to_yolo_seg(masks_dir="path/to/masks_dir", output_dir="path/to/output_dir", classes=80) ``` diff --git a/ultralytics/data/converter.py b/ultralytics/data/converter.py index 34b359d61..3464bc1b9 100644 --- a/ultralytics/data/converter.py +++ b/ultralytics/data/converter.py @@ -344,13 +344,13 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes): Args: masks_dir (str): The path to the directory where all mask images (png, jpg) are stored. output_dir (str): The path to the directory where the converted YOLO segmentation masks will be stored. - classes (int): Total classes in the dataset i.e for COCO classes=80 + classes (int): Total classes in the dataset i.e. for COCO classes=80 Example: ```python from ultralytics.data.converter import convert_segment_masks_to_yolo_seg - # for coco dataset, we have 80 classes + # The classes here is the total classes in the dataset, for COCO dataset we have 80 classes convert_segment_masks_to_yolo_seg('path/to/masks_directory', 'path/to/output/directory', classes=80) ``` @@ -373,7 +373,7 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes): """ import os - pixel_to_class_mapping = {i + 1: i for i in range(80)} + pixel_to_class_mapping = {i + 1: i for i in range(classes)} for mask_filename in os.listdir(masks_dir): if mask_filename.endswith(".png"): mask_path = os.path.join(masks_dir, mask_filename)