-
Notifications
You must be signed in to change notification settings - Fork 756
/
image.py
445 lines (356 loc) · 16.3 KB
/
image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
from __future__ import annotations
import io
import os
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.http import set_cookies
from ...exceptions import BadInput
from ...exceptions import InvalidArgument
from ...exceptions import InternalServerError
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.v1alpha1 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)
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]
@classmethod
def from_sample(
cls,
sample: ImageType | str,
pilmode: _Mode | None = DEFAULT_PIL_MODE,
*,
allowed_mime_types: t.Iterable[str] | None = None,
) -> Self:
from filetype.match import image_match
img_type = image_match(sample)
if img_type is None:
raise InvalidArgument(f"{sample} is not a valid image file type.")
kls = cls(
mime_type=img_type.mime,
pilmode=pilmode,
allowed_mime_types=allowed_mime_types,
)
if isinstance(sample, str) and os.path.exists(sample):
try:
with open(sample, "rb") as f:
kls.sample = PIL.Image.open(f)
except PIL.UnidentifiedImageError as err:
raise BadInput(f"Failed to parse sample image file: {err}") from None
elif LazyType["ext.NpNDArray"]("numpy.ndarray").isinstance(sample):
kls.sample = PIL.Image.fromarray(sample, mode=pilmode)
elif LazyType["PIL.Image.Image"]("PIL.Image.Image").isinstance(sample):
kls.sample = sample
else:
raise InvalidArgument(f"Unknown sample type: '{sample}'")
return kls
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_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(),
}
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
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},
)
async def from_proto(self, field: pb.File | bytes) -> ImageType:
from bentoml.grpc.utils import filetype_pb_to_mimetype_map
mapping = filetype_pb_to_mimetype_map()
# check if the request message has the correct field
if isinstance(field, bytes):
content = field
else:
assert isinstance(field, pb.File)
if field.kind:
try:
mime_type = mapping[field.kind]
if mime_type != self._mime_type:
raise BadInput(
f"Inferred mime_type from 'kind' is '{mime_type}', while '{self!r}' is expecting '{self._mime_type}'",
)
except KeyError:
raise BadInput(
f"{field.kind} is not a valid File kind. Accepted file kind: {[names for names,_ in pb.File.FileType.items()]}",
) from None
content = field.content
if not content:
raise BadInput("Content is empty!") from None
return PIL.Image.open(io.BytesIO(content))
async def to_proto(self, obj: ImageType) -> pb.File:
from bentoml.grpc.utils import mimetype_to_filetype_pb_map
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)
try:
kind = mimetype_to_filetype_pb_map()[self._mime_type]
except KeyError:
raise BadInput(
f"{self._mime_type} doesn't have a corresponding File 'kind'",
) from None
return pb.File(kind=kind, content=ret.getvalue())