MuseV/MMCM/mmcm/vision/data/video_dataset.py

import os
import itertools
from collections import namedtuple
from typing import Any, Iterator, List, Literal, Sequence
import math
from einops import rearrange

from moviepy.editor import VideoFileClip
import torchvision
from torch.utils.data.dataset import Dataset, IterableDataset
from PIL import Image
import cv2
import torch
import numpy as np

from ...utils.path_util import get_dir_file_map
from ...utils.itertools_util import generate_sample_idxs, overlap2step, step2overlap


VideoDatasetOutput = namedtuple("video_dataset_output", ["data", "index"])


def worker_init_fn(worker_id: int):
    worker_info = torch.utils.data.get_worker_info()
    dataset = worker_info.dataset  # the dataset copy in this worker process
    overall_start = 0
    overall_end = len(dataset)
    # configure the dataset to only process the split workload
    per_worker = int(
        math.ceil((overall_end - overall_start) / float(worker_info.num_workers))
    )
    worker_id = worker_info.id
    dataset_start = overall_start + worker_id * per_worker
    dataset_end = min(overall_start + per_worker, overall_end)
    dataset.sample_indexs = dataset.sample_indexs[dataset_start:dataset_end]


class SequentialDataset(IterableDataset):
    def __init__(
        self,
        raw_datas,
        time_size: int,
        step: int,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        max_num_per_batch: int = None,
        data_type: Literal["bgr", "rgb"] = "bgr",
        channels_order: str = "t h w c",
        sample_indexs: List[List[int]] = None,
    ) -> None:
        """_summary_

        Args:
            raw_datas (_type_): all original data
            time_size (int): frames number of a clip
            step (int): step of two windows
            overlap (int, optional): overlap of two windows. Defaults to None.
            sample_rate (int, optional): sample 1 evey sample_rate number. Defaults to 1.
            drop_last (bool, optional): whether drop the last if length of last batch < time_size. Defaults to False.
        """
        super().__init__()
        self.time_size = time_size
        if overlap is not None and step is None:
            step = overlap2step(overlap, time_size)
        if step is not None and overlap is None:
            overlap = step2overlap(step, time_size)
        self.overlap = overlap
        self.step = step
        self.sample_rate = sample_rate
        self.drop_last = drop_last
        self.raw_datas = raw_datas
        self.max_num_per_batch = max_num_per_batch
        if sample_indexs is not None:
            self.sample_indexs = sample_indexs
        else:
            self.generate_sample_idxs()
        self.current_pos = 0
        self.data_type = data_type
        self.channels_order = channels_order

    def generate_sample_idxs(
        self,
    ):
        self.sample_indexs = generate_sample_idxs(
            total=self.total_frames,
            window_size=self.time_size,
            step=self.step,
            sample_rate=self.sample_rate,
            drop_last=self.drop_last,
            max_num_per_window=self.max_num_per_batch,
        )

    def get_raw_datas(
        self,
    ):
        return self.raw_datas

    def get_raw_data(self, index: int):
        raise NotImplementedError

    def get_batch_raw_data(self, indexs: List[int]):
        datas = [self.get_raw_data(i) for i in indexs]
        datas = np.stack(datas, axis=0)
        return datas

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

    def __iter__(self) -> Iterator[Any]:
        return self

    def __getitem__(self, index):
        sample_indexs = self.sample_indexs[index]
        data = self.get_batch_raw_data(sample_indexs)
        if self.channels_order != "t h w c":
            data = rearrange(data, "t h w c -> {}".format(self.channels_order))
        sample_indexs = np.array(sample_indexs)
        return VideoDatasetOutput(data, sample_indexs)

    def get_data(self, index):
        return self.__getitem__(index)

    def __next__(self):
        while self.current_pos < len(self.sample_indexs):
            data = self.get_data(self.current_pos)
            self.current_pos += 1
            return data
        self.current_pos = 0
        raise StopIteration

    def preview(self, clip):
        """show data clip,
        play for image, video, and print for str list

        Args:
            clip (_type_): _description_
        """
        raise NotImplementedError

    def close(self):
        """
        close file handle if subclass open file
        """
        raise NotImplementedError

    @property
    def fps(self):
        raise NotImplementedError

    @property
    def total_frames(self):
        raise NotImplementedError

    @property
    def duration(self):
        raise NotImplementedError

    @property
    def width(self):
        raise NotImplementedError

    @property
    def height(self):
        raise NotImplementedError


class ItemsSequentialDataset(SequentialDataset):
    def __init__(
        self,
        raw_datas: Sequence,
        time_size: int,
        step: int,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        sample_indexs: List[List[int]] = None,
    ) -> None:
        super().__init__(
            raw_datas,
            time_size,
            step,
            overlap,
            sample_rate,
            drop_last,
            sample_indexs=sample_indexs,
        )

    def get_raw_data(self, index: int):
        return self.raw_datas[index]

    def prepare_raw_datas(self, raw_datas) -> Sequence:
        return raw_datas

    @property
    def total_frames(self):
        return len(self.raw_datas)


class ListSequentialDataset(ItemsSequentialDataset):
    def preview(self, clip):
        print(f"type is {self.__class__.__name__}, num is {len(clip)}")
        print(clip)


class ImagesSequentialDataset(ItemsSequentialDataset):
    def __init__(
        self,
        img_dir: Sequence,
        time_size: int,
        step: int,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        data_type: Literal["bgr", "rgb"] = "bgr",
        channels_order: str = "t h w c",
        sample_indexs: List[List[int]] = None,
    ) -> None:
        self.imgs_path = sorted(get_dir_file_map(img_dir).values())
        super().__init__(
            self.imgs_path,
            time_size,
            step,
            overlap,
            sample_rate,
            drop_last,
            data_ty=data_type,
            channels_order=channels_order,
            sample_indexs=sample_indexs,
        )


class PILImageSequentialDataset(ImagesSequentialDataset):
    def __getitem__(self, index: int) -> Image.Image:
        data, sample_indexs = super().__getitem__(index)
        data = [Image.open(x) for x in data]
        return VideoDatasetOutput(data, sample_indexs)


class MoviepyVideoDataset(SequentialDataset):
    def __init__(
        self,
        path,
        time_size: int,
        step: int,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        data_type: Literal["bgr", "rgb"] = "bgr",
        contenct_box: List[int] = None,
        sample_indexs: List[List[int]] = None,
    ) -> None:
        self.path = path
        self.f = self.prepare_raw_datas(self.path)
        super().__init__(
            self.f,
            time_size,
            step,
            overlap,
            sample_rate,
            drop_last,
            data_type=data_type,
            sample_indexs=sample_indexs,
        )
        self.contenct_box = contenct_box

    def prepare_raw_datas(self, path):
        f = VideoFileClip(path)
        return f

    def get_raw_data(self, index: int):
        return self.f.get_frame(index * 1 / self.f.fps)

    @property
    def fps(self):
        return self.f.fps

    @property
    def size(self):
        return self.f.size

    @property
    def total_frames(self):
        return int(self.duration * self.fps)

    @property
    def duration(self):
        return self.f.duration

    @property
    def width(self):
        return self.f.w

    @property
    def height(self):
        return self.f.h

    def __next__(
        self,
    ):
        video_clips = []
        cnt = 0
        frame_indexs = []

        for frame in itertools.islice(self.video.iter_frames(), step=self.step):
            if cnt >= self.total_frames:
                raise StopIteration
            else:
                frame_indexs.append(cnt)
                cnt += self.step
            if len(video_clips) < self.time_size:
                video_clips.append(frame)
            else:
                return_video_clips = video_clips
                return_frame_indexs = frame_indexs
                video_clips = []
                frame_indexs = []
                return VideoDatasetOutput(return_video_clips, return_frame_indexs)


class TorchVideoDataset(object):
    pass


class OpenCVVideoDataset(SequentialDataset):
    def __init__(
        self,
        path,
        time_size: int,
        step: int,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        data_type: Literal["bgr", "rgb"] = "bgr",
        channels_order: str = "t h w c",
        sample_indexs: List[List[int]] = None,
    ) -> None:
        self.path = path
        self.f = self.prepare_raw_datas(path)
        super().__init__(
            self.f,
            time_size,
            step,
            overlap,
            sample_rate,
            drop_last,
            data_type=data_type,
            channels_order=channels_order,
            sample_indexs=sample_indexs,
        )

    def prepare_raw_datas(self, path):
        f = cv2.VideoCapture(path)
        return f

    def get_raw_data(self, index: int):
        self.f.set(cv2.CAP_PROP_POS_FRAMES, index)
        if index < 0 or index >= self.total_frames:
            raise IndexError(
                f"index must in [0, {self.total_frames -1 }], but given index"
            )
        ret, frame = self.f.read()
        if self.data_type == "rgb":
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        return frame

    def get_raw_data_by_time(self, idx):
        raise NotImplementedError

    @property
    def total_frames(self):
        return int(self.f.get(cv2.CAP_PROP_FRAME_COUNT))

    @property
    def width(self):
        return int(self.f.get(cv2.CAP_PROP_FRAME_WIDTH))

    @property
    def height(self):
        return int(self.f.get(cv2.CAP_PROP_FRAME_HEIGHT))

    @property
    def durtion(self):
        return self.total_frames / self.fps

    @property
    def fps(self):
        return self.f.get(cv2.CAP_PROP_FPS)


class DecordVideoDataset(SequentialDataset):
    def __init__(
        self,
        path,
        time_size: int,
        step: int,
        device: str,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        device_id: int = 0,
        data_type: Literal["bgr", "rgb"] = "bgr",
        channels_order: str = "t h w c",
        sample_indexs: List[List[int]] = None,
    ) -> None:
        self.path = path
        self.device = device
        self.device_id = device_id
        self.f = self.prepare_raw_datas(path)
        super().__init__(
            self.f,
            time_size,
            step,
            overlap,
            sample_rate,
            drop_last,
            data_type=data_type,
            channels_order=channels_order,
            sample_indexs=sample_indexs,
        )

    def prepare_raw_datas(self, path):
        from decord import VideoReader
        from decord import cpu, gpu

        if self.device == "cpu":
            device = cpu(self.device_id)
        else:
            device = gpu(self.device_id)
        with open(path, "rb") as f:
            f = VideoReader(f, ctx=device)
        return f

    # decord 的 颜色通道 通道默认是 rgb
    def get_raw_data(self, index: int):
        data = self.f[index].asnumpy()
        if self.data_type == "bgr":
            data = data[:, :, ::-1]
        return data

    def get_batch_raw_data(self, indexs: List[int]):
        data = self.f.get_batch(indexs).asnumpy()

        if self.data_type == "bgr":
            data = data[:, :, :, ::-1]
        return data

    @property
    def total_frames(self):
        return len(self.f)

    @property
    def height(self):
        return self.f[0].shape[0]

    @property
    def width(self):
        return self.f[0].shape[1]

    @property
    def size(self):
        return self.f[0].shape[:2]

    @property
    def shape(self):
        return self.f[0].shape


class VideoMapClipDataset(SequentialDataset):
    def __init__(
        self,
        video_map: str,
        raw_datas,
        time_size: int,
        step: int,
        overlap: int = None,
        sample_rate: int = 1,
        drop_last: bool = False,
        max_num_per_batch: int = None,
    ) -> None:
        self.video_map = video_map
        super().__init__(
            raw_datas,
            time_size,
            step,
            overlap,
            sample_rate,
            drop_last,
            max_num_per_batch,
        )

    def generate_sample_idxs(self):
        # use video_map to generate matched sampled_index
        raise NotImplementedError