# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import tempfile import unittest import unittest.mock from contextlib import ExitStack from pathlib import Path import torch from parameterized import parameterized from pytorchvideo.data.dataset_manifest_utils import VideoClipInfo, VideoDatasetType from pytorchvideo.data.epic_kitchen import ActionData, EpicKitchenDataset from pytorchvideo.data.utils import save_dataclass_objs_to_headered_csv from utils import ( get_encoded_video_infos, get_flat_video_frames, MOCK_VIDEO_IDS, MOCK_VIDEO_INFOS, ) class TestEpicKitchenDataset(unittest.TestCase): ACTIONS_DATAS = { MOCK_VIDEO_IDS[0]: [ ActionData( "P01", "P01_01", "turn on light", "00:00:04.00", "00:00:06.00", 262, 370, "turn-on", 12, "light", 113, "['light']", "[113]", ), ActionData( "P01", "P01_01", "close door", "00:00:05.00", "00:00:07.00", 418, 569, "close", 3, "door", 8, "['door']", "[8]", ), ActionData( "P01", "P01_01", "close fridge", "00:01:1.91", "01:00:5.33", 1314, 1399, "close", 3, "fridge", 10, "['fridge']", "[10]", ), ], MOCK_VIDEO_IDS[1]: [ ActionData( "P02", "P02_002", "turn on light", "00:00:04.00", "00:00:06.00", 262, 370, "turn-on", 12, "light", 113, "['light']", "[113]", ) ], MOCK_VIDEO_IDS[2]: [ ActionData( "P02", "P02_005", "turn on light", "00:00:04.00", "00:00:06.00", 262, 370, "turn-on", 12, "light", 113, "['light']", "[113]", ) ], MOCK_VIDEO_IDS[3]: [ ActionData( "P07", "P07_002", "turn on light", "00:00:04.00", "00:00:06.00", 262, 370, "turn-on", 12, "light", 113, "['light']", "[113]", ) ], } def test_ActionData(self): action = ActionData( # This is a key-mapping as the underlying epic-kitchen # annotation files are of these string columns **{ "participant_id": "P07", "video_id": "P07_002", "narration": "turn on light", "start_timestamp": "00:00:04.00", "stop_timestamp": "00:00:06.50", "start_frame": "262", "stop_frame": "370", "verb": "turn-on", "verb_class": "12", "noun": "light", "noun_class": "113", "all_nouns": "['light', 'finger', 'wall']", "all_noun_classes": "[113, 1232, 1]", } ) self.assertEqual(action.video_id, "P07_002") self.assertEqual(action.start_time, 4.0) self.assertEqual(action.stop_time, 6.5) self.assertEqual(action.verb_class, 12) self.assertEqual(action.noun_class, 113) self.assertEqual(action.all_nouns, ["light", "finger", "wall"]) self.assertEqual(action.all_noun_classes, [113, 1232, 1]) @parameterized.expand([(VideoDatasetType.Frame,), (VideoDatasetType.EncodedVideo,)]) def test__len__(self, dataset_type): with tempfile.TemporaryDirectory(prefix=f"{TestEpicKitchenDataset}") as tempdir: tempdir = Path(tempdir) video_info_file = tempdir / "test_video_info.csv" save_dataclass_objs_to_headered_csv( list(MOCK_VIDEO_INFOS.values()), video_info_file ) action_file = tempdir / "action_video_info.csv" actions = [] for action_list in self.ACTIONS_DATAS.values(): for action in action_list: actions.append(action) save_dataclass_objs_to_headered_csv(actions, action_file) video_data_manifest_file_path = ( tempdir / "video_data_manifest_file_path.json" ) with ExitStack() as stack: if dataset_type == VideoDatasetType.Frame: video_data_dict = get_flat_video_frames(tempdir, "jpg") elif dataset_type == VideoDatasetType.EncodedVideo: video_data_dict = get_encoded_video_infos(tempdir, stack) save_dataclass_objs_to_headered_csv( list(video_data_dict.values()), video_data_manifest_file_path ) dataset = EpicKitchenDataset( video_info_file_path=str(video_info_file), actions_file_path=str(action_file), clip_sampler=lambda x, y: [ VideoClipInfo(str(i), i * 2.0, i * 2.0 + 0.9) for i in range(0, 7) ], video_data_manifest_file_path=str(video_data_manifest_file_path), dataset_type=dataset_type, ) self.assertEqual(len(dataset), 7) @parameterized.expand([(VideoDatasetType.Frame,), (VideoDatasetType.EncodedVideo,)]) def test__getitem__(self, dataset_type): with tempfile.TemporaryDirectory(prefix=f"{TestEpicKitchenDataset}") as tempdir: tempdir = Path(tempdir) video_info_file = tempdir / "test_video_info.csv" save_dataclass_objs_to_headered_csv( list(MOCK_VIDEO_INFOS.values()), video_info_file ) action_file = tempdir / "action_video_info.csv" actions = [] for action_list in self.ACTIONS_DATAS.values(): for action in action_list: actions.append(action) save_dataclass_objs_to_headered_csv(actions, action_file) video_data_manifest_file_path = ( tempdir / "video_data_manifest_file_path.json" ) with ExitStack() as stack: if dataset_type == VideoDatasetType.Frame: video_data_dict = get_flat_video_frames(tempdir, "jpg") elif dataset_type == VideoDatasetType.EncodedVideo: video_data_dict = get_encoded_video_infos(tempdir, stack) save_dataclass_objs_to_headered_csv( list(video_data_dict.values()), video_data_manifest_file_path ) video_ids = list(self.ACTIONS_DATAS) dataset = EpicKitchenDataset( video_info_file_path=str(video_info_file), actions_file_path=str(action_file), clip_sampler=lambda x, y: [ VideoClipInfo(video_ids[i // 2], i * 2.0, i * 2.0 + 0.9) for i in range(0, 7) ], video_data_manifest_file_path=str(video_data_manifest_file_path), dataset_type=dataset_type, ) get_clip_string = ( "pytorchvideo.data.frame_video.FrameVideo.get_clip" if dataset_type == VideoDatasetType.Frame else "pytorchvideo.data.encoded_video.EncodedVideo.get_clip" ) with unittest.mock.patch( get_clip_string, return_value=({"video": torch.rand(3, 5, 10, 20), "audio": []}), ) as _: clip_1 = dataset.__getitem__(1) for i, a in enumerate(clip_1["actions"]): self.assertEqual(a, self.ACTIONS_DATAS[video_ids[0]][i]) self.assertEqual(clip_1["start_time"], 2.0) self.assertEqual(clip_1["stop_time"], 2.9) self.assertEqual(clip_1["video_id"], MOCK_VIDEO_IDS[0]) clip_2 = dataset.__getitem__(2) for i, a in enumerate(clip_2["actions"]): self.assertEqual(a, self.ACTIONS_DATAS[video_ids[1]][i]) self.assertEqual(clip_2["start_time"], 4.0) self.assertEqual(clip_2["stop_time"], 4.9) self.assertEqual(clip_2["video_id"], MOCK_VIDEO_IDS[1])