| |
|
| | import base64
|
| | import os
|
| | import os.path as osp
|
| | import warnings
|
| |
|
| | import decord
|
| | import numpy as np
|
| | import torch
|
| |
|
| | from mmaction.apis import inference_recognizer, init_recognizer
|
| |
|
| | try:
|
| | from ts.torch_handler.base_handler import BaseHandler
|
| | except ImportError:
|
| | raise ImportError('`ts` is required. Try: pip install ts.')
|
| |
|
| |
|
| | class MMActionHandler(BaseHandler):
|
| |
|
| | def initialize(self, context):
|
| | properties = context.system_properties
|
| | self.map_location = 'cuda' if torch.cuda.is_available() else 'cpu'
|
| | self.device = torch.device(self.map_location + ':' +
|
| | str(properties.get('gpu_id')) if torch.cuda.
|
| | is_available() else self.map_location)
|
| | self.manifest = context.manifest
|
| |
|
| | model_dir = properties.get('model_dir')
|
| | serialized_file = self.manifest['model']['serializedFile']
|
| | checkpoint = os.path.join(model_dir, serialized_file)
|
| | self.config_file = os.path.join(model_dir, 'config.py')
|
| |
|
| | mapping_file_path = osp.join(model_dir, 'label_map.txt')
|
| | if not os.path.isfile(mapping_file_path):
|
| | warnings.warn('Missing the label_map.txt file. '
|
| | 'Inference output will not include class name.')
|
| | self.mapping = None
|
| | else:
|
| | lines = open(mapping_file_path).readlines()
|
| | self.mapping = [x.strip() for x in lines]
|
| |
|
| | self.model = init_recognizer(self.config_file, checkpoint, self.device)
|
| | self.initialized = True
|
| |
|
| | def preprocess(self, data):
|
| | videos = []
|
| |
|
| | for row in data:
|
| | video = row.get('data') or row.get('body')
|
| | if isinstance(video, str):
|
| | video = base64.b64decode(video)
|
| |
|
| | with open('/tmp/tmp.mp4', 'wb') as fout:
|
| | fout.write(video)
|
| |
|
| | video = decord.VideoReader('/tmp/tmp.mp4')
|
| | frames = [x.asnumpy() for x in video]
|
| | videos.append(np.stack(frames))
|
| |
|
| | return videos
|
| |
|
| | def inference(self, data, *args, **kwargs):
|
| | results = [inference_recognizer(self.model, item) for item in data]
|
| | return results
|
| |
|
| | def postprocess(self, data):
|
| |
|
| | output = []
|
| | for video_idx, video_result in enumerate(data):
|
| | output.append([])
|
| | assert isinstance(video_result, list)
|
| |
|
| | output[video_idx] = {
|
| | self.mapping[x[0]] if self.mapping else x[0]: float(x[1])
|
| | for x in video_result
|
| | }
|
| |
|
| | return output
|
| |
|
