""" Tahap 10: Training Angka Script ini melatih model TCN untuk mengenali angka dalam bahasa isyarat. Sesuai README: - Input: dataset/train/angka/ - Output: models/angka/model_angka.pt """ 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 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 AngkaDataset(Dataset): """Dataset untuk angka.""" 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: self.samples.append(np.load(npy_file)) 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_angka(data_dir: Path, model_dir: Path, epochs: int = EPOCHS, batch_size: int = BATCH_SIZE, lr: float = LEARNING_RATE): device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") print(f"\nLoading dataset from: {data_dir}") train_dataset = AngkaDataset(data_dir / "train" / "angka") test_dataset = AngkaDataset(data_dir / "test" / "angka") print(f"Train samples: {len(train_dataset)}, 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())}") 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) 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()):,}") criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=lr) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5) history = {"train_loss": [], "train_acc": [], "test_loss": [], "test_acc": []} best_acc = 0 print("\n" + "="*60) print("Training Angka Model") print("="*60) for epoch in range(epochs): print(f"\nEpoch {epoch+1}/{epochs}") 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_acc: best_acc = test_acc best_state = model.state_dict().copy() print(f"*** Best model! Acc: {best_acc:.2f}% ***") if best_state: model.load_state_dict(best_state) # Save model model_dir.mkdir(parents=True, exist_ok=True) angka_model_dir = model_dir / "angka" angka_model_dir.mkdir(exist_ok=True) model_path = angka_model_dir / "model_angka.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()} with open(angka_model_dir / "label_encoder.json", 'w') as f: json.dump({"label_to_idx": train_dataset.label_to_idx, "idx_to_label": idx_to_label}, f, indent=2) # Save history with open(angka_model_dir / "history.json", '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('Angka Model Accuracy') plt.xlabel('Epoch') plt.ylabel('Accuracy (%)') plt.legend() plt.grid(True) plt.savefig(ACCURACY_DIR / "angka_accuracy.png") plt.close() print(f"Accuracy plot saved to: {ACCURACY_DIR / 'angka_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('Angka Model Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend() plt.grid(True) plt.savefig(LOSS_DIR / "angka_loss.png") plt.close() print(f"Loss plot saved to: {LOSS_DIR / 'angka_loss.png'}") print(f"\nBest Test Accuracy: {best_acc:.2f}%") return history def main(): parser = argparse.ArgumentParser(description="Training model angka") parser.add_argument("--data-dir", type=str, default=None) parser.add_argument("--model-dir", type=str, default=None) 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 10: TRAINING ANGKA") print("="*60) train_angka(data_dir, model_dir, args.epochs, args.batch_size, args.lr) if __name__ == "__main__": main()