Coding-python-标签-核验txt中类别数并显示

代码

import os

def validate_labels(label_dir, num_classes):
    """
    Validates labels in a directory against a given number of classes.
    It checks if any label exceeds the number of classes defined and collects all unique labels.

    :param label_dir: Directory path where the label files (.txt) are stored.
    :param num_classes: Number of classes defined in the dataset.
    """
    found_labels = set()  # Use a set to store unique labels
    for filename in os.listdir(label_dir):
        if filename.endswith('.txt'):
            filepath = os.path.join(label_dir, filename)
            with open(filepath, 'r') as f:
                for line in f:
                    parts = line.split()
                    if len(parts) > 0:
                        label = int(parts[0])
                        if label==0:
                            print(filename)
                        found_labels.add(label)  # Add label to the set
                        if label >= num_classes:
                            print(f"ERROR: Label {label} in file {filename} exceeds the number of classes ({num_classes}).")
                            return False
    print("Unique labels found:", sorted(list(found_labels)))  # Print sorted list of unique labels
    if max(found_labels) + 1 == num_classes:
        print("All labels are within the defined class range.")
        return True
    else:
        print(f"Warning: Found labels up to {max(found_labels)}, but {num_classes} classes were expected.")
        return max(found_labels) + 1 == num_classes

# Example usage:
label_directory = 'train/labels'
number_of_classes = 17  # Replace with your actual number of classes

# Validate labels
is_valid = validate_labels(label_directory, number_of_classes)
if is_valid:
    print("Validation passed.")
else:
    print("Validation failed.")