toshas's picture
fix rotation due to invalid exif check order
59d338e
from __future__ import annotations
import warnings
from pathlib import Path
from typing import Literal, cast
import gradio
from gradio.image_utils import decode_base64_to_file, decode_base64_to_image, decode_base64_to_image_array, format_image
import numpy as np
import PIL.Image
from PIL import ImageOps
from gradio import processing_utils
from gradio.data_classes import ImageData
from gradio.exceptions import Error
PIL.Image.init() # fixes https://github.com/gradio-app/gradio/issues/2843 (remove when requiring Pillow 9.4+)
def patched_preprocess_image(
payload: ImageData | None,
cache_dir: str,
format: str,
image_mode: Literal[
"1", "L", "P", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"
]
| None,
type: Literal["numpy", "pil", "filepath"],
) -> np.ndarray | PIL.Image.Image | str | None:
if payload is None:
return payload
if payload.url and payload.url.startswith("data:"):
if type == "pil":
print("Preprocessing payload as PIL image")
return decode_base64_to_image(payload.url)
elif type == "numpy":
print("Preprocessing payload as numpy array")
return decode_base64_to_image_array(payload.url)
elif type == "filepath":
print("Preprocessing payload as file path")
return decode_base64_to_file(payload.url, cache_dir, format)
if payload.path is None:
raise ValueError("Image path is None.")
file_path = Path(payload.path)
if payload.orig_name:
p = Path(payload.orig_name)
name = p.stem
suffix = p.suffix.replace(".", "")
if suffix in ["jpg", "jpeg"]:
suffix = "jpeg"
else:
name = "image"
suffix = "webp"
if suffix.lower() == "svg":
if type == "filepath":
return str(file_path)
raise Error("SVG files are not supported as input images for this app.")
im = PIL.Image.open(file_path)
exif = im.getexif()
if suffix.lower() in ["heif", "heic"] and type == "filepath":
im = im.convert("RGB")
if exif.get(274, 1) != 1 and hasattr(ImageOps, "exif_transpose"):
# 274 is the code for image rotation and 1 means "correct orientation"
try:
im = ImageOps.exif_transpose(im)
Path(file_path).resolve().write_bytes(processing_utils.encode_pil_to_bytes(im, format="webp"))
except Exception:
warnings.warn(f"Failed to transpose image {file_path} based on EXIF data.")
max_dim = max(im.width, im.height)
if max_dim > 2048:
scale = min(1.0, 2048 / max_dim)
im = im.resize((round(im.width * scale), round(im.height * scale)), PIL.Image.BILINEAR)
Path(file_path).resolve().write_bytes(processing_utils.encode_pil_to_bytes(im, format="webp"))
if type == "filepath" and (image_mode in [None, im.mode]):
return str(file_path)
if suffix.lower() != "gif" and im is not None:
with warnings.catch_warnings():
warnings.simplefilter("ignore")
if image_mode is not None:
im = im.convert(image_mode)
return format_image(
im,
type=cast(Literal["numpy", "pil", "filepath"], type),
cache_dir=cache_dir,
name=name,
format=suffix,
)
gradio.image_utils.preprocess_image = patched_preprocess_image