TIF_E4253362/backend/scripts/03_split_dataset.py

388 lines
11 KiB
Python

"""
Tahap 3: Split Dataset
Script ini membagi dataset menjadi 80% training dan 20% testing
secara stratified untuk setiap kelas.
Sesuai README:
- Input: dataset/normalized/
- Output: dataset/train/ dan dataset/test/
- Pembagian dilakukan untuk setiap kelas
- Struktur folder dibuat sama persis dengan folder normalized/
Persentase:
- Train: 80%
- Test: 20%
"""
import sys
import json
import shutil
import numpy as np
from pathlib import Path
from tqdm import tqdm
from typing import List, Tuple, Dict
import argparse
from collections import defaultdict, Counter
# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))
from utils.config import (
NORMALIZED_DIR,
TRAIN_DIR,
TEST_DIR,
CATEGORIES,
TRAIN_RATIO
)
class DatasetSplitter:
"""
Splitter untuk membagi dataset menjadi train dan test sets.
Pembagian dilakukan untuk setiap kelas secara stratified.
"""
def __init__(self, train_ratio: float = TRAIN_RATIO, random_seed: int = 42):
"""
Inisialisasi splitter.
Args:
train_ratio: Rasio data untuk training (0.8 = 80%)
random_seed: Seed untuk random
"""
self.train_ratio = train_ratio
# Set random seed
np.random.seed(random_seed)
def split_samples(
self,
samples: List[Path],
labels: List[str]
) -> Tuple[List[Path], List[str], List[Path], List[str]]:
"""
Split samples menjadi train dan test sets secara stratified.
Args:
samples: List of sample paths
labels: List of corresponding labels
Returns:
Tuple (train_samples, train_labels, test_samples, test_labels)
"""
# Group samples by label
label_to_samples = defaultdict(list)
for sample, label in zip(samples, labels):
label_to_samples[label].append(sample)
train_samples = []
train_labels = []
test_samples = []
test_labels = []
# Split setiap kelas secara proporsional
for label, class_samples in label_to_samples.items():
# Shuffle samples
indices = np.arange(len(class_samples))
np.random.shuffle(indices)
# Calculate split point
n_train = int(len(class_samples) * self.train_ratio)
n_train = max(1, n_train) # Minimal 1 sample untuk train
# Split
train_indices = indices[:n_train]
test_indices = indices[n_train:]
# Jika test set kosong, ambil 1 dari train
if len(test_indices) == 0 and len(train_indices) > 1:
test_indices = [train_indices[-1]]
train_indices = train_indices[:-1]
train_samples.extend([class_samples[i] for i in train_indices])
train_labels.extend([label] * len(train_indices))
test_samples.extend([class_samples[i] for i in test_indices])
test_labels.extend([label] * len(test_indices))
return train_samples, train_labels, test_samples, test_labels
def copy_samples(
self,
samples: List[Path],
labels: List[str],
output_dir: Path,
category: str,
split: str
) -> int:
"""
Copy samples ke output directory.
Args:
samples: List of sample paths
labels: List of labels
output_dir: Root output directory
category: Category name
split: 'train' atau 'test'
Returns:
Number of copied files
"""
copied_count = 0
for sample_path, label in tqdm(
zip(samples, labels),
total=len(samples),
desc=f"Copying {split}"
):
try:
# Create output path (struktur sama persis dengan normalized/)
# Format: {train|test}/{category}/{label}/file.npy
# output_dir sudah berupa dataset/train atau dataset/test
output_path = output_dir / category / label / sample_path.name
# Skip jika sudah ada
if output_path.exists():
continue
# Create directory
output_path.parent.mkdir(parents=True, exist_ok=True)
# Copy file
shutil.copy2(sample_path, output_path)
copied_count += 1
except Exception as e:
print(f"Error copying {sample_path}: {e}")
return copied_count
def get_normalized_files(normalized_dir: Path, category: str) -> Tuple[List[Path], List[str]]:
"""
Get semua file normalized dan labels-nya.
Args:
normalized_dir: Root normalized directory
category: Category name
Returns:
Tuple (list of files, list of labels)
"""
category_dir = normalized_dir / category
if not category_dir.exists():
return [], []
files = []
labels = []
# Iterate semua subdirectories (labels)
for label_dir in category_dir.iterdir():
if label_dir.is_dir():
label = label_dir.name
# Get semua .npy files
for npy_file in label_dir.glob("*.npy"):
files.append(npy_file)
labels.append(label)
return files, labels
def split_category(
normalized_dir: Path,
train_dir: Path,
test_dir: Path,
category: str,
train_ratio: float = TRAIN_RATIO
) -> Dict:
"""
Split dataset untuk satu kategori.
Args:
normalized_dir: Directory normalized data
train_dir: Directory untuk data train
test_dir: Directory untuk data test
category: Category name
train_ratio: Train ratio
Returns:
Statistics dictionary
"""
print(f"\n{'='*60}")
print(f"Splitting dataset untuk kategori: {category}")
print(f"{'='*60}\n")
# Get files dan labels
files, labels = get_normalized_files(normalized_dir, category)
if not files:
print(f"Warning: Tidak ada file ditemukan untuk kategori {category}")
return {}
print(f"Total samples: {len(files)}")
# Count per class
class_counts = Counter(labels)
print(f"Classes: {len(class_counts)}")
for label, count in sorted(class_counts.items()):
print(f" - {label}: {count} samples")
# Split
splitter = DatasetSplitter(train_ratio=train_ratio)
train_files, train_labels, test_files, test_labels = splitter.split_samples(files, labels)
print(f"\nAfter split:")
print(f" Train: {len(train_files)} samples ({len(train_files)/len(files)*100:.1f}%)")
print(f" Test: {len(test_files)} samples ({len(test_files)/len(files)*100:.1f}%)")
# Copy files
print("\nCopying files...")
train_copied = splitter.copy_samples(
train_files, train_labels, train_dir, category, "train"
)
test_copied = splitter.copy_samples(
test_files, test_labels, test_dir, category, "test"
)
print(f"Copied: {train_copied} train, {test_copied} test files")
# Generate label encoder
unique_labels = sorted(set(labels))
label_encoder = {label: idx for idx, label in enumerate(unique_labels)}
# Save label encoder di folder models/
model_dir = Path(__file__).parent.parent / "models"
encoder_path = model_dir / f"{category}_label_encoder.json"
encoder_path.parent.mkdir(parents=True, exist_ok=True)
with open(encoder_path, 'w') as f:
json.dump(label_encoder, f, indent=2)
print(f"\nSaved label encoder to: {encoder_path}")
# Return statistics
return {
"category": category,
"train_ratio": train_ratio,
"total_samples": len(files),
"train_samples": len(train_files),
"test_samples": len(test_files),
"num_classes": len(unique_labels),
"classes": unique_labels
}
def split_all_categories(
normalized_dir: Path = NORMALIZED_DIR,
train_dir: Path = TRAIN_DIR,
test_dir: Path = TEST_DIR,
categories: List[str] = None,
train_ratio: float = TRAIN_RATIO
):
"""
Split semua kategori dataset.
Args:
normalized_dir: Directory normalized data
train_dir: Directory untuk data train
test_dir: Directory untuk data test
categories: List kategori
train_ratio: Train ratio
"""
if categories is None:
categories = CATEGORIES
print("="*60)
print("TAHAP 3: SPLIT DATASET")
print("="*60)
print(f"Input: {normalized_dir}")
print(f"Output:")
print(f" - Train: {train_dir}")
print(f" - Test: {test_dir}")
print(f"Categories: {categories}")
print(f"Train ratio: {train_ratio*100}%")
print(f"Test ratio: {(1-train_ratio)*100}%")
print("="*60)
all_split_info = {}
for category in categories:
split_info = split_category(
normalized_dir=normalized_dir,
train_dir=train_dir,
test_dir=test_dir,
category=category,
train_ratio=train_ratio
)
if split_info:
all_split_info[category] = split_info
# Print summary table
print("\n" + "="*60)
print("SPLIT SUMMARY")
print("="*60)
print(f"{'Kategori':<20} {'Total':<10} {'Train':<10} {'Test':<10} {'Kelas':<10}")
print("-"*60)
for cat, info in all_split_info.items():
print(f"{cat:<20} {info['total_samples']:<10} {info['train_samples']:<10} "
f"{info['test_samples']:<10} {info['num_classes']:<10}")
print("-"*60)
def main():
"""Main entry point."""
parser = argparse.ArgumentParser(
description="Split dataset untuk training dan testing"
)
parser.add_argument(
"--input",
type=str,
default=None,
help="Directory normalized data (default: dataset/normalized)"
)
parser.add_argument(
"--train-dir",
type=str,
default=None,
help="Directory untuk data train (default: dataset/train)"
)
parser.add_argument(
"--test-dir",
type=str,
default=None,
help="Directory untuk data test (default: dataset/test)"
)
parser.add_argument(
"--categories",
nargs="+",
choices=["abjad_bisindo", "abjad_sibi", "angka", "kata"],
default=None,
help="Kategori yang akan diproses"
)
parser.add_argument(
"--train-ratio",
type=float,
default=TRAIN_RATIO,
help=f"Rasio training (default: {TRAIN_RATIO})"
)
args = parser.parse_args()
# Setup paths
normalized_dir = Path(args.input) if args.input else NORMALIZED_DIR
train_dir = Path(args.train_dir) if args.train_dir else TRAIN_DIR
test_dir = Path(args.test_dir) if args.test_dir else TEST_DIR
# Split
split_all_categories(
normalized_dir=normalized_dir,
train_dir=train_dir,
test_dir=test_dir,
categories=args.categories,
train_ratio=args.train_ratio
)
if __name__ == "__main__":
main()