prepare_segmentation_data.py

import os

def convert_box_to_poly(label_path):
    with open(label_path, 'r') as f:
        lines = f.readlines()
    
    new_lines = []
    for line in lines:
        parts = list(map(float, line.strip().split()))
        if len(parts) == 5:
            # It's a box [cls, x, y, w, h]
            cls, x, y, w, h = parts
            # Convert to 4-point polygon: x1 y1, x2 y1, x2 y2, x1 y2
            x1, y1 = x - w/2, y - h/2
            x2, y2 = x + w/2, y + h/2
            # YOLO polygon format: cls x1 y1 x2 y2 x3 y3 x4 y4 ...
            # Using 4 points to represent the rectangle
            poly_parts = [int(cls), x1, y1, x2, y1, x2, y2, x1, y2]
            new_lines.append(" ".join(map(str, poly_parts)) + "\n")
        else:
            # It's already a polygon or other format, keep it
            new_lines.append(line)
            
    with open(label_path, 'w') as f:
        f.writelines(new_lines)

def prepare_data():
    label_dir = 'dataset/train/labels'
    for filename in os.listdir(label_dir):
        if filename.endswith('.txt'):
            convert_box_to_poly(os.path.join(label_dir, filename))
    print("Dataset converted to uniform polygon format.")

if __name__ == "__main__":
    prepare_data()