Add ten/five crop augmentation

configs/i3d.yml

@@ -6,9 +6,13 @@ streams: null # Streams to use for feature extraction (e.g. 'rgb' or 'flow'). Bo
 flow_type: 'pwc' # Flow to use in I3D. 'pwc' (PWCNet) is faster while 'raft' (RAFT) is more accurate.
 extraction_fps: null # For original video fps, leave as "null" (None)
 
+# Feature Augumentation
+augment: null # Augmentation to use for feature extraction. Can be ['ten_crop', 'five_crop']
+
 # Extraction Parameters
 device: 'cuda:0'  # device as in `torch`, can be 'cpu'
 on_extraction: 'print'  # what to do once the features are extracted. Can be ['print', 'save_numpy', 'save_pickle']
+save_option: null # choose what to save. If you only want rgb features, choose 'rgb_only' else null
 output_path: './output' # where to store results if saved
 tmp_path: './tmp' # folder to store the temporary files used for extraction (frames or aud files)
 keep_tmp_files: false # to keep temp files after feature extraction.

docs/models/i3d.md

@@ -44,6 +44,7 @@ You may test it yourself by providing `--show_pred` flag.
 | `video_paths`                           | `null`                                | A list of videos for feature extraction. E.g. `"[./sample/v_ZNVhz7ctTq0.mp4, ./sample/v_GGSY1Qvo990.mp4]"` or just one path `"./sample/v_GGSY1Qvo990.mp4"`.                      |
 | `file_with_video_paths`                 | `null`                                | A path to a text file with video paths (one path per line). Hint: given a folder `./dataset` with `.mp4` files one could use: `find ./dataset -name "*mp4" > ./video_paths.txt`. |
 | `on_extraction`                         | `print`                               | If `print`, the features are printed to the terminal. If `save_numpy` or `save_pickle`, the features are saved to either `.npy` file or `.pkl`.                                  |
+| `save_option`                         | `null`                               | If `rgb_only`, only the rgb features will be saved.                                  |
 | `output_path`                           | `"./output"`                          | A path to a folder for storing the extracted features (if `on_extraction` is either `save_numpy` or `save_pickle`).                                                              |
 | `keep_tmp_files`                        | `false`                               | If `true`, the reencoded videos will be kept in `tmp_path`.                                                                                                                      |
 | `tmp_path`                              | `"./tmp"`                             | A path to a folder for storing temporal files (e.g. reencoded videos).                                                                                                           |

models/_base/base_extractor.py

@@ -18,9 +18,11 @@ def __init__(self,
                  output_path: str,
                  keep_tmp_files: bool,
                  device: str,
+                 save_option=None,
                  ) -> None:
         self.feature_type = feature_type
         self.on_extraction = on_extraction
+        self.save_option = save_option
         self.tmp_path = tmp_path
         self.output_path = output_path
         self.keep_tmp_files = keep_tmp_files
@@ -76,6 +78,11 @@ def action_on_extraction(
             return
 
         for key, value in feats_dict.items():
+            if self.save_option == 'rgb_only':
+                if key != 'rgb':
+                    continue
+                else:
+                    key = None
             if self.on_extraction == 'print':
                 print(key)
                 print(value)
@@ -84,11 +91,18 @@ def action_on_extraction(
             elif self.on_extraction in ['save_numpy', 'save_pickle']:
                 # make dir if doesn't exist
                 os.makedirs(self.output_path, exist_ok=True)
-                fpath = make_path(self.output_path, video_path, key, action2ext[self.on_extraction])
-                if key != 'fps' and len(value) == 0:
-                    print(f'Warning: the value is empty for {key} @ {fpath}')
                 # save the info behind the each key
-                action2savefn[self.on_extraction](fpath, value)
+                if len(value.shape) < 3:
+                    fpath = make_path(self.output_path, video_path, key, action2ext[self.on_extraction])
+                    if key != 'fps' and len(value) == 0:
+                        print(f'Warning: the value is empty for {key} @ {fpath}')
+                    action2savefn[self.on_extraction](fpath, value)
+                else:
+                    for i in range(value.shape[0]):
+                        fpath = make_path(self.output_path, video_path, key, action2ext[self.on_extraction], i)
+                        if key != 'fps' and len(value) == 0:
+                            print(f'Warning: the value is empty for {key} @ {fpath}')
+                        action2savefn[self.on_extraction](fpath, value[i, :])
             else:
                 raise NotImplementedError(f'on_extraction: {self.on_extraction} is not implemented')
 

models/i3d/extract_i3d.py

@@ -24,6 +24,7 @@ def __init__(self, args) -> None:
         super().__init__(
             feature_type=args.feature_type,
             on_extraction=args.on_extraction,
+            save_option=args.save_option,
             tmp_path=args.tmp_path,
             output_path=args.output_path,
             keep_tmp_files=args.keep_tmp_files,
@@ -38,15 +39,24 @@ def __init__(self, args) -> None:
         self.extraction_fps = args.extraction_fps
         self.step_size = 64 if args.step_size is None else args.step_size
         self.stack_size = 64 if args.stack_size is None else args.stack_size
+        self.aug_type = args.augment
         self.resize_transforms = torchvision.transforms.Compose([
             torchvision.transforms.ToPILImage(),
             ResizeImproved(self.min_side_size),
             PILToTensor(),
             ToFloat(),
         ])
+        if self.aug_type is None:
+            aug_transform = TensorCenterCrop(self.central_crop_size)
+        elif self.aug_type == 'five_crop':
+            aug_transform = torchvision.transforms.FiveCrop(self.central_crop_size)
+            self.num_crop = 5
+        elif self.aug_type == 'ten_crop':
+            aug_transform = torchvision.transforms.TenCrop(self.central_crop_size)
+            self.num_crop = 10
         self.i3d_transforms = {
             'rgb': torchvision.transforms.Compose([
-                TensorCenterCrop(self.central_crop_size),
+                aug_transform,
                 ScaleTo1_1(),
                 PermuteAndUnsqueeze()
             ]),
@@ -82,8 +92,12 @@ def extract(self, video_path: str) -> Dict[str, np.ndarray]:
         # timestamp when the last frame in the stack begins (when the old frame of the last pair ends)
         timestamps_ms = []
         rgb_stack = []
-        feats_dict = {stream: [] for stream in self.streams}
-
+
+        if self.aug_type is not None:
+            feats_dict = {stream: [[] for _ in range(self.num_crop)] for stream in self.streams}
+        else:
+            feats_dict = {stream: [] for stream in self.streams}
+
         # sometimes when the target fps is 1 or 2, the first frame of the reencoded video is missing
         # and cap.read returns None but the rest of the frames are ok. timestep is 0.0 for the 2nd frame in
         # this case
@@ -113,7 +127,11 @@ def extract(self, video_path: str) -> Dict[str, np.ndarray]:
                 if len(rgb_stack) - 1 == self.stack_size:
                     batch_feats_dict = self.run_on_a_stack(rgb_stack, stack_counter, padder)
                     for stream in self.streams:
-                        feats_dict[stream].extend(batch_feats_dict[stream].tolist())
+                        if isinstance(batch_feats_dict[stream], tuple):
+                            for i in range(len(batch_feats_dict[stream])):
+                                feats_dict[stream][i].extend(batch_feats_dict[stream][i].tolist())
+                        else:
+                            feats_dict[stream].extend(batch_feats_dict[stream].tolist())
                     # leaving the elements if step_size < stack_size so they will not be loaded again
                     # if step_size == stack_size one element is left because the flow between the last element
                     # in the prev list and the first element in the current list
@@ -161,8 +179,11 @@ def run_on_a_stack(self, rgb_stack, stack_counter, padder=None) -> Dict[str, tor
                 raise NotImplementedError
             # apply transforms depending on the stream (flow or rgb)
             stream_slice = self.i3d_transforms[stream](stream_slice)
-            # extract features for a stream
-            batch_feats_dict[stream] = models[stream](stream_slice, features=True)  # (B, 1024)
+            if isinstance(stream_slice, tuple):
+                # extract features for a stream
+                batch_feats_dict[stream] = tuple([models[stream](stream_crop, features=True) for stream_crop in stream_slice])
+            else:    
+                batch_feats_dict[stream] = models[stream](stream_slice, features=True)  # (B, 1024)
             # add features to the output dict
             self.maybe_show_pred(stream_slice, self.name2module['model'][stream], stack_counter)
 

models/transforms.py

@@ -145,13 +145,17 @@ def __call__(self, tensor: torch.FloatTensor) -> torch.FloatTensor:
 
 class ScaleTo1_1(object):
 
-    def __call__(self, tensor: torch.FloatTensor) -> torch.FloatTensor:
+    def __call__(self, tensor):
+        if isinstance(tensor, tuple):
+            return tuple([(2 * t / 255) - 1 for t in tensor])
         return (2 * tensor / 255) - 1
 
 
 class PermuteAndUnsqueeze(object):
 
-    def __call__(self, tensor: torch.FloatTensor) -> torch.FloatTensor:
+    def __call__(self, tensor):
+        if isinstance(tensor, tuple):
+            return tuple([t.permute(1, 0, 2, 3).unsqueeze(0) for t in tensor])
         return tensor.permute(1, 0, 2, 3).unsqueeze(0)
 
 

utils/utils.py

@@ -50,9 +50,18 @@ def show_predictions_on_dataset(logits: torch.FloatTensor, dataset: Union[str, L
             print(f'{logit:8.3f} | {smax:.3f} | {cls}')
         print()
 
-def make_path(output_root, video_path, output_key, ext):
+def make_path(output_root, video_path, output_key, ext, idx=None):
     # extract file name and change the extention
-    fname = f'{Path(video_path).stem}_{output_key}{ext}'
+    if idx is not None:
+        if output_key is not None:
+            fname = f'{Path(video_path).stem}_{output_key}_{idx}{ext}'
+        else:
+            fname = f'{Path(video_path).stem}_{idx}{ext}'
+    else:
+        if output_key is not None:
+            fname = f'{Path(video_path).stem}_{output_key}{ext}'
+        else:
+            fname = f'{Path(video_path).stem}_{idx}{ext}'
     # construct the paths to save the features
     return os.path.join(output_root, fname)
 
@@ -131,8 +140,7 @@ def form_list_from_user_input(
         to_shuffle: bool = True,
     ) -> list:
     '''User specifies either list of videos in the cmd or a path to a file with video paths. This function
-       transforms the user input into a list of paths. Files are expected to be formatted with a single
-       video-path in each line.
+       transforms the user input into a list of paths.
 
     Args:
         video_paths (Union[str, ListConfig, None], optional): a list of video paths. Defaults to None.