File size: 3,118 Bytes
7f300d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import os
import sys
from pathlib import Path
def filter_illustrations_only(target_path: Path, inplace: bool = False) -> None:
"""
Filtre les annotations YOLO pour ne garder que la classe 'Illustration' (0).
Supprime toutes les autres classes (Initial, Ornament, Stamp, Table).
"""
# Mapping : seule la classe 0 (Illustration) est conservée
label_mapping = {
0: 0, # Illustration -> Gardée
1: None, # Initial -> Supprimée
2: None, # Ornament -> Supprimée
3: None, # Stamp -> Supprimée
4: None, # Table -> Supprimée
}
# Détection automatique du dossier 'labels'
if target_path.name == "labels":
input_dir = target_path
output_dir = target_path if inplace else target_path.parent / "labels_output"
else:
input_dir = target_path / "labels"
output_dir = input_dir if inplace else target_path / "labels_output"
if not input_dir.exists():
print(f"Erreur : Le dossier '{input_dir}' n'existe pas.")
return
if not inplace:
os.makedirs(output_dir, exist_ok=True)
total_files = 0
total_kept = 0
total_deleted = 0
empty_files_count = 0
for filename in os.listdir(input_dir):
if not filename.endswith(".txt"):
continue
total_files += 1
input_file = input_dir / filename
output_file = output_dir / filename
with open(input_file, "r", encoding="utf-8") as f:
lines = f.readlines()
new_lines = []
for line in lines:
parts = line.strip().split()
if not parts:
continue
try:
original_label = int(parts[0])
except ValueError:
continue # Ligne mal formée, ignorée
new_label = label_mapping.get(original_label)
if new_label is None:
total_deleted += 1
continue
parts[0] = str(new_label)
new_lines.append(" ".join(parts))
total_kept += 1
# Écriture du fichier (vide si aucune illustration)
with open(output_file, "w", encoding="utf-8") as f:
if new_lines:
f.write("\n".join(new_lines) + "\n")
else:
empty_files_count += 1
f.write("")
print("--- Résumé du traitement ---")
print(f"Fichiers traités : {total_files}")
print(f"Annotations conservées (Illu) : {total_kept}")
print(f"Annotations supprimées : {total_deleted}")
print(f"Fichiers .txt résultants vides : {empty_files_count}")
print(f"Dossier de sortie : {output_dir.resolve()}")
print("Traitement terminé avec succès !")
if __name__ == "__main__":
if len(sys.argv) < 2:
print("Usage: uv run filter_illustration_only.py <chemin_du_dossier>")
sys.exit(1)
dataset_path = Path(sys.argv[1])
# mets inplace=True si tu veux écraser les fichiers d'origine dans 'labels/'
filter_illustrations_only(dataset_path, inplace=False) |