datasets.py

import torch
import utils as utils
import torch.utils.data.dataset as Dataset
from torch.nn.utils.rnn import pad_sequence
from PIL import Image
import os
import random
import numpy as np
import copy
import pickle
from decord import VideoReader, cpu
import json
import pathlib
from torchvision import transforms 
from config import rgb_dirs, pose_dirs, ir_text_dirs
from typing import Dict, List, Tuple

# load sub-pose
def load_part_kp(skeletons, confs, force_ok=False):
    thr = 0.3
    kps_with_scores = {}
    scale = None
    
    for part in ['body', 'left', 'right', 'face_all']:
        kps = []
        confidences = []
        
        for skeleton, conf in zip(skeletons, confs):
            skeleton = skeleton[0]
            conf = conf[0]
            
            if part == 'body':
                hand_kp2d = skeleton[[0] + [i for i in range(3, 11)], :]
                confidence = conf[[0] + [i for i in range(3, 11)]]
            elif part == 'left':
                hand_kp2d = skeleton[91:112, :]
                hand_kp2d = hand_kp2d - hand_kp2d[0, :]
                confidence = conf[91:112]
            elif part == 'right':
                hand_kp2d = skeleton[112:133, :]
                hand_kp2d = hand_kp2d - hand_kp2d[0, :]
                confidence = conf[112:133]
            elif part == 'face_all':
                hand_kp2d = skeleton[[i for i in list(range(23,23+17))[::2]] + [i for i in range(83, 83 + 8)] + [53], :]
                hand_kp2d = hand_kp2d - hand_kp2d[-1, :]
                confidence = conf[[i for i in list(range(23,23+17))[::2]] + [i for i in range(83, 83 + 8)] + [53]]
            else:
                raise NotImplementedError
            
            kps.append(hand_kp2d)
            confidences.append(confidence)
            
        kps = np.stack(kps, axis=0)
        confidences = np.stack(confidences, axis=0)
        
        if part == 'body':
            if force_ok:
                result, scale, _ = crop_scale(np.concatenate([kps, confidences[...,None]], axis=-1), thr)
            else:
                result, scale, _ = crop_scale(np.concatenate([kps, confidences[...,None]], axis=-1), thr)
        else:
            assert not scale is None
            result = np.concatenate([kps, confidences[...,None]], axis=-1)
            if scale==0:
                result = np.zeros(result.shape)
            else:
                result[...,:2] = (result[..., :2]) / scale
                result = np.clip(result, -1, 1)
                # mask useless kp
                result[result[...,2]<=thr] = 0
            
        kps_with_scores[part] = torch.tensor(result)
        
    return kps_with_scores

# input: T, N, 3
# input is un-normed joints
def crop_scale(motion, thr):
    '''
        Motion: [(M), T, 17, 3].
        Normalize to [-1, 1]
    '''
    result = copy.deepcopy(motion)
    valid_coords = motion[motion[..., 2]>thr][:,:2]
    if len(valid_coords) < 4:
        return np.zeros(motion.shape), 0, None
    xmin = min(valid_coords[:,0])
    xmax = max(valid_coords[:,0])
    ymin = min(valid_coords[:,1])
    ymax = max(valid_coords[:,1])
    # ratio = np.random.uniform(low=scale_range[0], high=scale_range[1], size=1)[0]
    ratio = 1
    scale = max(xmax-xmin, ymax-ymin) * ratio
    if scale==0:
        return np.zeros(motion.shape), 0, None
    xs = (xmin+xmax-scale) / 2
    ys = (ymin+ymax-scale) / 2
    result[...,:2] = (motion[..., :2] - [xs,ys]) / scale
    result[...,:2] = (result[..., :2] - 0.5) * 2
    result = np.clip(result, -1, 1)
    # mask useless kp
    result[result[...,2]<=thr] = 0
    return result, scale, [xs,ys]

# bbox of hands
def bbox_4hands(left_keypoints, right_keypoints, hw):

    def compute_bbox(keypoints):
        min_x = np.min(keypoints[..., 0], axis=1)
        min_y = np.min(keypoints[..., 1], axis=1)
        max_x = np.max(keypoints[..., 0], axis=1)
        max_y = np.max(keypoints[..., 1], axis=1)
        
        return (max_x+min_x)/2, (max_y+min_y)/2, (max_x-min_x), (max_y-min_y)
    H,W = hw
    
    if left_keypoints is None:
        left_keypoints = np.zeros([1,21,2])
        
    if right_keypoints is None:
        right_keypoints = np.zeros([1,21,2])
    # [T, 21, 2]
    left_mean_x, left_mean_y, left_diff_x, left_diff_y = compute_bbox(left_keypoints)
    left_mean_x = W*left_mean_x
    left_mean_y = H*left_mean_y
    
    left_diff_x = W*left_diff_x
    left_diff_y = H*left_diff_y
    
    left_diff_x = max(left_diff_x)
    left_diff_y = max(left_diff_y)
    left_box_hw = max(left_diff_x,left_diff_y)
    
    right_mean_x, right_mean_y, right_diff_x, right_diff_y = compute_bbox(right_keypoints)
    right_mean_x = W*right_mean_x
    right_mean_y = H*right_mean_y
    
    right_diff_x = W*right_diff_x
    right_diff_y = H*right_diff_y
    
    right_diff_x = max(right_diff_x)
    right_diff_y = max(right_diff_y)
    right_box_hw = max(right_diff_x,right_diff_y)
    
    box_hw = int(max(left_box_hw, right_box_hw) * 1.2 / 2) * 2
    box_hw = max(box_hw, 0)

    left_new_box = np.stack([left_mean_x - box_hw/2, left_mean_y - box_hw/2, left_mean_x + box_hw/2, left_mean_y + box_hw/2]).astype(np.int16)
    right_new_box = np.stack([right_mean_x - box_hw/2, right_mean_y - box_hw/2, right_mean_x + box_hw/2, right_mean_y + box_hw/2]).astype(np.int16)
    
    return left_new_box.transpose(1,0), right_new_box.transpose(1,0), box_hw

def load_support_rgb_dict(tmp, skeletons, confs, full_path, data_transform):
    support_rgb_dict = {}
    
    confs = np.array(confs)
    skeletons = np.array(skeletons) 

    # sample index of low scores
    left_confs_filter = confs[:,0,91:112].mean(-1)
    left_confs_filter_indices = np.where(left_confs_filter > 0.3)[0]

    if len(left_confs_filter_indices) == 0:
        left_sampled_indices = None
        left_skeletons = None
    else:
        
        left_confs = confs[left_confs_filter_indices]
        left_confs = left_confs[:,0,[95,99,103,107,111]].min(-1)
        
        left_weights = np.max(left_confs) - left_confs + 1e-5
        left_probabilities = left_weights / np.sum(left_weights)
        
        left_sample_size = int(np.ceil(0.1 * len(left_confs_filter_indices)))
        
        left_sampled_indices = np.random.choice(left_confs_filter_indices.tolist(), 
                                                size=left_sample_size, 
                                                replace=False, 
                                                p=left_probabilities)
        # left_sampled_indices: values: 0-255(0,max_len)
        # tmp: values: 0-(end-start)
        left_sampled_indices = np.sort(left_sampled_indices)
        
        left_skeletons = skeletons[left_sampled_indices,0,91:112]

    right_confs_filter = confs[:,0,112:].mean(-1)
    right_confs_filter_indices = np.where(right_confs_filter > 0.3)[0]
    if len(right_confs_filter_indices) == 0:
        right_sampled_indices = None
        right_skeletons = None
        
    else:
        right_confs = confs[right_confs_filter_indices]
        right_confs = right_confs[:,0,[95+21,99+21,103+21,107+21,111+21]].min(-1)

        right_weights = np.max(right_confs) - right_confs + 1e-5
        right_probabilities = right_weights / np.sum(right_weights)
        
        right_sample_size = int(np.ceil(0.1 * len(right_confs_filter_indices)))
        
        right_sampled_indices = np.random.choice(right_confs_filter_indices.tolist(), 
                                                 size=right_sample_size, 
                                                 replace=False, 
                                                 p=right_probabilities)
        right_sampled_indices = np.sort(right_sampled_indices)
        
        right_skeletons = skeletons[right_sampled_indices,0,112:133]
        
    image_size = 112
    all_indices = []
    if not left_sampled_indices is None:
        all_indices.append(left_sampled_indices)
    if not right_sampled_indices is None:
        all_indices.append(right_sampled_indices)
    if len(all_indices) == 0:
        support_rgb_dict['left_sampled_indices'] = torch.tensor([-1])
        support_rgb_dict['left_hands'] = torch.zeros(1, 3, image_size, image_size)
        support_rgb_dict['left_skeletons_norm'] = torch.zeros(1, 21, 2)
        
        support_rgb_dict['right_sampled_indices'] = torch.tensor([-1])
        support_rgb_dict['right_hands'] = torch.zeros(1, 3, image_size, image_size)
        support_rgb_dict['right_skeletons_norm'] = torch.zeros(1, 21, 2)

        return support_rgb_dict

    sampled_indices = np.concatenate(all_indices)
    sampled_indices = np.unique(sampled_indices)
    sampled_indices_real = tmp[sampled_indices]

    # load image sample
    imgs = load_video_support_rgb(full_path, sampled_indices_real)

    # get hand bbox
    left_new_box, right_new_box, box_hw = bbox_4hands(left_skeletons,
                                                        right_skeletons,
                                                        imgs[0].shape[:2])
    
    # crop left and right hand
    image_size = 112
    if box_hw == 0:
        support_rgb_dict['left_sampled_indices'] = torch.tensor([-1])
        support_rgb_dict['left_hands'] = torch.zeros(1, 3, image_size, image_size)
        support_rgb_dict['left_skeletons_norm'] = torch.zeros(1, 21, 2)
        
        support_rgb_dict['right_sampled_indices'] = torch.tensor([-1])
        support_rgb_dict['right_hands'] = torch.zeros(1, 3, image_size, image_size)
        support_rgb_dict['right_skeletons_norm'] = torch.zeros(1, 21, 2)

        return support_rgb_dict

    factor = image_size / box_hw
    
    if left_sampled_indices is None:
        left_hands = torch.zeros(1, 3, image_size, image_size)
        left_skeletons_norm = torch.zeros(1, 21, 2)
        
    else:
        left_hands = torch.zeros(len(left_sampled_indices), 3, image_size, image_size)
            
        left_skeletons_norm = left_skeletons * imgs[0].shape[:2][::-1] - left_new_box[:, None, [0,1]]
        left_skeletons_norm = left_skeletons_norm / box_hw
        left_skeletons_norm = left_skeletons_norm.clip(0,1)

    if right_sampled_indices is None:
        right_hands = torch.zeros(1, 3, image_size, image_size)
        right_skeletons_norm = torch.zeros(1, 21, 2)
        
    else:
        right_hands = torch.zeros(len(right_sampled_indices), 3, image_size, image_size)
        
        right_skeletons_norm = right_skeletons * imgs[0].shape[:2][::-1] - right_new_box[:, None, [0,1]]
        right_skeletons_norm = right_skeletons_norm / box_hw
        right_skeletons_norm = right_skeletons_norm.clip(0,1)
    left_idx = 0
    right_idx = 0

    for idx, img in enumerate(imgs):
        mapping_idx = sampled_indices[idx]
        if not left_sampled_indices is None and left_idx < len(left_sampled_indices) and mapping_idx == left_sampled_indices[left_idx]:
            box = left_new_box[left_idx]
            
            img_draw = np.uint8(copy.deepcopy(img))[box[1]:box[3],box[0]:box[2],:]
            img_draw = np.pad(img_draw, ((0, max(0, box_hw-img_draw.shape[0])), (0, max(0, box_hw-img_draw.shape[1])), (0, 0)), mode='constant', constant_values=0)
            
            f_img = Image.fromarray(img_draw).convert('RGB').resize((image_size, image_size))
            f_img = data_transform(f_img).unsqueeze(0)
            left_hands[left_idx] = f_img
            left_idx += 1
            
        if not right_sampled_indices is None and right_idx < len(right_sampled_indices) and mapping_idx == right_sampled_indices[right_idx]:
            box = right_new_box[right_idx]
            
            img_draw = np.uint8(copy.deepcopy(img))[box[1]:box[3],box[0]:box[2],:]
            img_draw = np.pad(img_draw, ((0, max(0, box_hw-img_draw.shape[0])), (0, max(0, box_hw-img_draw.shape[1])), (0, 0)), mode='constant', constant_values=0)
            
            f_img = Image.fromarray(img_draw).convert('RGB').resize((image_size, image_size))
            f_img = data_transform(f_img).unsqueeze(0)
            right_hands[right_idx] = f_img
            right_idx += 1
   
    if left_sampled_indices is None:
        left_sampled_indices = np.array([-1])
        
    if right_sampled_indices is None:
        right_sampled_indices = np.array([-1])

    # get index, images and keypoints priors
    support_rgb_dict['left_sampled_indices'] = torch.tensor(left_sampled_indices)
    support_rgb_dict['left_hands'] = left_hands
    support_rgb_dict['left_skeletons_norm'] = torch.tensor(left_skeletons_norm)
    
    support_rgb_dict['right_sampled_indices'] = torch.tensor(right_sampled_indices)
    support_rgb_dict['right_hands'] = right_hands
    support_rgb_dict['right_skeletons_norm'] = torch.tensor(right_skeletons_norm)

    return support_rgb_dict

# use split rgb video for save time
def load_video_support_rgb(path, tmp):
    vr = VideoReader(path, num_threads=1, ctx=cpu(0))
    
    vr.seek(0)
    buffer = vr.get_batch(tmp).asnumpy()
    batch_image = buffer
    del vr

    return batch_image

# build base dataset
class Base_Dataset(Dataset.Dataset):
    def collate_fn(self, batch):
        tgt_batch,src_length_batch,name_batch,pose_tmp,gloss_batch,label_key_batch,ir_tgt_batch = [],[],[],[],[],[],[]
        
        for name_sample, pose_sample, text, gloss, _, label_key, ir_text in batch:
            name_batch.append(name_sample)
            pose_tmp.append(pose_sample)
            tgt_batch.append(text)
            gloss_batch.append(gloss)
            label_key_batch.append(label_key)
            ir_tgt_batch.append(ir_text)

        src_input = {}
        first_pose = pose_tmp[0]
        if isinstance(first_pose, dict):
            keys = first_pose.keys()
        else:
            keys = first_pose[0].keys()
            tmp = []
            for tup in pose_tmp:
                tmp.append(tup[0])
            pose_tmp = tmp

        for key in keys:
            max_len = max([len(vid[key]) for vid in pose_tmp])
            video_length = torch.LongTensor([len(vid[key]) for vid in pose_tmp])
            
            padded_video = [torch.cat(
                (
                    vid[key],
                    vid[key][-1][None].expand(max_len - len(vid[key]), -1, -1),
                )
                , dim=0)
                for vid in pose_tmp]
            
            img_batch = torch.stack(padded_video,0)
            
            src_input[key] = img_batch
            if 'attention_mask' not in src_input.keys():
                src_length_batch = video_length

                mask_gen = []
                for i in src_length_batch:
                    tmp = torch.ones([i]) + 7
                    mask_gen.append(tmp)
                mask_gen = pad_sequence(mask_gen, padding_value=0,batch_first=True)
                img_padding_mask = (mask_gen != 0).long()
                src_input['attention_mask'] = img_padding_mask

                src_input['name_batch'] = name_batch
                src_input['src_length_batch'] = src_length_batch

        if self.rgb_support:
            support_rgb_dicts = {key:[] for key in batch[0][-1].keys()}
            for _, _, _, _, support_rgb_dict in batch:
                for key in support_rgb_dict.keys():
                    support_rgb_dicts[key].append(support_rgb_dict[key])
            
            for part in ['left', 'right']:
                index_key = f'{part}_sampled_indices'
                skeletons_key = f'{part}_skeletons_norm'
                rgb_key = f'{part}_hands'
                len_key = f'{part}_rgb_len'

                index_batch = torch.cat(support_rgb_dicts[index_key], 0)
                skeletons_batch = torch.cat(support_rgb_dicts[skeletons_key], 0)
                img_batch = torch.cat(support_rgb_dicts[rgb_key], 0)
                
                src_input[index_key] = index_batch
                src_input[skeletons_key] = skeletons_batch
                src_input[rgb_key] = img_batch
                src_input[len_key] = [len(index) for index in support_rgb_dicts[index_key]]

        tgt_input = {}
        tgt_input['gt_sentence'] = tgt_batch
        tgt_input['gt_ir_sentence'] = ir_tgt_batch 
        tgt_input['gt_gloss'] = gloss_batch
        tgt_input['label_key'] = label_key_batch
        return src_input, tgt_input

def random_derangement(b: int, device=None) -> torch.Tensor:
        """
        Create a random derangement perm (perm[i] != i) over 0..b-1.
        Requires b >= 2.
        """
        if b <= 1:
            raise ValueError("Derangement requires b >= 2.")
        perm = torch.randperm(b, device=device)
        arange = torch.arange(b, device=device)

        fixed = perm == arange
        while fixed.any():
            idx = fixed.nonzero(as_tuple=False).squeeze(1)
            if idx.numel() == 1:
                i = idx.item()
                j = torch.randint(0, b, (1,), device=device).item()
                while j == i:
                    j = torch.randint(0, b, (1,), device=device).item()
                tmp = perm[i].clone()
                perm[i] = perm[j]
                perm[j] = tmp
            else:
                # rotate the fixed points to break fixed positions
                perm[idx] = perm[idx].roll(shifts=1, dims=0)
            fixed = perm == arange
        return perm

def find_target_window(frame_scores: torch.Tensor, clip_ratio: float) -> Tuple[int, int, int]:
    if frame_scores.dim() != 1:
        raise ValueError(f"frame_scores must be 1D [T], got {tuple(frame_scores.shape)}")
    T = frame_scores.shape[0]
    if T <= 0:
        raise ValueError("Empty frame_scores")
    if not (0 < float(clip_ratio) <= 1.0):
        raise ValueError("clip_ratio must be in (0, 1].")

    win = max(1, int(round(T * float(clip_ratio))))
    win = min(win, T) # the total num of frames 

    center = int(torch.argmax(frame_scores).item())
    half = win // 2
    tgt_start = max(0, min(center - half, T - win))
    tgt_end = tgt_start + win
    return tgt_start, tgt_end, win

def find_target_windows_batched(frame_scores: torch.Tensor, clip_ratio: float) -> List[Tuple[int, int, int]]:
    if frame_scores.dim() != 2:
        raise ValueError(f"frame_scores must be 2D [B, T], got {tuple(frame_scores.shape)}")
    B, _ = frame_scores.shape
    return [find_target_window(frame_scores[b], clip_ratio) for b in range(B)]

# Spatial
def apply_randomness(src_input: dict, local: bool = False, frame_scores=None, clip_ratio=0.2) -> dict:
    """
    Build negative samples by deranging along the batch dimension.
    """

    # Clone to avoid touching the originals
    body  = src_input['body'].clone()
    left  = src_input['left'].clone()
    right = src_input['right'].clone()
    face  = src_input['face_all'].clone()

    B, T = body.shape[0], body.shape[1]
    device = body.device
    if local:
        if frame_scores is None or frame_scores.shape[:2] != (B, T):
            raise ValueError("local=True requires frame_score of shape [B, T].")

        # 1) Use helper to locate the target window for replacement for each sample
        windows = find_target_windows_batched(frame_scores, clip_ratio)  # [(tgt_start, tgt_end, win), ...]
        for b in range(B):
            tgt_start, tgt_end, win = windows[b]

            # 2) If only one window is available (win == T), swap the first and second halves
            if T == win:
                mid = T // 2
                body[b]  = torch.cat([body[b, mid:],  body[b, :mid]],  dim=0)
                left[b]  = torch.cat([left[b, mid:],  left[b, :mid]],  dim=0)
                right[b] = torch.cat([right[b, mid:], right[b, :mid]], dim=0)
                face[b]  = torch.cat([face[b, mid:],  face[b, :mid]],  dim=0)
                continue

            # 3) Sample a donor window, avoiding overlap with the target window (up to 10 retries)
            donor_start = torch.randint(0, T - win + 1, (1,), device=device).item()
            tries = 0
            while donor_start == tgt_start and tries < 10:
                donor_start = torch.randint(0, T - win + 1, (1,), device=device).item()
                tries += 1
            if donor_start == tgt_start:
                # If still identical, fall back to swapping halves
                mid = T // 2
                body[b]  = torch.cat([body[b, mid:],  body[b, :mid]],  dim=0)
                left[b]  = torch.cat([left[b, mid:],  left[b, :mid]],  dim=0)
                right[b] = torch.cat([right[b, mid:], right[b, :mid]], dim=0)
                face[b]  = torch.cat([face[b, mid:],  face[b, :mid]],  dim=0)
                continue

            donor_end = donor_start + win

            # 4) Perform aligned replacement across all four components
            win_body  = body[b, donor_start:donor_end].clone()
            win_left  = left[b,  donor_start:donor_end].clone()
            win_right  = right[b, donor_start:donor_end].clone()
            win_face  = face[b,  donor_start:donor_end].clone()

            body[b,  tgt_start:tgt_end] = win_body
            left[b,  tgt_start:tgt_end] = win_left
            right[b, tgt_start:tgt_end] = win_right
            face[b,  tgt_start:tgt_end] = win_face

        return {'body': body, 'left': left, 'right': right, 'face_all': face}
    else:
        # ---- Global derangement (preserve cross-part alignment) ----
        sizes = {
            'body': body.size(2),       # 9
            'left': left.size(2),       # 21
            'right': right.size(2),     # 21
            'face_all': face.size(2),   # 18
        }
        order = ['body', 'left', 'right', 'face_all']

        # Record split ranges after concatenation on the joint axis
        starts = {}
        cur = 0
        for k in order:
            starts[k] = (cur, cur + sizes[k])
            cur += sizes[k]

        # Concatenate to [B, T, 69, 3]
        concat_all = torch.cat([body, left, right, face], dim=2)

        # Apply a derangement along batch
        perm = random_derangement(B, device=device)
        concat_shuf = concat_all.index_select(dim=0, index=perm)

        # Split back to four parts using recorded ranges
        b_s, b_e = starts['body']
        l_s, l_e = starts['left']
        r_s, r_e = starts['right']
        f_s, f_e = starts['face_all']

        body_shuf  = concat_shuf[:, :, b_s:b_e, :].contiguous()
        left_shuf  = concat_shuf[:, :, l_s:l_e, :].contiguous()
        right_shuf = concat_shuf[:, :, r_s:r_e, :].contiguous()
        face_shuf  = concat_shuf[:, :, f_s:f_e, :].contiguous()

        # Shape sanity checks
        assert body_shuf.shape  == body.shape
        assert left_shuf.shape  == left.shape
        assert right_shuf.shape == right.shape
        assert face_shuf.shape  == face.shape

        return {
            'body': body_shuf,
            'left': left_shuf,
            'right': right_shuf,
            'face_all': face_shuf,
            'perm': perm,
        }

def apply_blackness(src_input, local=False, frame_scores=None, clip_ratio=0.2):
    """
    Build negative samples by dark all the frames.
    """

    # Clone to avoid touching the originals
    body_black  = src_input['body'].clone()
    left_black  = src_input['left'].clone()
    right_black = src_input['right'].clone()
    face_black  = src_input['face_all'].clone()

    if local:
        if frame_scores is None:
            raise ValueError("local=True requires frame_score of shape [B, T].")
        B, T = body_black.shape[0], body_black.shape[1]
        if frame_scores.shape[:2] != (B, T):
            raise ValueError(f"frame_score must be [B, T], got {tuple(frame_scores.shape)}")

        windows = find_target_windows_batched(frame_scores, clip_ratio)  # [(tgt_start, tgt_end, win), ...]
        for b in range(B):
            tgt_start, tgt_end, win = windows[b]

            body_black[b,  tgt_start:tgt_end].zero_()
            left_black[b,  tgt_start:tgt_end].zero_()
            right_black[b, tgt_start:tgt_end].zero_()
            face_black[b,  tgt_start:tgt_end].zero_()

        return {'body': body_black, 'left': left_black, 'right': right_black, 'face_all': face_black}
    else:
        body_black.zero_()
        left_black.zero_()
        right_black.zero_()
        face_black.zero_()
        return {
            'body': body_black,
            'left': left_black,
            'right': right_black,
            'face_all': face_black
        }

def apply_random_sparse_mask(src_input, mask_ratio=0.2, local=False, frame_scores=None, clip_ratio=0.2) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
    """
    Generate 4 binary masks (0/1) for 'body','left','right','face_all'.
    For each sample b, the total number of zeroed joints across all 4 masks equals `num_global_mask_per_sample`.
    Zeros apply to all frames (T) and all coords (3) of those joints.
    """
    body  = src_input['body'].clone()      # [B, T,  9, 3]
    left  = src_input['left'].clone()      # [B, T, 21, 3]
    right = src_input['right'].clone()     # [B, T, 21, 3]
    face  = src_input['face_all'].clone()  # [B, T, 18, 3]

    B, T = body.shape[0], body.shape[1]
    body_keypoint  = body.shape[2]
    left_keypoint  = left.shape[2]
    right_keypoint = right.shape[2]
    face_keypoint  = face.shape[2]
    total_keypoint = body_keypoint + left_keypoint + right_keypoint + face_keypoint

    if not (0 < float(mask_ratio) <= 1.0):
        raise ValueError("mask_ratio must be in (0, 1].")

    # init 4 binary masks (1 keeps, 0 blacks)
    m_body  = torch.ones_like(body,  dtype=body.dtype)
    m_left  = torch.ones_like(left,  dtype=left.dtype)
    m_right = torch.ones_like(right, dtype=right.dtype)
    m_face  = torch.ones_like(face,  dtype=face.dtype)

    # global offsets
    off_body  = 0
    off_left  = off_body  + body_keypoint
    off_right = off_left  + left_keypoint
    off_face  = off_right + right_keypoint

    k_per_sample = max(1, min(int(total_keypoint * float(mask_ratio)), total_keypoint))

    if local:
        if frame_scores is None or frame_scores.shape[:2] != (B, T):
            raise ValueError(f"local=True requires frame_scores of shape [B, T], got {None if frame_scores is None else tuple(frame_scores.shape)}")

        windows = find_target_windows_batched(frame_scores, clip_ratio)

    device = body.device
    for b in range(B):
        gsel = torch.randperm(total_keypoint, device=device)[:k_per_sample]

        sel_body  = gsel[(gsel >= off_body)  & (gsel < off_left)]  - off_body
        sel_left  = gsel[(gsel >= off_left)  & (gsel < off_right)] - off_left
        sel_right = gsel[(gsel >= off_right) & (gsel < off_face)]  - off_right
        sel_face  = gsel[(gsel >= off_face)]                       - off_face

        if local:
            tgt_start, tgt_end, _ = windows[b]

            if sel_body.numel()  > 0: m_body[b,  tgt_start:tgt_end,  sel_body.long(),  :] = 0
            if sel_left.numel()  > 0: m_left[b,  tgt_start:tgt_end,  sel_left.long(),  :] = 0
            if sel_right.numel() > 0: m_right[b, tgt_start:tgt_end,  sel_right.long(), :] = 0
            if sel_face.numel()  > 0: m_face[b,  tgt_start:tgt_end,  sel_face.long(),  :] = 0
        else:
            # set those joints to 0 across all frames T and coords 3
            if sel_body.numel()  > 0: m_body[b, :, sel_body.long(),  :] = 0
            if sel_left.numel()  > 0: m_left[b, :, sel_left.long(),  :] = 0
            if sel_right.numel() > 0: m_right[b, :, sel_right.long(), :] = 0
            if sel_face.numel()  > 0: m_face[b, :, sel_face.long(),  :] = 0

    # apply
    out = {
        'body':      body  * m_body,
        'left':      left  * m_left,
        'right':     right * m_right,
        'face_all':  face  * m_face,
    }
    masks = {
        'body': m_body,
        'left': m_left,
        'right': m_right,
        'face_all': m_face,
    }
    return masks, out

def apply_ROI_remove(src_input, local=False, frame_scores=None, clip_ratio = 0.2, dominate_hand='right'):
    """
    Remove an ROI to create non-preference samples.
    Here the ROI is the *right hand* (key 'right': [B, T, 21, 3]).

    Args:
        src_input: dict with keys 'body','left','right','face_all'
                   shapes:
                     body     [B, T,  9, 3]
                     left     [B, T, 21, 3]
                     right    [B, T, 21, 3]
                     face_all [B, T, 18, 3]
        local: False -> remove right hand for the whole video (global)
               True  -> remove right hand only inside the target window found by frame_scores
        frame_scores: [B, T] importance scores (required when local=True)
        clip_ratio: ratio of frames to define the window length, e.g., 0.2 -> 20% of T

    Returns:
        A new dict with the same keys, where the right-hand ROI is zeroed
        (globally or inside the target window).
    """

    # Clone to avoid modifying originals
    body  = src_input['body'].clone()
    left  = src_input['left'].clone()
    right = src_input['right'].clone()
    face  = src_input['face_all'].clone()

    if local:
        # ---- Local mode: zero right hand only inside the target window ----
        if frame_scores is None:
            raise ValueError("local=True requires frame_scores of shape [B, T].")
        if frame_scores.dim() != 2:
            raise ValueError(f"frame_scores must be 2D [B, T], got {tuple(frame_scores.shape)}")

        B, T = body.shape[0], body.shape[1]
        if frame_scores.shape[:2] != (B, T):
            raise ValueError(f"frame_scores must match [B, T], got {tuple(frame_scores.shape)} vs ({B}, {T})")

        windows = find_target_windows_batched(frame_scores, clip_ratio)  # [(tgt_start, tgt_end, win), ...]

        for b in range(B):
            tgt_start, tgt_end, win = windows[b]
            if dominate_hand == 'right':
                right[b, tgt_start:tgt_end].zero_()
            elif dominate_hand == 'left':
                left[b, tgt_start:tgt_end].zero_()
            else:
                raise NotImplementedError
    else:
        # Global: zero the entire right hand across all frames
        if dominate_hand == 'right':
            right.zero_()
        elif dominate_hand == 'left':
            left.zero_()
        else:
            raise NotImplementedError
        
    return {'body': body, 'left': left, 'right': right, 'face_all': face}

# Temporal
def reverse(src_input, local=False, frame_scores=None, clip_ratio=0.2):
    """
    Reverse features along the time dimension (dim=1).
    Returns a new dict with time-reversed tensors.
    """
    # Clone to avoid in-place changes
    body  = src_input['body'].clone()
    left  = src_input['left'].clone()
    right = src_input['right'].clone()
    face  = src_input['face_all'].clone()
    
    if local:
        # --- local=True: only reverse the important window per sample ---
        if frame_scores is None:
            raise ValueError("local=True requires frame_scores of shape [B, T].")
        B, T = body.shape[0], body.shape[1]
        if frame_scores.shape[:2] != (B, T):
            raise ValueError(f"frame_scores must be [B, T], got {tuple(frame_scores.shape)}")

        windows = find_target_windows_batched(frame_scores, clip_ratio)

        # For each sample, reverse along the chosen window [tgt_start:tgt_end)
        for b in range(B):
            tgt_start, tgt_end, win = windows[b]
            # Slice [tgt_start:tgt_end] and flip along time dimension (dim=0 inside the slice)
            body[b,  tgt_start:tgt_end]  = torch.flip(body[b,  tgt_start:tgt_end],  dims=[0])
            left[b,  tgt_start:tgt_end]  = torch.flip(left[b,  tgt_start:tgt_end],  dims=[0])
            right[b, tgt_start:tgt_end]  = torch.flip(right[b, tgt_start:tgt_end], dims=[0])
            face[b,  tgt_start:tgt_end]  = torch.flip(face[b,  tgt_start:tgt_end],  dims=[0])

        return {
            'body': body.contiguous(),
            'left': left.contiguous(),
            'right': right.contiguous(),
            'face_all': face.contiguous(),
        }

    else:
        # Reverse along time dimension (dim=1)
        body_rev  = torch.flip(body,  dims=[1])
        left_rev  = torch.flip(left,  dims=[1])
        right_rev = torch.flip(right, dims=[1])
        face_rev  = torch.flip(face,  dims=[1])

        # Chcek
        assert body_rev.shape  == body.shape
        assert left_rev.shape  == left.shape
        assert right_rev.shape == right.shape
        assert face_rev.shape  == face.shape

        return {
            'body': body_rev.contiguous(),
            'left': left_rev.contiguous(),
            'right': right_rev.contiguous(),
            'face_all': face_rev.contiguous(),
        }

def shuffle(src_input, local=False, frame_scores=None, clip_ratio=0.2) -> dict:
    """
    Shuffle features along the time dimension (dim=1).
    Steps:
      1) record split positions for each part after concatenation
      2) concatenate parts on joint axis -> [B, T, 69, 3]
      3) shuffle along time axis with a single permutation shared by all parts
      4) split back using recorded positions
    """
    # Clone to avoid in-place edits
    body  = src_input['body'].clone()
    left  = src_input['left'].clone()
    right = src_input['right'].clone()
    face  = src_input['face_all'].clone()

    B, T = body.shape[0], body.shape[1]
    device = body.device

    # Record sizes and order
    sizes = {
        'body': body.size(2),       # 9
        'left': left.size(2),       # 21
        'right': right.size(2),     # 21
        'face_all': face.size(2),   # 18
    }
    order = ['body', 'left', 'right', 'face_all']

    # Record split positions
    starts = {}
    cur = 0
    for k in order:
        starts[k] = (cur, cur + sizes[k])
        cur += sizes[k]

    # Concatenate features of each part -> [B, T, 69, 3]
    concat_all = torch.cat([body, left, right, face], dim=2)

    t_perm = None
    if local:
        if frame_scores is None or frame_scores.shape[:2] != (B, T):
            raise ValueError("When local=True, frame_scores must be provided with shape [B, T].")
        # find per-sample target windows
        windows = find_target_windows_batched(frame_scores, clip_ratio)
        # apply in-window shuffle per sample
        for b in range(B):
            tgt_start, tgt_end, win = windows[b]
            if win <= 1:
                continue  # nothing to shuffle
            # build a permutation only for the window
            wperm = torch.randperm(win, device=device)
            # apply on the slice [tgt_start:tgt_end)
            window_slice = concat_all[b, tgt_start:tgt_end]                # [win, 69, 3]
            window_shuf  = window_slice.index_select(0, wperm)             # permute along time inside the window
            concat_all[b, tgt_start:tgt_end] = window_shuf
        concat_shuf = concat_all
    else:
        # Global time shuffle with a single permutation
        t_perm = torch.randperm(T, device=device)
        concat_shuf = concat_all.index_select(dim=1, index=t_perm)

    # Split back using recorded positions
    b_s, b_e = starts['body']
    l_s, l_e = starts['left']
    r_s, r_e = starts['right']
    f_s, f_e = starts['face_all']

    body_shuf  = concat_shuf[:, :, b_s:b_e, :].contiguous()
    left_shuf  = concat_shuf[:, :, l_s:l_e, :].contiguous()
    right_shuf = concat_shuf[:, :, r_s:r_e, :].contiguous()
    face_shuf  = concat_shuf[:, :, f_s:f_e, :].contiguous()

    # Sanity checks
    assert body_shuf.shape  == body.shape
    assert left_shuf.shape  == left.shape
    assert right_shuf.shape == right.shape
    assert face_shuf.shape  == face.shape

    return {
        'body': body_shuf,
        'left': left_shuf,
        'right': right_shuf,
        'face_all': face_shuf,
        'time_perm': t_perm,   # permutation applied on dim=1
    }

class S2T_Dataset(Base_Dataset):
    def __init__(self, path, args, phase, dataset_name=None):
        super(S2T_Dataset, self).__init__()

        self.args = copy.deepcopy(args)  # Make a local copy

        if dataset_name is not None:
            self.args.dataset = dataset_name  # Modify only the local copy

        self.rgb_support = self.args.rgb_support
        self.max_length = args.max_length
        self.raw_data = utils.load_dataset_file(path)
        self.phase = phase

        if self.args.dataset == "CSL_Daily" or self.args.dataset == "OpenASL":
            self.pose_dir = pose_dirs[self.args.dataset]
            self.rgb_dir = rgb_dirs[self.args.dataset]
        elif self.args.dataset in ["How2Sign"]:
            self.pose_dir = os.path.join(pose_dirs[self.args.dataset], phase)
            self.rgb_dir = os.path.join(rgb_dirs[self.args.dataset], phase)
        else:
            raise NotImplementedError

        self.list = list(self.raw_data.keys())

        self.data_transform = transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), 
                                    ])
        if args.language_loss_type != 'LM_Alone':
            dpo_ir_text_path = ir_text_dirs[self.args.dataset]
            with open(dpo_ir_text_path, "r", encoding="utf-8") as f:
                self.ir_text_dict = json.load(f) 
            if self.phase.lower() == 'train' or self.phase.lower() == 'check_fit':
                before = len(self.list)
                self.list = [k for k in self.list if k in self.ir_text_dict]
                after = len(self.list)
                if before != after:
                    print(f"[Dataset] Filtered {before - after} samples missing from ir_text_dict "
                        f"({after}/{before} retained)")
        else:
            self.ir_text_dict = {}

    def __len__(self):
        return len(self.list)
    
    def __getitem__(self, index):
        key = self.list[index]
        sample = self.raw_data[key]
        
        text = sample['text']
        
        ir_text = None
        if self.args.language_loss_type != 'LM_Alone':
            ir_text = self.ir_text_dict[key]['pre_text'] if self.phase.lower() == 'train' else None

        if "gloss" in sample.keys():
            gloss = " ".join(sample['gloss'])
        else:
            gloss = ''
        
        name_sample = sample['name']
        if (self.args.dataset == "How2Sign" or self.args.dataset == "OpenASL") and not self.args.rgb_support:
            pose_sample, support_rgb_dict = self.load_pose(sample['pose'])
        else:
            pose_sample, support_rgb_dict = self.load_pose(sample['video_path'])

        if isinstance(pose_sample, tuple):
            pose_sample = pose_sample[0]  

        return name_sample, pose_sample, text, gloss, support_rgb_dict, key, ir_text
    
    def load_pose(self, path):
        if self.args.dataset == 'OpenASL':
            path = path.replace(":", "_")

        pose = pickle.load(open(os.path.join(self.pose_dir, path.replace(".mp4", '.pkl')), 'rb'))
            
        if 'start' in pose.keys():
            assert pose['start'] < pose['end']
            duration = pose['end'] - pose['start']
            start = pose['start']
        else:
            duration = len(pose['scores'])
            start = 0
                
        if duration > self.max_length:
            tmp = sorted(random.sample(range(duration), k=self.max_length))
        else:
            tmp = list(range(duration))
        
        tmp = np.array(tmp) + start
            
        skeletons = pose['keypoints']
        confs = pose['scores']
        skeletons_tmp = []
        confs_tmp = []
        for index in tmp:
            skeletons_tmp.append(skeletons[index])
            confs_tmp.append(confs[index])

        skeletons = skeletons_tmp
        confs = confs_tmp
    
        kps_with_scores = load_part_kp(skeletons, confs, force_ok=True)

        support_rgb_dict = {}
        if self.rgb_support:
            full_path = os.path.join(self.rgb_dir, path)
            support_rgb_dict = load_support_rgb_dict(tmp, skeletons, confs, full_path, self.data_transform)
            
        return kps_with_scores, support_rgb_dict

    def __str__(self):
        return f'# {self.phase} with a total of {len(self)}'

class S2T_Dataset_news(Base_Dataset):
    def __init__(self, path, args, phase):
        super(S2T_Dataset_news, self).__init__()
        self.args = args
        self.rgb_support = self.args.rgb_support
        self.phase = phase
        self.max_length = args.max_length

        path = pathlib.Path(path)

        with path.open(encoding='utf-8') as f:
            self.annotation = json.load(f)
       
        if self.args.dataset == "CSL_News":
            self.pose_dir = pose_dirs[args.dataset]
            self.rgb_dir = rgb_dirs[args.dataset]
      
        else:
            raise NotImplementedError
        sum_sample = len(self.annotation)
        self.data_transform = transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), 
                                    ])

        if phase == 'train':
            self.start_idx = int(sum_sample * 0.0)
            self.end_idx = int(sum_sample * 0.99)
        else:
            self.start_idx = int(sum_sample * 0.99)
            self.end_idx = int(sum_sample)
        
    def __len__(self):
        return self.end_idx - self.start_idx
    
    def __getitem__(self, index):
        num_retries = 10  

        # skip some invalid video sample
        for _ in range(num_retries):
            sample = self.annotation[self.start_idx:self.end_idx][index]

            text = sample['text']
            name_sample = sample['video']
           
            try:
                pose_sample, support_rgb_dict = self.load_pose(sample['pose'], sample['video'])
            except:
                import traceback

                traceback.print_exc()
                print(f"Failed to load examples with video: {name_sample}. "
                            f"Will randomly sample an example as a replacement.")
                index = random.randint(0, len(self) - 1)
                continue

            break
           
        else:  
            raise RuntimeError(f"Failed to fetch video after {num_retries} retries.")
        
        return name_sample, pose_sample, text, _, support_rgb_dict
    
    def load_pose(self, pose_name, rgb_name):
        pose = pickle.load(open(os.path.join(self.pose_dir, pose_name), 'rb'))
        full_path = os.path.join(self.rgb_dir, rgb_name)
        
        duration = len(pose['scores'])

        if duration > self.max_length:
            tmp = sorted(random.sample(range(duration), k=self.max_length))
        else:
            tmp = list(range(duration))
        
        tmp = np.array(tmp)
            
        # dict_keys(['keypoints', 'scores'])
        # keypoints (1, 133, 2)
        # scores (1, 133)
        
        skeletons = pose['keypoints']
        confs = pose['scores']
        skeletons_tmp = []
        confs_tmp = []
        
        for index in tmp:
            skeletons_tmp.append(skeletons[index])
            confs_tmp.append(confs[index])

        skeletons = skeletons_tmp
        confs = confs_tmp
                
        kps_with_scores = load_part_kp(skeletons, confs)
        
        support_rgb_dict = {}
        if self.rgb_support:
            support_rgb_dict = load_support_rgb_dict(tmp, skeletons, confs, full_path, self.data_transform)

        return kps_with_scores, support_rgb_dict

    def __str__(self):
        return f'#total {len(self)}'