timesformer
copied
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
99 lines
3.4 KiB
99 lines
3.4 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 transform_video
|
|
from towhee.models import timesformer
|
|
|
|
log = logging.getLogger()
|
|
|
|
|
|
@register(output_schema=['labels', 'scores', 'features'])
|
|
class Timesformer(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:
|
|
- timesformer_k400_8x224
|
|
skip_preprocess (`str`):
|
|
Flag to skip video transforms.
|
|
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 = 'timesformer_k400_8x224',
|
|
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
|
|
if classmap is None:
|
|
class_file = os.path.join(str(Path(__file__).parent), 'kinetics_400.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 = timesformer.create_model(model_name=model_name, pretrained=True, device=self.device)
|
|
self.cfg = timesformer.get_configs(model_name=model_name)
|
|
self.cfg.update(side_size=self.cfg['img_size'], crop_size=self.cfg['img_size'])
|
|
|
|
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.cfg.update(num_frames=None)
|
|
|
|
data = transform_video(
|
|
video=video,
|
|
**self.cfg
|
|
)
|
|
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)
|
|
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
|