Spaces:
Running
Running
| from typing import Any, Dict, List, Optional, Union | |
| import numpy as np | |
| import PIL.Image | |
| import torch | |
| from ..utils.import_utils import is_kornia_available | |
| from .base import ProcessorMixin | |
| if is_kornia_available(): | |
| import kornia | |
| class CannyProcessor(ProcessorMixin): | |
| r""" | |
| Processor for obtaining the Canny edge detection of an image. | |
| Args: | |
| output_names (`List[str]`): | |
| The names of the outputs that the processor should return. The first output is the Canny edge detection of | |
| the input image. | |
| """ | |
| def __init__( | |
| self, | |
| output_names: List[str] = None, | |
| input_names: Optional[Dict[str, Any]] = None, | |
| device: Optional[torch.device] = None, | |
| ): | |
| super().__init__() | |
| self.output_names = output_names | |
| self.input_names = input_names | |
| self.device = device | |
| assert len(output_names) == 1 | |
| def forward(self, input: Union[torch.Tensor, PIL.Image.Image, List[PIL.Image.Image]]) -> torch.Tensor: | |
| r""" | |
| Obtain the Canny edge detection of the input image. | |
| Args: | |
| input (`torch.Tensor`, `PIL.Image.Image`, or `List[PIL.Image.Image]`): | |
| The input tensor, image or list of images for which the Canny edge detection should be obtained. | |
| If a tensor, must be a 3D (CHW) or 4D (BCHW) or 5D (BTCHW) tensor. The input tensor should have | |
| values in the range [0, 1]. | |
| Returns: | |
| torch.Tensor: | |
| The Canny edge detection of the input image. The output has the same shape as the input tensor. If | |
| the input is an image, the output is a 3D tensor. If the input is a list of images, the output is a 5D | |
| tensor. The output tensor has values in the range [0, 1]. | |
| """ | |
| if isinstance(input, PIL.Image.Image): | |
| input = kornia.utils.image.image_to_tensor(np.array(input)).unsqueeze(0) / 255.0 | |
| input = input.to(self.device) | |
| output = kornia.filters.canny(input)[1].repeat(1, 3, 1, 1).squeeze(0) | |
| elif isinstance(input, list): | |
| input = kornia.utils.image.image_list_to_tensor([np.array(img) for img in input]) / 255.0 | |
| output = kornia.filters.canny(input)[1].repeat(1, 3, 1, 1) | |
| else: | |
| ndim = input.ndim | |
| assert ndim in [3, 4, 5] | |
| batch_size = 1 if ndim == 3 else input.size(0) | |
| if ndim == 3: | |
| input = input.unsqueeze(0) # [C, H, W] -> [1, C, H, W] | |
| elif ndim == 5: | |
| input = input.flatten(0, 1) # [B, F, C, H, W] -> [B*F, C, H, W] | |
| output = kornia.filters.canny(input)[1].repeat(1, 3, 1, 1) | |
| output = output[0] if ndim == 3 else output.unflatten(0, (batch_size, -1)) if ndim == 5 else output | |
| # TODO(aryan): think about how one can pass parameters to the underlying function from | |
| # a UI perspective. It's important to think about ProcessorMixin in terms of a Graph-based | |
| # data processing pipeline. | |
| return {self.output_names[0]: output} | |