Spaces:
Runtime error
Runtime error
File size: 833 Bytes
944c93a 077dc3f 944c93a 4388025 077dc3f 4388025 077dc3f 4388025 bd2d69d 4388025 bd2d69d 077dc3f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
import os
from PIL import Image
import numpy as np
from tqdm import tqdm
class AbstractImageEmbedder:
def __init__(self, device: str = "cpu"):
self.device = device
def embed(self, image: Image) -> np.ndarray:
"""Embed an image
"""
raise NotImplementedError
def embed_folder(self, folder_path: str, output_path: str) -> None:
"""Embed all images in a folder and save them in a .npy file
"""
assert output_path.endswith(".npy"), "`output_path` must end with .npy"
embeddings = {}
for name in tqdm(os.listdir(folder_path)):
image_path = os.path.join(folder_path, name)
image = Image.open(image_path)
embedding = self.embed(image)
embeddings[name] = embedding
np.save(output_path, embeddings)
|