import torch import numpy as np import csv import os from tqdm import tqdm import argparse from pathlib import Path from typing import Dict, List, Optional, Tuple, Union, Any from examples.shubert.models.shubert import SHubertModel, SHubertConfig from transformers import ByT5Tokenizer, ByT5ForConditionalGeneration class SHubertProcessor: """ A class for processing multi-modal embeddings through SHubert model. """ def __init__(self, checkpoint_path: str, device: Optional[str] = None): """ Initialize the SHubertProcessor. Args: checkpoint_path: Path to the SHubert model checkpoint device: Device to use ('cuda' or 'cpu'). Auto-detected if None """ self.checkpoint_path = checkpoint_path self.device = device if device else torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load the model self.model = self._load_model() print(f"SHubertProcessor initialized on device: {self.device}") def _load_model(self) -> SHubertModel: """Load the SHubert model from checkpoint.""" # Initialize configuration cfg = SHubertConfig() # Initialize the model model = SHubertModel(cfg) # Load the checkpoint checkpoint = torch.load(self.checkpoint_path, map_location=self.device) # Extract state dict if 'model' in checkpoint: state_dict = checkpoint['model'] else: state_dict = checkpoint # Load the state dictionary into the model model.load_state_dict(state_dict, strict=False) model.eval() model.to(self.device) return model def process_embeddings(self, face_embeddings: np.ndarray, left_hand_embeddings: np.ndarray, right_hand_embeddings: np.ndarray, pose_embeddings: np.ndarray) -> np.ndarray: """ Process multi-modal embeddings through SHubert model. Args: face_embeddings: Face embeddings array of shape (num_frames, embedding_dim) left_hand_embeddings: Left hand embeddings array of shape (num_frames, embedding_dim) right_hand_embeddings: Right hand embeddings array of shape (num_frames, embedding_dim) pose_embeddings: Pose embeddings array of shape (num_frames, pose_dim) Returns: Numpy array of SHubert features with shape (num_layers, num_frames, feature_dim) """ # Convert to tensors and move to device face = torch.from_numpy(face_embeddings).float().to(self.device) left_hand = torch.from_numpy(left_hand_embeddings).float().to(self.device) right_hand = torch.from_numpy(right_hand_embeddings).float().to(self.device) body_posture = torch.from_numpy(pose_embeddings).float().to(self.device) length = face.shape[0] # Prepare input in the format expected by SHubert source = [{ "face": face, "left_hand": left_hand, "right_hand": right_hand, "body_posture": body_posture, # Add dummy labels to match the expected input format "label_face": torch.zeros((length, 1)).to(self.device), "label_left_hand": torch.zeros((length, 1)).to(self.device), "label_right_hand": torch.zeros((length, 1)).to(self.device), "label_body_posture": torch.zeros((length, 1)).to(self.device) }] # Extract features with torch.no_grad(): result = self.model.extract_features(source, padding_mask=None, kmeans_labels=None, mask=False) # Extract layer outputs layer_outputs = [] for layer in result['layer_results']: # layer_output has shape [T, B, D] # Since batch size B is 1, we can squeeze it layer_output = layer[-1] layer_output = layer_output.squeeze(1) # Shape: [T, D] layer_outputs.append(layer_output.cpu().numpy()) # Convert to NumPy array # Stack the outputs from all layers to get an array of shape [L, T, D] features = np.stack(layer_outputs, axis=0) # Shape: [L, T, D] return features def process_embeddings_from_files(self, face_path: str, left_hand_path: str, right_hand_path: str, pose_path: str) -> np.ndarray: """ Process embeddings loaded from files. Args: face_path: Path to face embeddings .npy file left_hand_path: Path to left hand embeddings .npy file right_hand_path: Path to right hand embeddings .npy file pose_path: Path to pose embeddings .npy file Returns: Numpy array of SHubert features with shape (num_layers, num_frames, feature_dim) """ # Load numpy arrays face_embeddings = np.load(face_path) left_hand_embeddings = np.load(left_hand_path) right_hand_embeddings = np.load(right_hand_path) pose_embeddings = np.load(pose_path) return self.process_embeddings(face_embeddings, left_hand_embeddings, right_hand_embeddings, pose_embeddings) def process_and_save_embeddings(self, face_embeddings: np.ndarray, left_hand_embeddings: np.ndarray, right_hand_embeddings: np.ndarray, pose_embeddings: np.ndarray, output_path: str) -> str: """ Process embeddings and save to file. Args: face_embeddings: Face embeddings array left_hand_embeddings: Left hand embeddings array right_hand_embeddings: Right hand embeddings array pose_embeddings: Pose embeddings array output_path: Path to save the output file Returns: Path to the saved file """ # Process embeddings features = self.process_embeddings(face_embeddings, left_hand_embeddings, right_hand_embeddings, pose_embeddings) # Create output directory if it doesn't exist output_dir = Path(output_path).parent output_dir.mkdir(parents=True, exist_ok=True) # Save features np.save(output_path, features) return str(output_path) def process_from_files_and_save(self, face_path: str, left_hand_path: str, right_hand_path: str, pose_path: str, output_path: str) -> str: """ Process embeddings from files and save results. Args: face_path: Path to face embeddings .npy file left_hand_path: Path to left hand embeddings .npy file right_hand_path: Path to right hand embeddings .npy file pose_path: Path to pose embeddings .npy file output_path: Path to save the output file Returns: Path to the saved file """ # Process embeddings features = self.process_embeddings_from_files(face_path, left_hand_path, right_hand_path, pose_path) # Create output directory if it doesn't exist output_dir = Path(output_path).parent output_dir.mkdir(parents=True, exist_ok=True) # Save features np.save(output_path, features) return str(output_path) class SHuBERTTextGenerator: """ A class that combines SHuBERT feature extraction with BYT5 text generation. """ def __init__(self, shubert_checkpoint: str, byt5_model_name: str = "google/byt5-base", device: Optional[str] = None): """ Initialize with SHuBERT and BYT5 models. Args: shubert_checkpoint: Path to SHuBERT model checkpoint byt5_model_name: Name of BYT5 model (default: "google/byt5-base") device: Device to use ('cuda' or 'cpu') """ self.device = device if device else torch.device("cuda" if torch.cuda.is_available() else "cpu") # Initialize SHuBERT processor self.shubert_processor = SHubertProcessor(shubert_checkpoint, self.device) # Initialize BYT5 model self.tokenizer = ByT5Tokenizer.from_pretrained(byt5_model_name) self.model = ByT5ForConditionalGeneration.from_pretrained(byt5_model_name).to(self.device) def generate_text(self, face_embeddings: np.ndarray, left_hand_embeddings: np.ndarray, right_hand_embeddings: np.ndarray, pose_embeddings: np.ndarray, max_length: int = 1024, num_beams: int = 5) -> str: """ Generate text from multi-modal embeddings. Args: face_embeddings: Face embeddings array left_hand_embeddings: Left hand embeddings array right_hand_embeddings: Right hand embeddings array pose_embeddings: Pose embeddings array max_length: Maximum length of generated text num_beams: Number of beams for beam search Returns: Generated text string """ # Get SHuBERT features features = self.shubert_processor.process_embeddings( face_embeddings, left_hand_embeddings, right_hand_embeddings, pose_embeddings) # Select features from specific layer (default: last layer) features = features[-1] # Shape: [T, D] # Convert to tensor and add batch dimension features = torch.from_numpy(features).float().unsqueeze(0).to(self.device) # Generate text generated_ids = self.model.generate( inputs_embeds=features, max_length=max_length, num_beams=num_beams, early_stopping=True ) # Decode generated tokens to text return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True) def generate_text_from_features(face_embeddings: np.ndarray, left_hand_embeddings: np.ndarray, right_hand_embeddings: np.ndarray, pose_embeddings: np.ndarray, shubert_checkpoint: str, byt5_model_name: str = "google/byt5-base", max_length: int = 1024, num_beams: int = 5) -> str: """ Convenience function to generate text from features. """ generator = SHuBERTTextGenerator(shubert_checkpoint, byt5_model_name) return generator.generate_text( face_embeddings, left_hand_embeddings, right_hand_embeddings, pose_embeddings, max_length=max_length, num_beams=num_beams ) # Convenience functions for backward compatibility def process_shubert_embeddings(face_embeddings: np.ndarray, left_hand_embeddings: np.ndarray, right_hand_embeddings: np.ndarray, pose_embeddings: np.ndarray, checkpoint_path: str) -> np.ndarray: """ Convenience function to process embeddings through SHubert. Args: face_embeddings: Face embeddings array left_hand_embeddings: Left hand embeddings array right_hand_embeddings: Right hand embeddings array pose_embeddings: Pose embeddings array checkpoint_path: Path to the SHubert model checkpoint Returns: Numpy array of SHubert features """ processor = SHubertProcessor(checkpoint_path) return processor.process_embeddings(face_embeddings, left_hand_embeddings, right_hand_embeddings, pose_embeddings) def process_sample(model: SHubertModel, face_path: str, left_hand_path: str, right_hand_path: str, body_posture_path: str) -> np.ndarray: """ Original function for backward compatibility with command-line usage. """ # Load numpy arrays face_np = np.load(face_path) left_hand_np = np.load(left_hand_path) right_hand_np = np.load(right_hand_path) body_posture_np = np.load(body_posture_path) face = torch.from_numpy(face_np).float().cuda() left_hand = torch.from_numpy(left_hand_np).float().cuda() right_hand = torch.from_numpy(right_hand_np).float().cuda() body_posture = torch.from_numpy(body_posture_np).float().cuda() length = face.shape[0] # Prepare input source = [{ "face": face, "left_hand": left_hand, "right_hand": right_hand, "body_posture": body_posture, # Add dummy labels to match the expected input format "label_face": torch.zeros((length, 1)).cuda(), "label_left_hand": torch.zeros((length, 1)).cuda(), "label_right_hand": torch.zeros((length, 1)).cuda(), "label_body_posture": torch.zeros((length, 1)).cuda() }] # Extract features with torch.no_grad(): result = model.extract_features(source, padding_mask=None, kmeans_labels=None, mask=False) # Extract layer outputs layer_outputs = [] for layer in result['layer_results']: # layer_output has shape [T, B, D] # Since batch size B is 1, we can squeeze it layer_output = layer[-1] layer_output = layer_output.squeeze(1) # Shape: [T, D] layer_outputs.append(layer_output.cpu().numpy()) # Convert to NumPy array # Stack the outputs from all layers to get an array of shape [L, T, D] features = np.stack(layer_outputs, axis=0) # Shape: [L, T, D] return features def load_model(checkpoint_path: str) -> SHubertModel: """ Original function for backward compatibility with command-line usage. """ cfg = SHubertConfig() # Initialize the model model = SHubertModel(cfg) # Load the checkpoint checkpoint = torch.load(checkpoint_path) # If the checkpoint is saved with a 'model' key if 'model' in checkpoint: state_dict = checkpoint['model'] else: state_dict = checkpoint # Load the state dictionary into the model model.load_state_dict(state_dict, strict=False) model.eval() model.cuda() # Move to GPU if available return model def main(csv_list: List[List[str]], checkpoint_path: str, output_dir: str, index: int): """ Original main function for backward compatibility with command-line usage. """ model = load_model(checkpoint_path) os.makedirs(output_dir, exist_ok=True) for row in csv_list: cues_list = row[0].split('\t') face_path, left_hand_path, right_hand_path, body_posture_path = cues_list[0], cues_list[1], cues_list[2], cues_list[3] output_filename = f"{os.path.basename(face_path).rsplit('.', 1)[0].rsplit('_', 1)[0]}.npy" output_path = os.path.join(output_dir, output_filename) # check if the output file already exists if os.path.exists(output_path): print(f"Skipping {output_path} as it already exists") continue # Process the sample features = process_sample(model, face_path, left_hand_path, right_hand_path, body_posture_path) np.save(output_path, features) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--index', type=int, required=True, help='index of the sub_list to work with') parser.add_argument('--csv_path', type=str, required=True, help='path to the CSV file') parser.add_argument('--checkpoint_path', type=str, required=True, help='path to the checkpoint file') parser.add_argument('--output_dir', type=str, required=True, help='directory to save output files') parser.add_argument('--batch_size', type=int, required=True, help='batch size for processing') args = parser.parse_args() index = args.index csv_path = args.csv_path checkpoint_path = args.checkpoint_path output_dir = args.output_dir batch_size = int(args.batch_size) # make output dir os.makedirs(output_dir, exist_ok=True) # Load CSV data fixed_list = [] with open(csv_path, 'r') as csvfile: reader = csv.reader(csvfile) for row in reader: fixed_list.append(row) # Process in batches video_batches = [fixed_list[i:i + batch_size] for i in range(0, len(fixed_list), batch_size)] csv_list = video_batches[index] main(csv_list, checkpoint_path, output_dir, index)