""" Tahap 4: Training Kata Script ini melatih model TCN untuk mengenali kata dalam bahasa isyarat. Sesuai README: - Input: dataset/train/kata/ - Output: models/kata/model_kata.pt - Output tambahan: Accuracy, Loss, Training History """ import sys import json import numpy as np from pathlib import Path from typing import Tuple import argparse import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset from tqdm import tqdm import matplotlib.pyplot as plt # Add parent directory to path sys.path.insert(0, str(Path(__file__).parent.parent)) from utils.tcn import TCNClassifier from utils.config import ( TRAIN_DIR, MODEL_DIR, FEATURE_SIZE, BATCH_SIZE, EPOCHS, LEARNING_RATE, ACCURACY_DIR, LOSS_DIR ) class KataDataset(Dataset): """Dataset untuk kata.""" def __init__(self, data_dir: Path): self.data_dir = data_dir self.samples = [] self.labels = [] self.label_to_idx = {} self._load_data() def _load_data(self): if not self.data_dir.exists(): return label_dirs = sorted([d for d in self.data_dir.iterdir() if d.is_dir()]) self.label_to_idx = {d.name: idx for idx, d in enumerate(label_dirs)} for label_dir in label_dirs: label_idx = self.label_to_idx[label_dir.name] for npy_file in label_dir.glob("*.npy"): try: sequence = np.load(npy_file) self.samples.append(sequence) self.labels.append(label_idx) except Exception as e: print(f"Error loading {npy_file}: {e}") def __len__(self): return len(self.samples) def __getitem__(self, idx) -> Tuple: return torch.FloatTensor(self.samples[idx]), torch.LongTensor([self.labels[idx]]).squeeze() def train_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss, correct, total = 0, 0, 0 for sequences, labels in tqdm(train_loader, desc="Training"): sequences, labels = sequences.to(device), labels.to(device) optimizer.zero_grad() outputs = model(sequences) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() return total_loss / len(train_loader), correct / total * 100 def evaluate(model, eval_loader, criterion, device): model.eval() total_loss, correct, total = 0, 0, 0 with torch.no_grad(): for sequences, labels in eval_loader: sequences, labels = sequences.to(device), labels.to(device) outputs = model(sequences) loss = criterion(outputs, labels) total_loss += loss.item() _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() return total_loss / len(eval_loader), correct / total * 100 def train_kata(data_dir: Path, model_dir: Path, epochs: int = EPOCHS, batch_size: int = BATCH_SIZE, learning_rate: float = LEARNING_RATE): """Training model untuk kata.""" # Setup device device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") # Create dataset print(f"\nLoading dataset from: {data_dir}") train_dataset = KataDataset(data_dir / "train" / "kata") test_dataset = KataDataset(data_dir / "test" / "kata") print(f"Train samples: {len(train_dataset)}") print(f"Test samples: {len(test_dataset)}") if len(train_dataset) == 0: raise ValueError("No training samples found!") num_classes = len(train_dataset.label_to_idx) print(f"Number of classes: {num_classes}") print(f"Classes: {list(train_dataset.label_to_idx.keys())}") # Create data loaders train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0) # Create model model = TCNClassifier( input_size=FEATURE_SIZE, num_classes=num_classes, num_channels=[64, 128, 256] ).to(device) print(f"\nTotal parameters: {sum(p.numel() for p in model.parameters()):,}") # Loss and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=5 ) # Training loop history = { "train_loss": [], "train_acc": [], "test_loss": [], "test_acc": [] } best_test_acc = 0 best_model_state = None print("\n" + "="*60) print("Training Kata Model") print("="*60) for epoch in range(epochs): print(f"\nEpoch {epoch+1}/{epochs}") print("-" * 40) train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device) test_loss, test_acc = evaluate(model, test_loader, criterion, device) scheduler.step(test_loss) history["train_loss"].append(train_loss) history["train_acc"].append(train_acc) history["test_loss"].append(test_loss) history["test_acc"].append(test_acc) print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%") print(f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%") if test_acc > best_test_acc: best_test_acc = test_acc best_model_state = model.state_dict().copy() print(f"*** New best model! Test Acc: {best_test_acc:.2f}% ***") # Load best model if best_model_state is not None: model.load_state_dict(best_model_state) # Save model (sesuai README: models/kata/model_kata.pt) model_dir.mkdir(parents=True, exist_ok=True) kata_model_dir = model_dir / "kata" kata_model_dir.mkdir(exist_ok=True) model_path = kata_model_dir / "model_kata.pt" torch.save(model, model_path) print(f"\nModel saved to: {model_path}") # Save label encoder idx_to_label = {v: k for k, v in train_dataset.label_to_idx.items()} encoder_path = kata_model_dir / "label_encoder.json" with open(encoder_path, 'w') as f: json.dump({"label_to_idx": train_dataset.label_to_idx, "idx_to_label": idx_to_label}, f, indent=2) # Save history history_path = kata_model_dir / "history.json" with open(history_path, 'w') as f: json.dump(history, f, indent=2) # Plot Accuracy plt.figure(figsize=(10, 6)) plt.plot(history['train_acc'], label='Train Accuracy') plt.plot(history['test_acc'], label='Test (Validation) Accuracy') plt.title('Kata Model Accuracy') plt.xlabel('Epoch') plt.ylabel('Accuracy (%)') plt.legend() plt.grid(True) plt.savefig(ACCURACY_DIR / "kata_accuracy.png") plt.close() print(f"Accuracy plot saved to: {ACCURACY_DIR / 'kata_accuracy.png'}") # Plot Loss plt.figure(figsize=(10, 6)) plt.plot(history['train_loss'], label='Train Loss') plt.plot(history['test_loss'], label='Test (Validation) Loss') plt.title('Kata Model Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend() plt.grid(True) plt.savefig(LOSS_DIR / "kata_loss.png") plt.close() print(f"Loss plot saved to: {LOSS_DIR / 'kata_loss.png'}") print("\n" + "="*60) print(f"Best Test Accuracy: {best_test_acc:.2f}%") print("="*60) return history def main(): """Main entry point.""" parser = argparse.ArgumentParser(description="Training model kata") parser.add_argument("--data-dir", type=str, default=None, help="Directory dataset (default: dataset/)") parser.add_argument("--model-dir", type=str, default=None, help="Directory model (default: models/)") parser.add_argument("--epochs", type=int, default=EPOCHS) parser.add_argument("--batch-size", type=int, default=BATCH_SIZE) parser.add_argument("--lr", type=float, default=LEARNING_RATE) args = parser.parse_args() data_dir = Path(args.data_dir) if args.data_dir else TRAIN_DIR.parent model_dir = Path(args.model_dir) if args.model_dir else MODEL_DIR print("="*60) print("TAHAP 4: TRAINING KATA") print("="*60) train_kata(data_dir, model_dir, args.epochs, args.batch_size, args.lr) if __name__ == "__main__": main()