Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ultralytics Assistant 2024-09-06 03:54:35 +08:00 committed by GitHub
parent 95d54828bb
commit ac2c2be8f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 45 additions and 62 deletions

View file

@ -370,13 +370,10 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes):
mask_yolo_03.txt
mask_yolo_04.txt
"""
import os
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)
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # Read the mask image in grayscale
for mask_path in Path(masks_dir).iterdir():
if mask_path.suffix == ".png":
mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) # Read the mask image in grayscale
img_height, img_width = mask.shape # Get image dimensions
LOGGER.info(f"Processing {mask_path} imgsz = {img_height} x {img_width}")
@ -406,7 +403,7 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes):
yolo_format.append(round(point[1] / img_height, 6))
yolo_format_data.append(yolo_format)
# Save Ultralytics YOLO format data to file
output_path = os.path.join(output_dir, os.path.splitext(mask_filename)[0] + ".txt")
output_path = Path(output_dir) / f"{Path(mask_filename).stem}.txt"
with open(output_path, "w") as file:
for item in yolo_format_data:
line = " ".join(map(str, item))