# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import pickle import traceback import warnings import numpy as np import torch from decord import VideoReader, cpu from torch.utils.data import Dataset from torchvision import transforms as T from torchvision.transforms.v2 import UniformTemporalSubsample from cosmos_predict2.data.dataset_utils import Resize_Preprocess, ToTensorVideo from imaginaire.utils import log """ Test the dataset with the following command: python -m cosmos_predict2.data.dataset_video """ class Dataset(Dataset): def __init__( self, dataset_dir, num_frames=93, video_size=(704, 1280), phase='train', params=None, ): """Dataset class for loading image-text-to-video generation data. Args: dataset_dir (str): Base path to the dataset directory num_frames (int): Number of frames to load per sequence video_size (list): Target size [H,W] for video frames Returns dict with: - video: RGB frames tensor [T,C,H,W] - video_name: Dict with episode/frame metadata """ super().__init__() self.dataset_dir = dataset_dir self.sequence_length = num_frames self.is_anydata = dataset_dir.endswith('.yaml') ''' NOTE(bvh): This file is obsolete for Any4D experiments. Use custom/dataset/anydata_dataset.py (AnyDataset) or custom/dataset/vidar_dataset.py (VidarDataset) directly instead. The experiment template (custom/experiment/template.py) no longer imports from here. This file is kept only for upstream cosmos configs (cosmos_nemo_assets, groot). ''' if self.is_anydata: # NOTE(dc): automatically detect if AnyData is used in support of legacy vidar dataset. # This is a temporary solution and should be removed in the future when we fully switch to AnyData. if any(dset in open(dataset_dir).read() for dset in ["[Unified]", "[Webbed]"]): from custom.dataset.anydata_dataset import AnyDataset self.anydata_dataset = AnyDataset(dataset_dir, num_frames, phase=phase, params=params) else: from custom.dataset.vidar_dataset import VidarDataset self.anydata_dataset = VidarDataset(dataset_dir, num_frames, phase=phase, params=params) return None video_dir = os.path.join(self.dataset_dir, "videos") self.t5_dir = os.path.join(self.dataset_dir, "t5_xxl") self.video_paths = [os.path.join(video_dir, f) for f in os.listdir(video_dir) if f.endswith(".mp4")] self.video_paths = sorted(self.video_paths) # remove video paths that does not have t5_embedding self.video_paths = [ path for path in self.video_paths if os.path.exists(os.path.join(self.t5_dir, os.path.basename(path).replace(".mp4", ".pickle"))) ] log.info(f"{len(self.video_paths)} videos in total") self.wrong_number = 0 self.preprocess = T.Compose([ToTensorVideo(), Resize_Preprocess(tuple(video_size))]) def __str__(self): if self.is_anydata: return self.anydata_dataset.__str__() return f"{len(self.video_paths)} samples from {self.dataset_dir}" def __len__(self): if self.is_anydata: return len(self.anydata_dataset) return len(self.video_paths) def _load_video(self, video_path): vr = VideoReader(video_path, ctx=cpu(0), num_threads=2) frame_ids = np.linspace(0, len(vr) - 1).astype(np.int32) vr.seek(0) frame_data = vr.get_batch(frame_ids).asnumpy() try: fps = vr.get_avg_fps() except Exception: # failed to read FPS fps = 24 return frame_data, fps def _get_frames(self, video_path): frames, fps = self._load_video(video_path) frames = frames.astype(np.uint8) frames = torch.from_numpy(frames).permute(0, 3, 1, 2) # (l, c, h, w) frames = UniformTemporalSubsample(self.sequence_length)(frames) frames = self.preprocess(frames) frames = torch.clamp(frames * 255.0, 0, 255).to(torch.uint8) return frames, fps def __getitem__(self, index): if self.is_anydata: return self.anydata_dataset[index] try: data = dict() video, fps = self._get_frames(self.video_paths[index]) video = video.permute(1, 0, 2, 3) # Rearrange from [T, C, H, W] to [C, T, H, W] video_path = self.video_paths[index] t5_embedding_path = os.path.join( self.t5_dir, os.path.basename(video_path).replace(".mp4", ".pickle"), ) data["video"] = video data["video_name"] = { "video_path": video_path, "t5_embedding_path": t5_embedding_path, # "start_frame_id": '0', } # Just add these to fit the interface # t5_embedding = np.load(sample["t5_embedding_path"])[0] with open(t5_embedding_path, "rb") as f: t5_embedding = pickle.load(f)[0] # [n_tokens, 1024] n_tokens = t5_embedding.shape[0] if n_tokens < 512: t5_embedding = np.concatenate( [t5_embedding, np.zeros((512 - n_tokens, 1024), dtype=np.float32)], axis=0 ) t5_text_mask = torch.zeros(512, dtype=torch.int64) t5_text_mask[:n_tokens] = 1 data["t5_text_embeddings"] = torch.from_numpy(t5_embedding) data["t5_text_mask"] = t5_text_mask data["fps"] = fps data["image_size"] = torch.tensor([704, 1280, 704, 1280]) data["num_frames"] = self.sequence_length data["padding_mask"] = torch.zeros(1, 704, 1280) return data except Exception: warnings.warn( f"Invalid data encountered: {self.video_paths[index]}. Skipped " f"(by randomly sampling another sample in the same dataset)." ) warnings.warn("FULL TRACEBACK:") warnings.warn(traceback.format_exc()) self.wrong_number += 1 log.info(self.wrong_number, rank0_only=False) return self[np.random.randint(len(self.samples))] if __name__ == "__main__": dataset = Dataset( dataset_dir="datasets/benchmark_train/gr1", num_frames=93, video_size=[480, 832], ) indices = [0, 13, -1] for idx in indices: data = dataset[idx] log.info( ( f"{idx=} " f"{data['video'].sum()=}\n" f"{data['video'].shape=}\n" f"{data['video_name']=}\n" f"{data['t5_text_embeddings'].shape=}\n" "---" ) )