总之我对 Paddle 自带的多语言不是太满意,毕竟我想一个配置直接无脑解决汉语+英语+日语多语言场景,这里就只能退一步先无脑用他们的汉英模型。模型下载部分可能需要注意下本地模型的加载:自己在开发机器时会发现模型名称按照文档写就可以,反正回自动下载;但如果要在非联网机器的话需要到他们的 model list 页面手动下载到本地解压,加载时也要注意目录结构。
在上面目录结构下,加载模型的代码如下
from paddleocr import PaddleOCR
ocr_model = PaddleOCR(
ocr_version="PP-OCRv3",
# det_model_dir="/home/andy/.paddleocr/whl/det/ml/Multilingual_PP-OCRv3_det_infer/", # multi-language
det_model_dir="/home/andy/.paddleocr/whl/det/ch/ch_PP-OCRv3_det_infer", # Chinese
# rec_model_dir="/home/andy/.paddleocr/whl/rec/ch/ch_ppocr_server_v2.0_rec_infer/", # chinese model on server
# rec_model_dir="/home/andy/.paddleocr/whl/rec/japan/japan_PP-OCRv3_rec_infer",
rec_model_dir="/home/andy/.paddleocr/whl/rec/ch/ch_PP-OCRv3_rec_infer", # Chinese model
# use_angle_cls=True,
# lang="ch",
) # need to run only once to download and load model into memory
使用模型的代码如下
try:
ocr_result = ocr_model.ocr(img=image_path, cls=False)
except:
print("Error: ", image_path)
continue
ocr_str = ""
ocr_full_result = []
for idx in range(len(ocr_result)):
for line in ocr_result[idx]:
ocr_str += (" " + line[1][0])
ocr_full_result.append(line)
ocr_str = ocr_str.strip()
注意在 CUDA 上跑的时候默认输出是 FP16,而 CPU 上的模型调用就只能输出 FP32 了,如果后面要在 CPU 上继续处理的话需要对齐一下类型
CLIP_MODEL = "ViT-L/14"
IMAGE_DIR = "/mnt/e/temp/iphone-import"
CLIP_MODEL_DOWNLOAD_ROOT = "/home/andy/dev/CLIP-dev/model"
def main():
IMAGE_PATH_LIST = sorted(glob(os.path.join(IMAGE_DIR, "*.*")))
IMAGE_PATH_LIST = [i for i in IMAGE_PATH_LIST if not (i.endswith(".GIF") or i.endswith(".MP4") or i.endswith("AAE"))]
print("Found {} images".format(len(IMAGE_PATH_LIST)))
device = "cuda"
model, preprocess = clip.load(CLIP_MODEL, device=device, download_root=CLIP_MODEL_DOWNLOAD_ROOT)
feature_dict = {}
for image_path in tqdm(IMAGE_PATH_LIST):
image_basename = os.path.basename(image_path)
image = preprocess(Image.open(image_path))
image = image.unsqueeze(0).to(device)
with torch.no_grad():
image_feat = model.encode_image(image)
image_feat = image_feat.detach().cpu().numpy() # It is still FP16 on CPU Here
feature_dict[image_basename] = image_feat
with gr.Blocks() as demo:
heading = gr.Markdown("# CLIP Image Search Demo")
# Use tabs
with gr.Tab("Using prompt and CLIP feature"):
prompt_textbox = gr.Textbox(lines=4, label="Prompt")
button_prompt = gr.Button("Search").style(size="lg")
with gr.Tab("Using image and CLIP feature"):
input_image = gr.Image(label="Image", type="pil")
button_image = gr.Button("Search").style(size="lg")
with gr.Accordion("Search options", open=False):
extension_choice = gr.CheckboxGroup(["jpg", "png", "gif"], label="extension", info="choose extension for search")
with gr.Row():
topn = gr.Number(value=64, label="topn")
minimum_width = gr.Number(value=0, label="minimum_width")
minimun_height = gr.Number(value=0, label="minimum_height")
minimun_importance = gr.Number(value=0, label="minimum_importance")
gr_gallery = gr.Gallery(label="results").style(grid=4, height=6)
button_prompt.click(submit, inputs=[prompt_textbox, topn, minimum_width, minimun_height, minimun_importance, extension_choice], outputs=[gr_gallery])
button_image.click(submit, inputs=[input_image, topn, minimum_width, minimun_height, minimun_importance, extension_choice], outputs=[gr_gallery])
demo.launch()
代码性能分析
这里分成两部分,一个是整个数据库初始化的操作,另一个是数据库建好之后的检索操作。
建库其实比较花时间,我在 Windows 10 的 WSL 下跑的 CUDA,具体数字不太记得,但大概是一小时跑 100k 数量图片的量级。我整个数据库全跑完大概要半天的时间。
检索的时间很大一部分来自于 MongoDB 查找(把百万量级的 feature 拿出来)。而算特征距离的速度,我实测每 8192 张图需要 20ms 左右,百万量级图片全算一遍距离需要 2~3s,估计还有不少优化空间。最后完整实测,一次全量搜索计算百万特征距然后找出 top N 在我 NAS 的 i3-10105T CPU 下需要 6~7s 的时间。
跑服务的话,因为 CLIP 模型要常驻内存,我用的 “ViT-L/14” 比较大,大概要占用 2~3GB 的内存,加上 mongodb 也要占用 1~2GB,不少 NAS 可能有些吃力,但 PC 一般问题不大。
效果展示与吐槽
首先是我发现 CLIP 模型是有一定 OCR 能力的。我先尝试搜索了一个 “anime screenshot with captions” 想搜一些动画截图。发现第一张图里面的英文文本是 “Your code’s buggy. We can’t sell this product.”,神奇的是竟然可以用这串文字也搜到同样的图。感觉 CLIP 模型也不算特别大,竟然已经训练出来还不错的 OCR 能力了。可惜就是 CLIP 对中文理解能力几乎为零完全没法用。