Module facetorch.analyzer.reader

class ImageReader(BaseReader):
    def __init__(
        self,
        transform: torchvision.transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
    ):
        """ImageReader is a wrapper around a functionality for reading images by Torchvision.

        Args:
            transform (torchvision.transforms.Compose): Transform compose object to be applied to the image, if fix_image_size is True.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transforms that are: resizing the image to a fixed size.

        """
        super().__init__(
            transform,
            device,
            optimize_transform,
        )

    @Timer("ImageReader.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(self, path_image: str, fix_img_size: bool = False) -> ImageData:
        """Reads an image from a path and returns a tensor of the image with values between 0-255 and shape (batch, channels, height, width). The order of color channels is RGB. PyTorch and Torchvision are used to read the image.

        Args:
            path_image (str): Path to the image.
            fix_img_size (bool): Whether to resize the image to a fixed size. If False, the size_portrait and size_landscape are ignored. Default is False.

        Returns:
            ImageData: ImageData object with image tensor and pil Image.
        """
        data = ImageData(path_input=path_image)
        data.img = torchvision.io.read_image(
            data.path_input, mode=torchvision.io.ImageReadMode.RGB
        )
        data.img = data.img.unsqueeze(0)
        data.img = data.img.to(self.device)

        if fix_img_size:
            data.img = self.transform(data.img)

        data.tensor = data.img.type(torch.float32)
        data.img = data.img.squeeze(0).cpu()
        data.set_dims()

        return data

class TensorReader(BaseReader):
    def __init__(
        self,
        transform: torchvision.transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
    ):
        """TensorReader is a wrapper around a functionality for reading tensors by Torchvision.

        Args:
            transform (torchvision.transforms.Compose): Transform compose object to be applied to the image, if fix_image_size is True.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transforms that are: resizing the image to a fixed size.

        """
        super().__init__(
            transform,
            device,
            optimize_transform,
        )

    @Timer("TensorReader.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(self, tensor: torch.Tensor, fix_img_size: bool = False) -> ImageData:
        """Reads a tensor and returns a tensor of the image with values between 0-255 and shape (batch, channels, height, width). The order of color channels is RGB. PyTorch and Torchvision are used to read the image.

        Args:
            tensor (torch.Tensor): Tensor of a single image with RGB values between 0-255 and shape (channels, height, width).
            fix_img_size (bool): Whether to resize the image to a fixed size. If False, the size_portrait and size_landscape are ignored. Default is False.

        Returns:
            ImageData: ImageData object with image tensor and pil Image.
        """
        return self.process_tensor(tensor, fix_img_size)

class UniversalReader(BaseReader):
    def __init__(
        self,
        transform: torchvision.transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
    ):
        """UniversalReader can read images from a path, URL, tensor, numpy array, bytes or PIL Image and return an ImageData object containing the image tensor.

        Args:
            transform (torchvision.transforms.Compose): Transform compose object to be applied to the image, if fix_image_size is True.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transforms that are: resizing the image to a fixed size.

        """
        super().__init__(transform, device, optimize_transform)

    @Timer("UniversalReader.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(
        self,
        image_source: Union[str, torch.Tensor, np.ndarray, bytes, Image.Image],
        fix_img_size: bool = False,
    ) -> ImageData:
        """Reads an image from a path, URL, tensor, numpy array, bytes or PIL Image and returns a tensor of the image with values between 0-255 and shape (batch, channels, height, width). The order of color channels is RGB. PyTorch and Torchvision are used to read the image.

        Args:
            image_source (Union[str, torch.Tensor, np.ndarray, bytes, Image.Image]): Image source to be read.
            fix_img_size (bool): Whether to resize the image to a fixed size. If False, the size_portrait and size_landscape are ignored. Default is False.

        Returns:
            ImageData: ImageData object with image tensor and pil Image.
        """
        if isinstance(image_source, str):
            if image_source.startswith("http"):
                return self.read_image_from_url(image_source, fix_img_size)
            else:
                return self.read_image_from_path(image_source, fix_img_size)
        elif isinstance(image_source, torch.Tensor):
            return self.read_tensor(image_source, fix_img_size)
        elif isinstance(image_source, np.ndarray):
            return self.read_numpy_array(image_source, fix_img_size)
        elif isinstance(image_source, bytes):
            return self.read_image_from_bytes(image_source, fix_img_size)
        elif isinstance(image_source, Image.Image):
            return self.read_pil_image(image_source, fix_img_size)
        else:
            raise ValueError("Unsupported data type")

    def read_tensor(self, tensor: torch.Tensor, fix_img_size: bool) -> ImageData:
        return self.process_tensor(tensor, fix_img_size)

    def read_pil_image(self, pil_image: Image.Image, fix_img_size: bool) -> ImageData:
        if pil_image.mode != "RGB":
            pil_image = pil_image.convert("RGB")
        tensor = torchvision.transforms.functional.pil_to_tensor(pil_image)
        return self.process_tensor(tensor, fix_img_size)

    def read_numpy_array(self, array: np.ndarray, fix_img_size: bool) -> ImageData:
        image_tensor = torch.from_numpy(array).float()
        if image_tensor.ndim == 3 and image_tensor.shape[2] == 3:
            image_tensor = image_tensor.permute(2, 0, 1).contiguous()
        else:
            raise ValueError(f"Unsupported numpy array shape: {image_tensor.shape}")
        return self.process_tensor(image_tensor, fix_img_size)

    def read_image_from_bytes(
        self, image_bytes: bytes, fix_img_size: bool
    ) -> ImageData:
        pil_image = Image.open(io.BytesIO(image_bytes))
        return self.read_pil_image(pil_image, fix_img_size)

    def read_image_from_path(self, path_image: str, fix_img_size: bool) -> ImageData:
        try:
            image_tensor = torchvision.io.read_image(path_image)
        except Exception as e:
            logger.error(f"Failed to read image from path {path_image}: {e}")
            raise ValueError(f"Could not read image from path {path_image}: {e}") from e

        return self.process_tensor(image_tensor, fix_img_size)

    def read_image_from_url(self, url: str, fix_img_size: bool) -> ImageData:
        try:
            response = requests.get(url, timeout=10)
            response.raise_for_status()
        except requests.RequestException as e:
            logger.error(f"Failed to fetch image from URL {url}: {e}")
            raise ValueError(f"Could not fetch image from URL {url}: {e}") from e

        image_bytes = response.content
        return self.read_image_from_bytes(image_bytes, fix_img_size)