GenHisDoc_model / filter_illustration_only.py

Jules Musquin

[ADD] creating metric and filer.py, updating README

7f300d9 3 months ago

3.12 kB

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)

	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)