stupdi/ASLTranslator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
ASLTranslator/training.py
Stupdi Go 2cbd178ba8 soy cooked pt 2
2026-01-19 23:21:35 -06:00

549 lines
18 KiB
Python
Raw Permalink Blame History

import os
import json
import math
import numpy as np
import polars as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from multiprocessing import Pool, cpu_count
from functools import partial
from tqdm import tqdm
from collections import Counter
# ===============================
# GPU CONFIGURATION
# ===============================
print("=" * 60)
print("GPU CONFIGURATION")
print("=" * 60)
if torch.cuda.is_available():
print(f"✓ CUDA available!")
print(f"✓ GPU: {torch.cuda.get_device_name(0)}")
device = torch.device('cuda:0')
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True
else:
print("✗ CUDA not available, using CPU")
device = torch.device('cpu')
print("=" * 60)
# ===============================
# SELECTED LANDMARK INDICES
# ===============================
IMPORTANT_FACE_INDICES = sorted(list(set([
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
55, 65, 66, 105, 107, 336, 296, 334,
33, 133, 160, 159, 158, 144, 145, 153,
362, 263, 387, 386, 385, 373, 374, 380,
1, 2, 98, 327,
61, 185, 40, 39, 37, 0, 267, 269, 270, 409,
291, 146, 91, 181, 84, 17, 314, 405, 321, 375,
78, 191, 80, 81, 82, 13, 312, 311, 310, 415,
308, 324, 318, 402, 317, 14, 87, 178, 88, 95
])))
NUM_FACE_POINTS = len(IMPORTANT_FACE_INDICES)
NUM_HAND_POINTS = 21 * 2
TOTAL_POINTS_PER_FRAME = NUM_HAND_POINTS + NUM_FACE_POINTS
# ===============================
# ENHANCED DATA EXTRACTION (POLARS)
# ===============================
def extract_multi_landmarks(path, min_valid_frames=3):
"""
Extract both hands + selected face landmarks with modality flags
Returns: dict with 'landmarks', 'left_hand_valid', 'right_hand_valid', 'face_valid'
"""
try:
df = pl.read_parquet(path)
seq = []
left_valid_frames = []
right_valid_frames = []
face_valid_frames = []
all_types = df.select("type").unique().to_series().to_list()
# Check if we have at least one of the required types
has_data = any(t in all_types for t in ["left_hand", "right_hand", "face"])
if not has_data:
return None
# Get all frames (might not start at 0)
frames = sorted(df.select("frame").unique().to_series().to_list())
if len(frames) < min_valid_frames:
return None
for frame in frames:
frame_df = df.filter(pl.col("frame") == frame)
frame_points = np.full((TOTAL_POINTS_PER_FRAME, 3), np.nan, dtype=np.float32)
pos = 0
left_valid = False
right_valid = False
face_valid = False
# Left hand (need at least 10 valid points)
left = frame_df.filter(pl.col("type") == "left_hand")
if left.height > 0:
valid_count = 0
for i in range(21):
row = left.filter(pl.col("landmark_index") == i)
if row.height > 0:
coords = row.select(["x", "y", "z"]).row(0)
if all(c is not None for c in coords):
frame_points[pos] = coords
valid_count += 1
pos += 1
left_valid = (valid_count >= 10)
else:
pos += 21
# Right hand (need at least 10 valid points)
right = frame_df.filter(pl.col("type") == "right_hand")
if right.height > 0:
valid_count = 0
for i in range(21):
row = right.filter(pl.col("landmark_index") == i)
if row.height > 0:
coords = row.select(["x", "y", "z"]).row(0)
if all(c is not None for c in coords):
frame_points[pos] = coords
valid_count += 1
pos += 1
right_valid = (valid_count >= 10)
else:
pos += 21
# Face (need at least 30% of selected landmarks)
face = frame_df.filter(pl.col("type") == "face")
if face.height > 0:
valid_count = 0
for idx in IMPORTANT_FACE_INDICES:
row = face.filter(pl.col("landmark_index") == idx)
if row.height > 0:
coords = row.select(["x", "y", "z"]).row(0)
if all(c is not None for c in coords):
frame_points[pos] = coords
valid_count += 1
pos += 1
face_valid = (valid_count >= len(IMPORTANT_FACE_INDICES) * 0.3)
# Accept frame if we have at least 20% valid data overall
valid_ratio = 1 - np.isnan(frame_points).mean()
if valid_ratio >= 0.20:
frame_points = np.nan_to_num(frame_points, nan=0.0)
seq.append(frame_points)
left_valid_frames.append(left_valid)
right_valid_frames.append(right_valid)
face_valid_frames.append(face_valid)
if len(seq) < min_valid_frames:
return None
return {
'landmarks': np.stack(seq),
'left_hand_valid': np.array(left_valid_frames),
'right_hand_valid': np.array(right_valid_frames),
'face_valid': np.array(face_valid_frames)
}
except Exception as e:
# Uncomment for debugging:
# print(f"Error processing {path}: {e}")
return None
def get_features_sequence(landmarks_data, max_frames=100):
"""
Enhanced feature extraction with separate modality processing
"""
if landmarks_data is None:
return None, None, None
landmarks_3d = landmarks_data['landmarks']
if len(landmarks_3d) == 0:
return None, None, None
T, N, _ = landmarks_3d.shape
# Separate modalities for independent normalization
left_hand = landmarks_3d[:, :21, :]
right_hand = landmarks_3d[:, 21:42, :]
face = landmarks_3d[:, 42:, :]
# Independent centering per modality
features_list = []
for modality, valid_mask in [
(left_hand, landmarks_data['left_hand_valid']),
(right_hand, landmarks_data['right_hand_valid']),
(face, landmarks_data['face_valid'])
]:
# Center on modality-specific mean
valid_frames = modality[valid_mask] if valid_mask.any() else modality
if len(valid_frames) > 0:
center = np.mean(valid_frames, axis=(0, 1), keepdims=True)
spread = np.std(valid_frames, axis=(0, 1), keepdims=True).max()
else:
center = 0
spread = 1
modality_norm = (modality - center) / max(spread, 1e-6)
flat = modality_norm.reshape(T, -1)
# Deltas
deltas = np.zeros_like(flat)
if T > 1:
deltas[1:] = flat[1:] - flat[:-1]
features_list.append(flat)
features_list.append(deltas)
# Combine all modalities
features = np.concatenate(features_list, axis=1)
# Create modality availability mask (which frames have which modalities)
modality_mask = np.stack([
landmarks_data['left_hand_valid'],
landmarks_data['right_hand_valid'],
landmarks_data['face_valid']
], axis=1).astype(np.float32) # (T, 3)
# Pad/truncate
if T < max_frames:
pad = np.zeros((max_frames - T, features.shape[1]), dtype=np.float32)
features = np.concatenate([features, pad], axis=0)
mask_pad = np.zeros((max_frames - T, 3), dtype=np.float32)
modality_mask = np.concatenate([modality_mask, mask_pad], axis=0)
frame_mask = np.zeros(max_frames, dtype=bool)
frame_mask[:T] = True
else:
features = features[:max_frames]
modality_mask = modality_mask[:max_frames]
frame_mask = np.ones(max_frames, dtype=bool)
features = np.nan_to_num(features, nan=0.0, posinf=0.0, neginf=0.0)
features = np.clip(features, -30, 30)
return features.astype(np.float32), frame_mask, modality_mask
def process_row(row_data, base_path, max_frames=100):
"""Process a single row - expects tuple of (path, sign)"""
path_rel, sign = row_data
path = os.path.join(base_path, path_rel)
if not os.path.exists(path):
return None, None, None, None
try:
lm_data = extract_multi_landmarks(path)
if lm_data is None:
return None, None, None, None
feat, frame_mask, modality_mask = get_features_sequence(lm_data, max_frames)
if feat is None:
return None, None, None, None
return feat, frame_mask, modality_mask, sign
except Exception:
return None, None, None, None
# ===============================
# ENHANCED MODEL WITH MODALITY AWARENESS
# ===============================
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=128):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
return x + self.pe[:, :x.size(1)]
class ModalityAwareTransformer(nn.Module):
def __init__(self, input_dim, num_classes, d_model=384, nhead=8, num_layers=5):
super().__init__()
# Main projection
self.proj = nn.Linear(input_dim, d_model)
# Modality embedding (3 modalities: left_hand, right_hand, face)
self.modality_embed = nn.Linear(3, d_model)
self.norm_in = nn.LayerNorm(d_model)
self.pos = PositionalEncoding(d_model)
enc_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=d_model * 4,
dropout=0.15,
activation='gelu',
batch_first=True,
norm_first=True
)
self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers)
self.head = nn.Sequential(
nn.LayerNorm(d_model),
nn.Dropout(0.25),
nn.Linear(d_model, num_classes)
)
def forward(self, x, modality_mask=None, key_padding_mask=None):
# Project features
x = self.proj(x)
# Add modality information if available
if modality_mask is not None:
mod_embed = self.modality_embed(modality_mask)
x = x + mod_embed
x = self.norm_in(x)
x = self.pos(x)
x = self.encoder(x, src_key_padding_mask=key_padding_mask)
# Weighted average (giving more weight to frames with more valid modalities)
if modality_mask is not None:
weights = modality_mask.sum(dim=-1, keepdim=True) + 1e-6 # (B, T, 1)
weights = weights / weights.sum(dim=1, keepdim=True)
x = (x * weights).sum(dim=1)
else:
x = x.mean(dim=1)
return self.head(x)
def load_kaggle_asl_data(base_path):
"""Load training metadata using Polars"""
train_path = os.path.join(base_path, "train.csv")
train_df = pl.read_csv(train_path)
return train_df, None
# ===============================
# MAIN
# ===============================
def main():
base_path = "asl_kaggle"
max_frames = 100
MIN_SAMPLES_PER_CLASS = 5
print("\nLoading metadata...")
train_df, _ = load_kaggle_asl_data(base_path)
print(f"Total samples in train.csv: {train_df.height}")
# Convert to simple tuples for multiprocessing compatibility
rows = [(row[0], row[1]) for row in train_df.select(["path", "sign"]).iter_rows()]
print("\nProcessing sequences with BOTH hands + FACE (enhanced)...")
print("This may take a few minutes...")
with Pool(cpu_count()) as pool:
results = list(tqdm(
pool.imap(
partial(process_row, base_path=base_path, max_frames=max_frames),
rows,
chunksize=80
),
total=len(rows),
desc="Landmarks extraction"
))
X_list, frame_masks_list, modality_masks_list, y_list = [], [], [], []
failed_count = 0
for feat, frame_mask, modality_mask, sign in results:
if feat is not None and frame_mask is not None:
X_list.append(feat)
frame_masks_list.append(frame_mask)
modality_masks_list.append(modality_mask)
y_list.append(sign)
else:
failed_count += 1
if not X_list:
print(f"\n❌ No valid sequences extracted!")
print(f"Failed to process: {failed_count}/{len(results)} files")
print("\nTroubleshooting tips:")
print("1. Check that parquet files contain 'left_hand', 'right_hand', or 'face' types")
print("2. Verify files have at least 3 frames")
print("3. Ensure landmark data is not all NaN")
return
print(f"\n✓ Successfully processed: {len(X_list)}/{len(results)} files")
print(f"✗ Failed: {failed_count}/{len(results)} files")
X = np.stack(X_list)
frame_masks = np.stack(frame_masks_list)
modality_masks = np.stack(modality_masks_list)
print(f"\nExtracted {len(X):,} sequences")
print(f"Feature shape: {X.shape[1:]} (input_dim = {X.shape[2]})")
print(f"Modality mask shape: {modality_masks.shape}")
# Global normalization
X = np.clip(X, -30, 30)
mean = X.mean(axis=(0, 1), keepdims=True)
std = X.std(axis=(0, 1), keepdims=True) + 1e-8
X = (X - mean) / std
# Labels
le = LabelEncoder()
y = le.fit_transform(y_list)
# Filter rare classes
counts = Counter(y)
valid = [k for k, v in counts.items() if v >= MIN_SAMPLES_PER_CLASS]
mask = np.isin(y, valid)
X = X[mask]
frame_masks = frame_masks[mask]
modality_masks = modality_masks[mask]
y = y[mask]
le = LabelEncoder()
y = le.fit_transform(y)
print(f"After filtering: {len(X):,} samples | {len(le.classes_)} classes")
# Analyze modality usage
print("\nModality statistics:")
print(f" Sequences with left hand: {(modality_masks[:, :, 0].sum(axis=1) > 0).mean() * 100:.1f}%")
print(f" Sequences with right hand: {(modality_masks[:, :, 1].sum(axis=1) > 0).mean() * 100:.1f}%")
print(f" Sequences with face: {(modality_masks[:, :, 2].sum(axis=1) > 0).mean() * 100:.1f}%")
# Split
X_tr, X_te, fm_tr, fm_te, mm_tr, mm_te, y_tr, y_te = train_test_split(
X, frame_masks, modality_masks, y, test_size=0.15, stratify=y, random_state=42
)
# Dataset
class ASLMultiDataset(Dataset):
def __init__(self, X, frame_masks, modality_masks, y):
self.X = torch.from_numpy(X).float()
self.frame_masks = torch.from_numpy(frame_masks).bool()
self.modality_masks = torch.from_numpy(modality_masks).float()
self.y = torch.from_numpy(y).long()
def __len__(self):
return len(self.X)
def __getitem__(self, idx):
return self.X[idx], self.frame_masks[idx], self.modality_masks[idx], self.y[idx]
batch_size = 64 if device.type == 'cuda' else 32
train_loader = DataLoader(
ASLMultiDataset(X_tr, fm_tr, mm_tr, y_tr),
batch_size=batch_size, shuffle=True,
num_workers=4, pin_memory=device.type == 'cuda'
)
test_loader = DataLoader(
ASLMultiDataset(X_te, fm_te, mm_te, y_te),
batch_size=batch_size * 2, shuffle=False,
num_workers=4, pin_memory=device.type == 'cuda'
)
# Enhanced model
model = ModalityAwareTransformer(
input_dim=X.shape[2],
num_classes=len(le.classes_),
d_model=384,
nhead=8,
num_layers=5
).to(device)
print(f"\nModel parameters: {sum(p.numel() for p in model.parameters()):,}")
criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
optimizer = optim.AdamW(model.parameters(), lr=4e-4, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10)
best_acc = 0.0
epochs = 70
print("\n" + "=" * 60)
print("TRAINING START")
print("=" * 60)
for epoch in range(epochs):
model.train()
total_loss = correct = total = 0
for x, frame_mask, modality_mask, yb in tqdm(train_loader, desc=f"Epoch {epoch + 1}", leave=False):
x = x.to(device)
frame_mask = frame_mask.to(device)
modality_mask = modality_mask.to(device)
yb = yb.to(device)
key_padding_mask = ~frame_mask
optimizer.zero_grad(set_to_none=True)
logits = model(x, modality_mask=modality_mask, key_padding_mask=key_padding_mask)
loss = criterion(logits, yb)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
correct += (logits.argmax(-1) == yb).sum().item()
total += yb.size(0)
train_acc = correct / total * 100
# Eval
model.eval()
correct = total = 0
with torch.no_grad():
for x, frame_mask, modality_mask, yb in test_loader:
x = x.to(device)
frame_mask = frame_mask.to(device)
modality_mask = modality_mask.to(device)
yb = yb.to(device)
key_padding_mask = ~frame_mask
logits = model(x, modality_mask=modality_mask, key_padding_mask=key_padding_mask)
correct += (logits.argmax(-1) == yb).sum().item()
total += yb.size(0)
test_acc = correct / total * 100
scheduler.step()
print(f"[{epoch + 1:2d}/{epochs}] Loss: {total_loss / len(train_loader):.4f} | "
f"Train: {train_acc:.2f}% | Test: {test_acc:.2f}%", end="")
if test_acc > best_acc:
best_acc = test_acc
torch.save(model.state_dict(), "best_asl_modality_aware.pth")
print(" → saved")
else:
print()
print("\n" + "=" * 60)
print(f"TRAINING COMPLETE - Best test accuracy: {best_acc:.2f}%")
print("=" * 60)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink