|
import argparse |
|
import csv |
|
import numpy as np |
|
import os |
|
|
|
from config import MODEL_LIST |
|
from utils import load_model |
|
|
|
|
|
def main(args): |
|
caption_txt_path = args.text_path |
|
f = open(caption_txt_path) |
|
captions = [sent.strip() for sent in f.readlines()] |
|
|
|
for model_name in MODEL_LIST: |
|
model, processor = load_model(f"koclip/{model_name}") |
|
captions_processed = [processor(sent,images=None,return_tensors='jax') for sent in captions] |
|
vec = [np.asarray(model.get_text_features(**c)) for c in captions_processed] |
|
|
|
with open(os.path.join(args.out_path, f"{model_name}.tsv"), "a+") as f: |
|
writer = csv.writer(f, delimiter="\t") |
|
for text, feature in zip(captions, vec): |
|
writer.writerow([text, ",".join(map(lambda x: str(x), feature))]) |
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--text_path", default="cifar100_captions_kr.txt") |
|
parser.add_argument("--out_path", default="features/text") |
|
args = parser.parse_args() |
|
main(args) |