Spaces:
Running
Running
| import yaml | |
| CLASSES_FILE = "backend/yolo_dataset/test_v1/classes.txt" | |
| DATA_YAML = "backend/yolo_dataset/test_v1/data.yaml" | |
| def main(): | |
| with open(CLASSES_FILE, 'r') as f: | |
| classes = [line.strip() for line in f.readlines() if line.strip()] | |
| with open(DATA_YAML, 'r') as f: | |
| data = yaml.safe_load(f) | |
| data['nc'] = len(classes) | |
| data['names'] = {i: name for i, name in enumerate(classes)} | |
| with open(DATA_YAML, 'w') as f: | |
| yaml.dump(data, f, default_flow_style=False) | |
| print(f"Updated data.yaml with {len(classes)} classes.") | |
| if __name__ == "__main__": | |
| main() | |