stgcn_runner.py

Go to the documentation of this file.

import torch
import torch.nn as nn
from smaction.builder.model_builder import build_action_model
 
from smaction.utils.transforms import *
import smaction.utils.action_utils as utils 
 
#skeleton, image, fusion 버전으로 분리 필요
class STGCNRunner(nn.Module):
    def __init__(self, backbone, head, device='cuda:0'):
        """
        input:
            model_cfg : 
            model_checkpoint : 
            device : 
        """
        super().__init__()
        self.backbone = build_action_model(backbone)
        self.head = build_action_model(head)
        self.device = device
        
 
    def make_action_input(self, pose_list, img_shape):
        pass
 
    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_input(pose_results, img_shape)
 
        return self.inference(data)
 
    def inference(self, x):
        feat = self.backbone(x)
        scores = self.head(feat)
        action_labels = scores.argmax()
        return {'action_labels':action_labels, 'scores':scores}
 
    def forward(self, x):
        return self.inference(x['keypoints'])