| from typing import List, Optional, Tuple, Dict, Callable | |
| import torch | |
| from torch import Tensor | |
| from torch.nn import Module | |
| from tha3.poser.poser import PoseParameterGroup, Poser | |
| from tha3.compute.cached_computation_func import TensorListCachedComputationFunc | |
| class GeneralPoser02(Poser): | |
| def __init__(self, | |
| module_loaders: Dict[str, Callable[[], Module]], | |
| device: torch.device, | |
| output_length: int, | |
| pose_parameters: List[PoseParameterGroup], | |
| output_list_func: TensorListCachedComputationFunc, | |
| subrect: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, | |
| default_output_index: int = 0, | |
| image_size: int = 256, | |
| dtype: torch.dtype = torch.float): | |
| self.dtype = dtype | |
| self.image_size = image_size | |
| self.default_output_index = default_output_index | |
| self.output_list_func = output_list_func | |
| self.subrect = subrect | |
| self.pose_parameters = pose_parameters | |
| self.device = device | |
| self.module_loaders = module_loaders | |
| self.modules = None | |
| self.num_parameters = 0 | |
| for pose_parameter in self.pose_parameters: | |
| self.num_parameters += pose_parameter.get_arity() | |
| self.output_length = output_length | |
| def get_image_size(self) -> int: | |
| return self.image_size | |
| def get_modules(self): | |
| if self.modules is None: | |
| self.modules = {} | |
| for key in self.module_loaders: | |
| module = self.module_loaders[key]() | |
| self.modules[key] = module | |
| module.to(self.device) | |
| module.train(False) | |
| return self.modules | |
| def get_pose_parameter_groups(self) -> List[PoseParameterGroup]: | |
| return self.pose_parameters | |
| def get_num_parameters(self) -> int: | |
| return self.num_parameters | |
| def pose(self, image: Tensor, pose: Tensor, output_index: Optional[int] = None) -> Tensor: | |
| if output_index is None: | |
| output_index = self.default_output_index | |
| output_list = self.get_posing_outputs(image, pose) | |
| return output_list[output_index] | |
| def get_posing_outputs(self, image: Tensor, pose: Tensor) -> List[Tensor]: | |
| modules = self.get_modules() | |
| if len(image.shape) == 3: | |
| image = image.unsqueeze(0) | |
| if len(pose.shape) == 1: | |
| pose = pose.unsqueeze(0) | |
| if self.subrect is not None: | |
| image = image[:, :, self.subrect[0][0]:self.subrect[0][1], self.subrect[1][0]:self.subrect[1][1]] | |
| batch = [image, pose] | |
| outputs = {} | |
| return self.output_list_func(modules, batch, outputs) | |
| def get_output_length(self) -> int: | |
| return self.output_length | |
| def free(self): | |
| self.modules = None | |
| def get_dtype(self) -> torch.dtype: | |
| return self.dtype | |