mmaction_pytorch_runner.py

Go to the documentation of this file.

from mmengine.dataset import pseudo_collate
from .mmaction_base_runner import MMActionBaseRunner
 
from smaction.utils.transforms import *
import smaction.utils.action_utils as utils 
import torch
 
#skeleton, image, fusion 버전으로 분리 필요
class MMActionPyTorchRunner(MMActionBaseRunner):
    def __init__(self, model_cfg, model_checkpoint, device='cuda', input_type='PoseC3D', k=60, img_shape=None, use_valid_check=False, use_dummy_pose=False):
        """
        input:
            model_cfg : 
            model_checkpoint : 
            device : 
        """
        super().__init__(model_cfg=model_cfg, model_checkpoint=model_checkpoint, device=device)
        self.input_func = get_action_input_func(input_type)
        self.img_shape = img_shape
        self.k = k #self.model.cfg.test_pipeline[0]['clip_len']
        self.use_valid_check = use_valid_check
        self.use_dummy_pose = use_dummy_pose
    
 
    def make_action_input(self, pose_list, img_shape):
        if img_shape is None:
            img_shape = self.img_shape
 
        data = self.input_func(pose_list, img_shape)
 
        data = self.test_pipeline(data)
        data = pseudo_collate([data])
 
        return data
 
    def run_recognizer(self, pose_results, img_shape=None):
        if self.use_valid_check and not utils.check_valid(pose_results, self.k):
            return None
 
        if self.use_dummy_pose:
            pose_results = utils.insert_dummy_pose(pose_results)
 
        data = self.make_action_inputmake_action_input(pose_results, img_shape)
 
        with torch.no_grad():
            result = self.model.test_step(data)[0]
 
        action_labels = result.pred_labels.item.cpu().numpy()
        scores = result.pred_scores.item.cpu().numpy()
 
        return {'action_labels':action_labels, 'scores':scores}