File size: 1,706 Bytes
2b07837 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | from abc import (
ABC,
abstractmethod,
)
from typing import Union
import numpy as np
import torch
class PoseInferenceBase(ABC):
@abstractmethod
def preprocess(
self, img: Union[np.ndarray, str], *args, **kwargs
) -> Union[torch.Tensor, np.ndarray]:
"""
Preprocesses the input image before inference.
Args:
img (Union[np.ndarray, str]): The input image as a NumPy array or a path to the image file.
Returns:
Union[torch.Tensor, np.ndarray]: The preprocessed image tensor or array ready for inference.
"""
raise NotImplementedError
@abstractmethod
def process(
self, img: Union[np.ndarray, str], *args, **kwargs
) -> Union[torch.Tensor, np.ndarray]:
"""
Performs inference on the input image.
Args:
img (Union[np.ndarray, str]): The input image as a NumPy array or a path to the image file.
Returns:
Union[torch.Tensor, np.ndarray]: The output of the inference process.
"""
raise NotImplementedError
@abstractmethod
def postprocess(
self, oriImg: np.ndarray, heatmaps: np.ndarray, pafs: np.ndarray, *args, **kwargs
) -> Union[torch.Tensor, np.ndarray]:
"""
Postprocesses the inference results.
Args:
oriImg (np.ndarray): The original input image.
heatmaps (np.ndarray): The heatmaps generated by the inference.
pafs (np.ndarray): The Part Affinity Fields (PAFs) generated by the inference.
Returns:
Union[torch.Tensor, np.ndarray]: The postprocessed results.
"""
raise NotImplementedError
|