Replaces uses of deprecated imghdr with Pillow to detect image MIME types

tensorboard/BUILD

@@ -397,6 +397,11 @@ tensorboard_zip_file(
 # `pip install numpy`
 py_library(name = "expect_numpy_installed")
 
+# This is a dummy rule used as a pillow dependency in open-source.
+# We expect pillow to already be installed on the system, e.g. via
+# `pip install pillow`
+py_library(name = "expect_pillow_installed")
+
 # This is a dummy rule used as a grpc dependency in open-source.
 # We expect grpc to already be installed on the system, e.g. via
 # `pip install grpcio`

tensorboard/pip_package/requirements.txt

@@ -25,6 +25,7 @@ numpy >= 1.12.0
 # (specifically the protobuf dependency). If we restrict protobuf >= 5.0.0 we
 # can get rid of the packaging dependency.
 packaging
+pillow
 # NOTE: this version must be >= the protoc version in our WORKSPACE file.
 # At the same time, any constraints we specify here must allow at least some
 # version to be installed that is also compatible with TensorFlow's constraints:

tensorboard/plugins/image/BUILD

@@ -25,6 +25,7 @@ py_library(
         "//tensorboard:plugin_util",
         "//tensorboard/backend:http_util",
         "//tensorboard/plugins:base_plugin",
+        "//tensorboard/util:img_mime_type_detector",
         "@org_pocoo_werkzeug",
     ],
 )

tensorboard/plugins/image/images_plugin.py

@@ -14,8 +14,6 @@
 # ==============================================================================
 """The TensorBoard Images plugin."""
 
-
-import imghdr
 import urllib.parse
 
 from werkzeug import wrappers
@@ -26,31 +24,12 @@
 from tensorboard.data import provider
 from tensorboard.plugins import base_plugin
 from tensorboard.plugins.image import metadata
+from tensorboard.util import img_mime_type_detector
 
 
-_IMGHDR_TO_MIMETYPE = {
-    "bmp": "image/bmp",
-    "gif": "image/gif",
-    "jpeg": "image/jpeg",
-    "png": "image/png",
-    "svg": "image/svg+xml",
-}
-
-_DEFAULT_IMAGE_MIMETYPE = "application/octet-stream"
 _DEFAULT_DOWNSAMPLING = 10  # images per time series
 
 
-# Extend imghdr.tests to include svg.
-def detect_svg(data, f):
-    del f  # Unused.
-    # Assume XML documents attached to image tag to be SVG.
-    if data.startswith(b"<?xml ") or data.startswith(b"<svg "):
-        return "svg"
-
-
-imghdr.tests.append(detect_svg)
-
-
 class ImagesPlugin(base_plugin.TBPlugin):
     """Images Plugin for TensorBoard."""
 
@@ -243,11 +222,8 @@ def _serve_individual_image(self, request):
                 "text/plain",
                 code=400,
             )
-        image_type = imghdr.what(None, data)
-        content_type = _IMGHDR_TO_MIMETYPE.get(
-            image_type, _DEFAULT_IMAGE_MIMETYPE
-        )
-        return http_util.Respond(request, data, content_type)
+        mime_type = img_mime_type_detector.from_bytes(data)
+        return http_util.Respond(request, data, mime_type)
 
     @wrappers.Request.application
     def _serve_tags(self, request):

tensorboard/plugins/metrics/BUILD

@@ -24,6 +24,7 @@ py_library(
         "//tensorboard/plugins/histogram:metadata",
         "//tensorboard/plugins/image:metadata",
         "//tensorboard/plugins/scalar:metadata",
+        "//tensorboard/util:img_mime_type_detector",
         "@org_pocoo_werkzeug",
     ],
 )

tensorboard/plugins/metrics/metrics_plugin.py

@@ -16,7 +16,6 @@
 
 
 import collections
-import imghdr
 import json
 
 from werkzeug import wrappers
@@ -30,18 +29,9 @@
 from tensorboard.plugins.image import metadata as image_metadata
 from tensorboard.plugins.metrics import metadata
 from tensorboard.plugins.scalar import metadata as scalar_metadata
+from tensorboard.util import img_mime_type_detector
 
 
-_IMGHDR_TO_MIMETYPE = {
-    "bmp": "image/bmp",
-    "gif": "image/gif",
-    "jpeg": "image/jpeg",
-    "png": "image/png",
-    "svg": "image/svg+xml",
-}
-
-_DEFAULT_IMAGE_MIMETYPE = "application/octet-stream"
-
 _SINGLE_RUN_PLUGINS = frozenset(
     [histogram_metadata.PLUGIN_NAME, image_metadata.PLUGIN_NAME]
 )
@@ -615,8 +605,8 @@ def _serve_image_data(self, request):
         if not blob_key:
             raise errors.InvalidArgumentError("Missing 'imageId' field")
 
-        (data, content_type) = self._image_data_impl(ctx, blob_key)
-        return http_util.Respond(request, data, content_type)
+        (data, mime_type) = self._image_data_impl(ctx, blob_key)
+        return http_util.Respond(request, data, mime_type)
 
     def _image_data_impl(self, ctx, blob_key):
         """Gets the image data for a blob key.
@@ -631,8 +621,5 @@ def _image_data_impl(self, ctx, blob_key):
               content_type: a string HTTP content type.
         """
         data = self._data_provider.read_blob(ctx, blob_key=blob_key)
-        image_type = imghdr.what(None, data)
-        content_type = _IMGHDR_TO_MIMETYPE.get(
-            image_type, _DEFAULT_IMAGE_MIMETYPE
-        )
-        return (data, content_type)
+        mime_type = img_mime_type_detector.from_bytes(data)
+        return (data, mime_type)

tensorboard/plugins/projector/BUILD

@@ -26,6 +26,7 @@ py_library(
         "//tensorboard/backend/event_processing:plugin_asset_util",
         "//tensorboard/compat:tensorflow",
         "//tensorboard/plugins:base_plugin",
+        "//tensorboard/util:img_mime_type_detector",
         "//tensorboard/util:tb_logging",
         "@org_pocoo_werkzeug",
     ],

tensorboard/plugins/projector/projector_plugin.py

@@ -17,7 +17,6 @@
 
 import collections
 import functools
-import imghdr
 import mimetypes
 import os
 import threading
@@ -35,7 +34,7 @@
 from tensorboard.plugins import base_plugin
 from tensorboard.plugins.projector import metadata
 from tensorboard.plugins.projector.projector_config_pb2 import ProjectorConfig
-from tensorboard.util import tb_logging
+from tensorboard.util import img_mime_type_detector, tb_logging
 
 logger = tb_logging.get_logger()
 
@@ -50,14 +49,6 @@
 BOOKMARKS_ROUTE = "/bookmarks"
 SPRITE_IMAGE_ROUTE = "/sprite_image"
 
-_IMGHDR_TO_MIMETYPE = {
-    "bmp": "image/bmp",
-    "gif": "image/gif",
-    "jpeg": "image/jpeg",
-    "png": "image/png",
-}
-_DEFAULT_IMAGE_MIMETYPE = "application/octet-stream"
-
 
 class LRUCache:
     """LRU cache.
@@ -786,8 +777,7 @@ def _serve_sprite_image(self, request):
         f = tf.io.gfile.GFile(fpath, "rb")
         encoded_image_string = f.read()
         f.close()
-        image_type = imghdr.what(None, encoded_image_string)
-        mime_type = _IMGHDR_TO_MIMETYPE.get(image_type, _DEFAULT_IMAGE_MIMETYPE)
+        mime_type = img_mime_type_detector.from_bytes(encoded_image_string)
         return Respond(request, encoded_image_string, mime_type)
 
 

tensorboard/util/BUILD

@@ -103,6 +103,14 @@ py_test(
     ],
 )
 
+py_library(
+    name = "img_mime_type_detector",
+    srcs = ["img_mime_type_detector.py"],
+    deps = [
+        "//tensorboard:expect_pillow_installed",
+    ],
+)
+
 py_library(
     name = "io_util",
     srcs = ["io_util.py"],

tensorboard/util/img_mime_type_detector.py

@@ -0,0 +1,40 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility to determine the MIME type of an image."""
+
+from PIL import Image
+import io
+
+_IMGHDR_TO_MIMETYPE = {
+    "bmp": "image/bmp",
+    "gif": "image/gif",
+    "jpeg": "image/jpeg",
+    "png": "image/png",
+}
+_DEFAULT_IMAGE_MIMETYPE = "application/octet-stream"
+
+
+def from_bytes(img_bytes: bytes) -> str:
+    """Returns the MIME type of an image from its bytes."""
+    format_lower = None
+    try:
+        img = Image.open(io.BytesIO(img_bytes))
+        format_lower = img.format.lower()
+        if format_lower == "jpg":
+            format_lower = "jpeg"
+    except:
+        # Let the default value be returned.
+        pass
+    return _IMGHDR_TO_MIMETYPE.get(format_lower, _DEFAULT_IMAGE_MIMETYPE)