aki-0421's picture
Upload folder using huggingface_hub
131c541 verified
import math
from PIL import Image
from sentence_transformers import SentenceTransformer
def resize_image_for_patch(image: Image.Image, patch_size: int = 14, max_patches: int = 400) -> Image.Image:
orig_width, orig_height = image.size
aspect_ratio = orig_width / orig_height
# Max width and height in pixels under the patch constraint
max_total_pixels = patch_size * math.sqrt(max_patches)
if aspect_ratio >= 1:
# Landscape or square orientation
target_width = patch_size * int(math.floor(math.sqrt(max_patches * aspect_ratio)))
target_height = int(target_width / aspect_ratio)
else:
# Portrait orientation
target_height = patch_size * int(math.floor(math.sqrt(max_patches / aspect_ratio)))
target_width = int(target_height * aspect_ratio)
# Ensure dimensions are multiples of patch_size
target_width -= target_width % patch_size
target_height -= target_height % patch_size
return image.resize((target_width, target_height), Image.BICUBIC)
# Init model
model = SentenceTransformer("./", device="cuda")
images = [
resize_image_for_patch(Image.open("/home/aki0421/Share/images/00085.png"))
]
image_embeddings = model.encode(images, convert_to_tensor=True)
sentences = [
"女の子が悲しんでいる。",
"落ち込んでる人",
"泣いている",
"笑っている",
"ピンクの髪の女の子",
"赤い髪の女の子",
"茶色の髪の女の子",
"赤い目",
"青い目",
"曇っている",
"雨が降っている",
"晴れている",
"キッチンにいます。",
"学校にいる",
"魔法少女のようだ",
"戦闘しますか?",
"男性ですか?",
"茶色い髪の女の子が悲しんでいるシーン",
"ピンクの髪の女の子が笑っているシーン"
]
text_embeddings = model.encode(sentences, convert_to_tensor=True)
similarities = model.similarity(text_embeddings, image_embeddings)
print(similarities)