def post_process_image()

in src/image_gen_aux/image_processor.py [0:0]

• 8 lines of code
• 3 McCabe index (conditional complexity)

    def post_process_image(self, image: torch.Tensor, return_type: str):
        """
        Post-process a PyTorch tensor image and convert it to the specified return type.

        Args:
            image (torch.Tensor): The input image as a PyTorch tensor.
            return_type (str): The desired return type, either "pt" for PyTorch tensor, "np" for NumPy array, or "pil" for PIL image.

        Returns:
            Union[torch.Tensor, np.ndarray, List[PIL.Image.Image]]: The post-processed image in the specified return type.
        """
        if return_type == "pt":
            return image

        image = self.pt_to_numpy(image)
        if return_type == "np":
            return image

        image = self.numpy_to_pil(image)
        return image