388 lines
11 KiB
Python
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()
|