127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
# Copyright 2018-2021 Streamlit Inc.
|
|
#
|
|
# 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.
|
|
|
|
import imghdr
|
|
import io
|
|
import mimetypes
|
|
from typing import cast
|
|
from urllib.parse import urlparse
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageFile
|
|
import hashlib
|
|
import base64
|
|
|
|
# Maximum content width
|
|
MAXIMUM_CONTENT_WIDTH = 1460 # 2 * 730
|
|
|
|
def _image_has_alpha_channel(image):
|
|
"""Check if the image has an alpha channel."""
|
|
if image.mode in ("RGBA", "LA") or (
|
|
image.mode == "P" and "transparency" in image.info
|
|
):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def _format_from_image_type(image, output_format):
|
|
"""Determine the output format based on image type."""
|
|
output_format = output_format.upper()
|
|
if output_format == "JPEG" or output_format == "PNG":
|
|
return output_format
|
|
|
|
# We are forgiving on the spelling of JPEG
|
|
if output_format == "JPG":
|
|
return "JPEG"
|
|
|
|
if _image_has_alpha_channel(image):
|
|
return "PNG"
|
|
|
|
return "JPEG"
|
|
|
|
def _PIL_to_bytes(image, format="JPEG", quality=100):
|
|
"""Convert PIL image to bytes."""
|
|
tmp = io.BytesIO()
|
|
|
|
# User must have specified JPEG, so we must convert it
|
|
if format == "JPEG" and _image_has_alpha_channel(image):
|
|
image = image.convert("RGB")
|
|
|
|
image.save(tmp, format=format, quality=quality)
|
|
|
|
return tmp.getvalue()
|
|
|
|
def _BytesIO_to_bytes(data):
|
|
"""Convert BytesIO to bytes."""
|
|
data.seek(0)
|
|
return data.getvalue()
|
|
|
|
def _normalize_to_bytes(data, width, output_format):
|
|
"""Normalize image data to bytes with proper format and size."""
|
|
image = Image.open(io.BytesIO(data))
|
|
actual_width, actual_height = image.size
|
|
format = _format_from_image_type(image, output_format)
|
|
if output_format.lower() == "auto":
|
|
ext = imghdr.what(None, data)
|
|
mimetype = mimetypes.guess_type("image.%s" % ext)[0]
|
|
else:
|
|
mimetype = "image/" + format.lower()
|
|
|
|
if width < 0 and actual_width > MAXIMUM_CONTENT_WIDTH:
|
|
width = MAXIMUM_CONTENT_WIDTH
|
|
|
|
if width > 0 and actual_width > width:
|
|
new_height = int(1.0 * actual_height * width / actual_width)
|
|
image = image.resize((width, new_height), resample=Image.BILINEAR)
|
|
data = _PIL_to_bytes(image, format=format, quality=90)
|
|
mimetype = "image/" + format.lower()
|
|
|
|
return data, mimetype
|
|
|
|
def generate_image_hash(image, mimetype):
|
|
"""Generate a SHA-224 hash for an image."""
|
|
hasher = hashlib.sha224()
|
|
hasher.update(image)
|
|
hasher.update(mimetype.encode())
|
|
return hasher.hexdigest()
|
|
|
|
def image_to_url(image, width=-1, output_format="auto"):
|
|
"""
|
|
Convert an image to a data URL.
|
|
|
|
Args:
|
|
image: The image data
|
|
width: Target width (negative means preserve original if under max width)
|
|
output_format: Output format (auto, jpeg, png)
|
|
|
|
Returns:
|
|
Data URL of the image
|
|
"""
|
|
image_data, mimetype = _normalize_to_bytes(image, width, output_format)
|
|
image_base64 = base64.b64encode(image_data).decode("utf-8")
|
|
return f"data:{mimetype};base64,{image_base64}"
|
|
|
|
def video_to_url(video_data):
|
|
"""
|
|
Convert video data to a data URL.
|
|
|
|
Args:
|
|
video_data: The video data
|
|
|
|
Returns:
|
|
Data URL of the video
|
|
"""
|
|
video_base64 = base64.b64encode(video_data).decode("utf-8")
|
|
return f"data:video/mp4;base64,{video_base64}"
|