import argparse import csv import os from PIL import Image from utils import load_model import jax.numpy as jnp from jax import jit from tqdm import tqdm def main(args): root = args.image_path files = list(os.listdir(root)) for f in files: assert(f[-4:] == ".jpg") for model_name in ["koclip-base", "koclip-large"]: model, processor = load_model(f"koclip/{model_name}") with tqdm(total=len(files)) as pbar: for counter in range(0, len(files), args.batch_size): images = [] image_ids = [] for idx in range(counter, min(len(files), counter + args.batch_size)): file_ = files[idx] image = Image.open(os.path.join(root, file_)).convert('RGB') images.append(image) image_ids.append(file_) pbar.update(args.batch_size) try: inputs = processor(text=[""], images=images, return_tensors="jax", padding=True) except: print(image_ids) break inputs['pixel_values'] = jnp.transpose(inputs['pixel_values'], axes=[0, 2, 3, 1]) features = model(**inputs).image_embeds with open(os.path.join(args.out_path, f"{model_name}.tsv"), "a+") as f: writer = csv.writer(f, delimiter="\t") for image_id, feature in zip(image_ids, features): writer.writerow([image_id, ",".join(map(lambda x: str(x), feature))]) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--batch_size", default=16) parser.add_argument("--image_path", default="images") parser.add_argument("--out_path", default="features") args = parser.parse_args() main(args)