"""Module contains models for work with images"""
import io
from enum import Enum
from typing import Optional, Union
import attr
import PIL.Image
from lunavl.sdk.estimators.body_estimators.bodywarper import BodyWarp
from lunavl.sdk.estimators.face_estimators.facewarper import FaceWarp
from lunavl.sdk.estimators.image_estimators.orientation_mode import OrientationType
from lunavl.sdk.image_utils.geometry import Rect
from lunavl.sdk.image_utils.image import RotationAngle, VLImage
from lunavl.sdk.image_utils.pil.np import pilToNumpy
from numpy import ndarray
from PIL.Image import Image as PilImage
from PIL.ImageOps import exif_transpose
from .sample import Sample
from ..errors import errors
from ..exif import Orientation, getExif
from ..utils.rotation import rotateBBoxes
[docs]def getAsPillow(data: Union[bytes, bytearray, PilImage, ndarray]) -> Union[PilImage, errors.LoopError]:
"""
Get image data as pillow object.
Returns:
PIL image if image was loaded success otherwise error
"""
if isinstance(data, (bytes, bytearray)):
try:
pillowImage = PIL.Image.open(io.BytesIO(data))
except OSError:
error = errors.InvalidType.format("Unsupported type")
return error
else:
return pillowImage
elif isinstance(data, PilImage):
pillowImage = data
elif isinstance(data, ndarray):
try:
pillowImage = PIL.Image.fromarray(data)
except TypeError:
error = errors.InvalidType.format("Unsupported type")
return error
else:
raise RuntimeError(f"Unsupported image format {type(data)}, use: bytes, bytearray, PilImage, ndarray")
return pillowImage
[docs]class ImageType(Enum):
"""Image type enum"""
IMAGE = 0 # usual image
FACE_WARP = 1 # vl face warp
BODY_WARP = 2 # vl body warp
[docs]def getTransposeParam(method: Orientation) -> tuple[RotationAngle, bool]:
"""
Get transpose param by exif orientation
Args:
method: method from a image
Returns:
tuple< rotation, neeed or not to flip the image>
"""
if method == Orientation.NOTHING:
rotationAngle = RotationAngle.ANGLE_0
flip = False
elif method == Orientation.ROTATED_90:
rotationAngle = RotationAngle.ANGLE_270
flip = False
elif method == Orientation.ROTATED_180:
rotationAngle = RotationAngle.ANGLE_180
flip = False
elif method == Orientation.ROTATED_270:
rotationAngle = RotationAngle.ANGLE_90
flip = False
elif method == Orientation.FLIPPED_LEFT_RIGHT:
flip = True
rotationAngle = RotationAngle.ANGLE_0
elif method == Orientation.FLIPPED_TOP_BOTTOM:
rotationAngle = RotationAngle.ANGLE_180
flip = True
elif method == Orientation.TRANSPOSED:
rotationAngle = RotationAngle.ANGLE_270
flip = True
else:
rotationAngle = RotationAngle.ANGLE_90
flip = True
return rotationAngle, flip
[docs]@attr.dataclass(slots=True)
class Image:
"""
A structure for work with image.
The image contains all temporary entities for estimation targets and provides conversation between them.
The image contains estimation results also.
"""
# origin image
origin: InputImage
# loaded sdk image
sdkImage: Optional[Union[VLImage, FaceWarp, BodyWarp]] = None
# image processing error
error: Optional[errors.LoopError] = None
# extracted exif
exif: Optional[dict] = None
# face bboxes for redect, may be different with origin.faceBoxes if client use exif info or autorotation
faceBoxes: Optional[list[Rect]] = None
# body bboxes for redect, may be different with origin.faceBoxes if client use exif info or autorotation
bodyBoxes: Optional[list[Rect]] = None
# estimated orientation (from image orientation estimator, exif orientation tag not influence on param)
orientation: Optional[OrientationType] = None
# estimated people count
peopleCount: Optional[int] = None
# list human samples
samples: list[Sample] = attr.field(factory=list)
# cached loaded pillow image
_pillowImage: Optional[PilImage] = None
# cached loaded numpy array
_ndarray: Optional[ndarray] = None
# exif orientation tag
exifOrientation: Optional[Orientation] = None
# cached raw image, may be different with origin.body in case of using exif transpose or autorotation
_rawImage: Optional[Union[bytes, bytearray, PilImage, ndarray]] = None
@property
def rawImage(self) -> Union[bytes, bytearray, PilImage, ndarray]:
"""Get data for load"""
if self._rawImage is None:
self._rawImage = self.origin.body
return self._rawImage
[docs] def getAsPillow(self) -> Optional[PilImage]:
"""
Get image data as pillow object.
Returns:
PIL image
"""
if self._pillowImage:
return self._pillowImage
img = getAsPillow(self.rawImage)
if isinstance(img, errors.LoopError):
self.error = img
return None
self._pillowImage = img
return img
[docs] def getAsNumpy(self) -> Optional[ndarray]:
"""
Convert image to numpy array.
If occured error function will set up in the image.
"""
if self._ndarray:
pass
elif isinstance(self.rawImage, (bytes, bytearray)):
pilImage = self.getAsPillow()
if not pilImage:
return
self._ndarray = pilToNumpy(pilImage)
elif isinstance(self.rawImage, PilImage):
self._ndarray = pilToNumpy(self.rawImage)
else:
self._ndarray = self.rawImage
return self._ndarray
[docs] @classmethod
def exifTranspose(cls, image: "Image") -> "Image":
"""
Based on genuine function from pillow
https://pillow.readthedocs.io/en/latest/_modules/PIL/ImageOps.html#exif_transpose
The only difference is that manipulation is being done inplace, allowing us to
leverage internal caching of exif data.
NOTE: Transposed image might have incorrect values in particular tags. For example,
orientation, length, width.
"""
pilImage = image.getAsPillow()
if not pilImage:
return image
method = Orientation.NOTHING
exif = image.extractExif()
if exifOrientation := exif.get("orientation"):
try:
method = Orientation(exifOrientation)
except ValueError:
# bad exif value, do nothing
pass
if method != Orientation.NOTHING and not (pilImage.format == "TIFF" and pilImage.tag_v2.get(0x0112)):
# pillow already transpose this image
# https://github.com/radarhere/Pillow/blob/main/src/PIL/TiffImagePlugin.py#L1136
# samples and np array is not rellevant
transposedImage = exif_transpose(pilImage)
newImage = cls(
origin=image.origin,
orientation=image.orientation,
exif=image.exif,
error=image.error,
faceBoxes=image.faceBoxes,
bodyBoxes=image.bodyBoxes,
)
newImage._pillowImage = transposedImage
newImage._rawImage = transposedImage
newImage.exifOrientation = exifOrientation
else:
newImage = image
if method != Orientation.NOTHING:
width = pilImage.width
height = pilImage.height
rotationAngle, flip = getTransposeParam(method)
if image.faceBoxes:
newImage.faceBoxes = rotateBBoxes(
image.faceBoxes, imageWidth=width, imageHeight=height, angle=rotationAngle, flip=flip
)
if image.bodyBoxes:
newImage.bodyBoxes = rotateBBoxes(
image.bodyBoxes, imageWidth=width, imageHeight=height, angle=rotationAngle, flip=flip
)
return newImage