Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import copy | |
| import random | |
| from collections import defaultdict | |
| from typing import Dict, List, Optional, Union | |
| from mmdet.datasets.transforms.frame_sampling import BaseFrameSample | |
| from mmdet.registry import TRANSFORMS | |
| class MixUniformRefFrameSample(BaseFrameSample): | |
| """Uniformly sample reference frames. | |
| Args: | |
| num_ref_imgs (int): Number of reference frames to be sampled. | |
| frame_range (int | list[int]): Range of frames to be sampled around | |
| key frame. If int, the range is [-frame_range, frame_range]. | |
| Defaults to 10. | |
| filter_key_img (bool): Whether to filter the key frame when | |
| sampling reference frames. Defaults to True. | |
| collect_video_keys (list[str]): The keys of video info to be | |
| collected. | |
| """ | |
| def __init__( | |
| self, | |
| num_ref_imgs: int = 1, | |
| frame_range: Union[int, List[int]] = 10, | |
| filter_key_img: bool = True, | |
| collect_video_keys: List[str] = ["video_id", "video_length"], | |
| ): | |
| self.num_ref_imgs = num_ref_imgs | |
| self.filter_key_img = filter_key_img | |
| if isinstance(frame_range, int): | |
| assert frame_range >= 0, "frame_range can not be a negative value." | |
| frame_range = [-frame_range, frame_range] | |
| elif isinstance(frame_range, list): | |
| assert len(frame_range) == 2, "The length must be 2." | |
| assert frame_range[0] <= 0 and frame_range[1] >= 0 | |
| for i in frame_range: | |
| assert isinstance(i, int), "Each element must be int." | |
| else: | |
| raise TypeError("The type of frame_range must be int or list.") | |
| self.frame_range = frame_range | |
| super().__init__(collect_video_keys=collect_video_keys) | |
| def sampling_frames(self, video_length: int, key_frame_id: int): | |
| """Sampling frames. | |
| Args: | |
| video_length (int): The length of the video. | |
| key_frame_id (int): The key frame id. | |
| Returns: | |
| list[int]: The sampled frame indices. | |
| """ | |
| if video_length > 1: | |
| left = max(0, key_frame_id + self.frame_range[0]) | |
| right = min(key_frame_id + self.frame_range[1], video_length - 1) | |
| frame_ids = list(range(0, video_length)) | |
| valid_ids = frame_ids[left : right + 1] | |
| if self.filter_key_img and key_frame_id in valid_ids: | |
| valid_ids.remove(key_frame_id) | |
| assert ( | |
| len(valid_ids) > 0 | |
| ), "After filtering key frame, there are no valid frames" | |
| if len(valid_ids) < self.num_ref_imgs: | |
| valid_ids = valid_ids * self.num_ref_imgs | |
| ref_frame_ids = random.sample(valid_ids, self.num_ref_imgs) | |
| else: | |
| ref_frame_ids = [key_frame_id] * self.num_ref_imgs | |
| sampled_frames_ids = [key_frame_id] + ref_frame_ids | |
| sampled_frames_ids = sorted(sampled_frames_ids) | |
| key_frames_ind = sampled_frames_ids.index(key_frame_id) | |
| key_frame_flags = [False] * len(sampled_frames_ids) | |
| key_frame_flags[key_frames_ind] = True | |
| return sampled_frames_ids, key_frame_flags | |
| def transform(self, video_infos: dict) -> Optional[Dict[str, List]]: | |
| """Transform the video information. | |
| Args: | |
| video_infos (dict): The whole video information. | |
| Returns: | |
| dict: The data information of the sampled frames. | |
| """ | |
| if "video_length" not in video_infos: | |
| generated_video_info = {} | |
| key_frame_id = 0 | |
| generated_video_info["video_id"] = video_infos["img_id"] | |
| generated_video_info["video_length"] = 1 | |
| generated_video_info["key_frame_id"] = key_frame_id | |
| generated_video_info["images"] = [video_infos] | |
| (sampled_frames_ids, key_frame_flags) = self.sampling_frames( | |
| generated_video_info["video_length"], key_frame_id=key_frame_id | |
| ) | |
| results = self.prepare_data(generated_video_info, sampled_frames_ids) | |
| results["key_frame_flags"] = key_frame_flags | |
| # results['is_image'] = True | |
| else: | |
| if "key_frame_id" in video_infos: | |
| key_frame_id = video_infos["key_frame_id"] | |
| assert isinstance(video_infos["key_frame_id"], int) | |
| else: | |
| key_frame_id = random.sample( | |
| list(range(video_infos["video_length"])), 1 | |
| )[0] | |
| (sampled_frames_ids, key_frame_flags) = self.sampling_frames( | |
| video_infos["video_length"], key_frame_id=key_frame_id | |
| ) | |
| results = self.prepare_data(video_infos, sampled_frames_ids) | |
| results["key_frame_flags"] = key_frame_flags | |
| # results['is_image'] = False | |
| return results | |
| def prepare_data( | |
| self, video_infos: dict, sampled_inds: List[int] | |
| ) -> Dict[str, List]: | |
| """Prepare data for the subsequent pipeline. | |
| Args: | |
| video_infos (dict): The whole video information. | |
| sampled_inds (list[int]): The sampled frame indices. | |
| Returns: | |
| dict: The processed data information. | |
| """ | |
| frames_anns = video_infos["images"] | |
| final_data_info = defaultdict(list) | |
| # for data in frames_anns: | |
| for index in sampled_inds: | |
| data = copy.deepcopy(frames_anns[index]) | |
| # copy the info in video-level into img-level | |
| for key in self.collect_video_keys: | |
| if key == "video_length": | |
| data["ori_video_length"] = video_infos[key] | |
| data["video_length"] = len(sampled_inds) | |
| else: | |
| data[key] = video_infos[key] | |
| # Collate data_list (list of dict to dict of list) | |
| for key, value in data.items(): | |
| final_data_info[key].append(value) | |
| return final_data_info | |
| def __repr__(self) -> str: | |
| repr_str = self.__class__.__name__ | |
| repr_str += f"(num_ref_imgs={self.num_ref_imgs}, " | |
| repr_str += f"frame_range={self.frame_range}, " | |
| repr_str += f"filter_key_img={self.filter_key_img}, " | |
| repr_str += f"collect_video_keys={self.collect_video_keys})" | |
| return repr_str | |