src/sata/sequence_dataset.py

import os, torch, random
import numpy as np
from pathlib import Path
from torch.utils.data import Dataset, Sampler, DataLoader
from functools import partial
from typing import List, Dict, Tuple, Optional, Union
from torch_geometric.data import Data, Batch

from mydataset import npz_2_data
from skel_pose_graph import SkelPoseGraph, find_feet

class SequenceDataset(Dataset):
    """
    Dataset that loads complete motion sequences with text descriptions.
    
    Each item is a complete motion sequence rather than a single frame.
    """
    
    copy_orig_contact = False

    def __init__(self, text_dir: Optional[str] = None, seq_length: Optional[int] = None):
        """
        Initialize the sequence dataset.
        
        Args:
            text_dir: Directory containing text description files (.txt)
            seq_length: Fixed sequence length. If None, load full sequences.
                       If specified, sequences will be truncated or padded to this length.
        """
        # Skeleton and pose data
        self.skel_list = []
        self.pose_list = []

        # File information
        self.filepaths = []
        self.frame_cnts = []
        self.start_frames = []
        self.end_frames = []

        # Motion set related info
        self.mi_ri_2_fi = []
        # mi: semantic motion index (same mi means semantically identical motion)
        # ri: 0<=ri<R, R: number of retargeted motions (including original data)
        # mi_ri_2_fi[mi][ri] = fi
        
        # Text related info
        self.text_dir = text_dir
        self.mi_2_text = []  # mi -> list of text descriptions for this motion
        
        # Sequence length configuration
        self.seq_length = seq_length

    def add_data(self, lo, go, qb, edges, q, p, qv, pv, pprev, c, r, filepath, mi, tf, text=None):
        """Add motion data to the dataset."""
        # Convert to SkelData and PoseData
        sd, pdl = npz_2_data(lo, go, qb, edges, q, p, qv, pv, pprev, c, r, tf)

        self.skel_list.append(sd)
        self.pose_list.extend(pdl)

        # Update file info
        fi = len(self.filepaths)
        nFrame = q.shape[0]
        start = sum(self.frame_cnts)
        end = start + nFrame
        self.filepaths.append(filepath)
        self.frame_cnts.append(nFrame)
        self.start_frames.append(start)
        self.end_frames.append(end)

        # Update motion set info
        assert len(self.mi_ri_2_fi) >= mi
        if len(self.mi_ri_2_fi) == mi:
            # New semantic motion set
            self.mi_ri_2_fi.append([])
            # Add text for new motion (only once per motion index)
            self.mi_2_text.append(text if text is not None else [])
        else:
            # Ensure consistent frame count among retargeted dataset
            orig_fi = self.mi_ri_2_fi[mi][0]
            orig_nFrame = self.frame_cnts[orig_fi]

            assert orig_nFrame == nFrame, f"Frame count mismatch: {orig_nFrame} vs {nFrame}"
            
            if self.copy_orig_contact:
                # Copy contact from the original motion (optional)
                lf, rf = find_feet(sd)
                orig_sd = self.skel_list[orig_fi]
                orig_start = self.start_frames[orig_fi]
                orig_pdl = self.pose_list[orig_start : orig_start + orig_nFrame]
                orig_lf, orig_rf = find_feet(orig_sd)
                for pdi, orig_pdi in zip(pdl, orig_pdl):
                    pdi.c[lf] = orig_pdi.c[orig_lf]
                    pdi.c[rf] = orig_pdi.c[orig_rf]

        self.mi_ri_2_fi[mi].append(fi)

    def add_data_from_npz(self, mi, npz_fp, bvh_fp=None, text=None):
        """Load data from an NPZ file."""
        data = np.load(npz_fp)
        if bvh_fp is None:
            bvh_fp = npz_fp  # placeholder
        self.add_data(**data, filepath=bvh_fp, mi=mi, text=text)

    def load_data_dir_pairs(self, data_dir):
        """
        Load paired motion data from a directory.
        
        The directory should contain:
        - pair.txt: List of (source, target) motion pairs
        - *.npz files: Motion data files
        """
        pair_path = os.path.join(data_dir, "pair.txt")
        assert os.path.exists(pair_path), f"{pair_path} does not exist"
        bvh_prefix = os.path.join(os.path.dirname(data_dir), "bvh")

        src_id_map = {}
        with open(pair_path, "r") as pair_file:
            for line in pair_file:
                if line.strip() == "":
                    continue
                src_rel_path, dst_rel_path = line.strip().split()

                if src_rel_path in src_id_map:
                    src_id = src_id_map[src_rel_path]
                else:  # New source
                    src_id = len(self.mi_ri_2_fi)
                    src_id_map[src_rel_path] = src_id
                    npz_fp = os.path.join(data_dir, src_rel_path)
                    bvh_fp = os.path.join(bvh_prefix, Path(src_rel_path).with_suffix(""))
                    # Load text if text_dir is provided
                    text = self._load_text_for_file(src_rel_path)
                    self.add_data_from_npz(src_id, npz_fp, bvh_fp, text=text)

                if dst_rel_path in src_id_map:
                    continue
                else:
                    npz_fp = os.path.join(data_dir, dst_rel_path)
                    bvh_fp = os.path.join(bvh_prefix, Path(dst_rel_path).with_suffix(""))
                    # No text for retargeted motions (they share text with source)
                    self.add_data_from_npz(src_id, npz_fp, bvh_fp, text=None)

    def _load_text_for_file(self, file_rel_path):
        """
        Load text descriptions for a motion file.
        
        Text files are stored in self.text_dir with the same name as motion file but .txt extension.
        Each line in the text file contains one text description in the format:
        caption#tokens#start_time#end_time
        
        Args:
            file_rel_path: Relative path to the motion file (e.g., 'test/006888.bvh.npz')
            
        Returns:
            List of text description dictionaries, or empty list if no text file found
        """
        if self.text_dir is None:
            return []
        
        # Extract filename without extensions
        basename = os.path.basename(file_rel_path)  # '006888.bvh.npz'
        filename = Path(basename).stem  # '006888.bvh'
        if filename.endswith('.bvh'):
            filename = filename[:-4]  # '006888'
        
        text_file_path = os.path.join(self.text_dir, filename + '.txt')
        
        if not os.path.exists(text_file_path):
            return []
        
        text_list = []
        try:
            with open(text_file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f.readlines(), 1):
                    line = line.strip()
                    if not line:
                        continue
                    
                    # Parse format: caption#tokens#start_time#end_time
                    parts = line.split('#')
                    if len(parts) < 4:
                        print(f"Warning: Malformed line {line_num} in {text_file_path}: expected 4 parts, got {len(parts)}")
                        continue
                    
                    try:
                        text_dict = {
                            'caption': parts[0].strip(),
                            'tokens': parts[1].strip().split(' ') if parts[1].strip() else [],
                            'start_time': float(parts[2].strip()) if parts[2].strip() else 0.0,
                            'end_time': float(parts[3].strip()) if parts[3].strip() else 0.0,
                        }
                        text_list.append(text_dict)
                    except ValueError as e:
                        print(f"Warning: Failed to parse line {line_num} in {text_file_path}: {e}")
                        continue
                        
        except Exception as e:
            print(f"Warning: Failed to load text from {text_file_path}: {e}")
            return []
        
        return text_list

    def get_mi_ri_fi_graph(self, mi, ri, frame):
        """Get skeleton pose graph for a specific motion, rig, and frame."""
        fi = self.mi_ri_2_fi[mi][ri]
        si = self.skel_list[fi]
        pi = self.pose_list[self.start_frames[fi] + frame]
        return SkelPoseGraph(si, pi)

    def get_sequence_graphs(self, mi, ri, start_frame=0, length=None):
        """
        Get a sequence of graphs for a specific motion and rig.
        
        Args:
            mi: Motion index
            ri: Rig index
            start_frame: Starting frame index
            length: Number of frames to load. If None, load all remaining frames.
            
        Returns:
            List of SkelPoseGraph objects
        """
        fi = self.mi_ri_2_fi[mi][ri]
        frame_cnt = self.frame_cnts[fi]
        
        if length is None:
            length = frame_cnt - start_frame
        
        end_frame = min(start_frame + length, frame_cnt)
        
        graphs = []
        for frame in range(start_frame, end_frame):
            graphs.append(self.get_mi_ri_fi_graph(mi, ri, frame))
        
        return graphs, end_frame - start_frame

    def __getitem__(self, idx):
        """
        Get a complete motion sequence with text.
        
        Args:
            idx: Tuple (mi, src_ri, tgt_ri) where:
                - mi: motion index
                - src_ri: source rig index
                - tgt_ri: target rig index
        
        Returns:
            Tuple (src_graphs, tgt_graphs, text, m_lens) where:
                - src_graphs: List of source skeleton pose graphs
                - tgt_graphs: List of target skeleton pose graphs
                - text: List of text descriptions for this motion
                - m_lens: Length of the motion sequence (number of frames)
        """
        mi, src_ri, tgt_ri = idx
        
        # Get frame count for this motion
        fi = self.mi_ri_2_fi[mi][src_ri]
        frame_cnt = self.frame_cnts[fi]
        
        # Handle sequence length
        if self.seq_length is not None:
            # Use specified sequence length
            actual_length = min(self.seq_length, frame_cnt)
            start_frame = 0
        else:
            # Use full sequence
            actual_length = frame_cnt
            start_frame = 0
        
        # Get sequence of graphs
        src_graphs, _ = self.get_sequence_graphs(mi, src_ri, start_frame, actual_length)
        tgt_graphs, _ = self.get_sequence_graphs(mi, tgt_ri, start_frame, actual_length)
        
        # Get text descriptions
        text = self.mi_2_text[mi] if mi < len(self.mi_2_text) else []
        
        # Actual sequence length
        m_lens = len(src_graphs)
        
        return src_graphs, tgt_graphs, text, m_lens

    def __len__(self):
        """Return the number of motion sequences in the dataset."""
        return len(self.mi_ri_2_fi)


class SequenceSampler(Sampler):
    """
    Sampler that yields motion sequence indices.
    
    Unlike frame-based samplers, this sampler yields complete motion sequences.
    """

    def __init__(self, dataset, batch_size, shuffle=True):
        """
        Initialize the sequence sampler.
        
        Args:
            dataset: SequenceDataset instance
            batch_size: Number of sequences per batch
            shuffle: Whether to shuffle motion indices
        """
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle

        # Motion indices
        self.motion_indices = list(range(len(self.dataset.mi_ri_2_fi)))
        
        # For specifying fixed src/tgt rig indices
        self.src_ri = None
        self.tgt_ri = None

    def __iter__(self):
        # Shuffle motion indices if required
        indices = self.motion_indices.copy()
        if self.shuffle:
            random.shuffle(indices)

        batch = []
        for mi in indices:
            # Random src/tgt skeletons (including the original ones)
            R = len(self.dataset.mi_ri_2_fi[mi])
            
            if R >= 2:
                src_ri, tgt_ri = random.sample(range(R), 2)
            else:
                # Only one rig available
                src_ri = tgt_ri = 0
            
            # Override if src/tgt ri is specified
            if self.src_ri is not None:
                src_ri = self.src_ri
            if self.tgt_ri is not None:
                tgt_ri = self.tgt_ri

            batch.append((mi, src_ri, tgt_ri))

            # Yield batch when full
            if len(batch) == self.batch_size:
                yield batch
                batch = []

        # Yield remaining items
        if len(batch) > 0:
            yield batch

    def __len__(self):
        """Return the number of batches."""
        return (len(self.motion_indices) + self.batch_size - 1) // self.batch_size


def sequence_collate_fn(batch, mask_option=[], device="cpu", pad_to_max=False, return_list=False):
    """
    Collate function for sequence-based batches.
    
    Args:
        batch: List of tuples (src_graphs, tgt_graphs, text, m_lens) from __getitem__
        mask_option: List of strings indicating which graphs to apply random masking to
        device: Device to move tensors to
        pad_to_max: Default false. If True, pad all sequences to the maximum length in the batch
        return_list: If False (default), return aggregated Batch objects.
                    If True, return list of Batch objects (one per sequence).
        
    Returns:
        If return_list=False (default):
            Tuple (src_batch, tgt_batch, text_list, m_lens, masks, batch_info) where:
                - src_batch: Single aggregated Batch containing all sequences
                - tgt_batch: Single aggregated Batch containing all sequences
                - text_list: List of text descriptions
                - m_lens: Numpy array of sequence lengths [T_0, T_1, ..., T_{B-1}]
                - masks: Tensor of shape (batch_size, max_len) indicating valid frames (if pad_to_max)
                - batch_info: Dict with keys:
                    - 'cumsum_lens': Cumulative sum of lengths for splitting [0, T_0, T_0+T_1, ...]
                    - 'num_sequences': Number of sequences in batch (batch_size)
                    - 'total_frames': Total number of frames across all sequences
        
        If return_list=True:
            Tuple (src_batch_list, tgt_batch_list, text_list, m_lens, masks, batch_info) where:
                - src_batch_list: List of batched source graphs (one per sequence)
                - tgt_batch_list: List of batched target graphs (one per sequence)
                - text_list: List of text descriptions
                - m_lens: Numpy array of sequence lengths
                - masks: Tensor of shape (batch_size, max_len) indicating valid frames (if pad_to_max)
                - batch_info: Dict with metadata
    """
    # Unpack batch
    src_graphs_list = [item[0] for item in batch]
    tgt_graphs_list = [item[1] for item in batch]
    text_list = [item[2] for item in batch]
    m_lens = [item[3] for item in batch]
    
    if pad_to_max and len(set(m_lens)) > 1:
        # Sequences have different lengths, need to pad
        max_len = max(m_lens)
        
        # Create masks (True for valid frames, False for padding)
        masks = []
        padded_src_graphs_list = []
        padded_tgt_graphs_list = []
        
        for src_graphs, tgt_graphs, seq_len in zip(src_graphs_list, tgt_graphs_list, m_lens):
            # Create mask for this sequence
            mask = [True] * seq_len + [False] * (max_len - seq_len)
            masks.append(mask)
            
            # Pad sequences by repeating the last frame
            if seq_len < max_len:
                padding_needed = max_len - seq_len
                # Repeat last frame for padding
                src_graphs_padded = src_graphs + [src_graphs[-1]] * padding_needed
                tgt_graphs_padded = tgt_graphs + [tgt_graphs[-1]] * padding_needed
            else:
                src_graphs_padded = src_graphs
                tgt_graphs_padded = tgt_graphs
            
            padded_src_graphs_list.append(src_graphs_padded)
            padded_tgt_graphs_list.append(tgt_graphs_padded)
        
        src_graphs_list = padded_src_graphs_list
        tgt_graphs_list = padded_tgt_graphs_list
        m_lens = [max_len] * len(m_lens)
        masks = torch.BoolTensor(masks).to(device)
    else:
        # All sequences have the same length, no padding needed
        masks = None
    
    # Batch graphs for each sequence
    src_batch_list = []
    tgt_batch_list = []
    
    for src_graphs, tgt_graphs in zip(src_graphs_list, tgt_graphs_list):
        src_batch = Batch.from_data_list(src_graphs).to(device)
        tgt_batch = Batch.from_data_list(tgt_graphs).to(device)
        
        # Apply masking if requested
        if "src" in mask_option:
            src_batch.mask = rnd_mask(src_batch, consq_n=len(src_graphs))
        if "tgt" in mask_option:
            tgt_batch.mask = rnd_mask(tgt_batch, consq_n=len(tgt_graphs))
        
        src_batch_list.append(src_batch)
        tgt_batch_list.append(tgt_batch)
    
    # Random select one text for each sequence if multiple texts are available
    for i, texts in enumerate(text_list):
        text_list[i] = random.choice(texts) if texts else {}

    # Create batch info for splitting later
    cumsum_lens = np.concatenate([[0], np.cumsum(m_lens)])
    batch_info = {
        'cumsum_lens': cumsum_lens,  # [0, T_0, T_0+T_1, ..., sum(T_i)]
        'num_sequences': len(m_lens),  # Batch size
        'total_frames': int(np.sum(m_lens)),  # Total frames
    }
    
    # Return list format if requested
    if return_list:
        return src_batch_list, tgt_batch_list, text_list, m_lens, masks, batch_info
    
    # Aggregate all sequences into a single Batch (default behavior)
    # Flatten all graphs from all sequences
    all_src_graphs = [graph for graphs in src_graphs_list for graph in graphs]
    all_tgt_graphs = [graph for graphs in tgt_graphs_list for graph in graphs]
    
    # Create single aggregated batch
    src_batch = Batch.from_data_list(all_src_graphs).to(device)
    tgt_batch = Batch.from_data_list(all_tgt_graphs).to(device)
    
    # Apply masking if requested (on the aggregated batch)
    if "src" in mask_option:
        src_batch.mask = rnd_mask(src_batch, consq_n=sum(m_lens))
    if "tgt" in mask_option:
        tgt_batch.mask = rnd_mask(tgt_batch, consq_n=sum(m_lens))
    
    
    return src_batch, tgt_batch, text_list, m_lens, masks, batch_info


def get_sequence_data_loader(
    data_dir, 
    batch_size, 
    shuffle=True, 
    mask_option=[], 
    device="cpu", 
    text_dir=None,
    seq_length=None,
    pad_to_max=True,
    return_list=False
):
    """
    Create a DataLoader for sequence-based motion data.
    
    Args:
        data_dir: Directory containing motion data
        batch_size: Number of sequences per batch
        shuffle: Whether to shuffle the data
        mask_option: List of strings for masking ('src', 'tgt')
        device: Device to load data to
        text_dir: Directory containing text descriptions
        seq_length: Fixed sequence length (None for variable length)
        pad_to_max: Whether to pad sequences to max length in batch
        return_list: If False (default), return aggregated Batch objects.
                    If True, return list of Batch objects (one per sequence).
        
    Returns:
        DataLoader instance
    """
    ds = SequenceDataset(text_dir=text_dir, seq_length=seq_length)
    ds.load_data_dir_pairs(data_dir)
    
    sampler = SequenceSampler(ds, batch_size=batch_size, shuffle=shuffle)
    
    dl = DataLoader(
        ds,
        batch_sampler=sampler,
        collate_fn=partial(
            sequence_collate_fn,
            mask_option=mask_option,
            device=device,
            pad_to_max=pad_to_max,
            return_list=return_list,
        ),
        pin_memory=True,
    )
    return dl