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', input_type: str = 'video', skip_preprocess: bool = False, classmap: dict = None, topk: int = 5, ): super().__init__(framework=framework) self.model_name = model_name self.input_type = input_type 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) features = feats.to('cpu').squeeze(0).detach().numpy() outs = self.model.head(feats, input_type = self.input_type) 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