INNER CODE UNIT · Python
encode_images
patrickjohncyh/fashion-clip · fashion_clip/fashion_clip.py:177
def encode_images(self, images: Union[List[str], List[PIL.Image.Image]], batch_size: int):
def transform_fn(el):
imgs = el['image'] if isinstance(el['image'][0], PIL.Image.Image) else [Image().decode_example(_) for _ in el['image']]
return self.preprocess(images=imgs, return_tensors='pt')
dataset = Dataset.from_dict({'image': images})
dataset = dataset.cast_column('image',Image(decode=False)) if isinstance(images[0], str) else dataset
# dataset = dataset.map(map_fn,
# batched=True,
# remove_columns=['image'])
dataset.set_format('torch')
dataset.set_transform(transform_fn)
dataloader = DataLoader(dataset, batch_size=batch_size)
image_embeddings = []
pbar = tqdm(total=len(images) // batch_size, position=0)
with torch.no_grad():
for batch in dataloader:
batch = {k:v.to(self.device) for k,v in batch.items()}