|
import os |
|
from pyclip import Clip, enum_devices, sys_init, sys_deinit, ClipDeviceType |
|
import cv2 |
|
import glob |
|
import argparse |
|
import tqdm |
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--ienc', type=str, default='cnclip/cnclip_vit_l14_336px_vision_u16u8.axmodel') |
|
parser.add_argument('--tenc', type=str, default='cnclip/cnclip_vit_l14_336px_text_u16.axmodel') |
|
parser.add_argument('--vocab', type=str, default='cnclip/cn_vocab.txt') |
|
parser.add_argument('--isCN', type=int, default=1) |
|
parser.add_argument('--db_path', type=str, default='clip_feat_db_coco') |
|
parser.add_argument('--image_folder', type=str, default='coco_1000') |
|
args = parser.parse_args() |
|
|
|
image_folder = args.image_folder |
|
|
|
|
|
print("可用设备:", enum_devices()) |
|
|
|
|
|
sys_init(ClipDeviceType.axcl_device, 0) |
|
|
|
try: |
|
|
|
clip = Clip({ |
|
'text_encoder_path': args.tenc, |
|
'image_encoder_path': args.ienc, |
|
'tokenizer_path': args.vocab, |
|
'db_path': args.db_path, |
|
'isCN': args.isCN |
|
}) |
|
|
|
|
|
|
|
image_files = glob.glob(os.path.join(image_folder, '*.jpg')) |
|
for image_file in tqdm.tqdm(image_files): |
|
img = cv2.imread(image_file) |
|
cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) |
|
filename = os.path.basename(image_file) |
|
clip.add_image(filename, img) |
|
|
|
|
|
results = clip.match_text('dog', top_k=10) |
|
print("匹配结果:", results) |
|
|
|
finally: |
|
|
|
sys_deinit(ClipDeviceType.axcl_device, 0) |