Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading