preprocessing.py

import os
import shutil
import pandas as pd
from pathlib import Path

# ------------------ Preprocessing function ------------------
def reorganize_dataset(split_dir: str, csv_file: str, out_dir: str):
    """
    Reorganizes the dataset splitting it into as many folders as there are characters (labels).
    split_dir: directory with original images
    csv_file: CSV file path with labels
    out_dir: output directory
    """

    df = pd.read_csv(csv_file)

    # Characters columns
    label_cols = [c for c in df.columns if c not in ["filename", "Unlabeled"]]

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Images with more than one label
    ambiguous_images = []

    for _, row in df.iterrows():
        filename = row["filename"]
        # Find the character (column with label 1)
        # Case 1: one label
        labels = [col for col in label_cols if row[col] == 1]
        # Case 2: unlabeled
        if len(labels) == 0:
            print(f"No labels found for {filename}.")
            # Skip
            continue
        # Case 3: more labels (ambiguous)
        elif len(labels) > 1:
            ambiguous_images.append((filename, labels))

        # Use the first label found
        label = labels[0].strip()
        src = Path(split_dir) / filename
        dst_dir = out_dir / label
        dst_dir.mkdir(parents=True, exist_ok=True)
        dst = dst_dir / filename

        # Copy original images in the new folders
        if src.exists():
            shutil.copy(src, dst)
        else:
            print(f"[WARN] File not found: {src}")

    # Report ambiguous images
    if ambiguous_images:
        print(f"\n⚠️ Images with more labels found in {out_dir}:")
        for fname, labels in ambiguous_images:
            print(f"  - {fname}: {labels} (copied in '{labels[0].strip()}')")
    else:
        print(f"\n✅ No ambiguous images found in {out_dir}.")

    print(f"Images copied in {out_dir}.")


# -------------- Execution function --------------
if __name__ == "__main__":
    # train
    reorganize_dataset(
        split_dir="Anime -Naruto-.v3i.multiclass/train",
        csv_file="Anime -Naruto-.v3i.multiclass/train/_classes.csv",
        out_dir="dataset/train"
    )

    # valid
    reorganize_dataset(
        split_dir="Anime -Naruto-.v3i.multiclass/valid",
        csv_file="Anime -Naruto-.v3i.multiclass/valid/_classes.csv",
        out_dir="dataset/valid"
    )

    # test
    reorganize_dataset(
        split_dir="Anime -Naruto-.v3i.multiclass/test",
        csv_file="Anime -Naruto-.v3i.multiclass/test/_classes.csv",
        out_dir="dataset/test"
    )