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