95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
|
import os.path
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Callable, Optional, Union
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torchvision.datasets.utils import download_url, verify_str_arg
|
||
|
|
from torchvision.datasets.vision import VisionDataset
|
||
|
|
|
||
|
|
|
||
|
|
class MovingMNIST(VisionDataset):
|
||
|
|
"""`MovingMNIST <http://www.cs.toronto.edu/~nitish/unsupervised_video/>`_ Dataset.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
root (str or ``pathlib.Path``): Root directory of dataset where ``MovingMNIST/mnist_test_seq.npy`` exists.
|
||
|
|
split (string, optional): The dataset split, supports ``None`` (default), ``"train"`` and ``"test"``.
|
||
|
|
If ``split=None``, the full data is returned.
|
||
|
|
split_ratio (int, optional): The split ratio of number of frames. If ``split="train"``, the first split
|
||
|
|
frames ``data[:, :split_ratio]`` is returned. If ``split="test"``, the last split frames ``data[:, split_ratio:]``
|
||
|
|
is returned. If ``split=None``, this parameter is ignored and the all frames data is returned.
|
||
|
|
transform (callable, optional): A function/transform that takes in a torch Tensor
|
||
|
|
and returns a transformed version. E.g, ``transforms.RandomCrop``
|
||
|
|
download (bool, optional): If true, downloads the dataset from the internet and
|
||
|
|
puts it in root directory. If dataset is already downloaded, it is not
|
||
|
|
downloaded again.
|
||
|
|
"""
|
||
|
|
|
||
|
|
_URL = "http://www.cs.toronto.edu/~nitish/unsupervised_video/mnist_test_seq.npy"
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
root: Union[str, Path],
|
||
|
|
split: Optional[str] = None,
|
||
|
|
split_ratio: int = 10,
|
||
|
|
download: bool = False,
|
||
|
|
transform: Optional[Callable] = None,
|
||
|
|
) -> None:
|
||
|
|
super().__init__(root, transform=transform)
|
||
|
|
|
||
|
|
self._base_folder = os.path.join(self.root, self.__class__.__name__)
|
||
|
|
self._filename = self._URL.split("/")[-1]
|
||
|
|
|
||
|
|
if split is not None:
|
||
|
|
verify_str_arg(split, "split", ("train", "test"))
|
||
|
|
self.split = split
|
||
|
|
|
||
|
|
if not isinstance(split_ratio, int):
|
||
|
|
raise TypeError(f"`split_ratio` should be an integer, but got {type(split_ratio)}")
|
||
|
|
elif not (1 <= split_ratio <= 19):
|
||
|
|
raise ValueError(f"`split_ratio` should be `1 <= split_ratio <= 19`, but got {split_ratio} instead.")
|
||
|
|
self.split_ratio = split_ratio
|
||
|
|
|
||
|
|
if download:
|
||
|
|
self.download()
|
||
|
|
|
||
|
|
if not self._check_exists():
|
||
|
|
raise RuntimeError("Dataset not found. You can use download=True to download it.")
|
||
|
|
|
||
|
|
data = torch.from_numpy(np.load(os.path.join(self._base_folder, self._filename)))
|
||
|
|
if self.split == "train":
|
||
|
|
data = data[: self.split_ratio]
|
||
|
|
elif self.split == "test":
|
||
|
|
data = data[self.split_ratio :]
|
||
|
|
self.data = data.transpose(0, 1).unsqueeze(2).contiguous()
|
||
|
|
|
||
|
|
def __getitem__(self, idx: int) -> torch.Tensor:
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
index (int): Index
|
||
|
|
Returns:
|
||
|
|
torch.Tensor: Video frames (torch Tensor[T, C, H, W]). The `T` is the number of frames.
|
||
|
|
"""
|
||
|
|
data = self.data[idx]
|
||
|
|
if self.transform is not None:
|
||
|
|
data = self.transform(data)
|
||
|
|
|
||
|
|
return data
|
||
|
|
|
||
|
|
def __len__(self) -> int:
|
||
|
|
return len(self.data)
|
||
|
|
|
||
|
|
def _check_exists(self) -> bool:
|
||
|
|
return os.path.exists(os.path.join(self._base_folder, self._filename))
|
||
|
|
|
||
|
|
def download(self) -> None:
|
||
|
|
if self._check_exists():
|
||
|
|
return
|
||
|
|
|
||
|
|
download_url(
|
||
|
|
url=self._URL,
|
||
|
|
root=self._base_folder,
|
||
|
|
filename=self._filename,
|
||
|
|
md5="be083ec986bfe91a449d63653c411eb2",
|
||
|
|
)
|