Skip to content

Commit

Permalink
fix: fix pad to square for some rectangular images (#421)
Browse files Browse the repository at this point in the history
  • Loading branch information
joein authored Dec 12, 2024
1 parent 7627d78 commit 3a847f6
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 9 deletions.
22 changes: 14 additions & 8 deletions fastembed/image/transform/functional.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Sized, Union, Optional
from typing import Sized, Union

import numpy as np
from PIL import Image
Expand Down Expand Up @@ -127,18 +127,24 @@ def pil2ndarray(image: Union[Image.Image, np.ndarray]):
def pad2square(
image: Image.Image,
size: int,
fill_color: Optional[Union[str, int, tuple[int, ...]]] = None,
fill_color: Union[str, int, tuple[int, ...]] = 0,
) -> Image.Image:
height, width = image.height, image.width

# if the size is larger than the new canvas
if width > size or height > size:
left, right = 0, width
top, bottom = 0, height

crop_required = False
if width > size:
left = (width - size) // 2
top = (height - size) // 2
right = left + size
crop_required = True

if height > size:
top = (height - size) // 2
bottom = top + size
image = image.crop((left, top, right, bottom))
crop_required = True

new_image = Image.new(mode="RGB", size=(size, size), color=fill_color or 0)
new_image.paste(image)
new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
return new_image
2 changes: 1 addition & 1 deletion fastembed/text/text_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def __init__(
return

raise ValueError(
f"Model {model_name} is not supported in TextEmbedding."
f"Model {model_name} is not supported in TextEmbedding. "
"Please check the supported models using `TextEmbedding.list_supported_models()`"
)

Expand Down

0 comments on commit 3a847f6

Please sign in to comment.