logo
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions

114 lines
4.0 KiB

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
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=224,
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