Source code for luna_handlers.sdk.sdk_loop.models.image

"""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]@attr.define(slots=True) class InputImage: """ Container for input image """ # body body: Union[bytes, bytearray, PilImage, ndarray] # image type (0-image, 1 - face warp, 2 - body warp) imageType: ImageType # image filename filename: str = "" # face bboxes faceBoxes: Optional[list[Rect]] = None # body bboxes bodyBoxes: Optional[list[Rect]] = None # cached pillow image _pillowImage: Optional[PilImage] = None
[docs] def getAsPillow(self) -> Union[PilImage, errors.LoopError]: """ Get image data as pillow object. Returns: PIL image if image was loaded success otherwise error """ if self._pillowImage is None: res = getAsPillow(self.body) if isinstance(res, errors.LoopError): return res self._pillowImage = res return self._pillowImage
[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
[docs] def extractExif(self) -> dict: """ Extract exif from image. Set an error if load image is failed. If image contains an error will be return empty dict Returns: extracted exif tags """ exif = {} if self.error: # some trouble pass else: if not self.exif: pilImage = self.origin.getAsPillow() if isinstance(pilImage, errors.LoopError): self.error = pilImage else: exif = getExif(pilImage) self.exif = exif return exif