HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Update convert_segment_masks_to_yolo_seg to support custom datasets (#15176)

Browse source
This commit is contained in:
Muhammad Rizwan Munawar 2024-08-11 22:53:20 +05:00 • committed by GitHub
parent 3e4a581c35
commit 4d35458e7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

2
docs/en/usage/simple-utilities.md
View file

@ -63,7 +63,7 @@ The converted masks will be saved in the specified output directory.
```python ```python
from ultralytics.data.converter import convert_segment_masks_to_yolo_seg 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) convert_segment_masks_to_yolo_seg(masks_dir="path/to/masks_dir", output_dir="path/to/output_dir", classes=80)
``` ```

6
ultralytics/data/converter.py
View file

@ -344,13 +344,13 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes):
Args: Args:
masks_dir (str): The path to the directory where all mask images (png, jpg) are stored. 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. 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: Example:
```python ```python
from ultralytics.data.converter import convert_segment_masks_to_yolo_seg 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) 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 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): for mask_filename in os.listdir(masks_dir):
if mask_filename.endswith(".png"): if mask_filename.endswith(".png"):
mask_path = os.path.join(masks_dir, mask_filename) mask_path = os.path.join(masks_dir, mask_filename)