File size: 3,371 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | from typing import Any, Dict, Optional
import io
import tarfile
import torch
import numpy as np
from omegaconf import DictConfig
from tqdm import tqdm
from torchvision.transforms import InterpolationMode
from torchvision import transforms
from .base_video import (
BaseVideoDataset,
BaseSimpleVideoDataset,
BaseAdvancedVideoDataset,
)
class MinecraftBaseVideoDataset(BaseVideoDataset):
_ALL_SPLITS = ["training", "validation"]
def download_dataset(self):
from internetarchive import download
part_suffixes = [
"aa",
"ab",
"ac",
"ad",
"ae",
"af",
"ag",
"ah",
"ai",
"aj",
"ak",
]
for part_suffix in part_suffixes:
identifier = f"minecraft_marsh_dataset_{part_suffix}"
file_name = f"minecraft.tar.part{part_suffix}"
download(identifier, file_name, destdir=self.save_dir, verbose=True)
combined_bytes = io.BytesIO()
for part_suffix in part_suffixes:
identifier = f"minecraft_marsh_dataset_{part_suffix}"
file_name = f"minecraft.tar.part{part_suffix}"
part_file = self.save_dir / identifier / file_name
with open(part_file, "rb") as part:
combined_bytes.write(part.read())
combined_bytes.seek(0)
with tarfile.open(fileobj=combined_bytes, mode="r") as combined_archive:
combined_archive.extractall(self.save_dir)
(self.save_dir / "minecraft/test").rename(self.save_dir / "validation")
(self.save_dir / "minecraft/train").rename(self.save_dir / "training")
(self.save_dir / "minecraft").rmdir()
for part_suffix in part_suffixes:
identifier = f"minecraft_marsh_dataset_{part_suffix}"
file_name = f"minecraft.tar.part{part_suffix}"
part_file = self.save_dir / identifier / file_name
part_file.rmdir()
def video_length(self, video_metadata: Dict[str, Any]) -> int:
return 300
def build_transform(self):
return transforms.Resize(
(self.resolution, self.resolution),
interpolation=InterpolationMode.NEAREST_EXACT,
antialias=True,
)
class MinecraftSimpleVideoDataset(MinecraftBaseVideoDataset, BaseSimpleVideoDataset):
"""
Minecraft simple video dataset
"""
def __init__(self, cfg: DictConfig, split: str = "training"):
if split == "test":
split = "validation"
BaseSimpleVideoDataset.__init__(self, cfg, split)
class MinecraftAdvancedVideoDataset(
MinecraftBaseVideoDataset, BaseAdvancedVideoDataset
):
"""
Minecraft advanced video dataset
"""
def __init__(
self,
cfg: DictConfig,
split: str = "training",
current_epoch: Optional[int] = None,
):
if split == "test":
split = "validation"
BaseAdvancedVideoDataset.__init__(self, cfg, split, current_epoch)
def load_cond(
self, video_metadata: Dict[str, Any], start_frame: int, end_frame: int
) -> torch.Tensor:
path = video_metadata["video_paths"].with_suffix(".npz")
actions = np.load(path)["actions"][start_frame:end_frame]
return torch.from_numpy(np.eye(4)[actions]).float()
|