add omnivore

Signed-off-by: gexy5 <xinyu.ge@zilliz.com>

omnivore.py

@@ -0,0 +1,118 @@
+import logging
+import os
+import json
+from pathlib import Path
+from typing import List
+
+import torch
+import numpy
+
+from towhee import register
+from towhee.operator.base import NNOperator
+from towhee.types.video_frame import VideoFrame
+from towhee.models.utils.video_transforms import get_configs, transform_video
+from towhee.models.omnivore.omnivore import create_model
+
+log = logging.getLogger()
+
+@register(output_schema=['labels', 'scores', 'features'])
+class Omnivore(NNOperator):
+    """
+    Generate a list of class labels given a video input data.
+    Default labels are from [Kinetics400 Dataset](https://deepmind.com/research/open-source/kinetics).
+    Args:
+        model_name (`str`):
+            Supported model names:
+            - omnivore_swinT
+            - omnivore_swinS
+            - omnivore_swinB
+            - omnivore_swinB_in21k
+            - omnivore_swinL_in21k
+            - omnivore_swinB_epic
+        skip_preprocess (`str`):
+            Flag to skip video transforms.
+        predict (`bool`):
+            Flag to control whether predict labels. If False, then return video embedding.
+        classmap (`dict=None`):
+            The dictionary maps classes to integers.
+        topk (`int=5`):
+            The number of classification labels to be returned (ordered by possibility from high to low).
+    """
+    def __init__(self,
+                 model_name: str = 'omnivore_swinT',
+                 framework: str = 'pytorch',
+                 skip_preprocess: bool = False,
+                 classmap: dict = None,
+                 topk: int = 5,
+                 ):
+        super().__init__(framework=framework)
+        self.model_name = model_name
+        self.skip_preprocess = skip_preprocess
+        self.topk = topk
+        self.dataset_name = 'kinetics_400'
+        if classmap is None:
+            class_file = os.path.join(str(Path(__file__).parent), self.dataset_name+'.json')
+            with open(class_file, "r") as f:
+                kinetics_classes = json.load(f)
+            self.classmap = {}
+            for k, v in kinetics_classes.items():
+                self.classmap[v] = str(k).replace('"', '')
+        else:
+            self.classmap = classmap
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model = create_model(model_name=model_name, pretrained=True, device=self.device)
+        self.input_mean=[0.485, 0.456, 0.406]
+        self.input_std=[0.229, 0.224, 0.225]
+        self.transform_cfgs = get_configs(
+                side_size=256,
+                crop_size=224,
+                num_frames=24,
+                mean=self.input_mean,
+                std=self.input_std,
+                )
+        self.model.eval()
+
+    def __call__(self, video: List[VideoFrame]):
+        """
+        Args:
+            video (`List[VideoFrame]`):
+                Video path in string.
+
+        Returns:
+            (labels, scores)
+                A tuple of lists (labels, scores).
+            OR emb
+                Video embedding.
+        """
+        # Convert list of towhee.types.Image to numpy.ndarray in float32
+        video = numpy.stack([img.astype(numpy.float32)/255. for img in video], axis=0)
+        assert len(video.shape) == 4
+        video = video.transpose(3, 0, 1, 2)  # twhc -> ctwh
+
+        # Transform video data given configs
+        if self.skip_preprocess:
+            self.transform_cfgs.update(num_frames=None)
+
+        data = transform_video(
+                    video=video,
+                    **self.transform_cfgs
+                    )
+        inputs = data.to(self.device)[None, ...]
+
+        feats = self.model.forward_features(inputs)
+        if self.model.reshape:
+            if self.model.is_shift and self.model.temporal_pool:
+                base_out = feats.view((-1, self.model.num_segments // 2) + feats.size()[1:])
+            else:
+                base_out = feats.view((-1, self.model.num_segments) + feats.size()[1:])
+            output = self.model.consensus(base_out)
+        features = output.to('cpu').squeeze(0).detach().numpy()
+
+        outs = self.model.head(feats)
+        post_act = torch.nn.Softmax(dim=1)
+        preds = post_act(outs)
+        pred_scores, pred_classes = preds.topk(k=self.topk)
+        labels = [self.classmap[int(i)] for i in pred_classes[0]]
+        scores = [round(float(x), 5) for x in pred_scores[0]]
+
+        return labels, scores, features