musev/data/data_util.py

from typing import List, Dict, Literal, Union, Tuple
import os
import string
import logging

import torch
import numpy as np
from einops import rearrange, repeat

logger = logging.getLogger(__name__)


def generate_tasks_of_dir(
    path: str,
    output_dir: str,
    exts: Tuple[str],
    same_dir_name: bool = False,
    **kwargs,
) -> List[Dict]:
    """covert video directory into tasks

    Args:
        path (str): _description_
        output_dir (str): _description_
        exts (Tuple[str]): _description_
        same_dir_name (bool, optional): 存储路径是否保留和源视频相同的父文件名. Defaults to False.
            whether keep the same parent dir name as the source video
    Returns:
        List[Dict]: _description_
    """
    tasks = []
    for rootdir, dirs, files in os.walk(path):
        for basename in files:
            if basename.lower().endswith(exts):
                video_path = os.path.join(rootdir, basename)
                filename, ext = basename.split(".")
                rootdir_name = os.path.basename(rootdir)
                if same_dir_name:
                    save_path = os.path.join(
                        output_dir, rootdir_name, f"{filename}.h5py"
                    )
                    save_dir = os.path.join(output_dir, rootdir_name)
                else:
                    save_path = os.path.join(output_dir, f"{filename}.h5py")
                    save_dir = output_dir
                task = {
                    "video_path": video_path,
                    "output_path": save_path,
                    "output_dir": save_dir,
                    "filename": filename,
                    "ext": ext,
                }
                task.update(kwargs)
                tasks.append(task)
    return tasks


def sample_by_idx(
    T: int,
    n_sample: int,
    sample_rate: int,
    sample_start_idx: int = None,
    change_sample_rate: bool = False,
    seed: int = None,
    whether_random: bool = True,
    n_independent: int = 0,
) -> List[int]:
    """given a int to represent candidate list, sample n_sample with sample_rate from the candidate list

    Args:
        T (int): _description_
        n_sample (int): 目标采样数目. sample number
        sample_rate (int): 采样率, 每隔sample_rate个采样一个. sample interval, pick one per sample_rate number
        sample_start_idx (int, optional): 采样开始位置的选择. start position to sample . Defaults to 0.
        change_sample_rate (bool, optional): 是否可以通过降低sample_rate的方式来完成采样. whether allow changing sample_rate to finish sample process. Defaults to False.
        whether_random (bool, optional): 是否最后随机选择开始点. whether randomly choose sample start position. Defaults to False.

    Raises:
        ValueError: T / sample_rate should be larger than n_sample
    Returns:
        List[int]: 采样的索引位置. sampled index position
    """
    if T < n_sample:
        raise ValueError(f"T({T}) < n_sample({n_sample})")
    else:
        if T / sample_rate < n_sample:
            if not change_sample_rate:
                raise ValueError(
                    f"T({T}) / sample_rate({sample_rate}) < n_sample({n_sample})"
                )
            else:
                while T / sample_rate < n_sample:
                    sample_rate -= 1
                    logger.error(
                        f"sample_rate{sample_rate+1} is too large, decrease to {sample_rate}"
                    )
                    if sample_rate == 0:
                        raise ValueError("T / sample_rate < n_sample")

    if sample_start_idx is None:
        if whether_random:
            sample_start_idx_candidates = np.arange(T - n_sample * sample_rate)
            if seed is not None:
                np.random.seed(seed)
            sample_start_idx = np.random.choice(sample_start_idx_candidates, 1)[0]

        else:
            sample_start_idx = 0
    sample_end_idx = sample_start_idx + sample_rate * n_sample
    sample = list(range(sample_start_idx, sample_end_idx, sample_rate))
    if n_independent == 0:
        n_independent_sample = None
    else:
        left_candidate = np.array(
            list(range(0, sample_start_idx)) + list(range(sample_end_idx, T))
        )
        if len(left_candidate) >= n_independent:
            # 使用两端的剩余空间采样, use the left space to sample
            n_independent_sample = np.random.choice(left_candidate, n_independent)
        else:
            # 当两端没有剩余采样空间时，使用任意不是sample中的帧
            # if no enough space to sample, use any frame not in sample
            left_candidate = np.array(list(set(range(T) - set(sample))))
            n_independent_sample = np.random.choice(left_candidate, n_independent)

    return sample, sample_rate, n_independent_sample


def sample_tensor_by_idx(
    tensor: Union[torch.Tensor, np.ndarray],
    n_sample: int,
    sample_rate: int,
    sample_start_idx: int = 0,
    change_sample_rate: bool = False,
    seed: int = None,
    dim: int = 0,
    return_type: Literal["numpy", "torch"] = "torch",
    whether_random: bool = True,
    n_independent: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor, int, torch.Tensor, torch.Tensor]:
    """sample sub_tensor

    Args:
        tensor (Union[torch.Tensor, np.ndarray]): _description_
        n_sample (int): _description_
        sample_rate (int): _description_
        sample_start_idx (int, optional): _description_. Defaults to 0.
        change_sample_rate (bool, optional): _description_. Defaults to False.
        seed (int, optional): _description_. Defaults to None.
        dim (int, optional): _description_. Defaults to 0.
        return_type (Literal[&quot;numpy&quot;, &quot;torch&quot;], optional): _description_. Defaults to "torch".
        whether_random (bool, optional): _description_. Defaults to True.
        n_independent (int, optional): 独立于n_sample的采样数量. Defaults to 0.
            n_independent sample number that is independent of n_sample

    Returns:
        Tuple[torch.Tensor, torch.Tensor, int, torch.Tensor, torch.Tensor]: sampled tensor
    """
    if isinstance(tensor, np.ndarray):
        tensor = torch.from_numpy(tensor)
    T = tensor.shape[dim]
    sample_idx, sample_rate, independent_sample_idx = sample_by_idx(
        T,
        n_sample,
        sample_rate,
        sample_start_idx,
        change_sample_rate,
        seed,
        whether_random=whether_random,
        n_independent=n_independent,
    )
    sample_idx = torch.LongTensor(sample_idx)
    sample = torch.index_select(tensor, dim, sample_idx)
    if independent_sample_idx is not None:
        independent_sample_idx = torch.LongTensor(independent_sample_idx)
        independent_sample = torch.index_select(tensor, dim, independent_sample_idx)
    else:
        independent_sample = None
        independent_sample_idx = None
    if return_type == "numpy":
        sample = sample.cpu().numpy()
    return sample, sample_idx, sample_rate, independent_sample, independent_sample_idx


def concat_two_tensor(
    data1: torch.Tensor,
    data2: torch.Tensor,
    dim: int,
    method: Literal[
        "first_in_first_out", "first_in_last_out", "intertwine", "index"
    ] = "first_in_first_out",
    data1_index: torch.long = None,
    data2_index: torch.long = None,
    return_index: bool = False,
):
    """concat two tensor along dim with given method

    Args:
        data1 (torch.Tensor): first in data
        data2 (torch.Tensor): last in  data
        dim (int): _description_
        method (Literal[ &quot;first_in_first_out&quot;, &quot;first_in_last_out&quot;, &quot;intertwine&quot; ], optional): _description_. Defaults to "first_in_first_out".

    Raises:
        NotImplementedError: unsupported method
        ValueError: unsupported method

    Returns:
        _type_: _description_
    """
    len_data1 = data1.shape[dim]
    len_data2 = data2.shape[dim]

    if method == "first_in_first_out":
        res = torch.concat([data1, data2], dim=dim)
        data1_index = range(len_data1)
        data2_index = [len_data1 + x for x in range(len_data2)]
    elif method == "first_in_last_out":
        res = torch.concat([data2, data1], dim=dim)
        data2_index = range(len_data2)
        data1_index = [len_data2 + x for x in range(len_data1)]
    elif method == "intertwine":
        raise NotImplementedError("intertwine")
    elif method == "index":
        res = concat_two_tensor_with_index(
            data1=data1,
            data1_index=data1_index,
            data2=data2,
            data2_index=data2_index,
            dim=dim,
        )
    else:
        raise ValueError(
            "only support first_in_first_out, first_in_last_out, intertwine, index"
        )
    if return_index:
        return res, data1_index, data2_index
    else:
        return res


def concat_two_tensor_with_index(
    data1: torch.Tensor,
    data1_index: torch.LongTensor,
    data2: torch.Tensor,
    data2_index: torch.LongTensor,
    dim: int,
) -> torch.Tensor:
    """_summary_

    Args:
        data1 (torch.Tensor): b1*c1*h1*w1*...
        data1_index (torch.LongTensor): N, if dim=1, N=c1
        data2 (torch.Tensor): b2*c2*h2*w2*...
        data2_index (torch.LongTensor): M, if dim=1, M=c2
        dim (int): int

    Returns:
        torch.Tensor: b*c*h*w*..., if dim=1, b=b1=b2, c=c1+c2, h=h1=h2, w=w1=w2,...
    """
    shape1 = list(data1.shape)
    shape2 = list(data2.shape)
    target_shape = list(shape1)
    target_shape[dim] = shape1[dim] + shape2[dim]
    target = torch.zeros(target_shape, device=data1.device, dtype=data1.dtype)
    target = batch_index_copy(target, dim=dim, index=data1_index, source=data1)
    target = batch_index_copy(target, dim=dim, index=data2_index, source=data2)
    return target


def repeat_index_to_target_size(
    index: torch.LongTensor, target_size: int
) -> torch.LongTensor:
    if len(index.shape) == 1:
        index = repeat(index, "n -> b n", b=target_size)
    if len(index.shape) == 2:
        remainder = target_size % index.shape[0]
        assert (
            remainder == 0
        ), f"target_size % index.shape[0] must be zero, but give {target_size % index.shape[0]}"
        index = repeat(index, "b n -> (b c) n", c=int(target_size / index.shape[0]))
    return index


def batch_concat_two_tensor_with_index(
    data1: torch.Tensor,
    data1_index: torch.LongTensor,
    data2: torch.Tensor,
    data2_index: torch.LongTensor,
    dim: int,
) -> torch.Tensor:
    return concat_two_tensor_with_index(data1, data1_index, data2, data2_index, dim)


def interwine_two_tensor(
    data1: torch.Tensor,
    data2: torch.Tensor,
    dim: int,
    return_index: bool = False,
) -> torch.Tensor:
    shape1 = list(data1.shape)
    shape2 = list(data2.shape)
    target_shape = list(shape1)
    target_shape[dim] = shape1[dim] + shape2[dim]
    target = torch.zeros(target_shape, device=data1.device, dtype=data1.dtype)
    data1_reshape = torch.swapaxes(data1, 0, dim)
    data2_reshape = torch.swapaxes(data2, 0, dim)
    target = torch.swapaxes(target, 0, dim)
    total_index = set(range(target_shape[dim]))
    data1_index = range(0, 2 * shape1[dim], 2)
    data2_index = sorted(list(set(total_index) - set(data1_index)))
    data1_index = torch.LongTensor(data1_index)
    data2_index = torch.LongTensor(data2_index)
    target[data1_index, ...] = data1_reshape
    target[data2_index, ...] = data2_reshape
    target = torch.swapaxes(target, 0, dim)
    if return_index:
        return target, data1_index, data2_index
    else:
        return target


def split_index(
    indexs: torch.Tensor,
    n_first: int = None,
    n_last: int = None,
    method: Literal[
        "first_in_first_out", "first_in_last_out", "intertwine", "index", "random"
    ] = "first_in_first_out",
):
    """_summary_

    Args:
        indexs (List): _description_
        n_first (int): _description_
        n_last (int): _description_
        method (Literal[ &quot;first_in_first_out&quot;, &quot;first_in_last_out&quot;, &quot;intertwine&quot;, &quot;index&quot; ], optional): _description_. Defaults to "first_in_first_out".

    Raises:
        NotImplementedError: _description_

    Returns:
        first_index: _description_
        last_index:
    """
    # assert (
    #     n_first is None and n_last is None
    # ), "must assign one value for n_first or n_last"
    n_total = len(indexs)
    if n_first is None:
        n_first = n_total - n_last
    if n_last is None:
        n_last = n_total - n_first
    assert len(indexs) == n_first + n_last
    if method == "first_in_first_out":
        first_index = indexs[:n_first]
        last_index = indexs[n_first:]
    elif method == "first_in_last_out":
        first_index = indexs[n_last:]
        last_index = indexs[:n_last]
    elif method == "intertwine":
        raise NotImplementedError
    elif method == "random":
        idx_ = torch.randperm(len(indexs))
        first_index = indexs[idx_[:n_first]]
        last_index = indexs[idx_[n_first:]]
    return first_index, last_index


def split_tensor(
    tensor: torch.Tensor,
    dim: int,
    n_first=None,
    n_last=None,
    method: Literal[
        "first_in_first_out", "first_in_last_out", "intertwine", "index", "random"
    ] = "first_in_first_out",
    need_return_index: bool = False,
):
    device = tensor.device
    total = tensor.shape[dim]
    if n_first is None:
        n_first = total - n_last
    if n_last is None:
        n_last = total - n_first
    indexs = torch.arange(
        total,
        dtype=torch.long,
        device=device,
    )
    (
        first_index,
        last_index,
    ) = split_index(
        indexs=indexs,
        n_first=n_first,
        method=method,
    )
    first_tensor = torch.index_select(tensor, dim=dim, index=first_index)
    last_tensor = torch.index_select(tensor, dim=dim, index=last_index)
    if need_return_index:
        return (
            first_tensor,
            last_tensor,
            first_index,
            last_index,
        )
    else:
        return (first_tensor, last_tensor)


# TODO: 待确定batch_index_select的优化
def batch_index_select(
    tensor: torch.Tensor, index: torch.LongTensor, dim: int
) -> torch.Tensor:
    """_summary_

    Args:
        tensor (torch.Tensor): D1*D2*D3*D4...
        index (torch.LongTensor): D1*N or N, N<= tensor.shape[dim]
        dim (int): dim to select

    Returns:
        torch.Tensor: D1*...*N*...
    """
    # TODO: now only support N same for every d1
    if len(index.shape) == 1:
        return torch.index_select(tensor, dim=dim, index=index)
    else:
        index = repeat_index_to_target_size(index, tensor.shape[0])
        out = []
        for i in torch.arange(tensor.shape[0]):
            sub_tensor = tensor[i]
            sub_index = index[i]
            d = torch.index_select(sub_tensor, dim=dim - 1, index=sub_index)
            out.append(d)
        return torch.stack(out).to(dtype=tensor.dtype)


def batch_index_copy(
    tensor: torch.Tensor, dim: int, index: torch.LongTensor, source: torch.Tensor
) -> torch.Tensor:
    """_summary_

    Args:
        tensor (torch.Tensor): b*c*h
        dim (int):
        index (torch.LongTensor): b*d,
        source (torch.Tensor):
            b*d*h*..., if dim=1
            b*c*d*..., if dim=2

    Returns:
        torch.Tensor: b*c*d*...
    """
    if len(index.shape) == 1:
        tensor.index_copy_(dim=dim, index=index, source=source)
    else:
        index = repeat_index_to_target_size(index, tensor.shape[0])

        batch_size = tensor.shape[0]
        for b in torch.arange(batch_size):
            sub_index = index[b]
            sub_source = source[b]
            sub_tensor = tensor[b]
            sub_tensor.index_copy_(dim=dim - 1, index=sub_index, source=sub_source)
            tensor[b] = sub_tensor
    return tensor


def batch_index_fill(
    tensor: torch.Tensor,
    dim: int,
    index: torch.LongTensor,
    value: Literal[torch.Tensor, torch.float],
) -> torch.Tensor:
    """_summary_

    Args:
        tensor (torch.Tensor): b*c*h
        dim (int):
        index (torch.LongTensor): b*d,
        value (torch.Tensor): b

    Returns:
        torch.Tensor: b*c*d*...
    """
    index = repeat_index_to_target_size(index, tensor.shape[0])
    batch_size = tensor.shape[0]
    for b in torch.arange(batch_size):
        sub_index = index[b]
        sub_value = value[b] if isinstance(value, torch.Tensor) else value
        sub_tensor = tensor[b]
        sub_tensor.index_fill_(dim - 1, sub_index, sub_value)
        tensor[b] = sub_tensor
    return tensor


def adaptive_instance_normalization(
    src: torch.Tensor,
    dst: torch.Tensor,
    eps: float = 1e-6,
):
    """
    Args:
        src (torch.Tensor): b c t h w
        dst (torch.Tensor): b c t h w
    """
    ndim = src.ndim
    if ndim == 5:
        dim = (2, 3, 4)
    elif ndim == 4:
        dim = (2, 3)
    elif ndim == 3:
        dim = 2
    else:
        raise ValueError("only support ndim in [3,4,5], but given {ndim}")
    var, mean = torch.var_mean(src, dim=dim, keepdim=True, correction=0)
    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
    dst = align_repeat_tensor_single_dim(dst, src.shape[0], dim=0)
    mean_acc, var_acc = torch.var_mean(dst, dim=dim, keepdim=True, correction=0)
    # mean_acc = sum(mean_acc) / float(len(mean_acc))
    # var_acc = sum(var_acc) / float(len(var_acc))
    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5
    src = (((src - mean) / std) * std_acc) + mean_acc
    return src


def adaptive_instance_normalization_with_ref(
    src: torch.LongTensor,
    dst: torch.LongTensor,
    style_fidelity: float = 0.5,
    do_classifier_free_guidance: bool = True,
):
    # logger.debug(
    #     f"src={src.shape}, min={src.min()}, max={src.max()}, mean={src.mean()}, \n"
    #     f"dst={src.shape}, min={dst.min()}, max={dst.max()}, mean={dst.mean()}"
    # )
    batch_size = src.shape[0] // 2
    uc_mask = torch.Tensor([1] * batch_size + [0] * batch_size).type_as(src).bool()
    src_uc = adaptive_instance_normalization(src, dst)
    src_c = src_uc.clone()
    # TODO: 该部分默认 do_classifier_free_guidance and style_fidelity > 0 = True
    if do_classifier_free_guidance and style_fidelity > 0:
        src_c[uc_mask] = src[uc_mask]
    src = style_fidelity * src_c + (1.0 - style_fidelity) * src_uc
    return src


def batch_adain_conditioned_tensor(
    tensor: torch.Tensor,
    src_index: torch.LongTensor,
    dst_index: torch.LongTensor,
    keep_dim: bool = True,
    num_frames: int = None,
    dim: int = 2,
    style_fidelity: float = 0.5,
    do_classifier_free_guidance: bool = True,
    need_style_fidelity: bool = False,
):
    """_summary_

    Args:
        tensor (torch.Tensor): b c t h w
        src_index (torch.LongTensor): _description_
        dst_index (torch.LongTensor): _description_
        keep_dim (bool, optional): _description_. Defaults to True.

    Returns:
        _type_: _description_
    """
    ndim = tensor.ndim
    dtype = tensor.dtype
    if ndim == 4 and num_frames is not None:
        tensor = rearrange(tensor, "(b t) c h w->  b c t h w ", t=num_frames)
    src = batch_index_select(tensor, dim=dim, index=src_index).contiguous()
    dst = batch_index_select(tensor, dim=dim, index=dst_index).contiguous()
    if need_style_fidelity:
        src = adaptive_instance_normalization_with_ref(
            src=src,
            dst=dst,
            style_fidelity=style_fidelity,
            do_classifier_free_guidance=do_classifier_free_guidance,
            need_style_fidelity=need_style_fidelity,
        )
    else:
        src = adaptive_instance_normalization(
            src=src,
            dst=dst,
        )
    if keep_dim:
        src = batch_concat_two_tensor_with_index(
            src.to(dtype=dtype),
            src_index,
            dst.to(dtype=dtype),
            dst_index,
            dim=dim,
        )

    if ndim == 4 and num_frames is not None:
        src = rearrange(tensor, "b c t h w ->(b t) c h w")
    return src


def align_repeat_tensor_single_dim(
    src: torch.Tensor,
    target_length: int,
    dim: int = 0,
    n_src_base_length: int = 1,
    src_base_index: List[int] = None,
) -> torch.Tensor:
    """沿着 dim 纬度， 补齐 src 的长度到目标 target_length。
    当 src 长度不如 target_length 时， 取其中 前 n_src_base_length 然后 repeat 到 target_length

    align length of src to target_length along dim
    when src length is less than target_length, take the first n_src_base_length and repeat to target_length

    Args:
        src (torch.Tensor): 输入 tensor, input tensor
        target_length (int): 目标长度, target_length
        dim (int, optional): 处理纬度, target dim . Defaults to 0.
        n_src_base_length (int, optional): src 的基本单元长度, basic length of src. Defaults to 1.

    Returns:
        torch.Tensor: _description_
    """
    src_dim_length = src.shape[dim]
    if target_length > src_dim_length:
        if target_length % src_dim_length == 0:
            new = src.repeat_interleave(
                repeats=target_length // src_dim_length, dim=dim
            )
        else:
            if src_base_index is None and n_src_base_length is not None:
                src_base_index = torch.arange(n_src_base_length)

            new = src.index_select(
                dim=dim,
                index=torch.LongTensor(src_base_index).to(device=src.device),
            )
            new = new.repeat_interleave(
                repeats=target_length // len(src_base_index),
                dim=dim,
            )
    elif target_length < src_dim_length:
        new = src.index_select(
            dim=dim,
            index=torch.LongTensor(torch.arange(target_length)).to(device=src.device),
        )
    else:
        new = src
    return new


def fuse_part_tensor(
    src: torch.Tensor,
    dst: torch.Tensor,
    overlap: int,
    weight: float = 0.5,
    skip_step: int = 0,
) -> torch.Tensor:
    """fuse overstep tensor with weight of src into dst
    out = src_fused_part * weight + dst * (1-weight) for overlap

    Args:
        src (torch.Tensor): b c t h w
        dst (torch.Tensor): b c t h w
        overlap (int): 1
        weight (float, optional): weight of src tensor part. Defaults to 0.5.

    Returns:
        torch.Tensor: fused tensor
    """
    if overlap == 0:
        return dst
    else:
        dst[:, :, skip_step : skip_step + overlap] = (
            weight * src[:, :, -overlap:]
            + (1 - weight) * dst[:, :, skip_step : skip_step + overlap]
        )
        return dst