| from typing import ( |
| Optional, |
| Tuple, |
| Union, |
| ) |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
| from loguru import logger |
|
|
| from src.models.networks import OpenPoseNet |
|
|
| from .base import PoseInferenceBase |
| from .decode import decode_pose |
|
|
|
|
| class PoseInference(PoseInferenceBase): |
| def __init__(self, model_weight_path: str, device: Optional[str] = None): |
| super().__init__() |
| self.net = OpenPoseNet() |
| if not device: |
| self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
| else: |
| self.device = device |
|
|
| net_weights = torch.load(model_weight_path, map_location=self.device) |
| keys = list(net_weights.keys()) |
|
|
| weights_load = {} |
| for i in range(len(keys)): |
| weights_load[list(self.net.state_dict().keys())[i]] = net_weights[list(keys)[i]] |
|
|
| state = self.net.state_dict() |
| state.update(weights_load) |
| self.net.load_state_dict(state) |
| self.net.eval() |
|
|
| logger.info(f"Load model successfully to device '{self.device}' for inference") |
|
|
| def preprocess( |
| self, |
| img: Union[np.ndarray, str], |
| size: Tuple[int] = (368, 368), |
| color_mean: Tuple[float] = (0.485, 0.456, 0.406), |
| color_std: Tuple[float] = (0.229, 0.224, 0.225), |
| ) -> torch.Tensor: |
| if isinstance(img, str): |
| original_img = cv2.imread(img) |
| elif isinstance(img, np.ndarray): |
| original_img = img |
| else: |
| raise ValueError("'img' parameter must be of type string or numpy array") |
|
|
| original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB) |
| img = cv2.resize(original_img, size, interpolation=cv2.INTER_CUBIC) |
| img = img.astype(np.float32) / 255.0 |
|
|
| preprocessed_img = img.copy() |
|
|
| for i in range(3): |
| preprocessed_img[:, :, i] = preprocessed_img[:, :, i] - color_mean[i] |
| preprocessed_img[:, :, i] = preprocessed_img[:, :, i] / color_std[i] |
|
|
| img = preprocessed_img.transpose((2, 0, 1)).astype(np.float32) |
| img = torch.from_numpy(img) |
| img = img.unsqueeze(0) |
|
|
| return img, original_img |
|
|
| def process( |
| self, |
| img: Union[np.ndarray, str], |
| size: Tuple[int] = (368, 368), |
| color_mean: Tuple[float] = (0.485, 0.456, 0.406), |
| color_std: Tuple[float] = (0.229, 0.224, 0.225), |
| ): |
| preprocessed_img, original_img = self.preprocess( |
| img=img, size=size, color_mean=color_mean, color_std=color_std |
| ) |
| |
| predicted_outputs, _ = self.net(preprocessed_img) |
|
|
| shape = original_img.shape |
| heatmaps = PoseInference._generate_heatmap(predicted_outputs, size, shape) |
| pafs = PoseInference._generate_part_affinity_fields(predicted_outputs, size, shape) |
|
|
| result_img = self.postprocess(original_img, heatmaps, pafs) |
| result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR) |
| return result_img |
|
|
| def postprocess( |
| self, |
| oriImg: np.ndarray, |
| heatmaps: np.ndarray, |
| pafs: np.ndarray, |
| ) -> np.ndarray: |
| _, result_img, _, _ = decode_pose(oriImg, heatmaps, pafs) |
| return result_img |
|
|
| @staticmethod |
| def _generate_heatmap( |
| predicted_outputs: torch.Tensor, |
| size: Tuple[int], |
| oriImg_shape: Tuple[int], |
| ) -> np.ndarray: |
| _heatmaps = predicted_outputs[1][0].detach().numpy().transpose(1, 2, 0) |
| _heatmaps = cv2.resize(_heatmaps, size, interpolation=cv2.INTER_CUBIC) |
| _heatmaps = cv2.resize( |
| _heatmaps, |
| (oriImg_shape[1], oriImg_shape[0]), |
| interpolation=cv2.INTER_CUBIC, |
| ) |
| logger.info("Generate heatmap ...") |
| return _heatmaps |
|
|
| @staticmethod |
| def _generate_part_affinity_fields( |
| predicted_outputs: torch.Tensor, |
| size: Tuple[int], |
| oriImg_shape: Tuple[int], |
| ) -> np.ndarray: |
| _pafs = predicted_outputs[0][0].detach().numpy().transpose(1, 2, 0) |
| _pafs = cv2.resize(_pafs, size, interpolation=cv2.INTER_CUBIC) |
| _pafs = cv2.resize( |
| _pafs, |
| (oriImg_shape[1], oriImg_shape[0]), |
| interpolation=cv2.INTER_CUBIC, |
| ) |
| logger.info("Generate part affinity fields ...") |
| return _pafs |
|
|
