tsm
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
108 lines
4.1 KiB
108 lines
4.1 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.tsm.tsm import create_model
|
|
|
|
log = logging.getLogger()
|
|
|
|
|
|
@register(output_schema=['labels', 'scores', 'features'])
|
|
class Tsm(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:
|
|
- tsm_k400_r50_seg8
|
|
skip_preprocess (`str`):
|
|
Flag to skip video transforms.
|
|
predict (`bool`):
|
|
Flag to control whether predict labels. If False, then return video embedding.
|
|
classmap (`str=None`):
|
|
Path of the json file to match class names.
|
|
topk (`int=5`):
|
|
The number of classification labels to be returned (ordered by possibility from high to low).
|
|
"""
|
|
def __init__(self,
|
|
model_name: str = 'tsm_k400_r50_seg8',
|
|
framework: str = 'pytorch',
|
|
skip_preprocess: bool = False,
|
|
classmap: str = None,
|
|
topk: int = 5,
|
|
):
|
|
super().__init__(framework=framework)
|
|
self.model_name = model_name
|
|
self.skip_preprocess = skip_preprocess
|
|
self.topk = topk
|
|
if 'k400' in model_name:
|
|
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'
|
|
if model_name == 'tsm_k400_r50_seg8':
|
|
self.weights_path = os.path.join(str(Path(__file__).parent), 'TSM_kinetics_RGB_resnet50_shift8_blockres_avg_segment8_e50.pth')
|
|
self.model = create_model(model_name=model_name, pretrained=True, weights_path=self.weights_path, device=self.device)
|
|
self.transform_cfgs = get_configs(
|
|
side_size=224,
|
|
crop_size=224,
|
|
num_frames=8,
|
|
mean=self.model.input_mean,
|
|
std=self.model.input_std,
|
|
)
|
|
|
|
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)
|
|
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
|