feat: Change OpenAI Image Format to JPG (#4117)

This commit is contained in:
Michael Genson 2024-08-30 16:24:25 -05:00 committed by GitHub
parent 2ad6e1b198
commit a3f474e088
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 34 additions and 13 deletions

View File

@ -9,9 +9,17 @@ from pillow_heif import register_avif_opener, register_heif_opener
register_heif_opener()
register_avif_opener()
WEBP = ".webp"
FORMAT = "WEBP"
@dataclass
class ImageFormat:
suffix: str
format: str
modes: list[str]
"""If the image is not in the correct mode, it will be converted to the first mode in the list"""
JPG = ImageFormat(".jpg", "JPEG", ["RGB"])
WEBP = ImageFormat(".webp", "WEBP", ["RGB", "RGBA"])
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".heic", ".avif"}
@ -57,24 +65,37 @@ class ABCMinifier(ABC):
return
for file in image.parent.glob("*.*"):
if file.suffix != WEBP:
if file.suffix != WEBP.suffix:
file.unlink()
class PillowMinifier(ABCMinifier):
@staticmethod
def to_webp(image_file: Path, dest: Path | None = None, quality: int = 100) -> Path:
def _convert_image(
image_file: Path, image_format: ImageFormat, dest: Path | None = None, quality: int = 100
) -> Path:
"""
Converts an image to the webp format in-place. The original image is not
removed By default, the quality is set to 100.
Converts an image to the specified format in-place. The original image is not
removed. By default, the quality is set to 100.
"""
img = Image.open(image_file)
dest = dest or image_file.with_suffix(WEBP)
img.save(dest, FORMAT, quality=quality)
img = Image.open(image_file)
if img.mode not in image_format.modes:
img = img.convert(image_format.modes[0])
dest = dest or image_file.with_suffix(image_format.suffix)
img.save(dest, image_format.format, quality=quality)
return dest
@staticmethod
def to_jpg(image_file: Path, dest: Path | None = None, quality: int = 100) -> Path:
return PillowMinifier._convert_image(image_file, JPG, dest, quality)
@staticmethod
def to_webp(image_file: Path, dest: Path | None = None, quality: int = 100) -> Path:
return PillowMinifier._convert_image(image_file, WEBP, dest, quality)
@staticmethod
def crop_center(pil_img: Image, crop_width=300, crop_height=300):
img_width, img_height = pil_img.size
@ -122,7 +143,7 @@ class PillowMinifier(ABCMinifier):
else:
img = Image.open(image_file)
tiny_image = PillowMinifier.crop_center(img)
tiny_image.save(tiny_dest, FORMAT, quality=70)
tiny_image.save(tiny_dest, WEBP.format, quality=70)
self._logger.info("Tiny image saved")
success = True

View File

@ -65,12 +65,12 @@ class OpenAILocalImage(OpenAIImageBase):
path: Path
def get_image_url(self) -> str:
image = img.PillowMinifier.to_webp(
self.path, dest=self.path.parent.joinpath(f"{self.filename}-min-original.webp")
image = img.PillowMinifier.to_jpg(
self.path, dest=self.path.parent.joinpath(f"{self.filename}-min-original.jpg")
)
with open(image, "rb") as f:
b64content = base64.b64encode(f.read()).decode("utf-8")
return f"data:image/webp;base64,{b64content}"
return f"data:image/jpg;base64,{b64content}"
class OpenAIService(BaseService):