MuseV/MMCM/mmcm/data/media_map/media_map.py

from __future__ import annotations
import bisect

import logging
from copy import deepcopy

from functools import partial
from typing import Any, Callable, Iterable, List, Union, Tuple, Dict

import numpy as np
from ..clip.clip_process import get_subseq_by_time
from ..clip.clip_stat import stat_clipseq_duration
from ..clip import Clip, ClipSeq, ClipIds, MatchedClipIds, MatchedClipIdsSeq
from .media_map_process import get_sub_mediamap_by_time
from ..emb import MediaMapEmb, H5pyMediaMapEmb
from ..general.items import Item, Items
from ...utils.data_util import pick_subdct
from ...utils.util import convert_class_attr_to_dict, load_dct_from_file

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name


__all__ = ["MetaInfo", "MetaInfoList", "MediaMap", "MediaMapSeq"]


class MetaInfo(Item):
    """歌曲、视频等媒体文件级别的元信息"""

    def __init__(
        self,
        mediaid=None,
        media_name=None,
        media_duration=None,
        signature=None,
        media_path: str = None,
        media_map_path: str = None,
        start: float = None,
        end: float = None,
        ext=None,
        **kwargs,
    ):
        super(MetaInfo).__init__()
        self.mediaid = mediaid
        self.media_name = media_name
        self.media_duration = media_duration
        self.signature = signature
        self.media_path = media_path
        self.media_map_path = media_map_path
        self.start = start
        self.end = end
        self.ext = ext
        self.__dict__.update(**kwargs)
        self.preprocess()

    def preprocess(self):
        self.set_start_end()

    def set_start_end(self):
        if self.start is None:
            self.start = 0
        elif self.start >= 0 and self.start <= 1:
            self.start = self.start * self.media_duration

        if self.end is None:
            self.end = self.media_duration
        elif self.end >= 0 and self.end <= 1:
            self.end = self.end * self.media_duration


class MetaInfoList(Items):
    """媒体元数据列表，主要用于多歌曲、多视频剪辑时存储原单一媒体文件的元信息"""

    def __init__(self, items: Union[MetaInfo, List[MetaInfo]] = None):
        """
        Args:
            meta_info_list (list, optional): MetaInfo 列表. Defaults to None.
        """
        if items is None:
            items = []
        else:
            items = items if isinstance(items, list) else [items]
        super().__init__(items)
        self.meta_info_list = self.items
        if len(self.items) > 1:
            self.reset()

    def __len__(self):
        return len(self.meta_info_list)

    def __getitem__(self, i) -> MetaInfo:
        return self.meta_info_list[i]

    @property
    def groupnum(self) -> int:
        return len(self.meta_info_list)


class MediaMap(object):
    """媒体信息基类，也可以理解为音乐谱面、视觉谱面、音游谱面基类。主要有 MetaInfo、MetaInfoList、ClipSeq 属性。
    不同的媒体信息的 属性 类会有不同，所以在类变量里做定义。如有变化，可以定义自己的属性类。
    """

    def __init__(
        self,
        meta_info: MetaInfo = None,
        clipseq: ClipSeq = None,
        stageseq: ClipSeq = None,
        frameseq: ClipSeq = None,
        emb: H5pyMediaMapEmb = None,
        **kwargs,
    ):
        """用于存储media的相关信息，media_info是json或直接字典

        Args:
            meta_info (MetaInfo): 当sub_meta_info不为None时, meta_info由sub_meta_info整合而成
            sub_meta_info (None or [MetaInfo]): 当多个MediaInfo拼在一起时,用于保留子MediaInfo的信息
            clipseq (ClipSeq): # 按照clipidx排序;
            stageseq (ClipSeq): # 比 clipseq 更高纬度的片段划分，例如clips是镜头分割，stages是scenes分割；clips是关键点分割，stages是结构分割；
            frameseq (ClipSeq): # 比 clipseq 更低纬度的片段划分
            kwargs (dict, optional): 所有相关信息都会作为 meta_info 的补充，赋值到 meta_info 中
        """
        self.meta_info = meta_info
        self.clipseq = clipseq
        self.frameseq = frameseq
        self.stageseq = stageseq
        self.emb = emb
        self.meta_info.__dict__.update(**kwargs)
        self.preprocess()

    def preprocess(
        self,
    ):
        if (self.meta_info.start != 0 and self.meta_info.start is not None) or (
            self.meta_info.end is not None and self.meta_info.end == 1
        ):
            self.drop_head_and_tail()
        self.meta_info.preprocess()
        if self.clipseq is not None:
            self.clipseq.preprocess()
        if self.frameseq is not None:
            self.frameseq.preprocess()
        if self.stageseq is not None:
            self.stageseq.preprocess()
        self.clip_start_idx = self.clipseq[0].clipid
        self.clip_end_idx = self.clipseq[-1].clipid

    def drop_head_and_tail(self) -> MediaMap:
        self.clipseq = get_subseq_by_time(
            self.clipseq,
            start=self.meta_info.start,
            end=self.meta_info.end,
            duration=self.meta_info.media_duration,
        )
        if self.stageseq is not None:
            self.stageseq = get_subseq_by_time(
                self.clipseq,
                start=self.meta_info.start,
                end=self.meta_info.end,
                duration=self.meta_info.media_duration,
            )

    def set_clip_value(self, k, v):
        """为clipseq中的每个clip赋值，

        Args:
            k (str): Clip中字段名
            v (any): Clip中字段值
        """
        self.clipseq.set_clip_value(k, v)

    def spread_metainfo_2_clip(
        self, target_keys: List = None, ignored_keys: List = None
    ) -> None:
        """将metainfo中的信息赋值到clip中，便于clip后面做相关处理。

        Args:
            target_keys ([str]): 待赋值的目标字段
        """
        dst = pick_subdct(
            self.meta_info.__dict__, target_keys=target_keys, ignored_keys=ignored_keys
        )
        for k, v in dst.items():
            self.set_clip_value(k, v)

    def spread_parameters(self, target_keys: list, ignored_keys) -> None:
        """元数据广播，将 media_info 的元数据广播到 clip 中，以及调用 clip 自己的参数传播。"""
        self.spread_metainfo_2_clip(target_keys=target_keys, ignored_keys=ignored_keys)
        for clip in self.clipseq:
            clip.spread_parameters()

    def stat(
        self,
    ):
        """统计 media_info 相关信息，便于了解，目前统计内容有
        1. 片段长度
        """
        self.stat_clipseq_duration()

    def stat_clipseq_duration(
        self,
    ):
        hist, bin_edges = stat_clipseq_duration(self.clipseq)
        print(self.media_name, "bin_edges", bin_edges)
        print(self.media_name, "hist", hist)

    def to_dct(self, target_keys: list = None, ignored_keys: list = None):
        raise NotImplementedError

    @property
    def duration(
        self,
    ):
        return self.clipseq.duration

    @property
    def mediaid(
        self,
    ):
        return self.meta_info.mediaid

    @property
    def media_name(
        self,
    ):
        return self.meta_info.media_name

    @property
    def duration_seq_emb(self):
        return self.clipseq.duration_seq_emb

    @property
    def timestamp_seq_emb(self):
        return self.clipseq.timestamp_seq_emb

    @property
    def rela_timestamp_seq_emb(self):
        return self.clipseq.rela_timestamp_seq_emb

    def get_emb(self, key, idx=None):
        # TODO: 待修改为更通用的形式
        if idx is None:
            idx = range(self.clip_start_idx, self.clip_end_idx + 1)
        elif isinstance(idx, int):
            idx += self.clip_start_idx
        elif isinstance(idx, Iterable):
            idx = [x + self.clip_start_idx for x in idx]
        else:
            raise ValueError(
                f"idx only support None, int, Iterable, but given {idx},type is {type(idx)}"
            )
        return self.emb.get_value(key, idx=idx)

    def get_meta_info_attr(self, key: str) -> Any:
        return getattr(self.meta_info, key)

    @classmethod
    def from_json_path(
        cls, path: Dict, emb_path: str, media_path: str = None, **kwargs
    ) -> MediaMap:
        media_map = load_dct_from_file(path)
        emb = H5pyMediaMapEmb(emb_path)
        return cls.from_data(media_map, emb=emb, media_path=media_path, **kwargs)


class MediaMapSeq(Items):
    def __init__(self, maps: List[MediaMap]) -> None:
        super().__init__(maps)
        self.maps = self.data
        self.preprocess()
        self.each_map_clipseq_num = [len(m.clipseq) for m in self.maps]
        self.each_map_clipseq_num_cumsum = np.cumsum([0] + self.each_map_clipseq_num)

    @property
    def clipseq(self):
        clipseq = []
        for m in self.maps:
            clipseq.extend(m.clipseq.data)
        return type(self.maps[0].clipseq)(clipseq)

    @property
    def stagesseq(self):
        stagesseq = []
        for m in self.maps:
            stagesseq.extend(m.stagesseq.data)
        return type(self.maps[0].stagesseq)(stagesseq)

    @property
    def frameseq(self):
        frameseq = []
        for m in self.maps:
            frameseq.extend(m.frameseq.data)
        return type(self.maps[0].frameseq)(frameseq)

    def preprocess(self):
        for m in self.maps:
            m.preprocess()

    def _combine_str(
        self,
        attrs: List[str],
        sep: str = "|",
        single_maxlen: int = 10,
        total_max_length: int = 60,
    ) -> str:
        return sep.join([str(attr)[:single_maxlen] for attr in attrs])[
            :total_max_length
        ]

    def get_meta_info_attr(self, key: str, func: Callable) -> Any:
        attrs = [m.get_meta_info_attr(key) for m in self.maps]
        return func(attrs)

    @property
    def mediaid(self) -> str:
        return self.get_meta_info_attr(key="mediaid", func=self._combine_str)

    @property
    def media_name(self) -> str:
        return self.get_meta_info_attr(key="media_name", func=self._combine_str)

    @property
    def duration(self) -> float:
        return sum([m.duration for m in self.maps])

    @property
    def media_duration(self) -> float:
        return self.get_meta_info_attr(key="media_duration", func=sum)

    @classmethod
    def from_json_paths(
        cls,
        media_map_class: MediaMap,
        media_paths: str,
        media_map_paths: str,
        emb_paths: str,
        **kwargs,
    ) -> MediaMapSeq:
        map_seq = [
            media_map_class.from_json_path(
                path=media_map_paths[i],
                emb_path=emb_paths[i],
                media_path=media_paths[i],
                **kwargs,
            )
            for i in range(len(media_map_paths))
        ]
        return cls(map_seq)

    # TODO: implement mapseq stat func
    def stat(self):
        for m in self.maps:
            m.stat()

    def _combine_embs(self, embs):
        return np.concatenate(embs, axis=0)

    @property
    def duration_seq_emb(self):
        embs = [m.duration_seq_emb for m in self.maps]
        return self._combine_embs(embs)

    @property
    def timestamp_seq_emb(self):
        embs = [m.timestamp_seq_emb for m in self.maps]
        return self._combine_embs(embs)

    @property
    def rela_timestamp_seq_emb(self):
        embs = [m.rela_timestamp_seq_emb for m in self.maps]
        return self._combine_embs(embs)

    def clip_idx_2_map_idx(self, idx):
        target_map_idx = bisect.bisect_right(self.each_map_clipseq_num_cumsum, idx)
        target_map_idx = min(max(0, target_map_idx - 1), len(self.maps) - 1)
        target_map_clip_idx = idx - self.each_map_clipseq_num_cumsum[target_map_idx]
        return target_map_idx, target_map_clip_idx

    def get_emb(self, key: str, idx: Union[None, int, List[int]] = None) -> np.array:
        if idx is None:
            embs = [m.get_emb(key, idx=idx) for m in self.maps]
        else:
            if not isinstance(idx, list):
                idx = [idx]
            embs = []
            for c_idx in idx:
                target_map_idx, target_map_clip_idx = self.clip_idx_2_map_idx(c_idx)
                embs.append(
                    self.maps[target_map_idx].get_emb(key, int(target_map_clip_idx))
                )
        if len(embs) == 1:
            return embs[0]
        else:
            return self._combine_embs(embs)