Add Image API

docs/source/core.rst

@@ -0,0 +1,8 @@
+kornia.core
+===========
+
+.. currentmodule:: kornia.core
+
+.. autoclass:: TensorWrapper
+    :members:
+    :undoc-members:

docs/source/image.rst

@@ -0,0 +1,26 @@
+kornia.image
+============
+
+Module to provide a high level API to process images.
+
+.. currentmodule:: kornia.image
+
+.. autoclass:: ImageSize
+    :members:
+    :undoc-members:
+
+.. autoclass:: PixelFormat
+    :members:
+    :undoc-members:
+
+.. autoclass:: ChannelsOrder
+    :members:
+    :undoc-members:
+
+.. autoclass:: ImageLayout
+   :members:
+   :undoc-members:
+
+.. autoclass:: Image
+    :members:
+    :undoc-members:

docs/source/index.rst

@@ -63,6 +63,7 @@ Join the community
    geometry
    sensors
    io
+   image
    losses
    metrics
    morphology

kornia/color/gray.py

@@ -15,18 +15,19 @@ def grayscale_to_rgb(image: Tensor) -> Tensor:
     The image data is assumed to be in the range of (0, 1).
 
     Args:
-        image: grayscale image to be converted to RGB with shape :math:`(*,1,H,W)`.
+        image: grayscale image tensor to be converted to RGB with shape :math:`(*,1,H,W)`.
 
     Returns:
         RGB version of the image with shape :math:`(*,3,H,W)`.
 
     Example:
+        >>> import torch
         >>> input = torch.randn(2, 1, 4, 5)
         >>> gray = grayscale_to_rgb(input) # 2x3x4x5
     """
     KORNIA_CHECK_IS_TENSOR(image)
 
-    if image.dim() < 3 or image.size(-3) != 1:
+    if len(image.shape) < 3 or image.shape[-3] != 1:
         raise ValueError(f"Input size must have a shape of (*, 1, H, W). " f"Got {image.shape}.")
 
     return concatenate([image, image, image], -3)
@@ -51,6 +52,7 @@ def rgb_to_grayscale(image: Tensor, rgb_weights: Optional[Tensor] = None) -> Ten
        color_conversions.html>`__.
 
     Example:
+        >>> import torch
         >>> input = torch.rand(2, 3, 4, 5)
         >>> gray = rgb_to_grayscale(input) # 2x1x4x5
     """
@@ -81,7 +83,7 @@ def rgb_to_grayscale(image: Tensor, rgb_weights: Optional[Tensor] = None) -> Ten
     return w_r * r + w_g * g + w_b * b
 
 
-def bgr_to_grayscale(image: torch.Tensor) -> torch.Tensor:
+def bgr_to_grayscale(image: Tensor) -> Tensor:
     r"""Convert a BGR image to grayscale.
 
     The image data is assumed to be in the range of (0, 1). First flips to RGB, then converts.
@@ -93,6 +95,7 @@ def bgr_to_grayscale(image: torch.Tensor) -> torch.Tensor:
         grayscale version of the image with shape :math:`(*,1,H,W)`.
 
     Example:
+        >>> import torch
         >>> input = torch.rand(2, 3, 4, 5)
         >>> gray = bgr_to_grayscale(input) # 2x1x4x5
     """
@@ -118,6 +121,7 @@ class GrayscaleToRgb(Module):
         https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html
 
     Example:
+        >>> import torch
         >>> input = torch.rand(2, 1, 4, 5)
         >>> rgb = GrayscaleToRgb()
         >>> output = rgb(input)  # 2x3x4x5
@@ -140,13 +144,16 @@ class RgbToGrayscale(Module):
         https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html
 
     Example:
+        >>> import torch
         >>> input = torch.rand(2, 3, 4, 5)
         >>> gray = RgbToGrayscale()
         >>> output = gray(input)  # 2x1x4x5
     """
 
     def __init__(self, rgb_weights: Optional[Tensor] = None) -> None:
         super().__init__()
+        if rgb_weights is None:
+            rgb_weights = Tensor([0.299, 0.587, 0.114])
         self.rgb_weights = rgb_weights
 
     def forward(self, image: Tensor) -> Tensor:
@@ -166,6 +173,7 @@ class BgrToGrayscale(Module):
         https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html
 
     Example:
+        >>> import torch
         >>> input = torch.rand(2, 3, 4, 5)
         >>> gray = BgrToGrayscale()
         >>> output = gray(input)  # 2x1x4x5

kornia/core/__init__.py

@@ -19,6 +19,7 @@
     zeros,
     zeros_like,
 )
+from .tensor_wrapper import TensorWrapper
 
 __all__ = [
     "concatenate",
@@ -40,4 +41,5 @@
     "zeros_like",
     "linspace",
     "diag",
+    "TensorWrapper",
 ]

kornia/core/_backend.py

@@ -1,8 +1,6 @@
-from typing import Union
+from __future__ import annotations
 
 import torch
-import torch.nn.functional as F
-from torch import device
 
 # classes
 Tensor = torch.Tensor
@@ -15,8 +13,8 @@
 concatenate = torch.cat
 stack = torch.stack
 linspace = torch.linspace
-normalize = F.normalize
-pad = F.pad
+normalize = torch.nn.functional.normalize
+pad = torch.nn.functional.pad
 eye = torch.eye
 zeros = torch.zeros
 zeros_like = torch.zeros_like
@@ -31,5 +29,5 @@
 rand = torch.rand
 
 # type alias
-Device = Union[str, device, None]
-Dtype = Union[torch.dtype, None]
+Device = str | torch.device | None
+Dtype = torch.dtype | None

kornia/image/__init__.py

@@ -1 +1,4 @@
-from .base import ImageSize
+from .base import ChannelsOrder, ImageLayout, ImageSize, PixelFormat
+from .image import Image
+
+__all__ = ["ImageSize", "PixelFormat", "ChannelsOrder", "ImageLayout", "Image"]

kornia/image/base.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-
-from kornia.core import Tensor
+from enum import Enum
 
 
 @dataclass
@@ -11,13 +10,58 @@ class ImageSize:
 
     Args:
         height: image height.
-        width: image width
+        width: image width.
+
     Example:
         >>> size = ImageSize(3, 4)
         >>> size.height
         3
         >>> size.width
         4
     """
-    height: int | Tensor
-    width: int | Tensor
+    height: int
+    width: int
+
+
+class PixelFormat(Enum):
+    r"""Enum that represents the pixel format of an image."""
+    GRAY = 0
+    RGB = 1
+    BGR = 2
+
+
+class ChannelsOrder(Enum):
+    r"""Enum that represents the channels order of an image."""
+    CHANNELS_FIRST = 0
+    CHANNELS_LAST = 1
+
+
+@dataclass
+class ImageLayout:
+    """Data class to represent the layout of an image.
+
+    Args:
+        image_size: image size.
+        channels: number of channels.
+        pixel_format: pixel format.
+        channels_order: channels order.
+
+    Example:
+        >>> layout = ImageLayout(ImageSize(3, 4), 3, PixelFormat.RGB, ChannelsOrder.CHANNEL_LAST)
+        >>> layout.image_size
+        ImageSize(height=3, width=4)
+        >>> layout.channels
+        3
+        >>> layout.pixel_format
+        <PixelFormat.RGB: 1>
+        >>> layout.channels_order
+        <ChannelsOrder.CHANNEL_LAST: 1>
+    """
+
+    image_size: ImageSize
+    channels: int
+    pixel_format: PixelFormat
+    channels_order: ChannelsOrder
+
+
+# TODO: define CompressedImage