base.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  1. # LICENSE HEADER MANAGED BY add-license-header
  2. #
  3. # Copyright 2018 Kornia Team
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. #
  17. from __future__ import annotations
  18. from dataclasses import dataclass
  19. from enum import Enum
  20. from kornia.core import Tensor
  21. @dataclass(frozen=True)
  22. class ImageSize:
  23. r"""Data class to represent image shape.
  24. Args:
  25. height: image height.
  26. width: image width.
  27. Example:
  28. >>> size = ImageSize(3, 4)
  29. >>> size.height
  30. 3
  31. >>> size.width
  32. 4
  33. """
  34. height: int | Tensor
  35. width: int | Tensor
  36. class ColorSpace(Enum):
  37. r"""Enum that represents the color space of an image."""
  38. UNKNOWN = 0 # for now, in case of multi band images
  39. GRAY = 1
  40. RGB = 2
  41. BGR = 3
  42. @dataclass(frozen=True)
  43. class PixelFormat:
  44. r"""Data class to represent the pixel format of an image.
  45. Args:
  46. color_space: color space.
  47. bit_depth: the number of bits per channel.
  48. Example:
  49. >>> pixel_format = PixelFormat(ColorSpace.RGB, 8)
  50. >>> pixel_format.color_space
  51. <ColorSpace.RGB: 2>
  52. >>> pixel_format.bit_depth
  53. 8
  54. """
  55. color_space: ColorSpace
  56. bit_depth: int
  57. class ChannelsOrder(Enum):
  58. r"""Enum that represents the channels order of an image."""
  59. CHANNELS_FIRST = 0
  60. CHANNELS_LAST = 1
  61. @dataclass(frozen=True)
  62. class ImageLayout:
  63. """Data class to represent the layout of an image.
  64. Args:
  65. image_size: image size.
  66. channels: number of channels.
  67. channels_order: channels order.
  68. Example:
  69. >>> layout = ImageLayout(ImageSize(3, 4), 3, ChannelsOrder.CHANNELS_LAST)
  70. >>> layout.image_size
  71. ImageSize(height=3, width=4)
  72. >>> layout.channels
  73. 3
  74. >>> layout.channels_order
  75. <ChannelsOrder.CHANNELS_LAST: 1>
  76. """
  77. image_size: ImageSize
  78. channels: int
  79. channels_order: ChannelsOrder
  80. # TODO: define CompressedImage