60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
import os
|
|
|
|
# ----------------------------
|
|
# Configuration
|
|
# ----------------------------
|
|
|
|
src_dir = "datasets/visiope/test/labels"
|
|
dest_dir = "datasets/_unified/test/labels"
|
|
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
|
|
# Reference class order you want to follow
|
|
"""[
|
|
'w_pawn','w_knight','w_bishop','w_rook','w_queen','w_king',
|
|
'b_pawn','b_knight','b_bishop','b_rook','b_queen','b_king'
|
|
]"""
|
|
|
|
reference_classes = [
|
|
'w_pawn','w_knight','w_bishop','w_rook','w_queen','w_king',
|
|
'b_pawn','b_knight','b_bishop','b_rook','b_queen','b_king'
|
|
]
|
|
|
|
# Current class order in your dataset (change this to match your dataset!)
|
|
current_classes = ['bishop', 'black-bishop', 'black-king', 'black-knight', 'black-pawn', 'black-queen', 'black-rook',
|
|
'white-bishop', 'white-king', 'white-knight', 'white-pawn', 'white-queen', 'white-rook']
|
|
|
|
# ----------------------------
|
|
# Build index mapping
|
|
# ----------------------------
|
|
index_map = {current_classes.index(cls): reference_classes.index(cls) for cls in current_classes}
|
|
|
|
# ----------------------------
|
|
# Process each label file
|
|
# ----------------------------
|
|
|
|
count = 0
|
|
for filename in os.listdir(src_dir):
|
|
if filename.endswith(".txt"):
|
|
input_path = os.path.join(src_dir, filename)
|
|
output_path = os.path.join(dest_dir, filename)
|
|
|
|
with open(input_path, "r") as f:
|
|
lines = f.readlines()
|
|
|
|
new_lines = []
|
|
for line in lines:
|
|
parts = line.strip().split()
|
|
old_idx = int(parts[0])
|
|
new_idx = index_map[old_idx]
|
|
new_lines.append(" ".join([str(new_idx)] + parts[1:]))
|
|
|
|
with open(output_path, "w") as f:
|
|
f.write("\n".join(new_lines))
|
|
|
|
if count%100 == 0:
|
|
print(count)
|
|
count += 1
|
|
|
|
print(f"All labels remapped and saved to '{dest_dir}'")
|