Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import math | |
| import random | |
| from typing import Iterator, Optional, Sized | |
| import numpy as np | |
| from mmdet.datasets.base_det_dataset import BaseDetDataset | |
| from mmdet.datasets.base_video_dataset import BaseVideoDataset | |
| from mmdet.registry import DATA_SAMPLERS | |
| from mmengine.dataset import ClassBalancedDataset, ConcatDataset | |
| from mmengine.dist import get_dist_info, sync_random_seed | |
| from torch.utils.data import Sampler | |
| from torch.utils.data.dataset import ConcatDataset as _ConcatDataset | |
| from ..dataset_wrappers import SeqMultiImageMixDataset | |
| class HybridVideoImgSampler(Sampler): | |
| """Sampler that providing image-level sampling outputs for video datasets | |
| in tracking tasks. It could be both used in both distributed and | |
| non-distributed environment. | |
| If using the default sampler in pytorch, the subsequent data receiver will | |
| get one video, which is not desired in some cases: | |
| (Take a non-distributed environment as an example) | |
| 1. In test mode, we want only one image is fed into the data pipeline. This | |
| is in consideration of memory usage since feeding the whole video commonly | |
| requires a large amount of memory (>=20G on MOTChallenge17 dataset), which | |
| is not available in some machines. | |
| 2. In training mode, we may want to make sure all the images in one video | |
| are randomly sampled once in one epoch and this can not be guaranteed in | |
| the default sampler in pytorch. | |
| Args: | |
| dataset (Sized): Dataset used for sampling. | |
| seed (int, optional): random seed used to shuffle the sampler. This | |
| number should be identical across all processes in the distributed | |
| group. Defaults to None. | |
| """ | |
| def __init__(self, dataset: Sized, seed: Optional[int] = None,) -> None: | |
| rank, world_size = get_dist_info() | |
| self.rank = rank | |
| self.world_size = world_size | |
| self.epoch = 0 | |
| if seed is None: | |
| self.seed = sync_random_seed() | |
| else: | |
| self.seed = seed | |
| self.dataset = dataset | |
| self.indices = [] | |
| # Hard code here to handle different dataset wrapper | |
| if isinstance(self.dataset, ConcatDataset): | |
| cat_datasets = self.dataset.datasets | |
| assert isinstance( | |
| cat_datasets[0], BaseVideoDataset | |
| ), f"expected BaseVideoDataset, but got {type(cat_datasets[0])}" | |
| self.test_mode = cat_datasets[0].test_mode | |
| assert not self.test_mode, "'ConcatDataset' should not exist in " | |
| "test mode" | |
| for dataset in cat_datasets: | |
| num_videos = len(dataset) | |
| for video_ind in range(num_videos): | |
| self.indices.extend( | |
| [ | |
| (video_ind, frame_ind) | |
| for frame_ind in range(dataset.get_len_per_video(video_ind)) | |
| ] | |
| ) | |
| elif isinstance(self.dataset, ClassBalancedDataset): | |
| ori_dataset = self.dataset.dataset | |
| assert isinstance( | |
| ori_dataset, BaseVideoDataset | |
| ), f"expected BaseVideoDataset, but got {type(ori_dataset)}" | |
| self.test_mode = ori_dataset.test_mode | |
| assert not self.test_mode, "'ClassBalancedDataset' should not " | |
| "exist in test mode" | |
| video_indices = self.dataset.repeat_indices | |
| for index in video_indices: | |
| self.indices.extend( | |
| [ | |
| (index, frame_ind) | |
| for frame_ind in range(ori_dataset.get_len_per_video(index)) | |
| ] | |
| ) | |
| elif isinstance(self.dataset, BaseVideoDataset): | |
| self.test_mode = self.dataset.test_mode | |
| num_videos = len(self.dataset) | |
| if self.test_mode: | |
| # in test mode, the images belong to the same video must be put | |
| # on the same device. | |
| if num_videos < self.world_size: | |
| raise ValueError( | |
| f"only {num_videos} videos loaded," | |
| f"but {self.world_size} gpus were given." | |
| ) | |
| chunks = np.array_split(list(range(num_videos)), self.world_size) | |
| for videos_inds in chunks: | |
| indices_chunk = [] | |
| for video_ind in videos_inds: | |
| indices_chunk.extend( | |
| [ | |
| (video_ind, frame_ind) | |
| for frame_ind in range( | |
| self.dataset.get_len_per_video(video_ind) | |
| ) | |
| ] | |
| ) | |
| self.indices.append(indices_chunk) | |
| else: | |
| for video_ind in range(num_videos): | |
| self.indices.extend( | |
| [ | |
| (video_ind, frame_ind) | |
| for frame_ind in range( | |
| self.dataset.get_len_per_video(video_ind) | |
| ) | |
| ] | |
| ) | |
| else: | |
| assert isinstance(self.dataset, SeqMultiImageMixDataset), ( | |
| "HybridVideoImgSampler is only supported in BaseVideoDataset or " | |
| "dataset wrapper: ClassBalancedDataset and ConcatDataset,SeqMultiImageMixDataset, but " | |
| f"got {type(self.dataset)} " | |
| ) | |
| self.test_mode = self.dataset.test_mode | |
| # num_videos = len(self.dataset) | |
| if self.test_mode: | |
| print("Not support test mode") | |
| raise NotImplementedError | |
| else: | |
| assert isinstance( | |
| self.dataset.dataset, _ConcatDataset | |
| ), "HybridVideoImgSampler is only supported in _ConcatDataset" | |
| cat_datasets = self.dataset.dataset.datasets | |
| for dataset in cat_datasets: | |
| self.test_mode = dataset.test_mode | |
| assert not self.test_mode, "'ConcatDataset' should not exist in " | |
| "test mode" | |
| if isinstance(dataset, BaseVideoDataset): | |
| num_videos = len(dataset) | |
| video_indices = [] | |
| for video_ind in range(num_videos): | |
| video_indices.extend( | |
| [ | |
| (video_ind, frame_ind) | |
| for frame_ind in range( | |
| dataset.get_len_per_video(video_ind) | |
| ) | |
| ] | |
| ) | |
| elif isinstance(dataset, BaseDetDataset): | |
| img_indices = [] | |
| num_imgs = len(dataset) | |
| for img_ind in range(num_imgs): | |
| img_indices.extend([img_ind]) | |
| ###### special process to make debug task easier ##### | |
| def alternate_merge(list1, list2): | |
| # Create a new list to hold the merged elements | |
| merged_list = [] | |
| # Get the length of the shorter list | |
| min_length = min(len(list1), len(list2)) | |
| # Append elements alternately from both lists | |
| for i in range(min_length): | |
| merged_list.append(list1[i]) | |
| merged_list.append(list2[i]) | |
| # Append the remaining elements from the longer list | |
| if len(list1) > len(list2): | |
| merged_list.extend(list1[min_length:]) | |
| else: | |
| merged_list.extend(list2[min_length:]) | |
| return merged_list | |
| self.indices = alternate_merge(img_indices, video_indices) | |
| if self.test_mode: | |
| self.num_samples = len(self.indices[self.rank]) | |
| self.total_size = sum([len(index_list) for index_list in self.indices]) | |
| else: | |
| self.num_samples = int(math.ceil(len(self.indices) * 1.0 / self.world_size)) | |
| self.total_size = self.num_samples * self.world_size | |
| def __iter__(self) -> Iterator: | |
| if self.test_mode: | |
| # in test mode, the order of frames can not be shuffled. | |
| indices = self.indices[self.rank] | |
| else: | |
| # deterministically shuffle based on epoch | |
| rng = random.Random(self.epoch + self.seed) | |
| indices = rng.sample(self.indices, len(self.indices)) | |
| # add extra samples to make it evenly divisible | |
| indices += indices[: (self.total_size - len(indices))] | |
| assert len(indices) == self.total_size | |
| # subsample | |
| indices = indices[self.rank : self.total_size : self.world_size] | |
| assert len(indices) == self.num_samples | |
| return iter(indices) | |
| def __len__(self): | |
| return self.num_samples | |
| def set_epoch(self, epoch): | |
| self.epoch = epoch | |