Source code for bentoml._internal.io_descriptors.image

from __future__ import annotations

import io
import typing as t
import functools
from typing import TYPE_CHECKING
from urllib.parse import quote

from starlette.requests import Request
from multipart.multipart import parse_options_header
from starlette.responses import Response
from starlette.datastructures import UploadFile

from .base import IODescriptor
from ..types import LazyType
from ..utils import LazyLoader
from ..utils import resolve_user_filepath
from ..utils.http import set_cookies
from ...exceptions import BadInput
from ...exceptions import InvalidArgument
from ...exceptions import InternalServerError
from ...exceptions import MissingDependencyException
from ..service.openapi import SUCCESS_DESCRIPTION
from ..service.openapi.specification import Schema
from ..service.openapi.specification import MediaType

PIL_EXC_MSG = "'Pillow' is required to use the Image IO descriptor. Install with 'pip install bentoml[io-image]'."

if TYPE_CHECKING:
    from types import UnionType

    import PIL
    import PIL.Image
    from typing_extensions import Self

    from bentoml.grpc.v1 import service_pb2 as pb

    from .. import external_typing as ext
    from .base import OpenAPIResponse
    from ..context import InferenceApiContext as Context

    _Mode = t.Literal[
        "1", "CMYK", "F", "HSV", "I", "L", "LAB", "P", "RGB", "RGBA", "RGBX", "YCbCr"
    ]
else:
    from bentoml.grpc.utils import import_generated_stubs

    # NOTE: pillow-simd only benefits users who want to do preprocessing
    # TODO: add options for users to choose between simd and native mode
    PIL = LazyLoader("PIL", globals(), "PIL", exc_msg=PIL_EXC_MSG)
    PIL.Image = LazyLoader("PIL.Image", globals(), "PIL.Image", exc_msg=PIL_EXC_MSG)

    pb, _ = import_generated_stubs()

# NOTES: we will keep type in quotation to avoid backward compatibility
#  with numpy < 1.20, since we will use the latest stubs from the main branch of numpy.
#  that enable a new way to type hint an ndarray.
ImageType = t.Union["PIL.Image.Image", "ext.NpNDArray"]

DEFAULT_PIL_MODE = "RGB"


PIL_WRITE_ONLY_FORMATS = {"PALM", "PDF"}
READABLE_MIMES: set[str] = None  # type: ignore (lazy constant)
MIME_EXT_MAPPING: dict[str, str] = None  # type: ignore (lazy constant)


@functools.lru_cache(maxsize=1)
def initialize_pillow():
    global MIME_EXT_MAPPING  # pylint: disable=global-statement
    global READABLE_MIMES  # pylint: disable=global-statement

    try:
        import PIL.Image
    except ImportError:
        raise InternalServerError(PIL_EXC_MSG)

    PIL.Image.init()
    MIME_EXT_MAPPING = {v: k for k, v in PIL.Image.MIME.items()}  # type: ignore (lazy constant)
    READABLE_MIMES = {k for k, v in MIME_EXT_MAPPING.items() if v not in PIL_WRITE_ONLY_FORMATS}  # type: ignore (lazy constant)


[docs]class Image(IODescriptor[ImageType], descriptor_id="bentoml.io.Image"): """ :obj:`Image` defines API specification for the inputs/outputs of a Service, where either inputs will be converted to or outputs will be converted from images as specified in your API function signature. A sample object detection service: .. code-block:: python :caption: `service.py` from __future__ import annotations from typing import TYPE_CHECKING from typing import Any import bentoml from bentoml.io import Image from bentoml.io import NumpyNdarray if TYPE_CHECKING: from PIL.Image import Image from numpy.typing import NDArray runner = bentoml.tensorflow.get('image-classification:latest').to_runner() svc = bentoml.Service("vit-object-detection", runners=[runner]) @svc.api(input=Image(), output=NumpyNdarray(dtype="float32")) async def predict_image(f: Image) -> NDArray[Any]: assert isinstance(f, Image) arr = np.array(f) / 255.0 assert arr.shape == (28, 28) # We are using greyscale image and our PyTorch model expect one # extra channel dimension arr = np.expand_dims(arr, (0, 3)).astype("float32") # reshape to [1, 28, 28, 1] return await runner.async_run(arr) Users then can then serve this service with :code:`bentoml serve`: .. code-block:: bash % bentoml serve ./service.py:svc --reload Users can then send requests to the newly started services with any client: .. tab-set:: .. tab-item:: Bash .. code-block:: bash # we will run on our input image test.png # image can get from http://images.cocodataset.org/val2017/000000039769.jpg % curl -H "Content-Type: multipart/form-data" \\ -F 'fileobj=@test.jpg;type=image/jpeg' \\ http://0.0.0.0:3000/predict_image # [{"score":0.8610631227493286,"label":"Egyptian cat"}, # {"score":0.08770329505205154,"label":"tabby, tabby cat"}, # {"score":0.03540956228971481,"label":"tiger cat"}, # {"score":0.004140055272728205,"label":"lynx, catamount"}, # {"score":0.0009498853469267488,"label":"Siamese cat, Siamese"}]% .. tab-item:: Python .. code-block:: python :caption: `request.py` import requests requests.post( "http://0.0.0.0:3000/predict_image", files = {"upload_file": open('test.jpg', 'rb')}, headers = {"content-type": "multipart/form-data"} ).text Args: pilmode: Color mode for PIL. Default to ``RGB``. mime_type: The MIME type of the file type that this descriptor should return. Only relevant when used as an output descriptor. allowed_mime_types: A list of MIME types to restrict input to. Returns: :obj:`Image`: IO Descriptor that either a :code:`PIL.Image.Image` or a :code:`np.ndarray` representing an image. """ _proto_fields = ("file",) def __init__( self, pilmode: _Mode | None = DEFAULT_PIL_MODE, mime_type: str = "image/jpeg", *, allowed_mime_types: t.Iterable[str] | None = None, ): initialize_pillow() if pilmode is not None and pilmode not in PIL.Image.MODES: # pragma: no cover raise InvalidArgument( f"Invalid Image pilmode '{pilmode}'. Supported PIL modes are {', '.join(PIL.Image.MODES)}." ) from None self._mime_type = mime_type.lower() self._allowed_mimes: set[str] = ( READABLE_MIMES if allowed_mime_types is None else {mtype.lower() for mtype in allowed_mime_types} ) self._allow_all_images = allowed_mime_types is None if self._mime_type not in MIME_EXT_MAPPING: # pragma: no cover raise InvalidArgument( f"Invalid Image mime_type '{mime_type}'; supported mime types are {', '.join(PIL.Image.MIME.values())} " ) for mtype in self._allowed_mimes: if mtype not in MIME_EXT_MAPPING: # pragma: no cover raise InvalidArgument( f"Invalid Image MIME in allowed_mime_types: '{mtype}'; supported mime types are {', '.join(PIL.Image.MIME.values())} " ) if mtype not in READABLE_MIMES: raise InvalidArgument( f"Pillow does not support reading '{mtype}' files." ) self._pilmode: _Mode | None = pilmode self._format: str = MIME_EXT_MAPPING[self._mime_type] def _from_sample(self, sample: ImageType | str) -> ImageType: try: from filetype.match import image_match except ModuleNotFoundError: raise MissingDependencyException( "'filetype' is required to use 'from_sample'. Install it with 'pip install bentoml[io-image]'." ) img_type = image_match(sample) if img_type is None: raise InvalidArgument(f"{sample} is not a valid image file type.") if LazyType["ext.NpNDArray"]("numpy.ndarray").isinstance(sample): sample = PIL.Image.fromarray(sample) elif isinstance(sample, str): p = resolve_user_filepath(sample, ctx=None) try: with open(p, "rb") as f: sample = PIL.Image.open(f) except PIL.UnidentifiedImageError as err: raise BadInput(f"Failed to parse sample image file: {err}") from None self._mime_type = img_type.mime return sample def to_spec(self) -> dict[str, t.Any]: return { "id": self.descriptor_id, "args": { "pilmode": self._pilmode, "mime_type": self._mime_type, "allowed_mime_types": list(self._allowed_mimes), }, } @classmethod def from_spec(cls, spec: dict[str, t.Any]) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in Image spec: {spec}") return cls(**spec["args"]) def input_type(self) -> UnionType: return ImageType def openapi_schema(self) -> Schema: return Schema(type="string", format="binary") def openapi_components(self) -> dict[str, t.Any] | None: pass def openapi_example(self): pass def openapi_request_body(self) -> dict[str, t.Any]: return { "content": { mtype: MediaType(schema=self.openapi_schema()) for mtype in self._allowed_mimes }, "required": True, "x-bentoml-io-descriptor": self.to_spec(), } def openapi_responses(self) -> OpenAPIResponse: return { "description": SUCCESS_DESCRIPTION, "content": {self._mime_type: MediaType(schema=self.openapi_schema())}, "x-bentoml-io-descriptor": self.to_spec(), }
[docs] async def from_http_request(self, request: Request) -> ImageType: content_type, _ = parse_options_header(request.headers["content-type"]) mime_type = content_type.decode().lower() bytes_: bytes | str | None = None if mime_type == "multipart/form-data": form = await request.form() found_mimes: list[str] = [] for val in form.values(): val_content_type = val.content_type # type: ignore (bad starlette types) if isinstance(val, UploadFile): found_mimes.append(val_content_type) if self._allowed_mimes is None: if ( val_content_type in MIME_EXT_MAPPING or val_content_type.startswith("image/") ): bytes_ = await val.read() break elif val_content_type in self._allowed_mimes: bytes_ = await val.read() break else: if len(found_mimes) == 0: raise BadInput("no image file found in multipart form") else: if self._allowed_mimes is None: raise BadInput( f"no multipart image file (supported images are: {', '.join(MIME_EXT_MAPPING.keys())}, or 'image/*'), got files with content types {', '.join(found_mimes)}" ) else: raise BadInput( f"no multipart image file (allowed mime types are: {', '.join(self._allowed_mimes)}), got files with content types {', '.join(found_mimes)}" ) elif self._allowed_mimes is None: if mime_type in MIME_EXT_MAPPING or mime_type.startswith("image/"): bytes_ = await request.body() elif mime_type in self._allowed_mimes: bytes_ = await request.body() else: if self._allowed_mimes is None: raise BadInput( f"unsupported mime type {mime_type}; supported mime types are: {', '.join(MIME_EXT_MAPPING.keys())}, or 'image/*'" ) else: raise BadInput( f"mime type {mime_type} is not allowed, allowed mime types are: {', '.join(self._allowed_mimes)}" ) assert bytes_ is not None if isinstance(bytes_, str): bytes_ = bytes(bytes_, "UTF-8") try: return PIL.Image.open(io.BytesIO(bytes_)) except PIL.UnidentifiedImageError as err: raise BadInput(f"Failed to parse uploaded image file: {err}") from None
[docs] async def to_http_response( self, obj: ImageType, ctx: Context | None = None ) -> Response: if LazyType["ext.NpNDArray"]("numpy.ndarray").isinstance(obj): image = PIL.Image.fromarray(obj, mode=self._pilmode) elif LazyType["PIL.Image.Image"]("PIL.Image.Image").isinstance(obj): image = obj else: raise BadInput( f"Unsupported Image type received: '{type(obj)}', the Image IO descriptor only supports 'np.ndarray' and 'PIL.Image'." ) from None filename = f"output.{self._format.lower()}" ret = io.BytesIO() image.save(ret, format=self._format) # rfc2183 content_disposition_filename = quote(filename) if content_disposition_filename != filename: content_disposition = "attachment; filename*=utf-8''{}".format( content_disposition_filename ) else: content_disposition = f'attachment; filename="{filename}"' if ctx is not None: if "content-disposition" not in ctx.response.headers: ctx.response.headers["content-disposition"] = content_disposition res = Response( ret.getvalue(), media_type=self._mime_type, headers=ctx.response.headers, # type: ignore (bad starlette types) status_code=ctx.response.status_code, ) set_cookies(res, ctx.response.cookies) return res else: return Response( ret.getvalue(), media_type=self._mime_type, headers={"content-disposition": content_disposition}, )
[docs] async def from_proto(self, field: pb.File | bytes) -> ImageType: if isinstance(field, bytes): content = field else: assert isinstance(field, pb.File) if field.kind and field.kind != self._mime_type: raise BadInput( f"MIME type from 'kind' is '{field.kind}', while '{self!r}' is expecting '{self._mime_type}'", ) content = field.content if not content: raise BadInput("Content is empty!") from None return PIL.Image.open(io.BytesIO(content))
[docs] async def to_proto(self, obj: ImageType) -> pb.File: if LazyType["ext.NpNDArray"]("numpy.ndarray").isinstance(obj): image = PIL.Image.fromarray(obj, mode=self._pilmode) elif LazyType["PIL.Image.Image"]("PIL.Image.Image").isinstance(obj): image = obj else: raise BadInput( f"Unsupported Image type received: '{type(obj)}', the Image IO descriptor only supports 'np.ndarray' and 'PIL.Image'.", ) from None ret = io.BytesIO() image.save(ret, format=self._format) return pb.File(kind=self._mime_type, content=ret.getvalue())