Add inference_test.py

inference_test.py

@@ -0,0 +1,83 @@
+""" 
+---------------------------------------------------------------
+Minimal Python Script for Testing Manga-OCR Model (NCNN Version)
+© All Rights Reserved by MW Enterprise and Affiliates
+
+Credits:
+Model conversion courtesy of AJATiX
+
+漫画OCR模型测试脚本(NCNN转换版)
+腾讯模型转换 - 薄荷糖企业及关联公司版权所有
+
+致谢:
+模型转换由AJATiX提供
+---------------------------------------------------------------
+
+
+Development Roadmap:
+
+    Implement our own token decoder (dependency reduction)
+
+    Optimize image preprocessing to get better results
+
+    C++ migration with cross-platform clipboard image detection"
+
+"""
+import numpy as np
+import ncnn
+import cv2
+import re
+import jaconv
+from transformers import AutoTokenizer
+
+
+def preprocess_image(image_path):
+    img = cv2.imread(image_path)
+    mat_in = ncnn.Mat.from_pixels_resize(
+        img, ncnn.Mat.PixelType.PIXEL_RGB, img.shape[1], img.shape[0], 224, 224
+    )
+    mean_vals, norm_vals = [0.5] * 3, [1 / 255.0] * 3
+    mat_in.substract_mean_normalize(mean_vals, norm_vals)
+    return mat_in
+
+
+def run_inference(mat_in, model_path="model"):
+    net = ncnn.Net()
+    net.opt.num_threads = 16
+    net.opt.use_vulkan_compute = False
+    net.load_param(f"{model_path}/ncnn_model.param")
+    net.load_model(f"{model_path}/ncnn_model.bin")
+
+    ex = net.create_extractor()
+    ex.input("in0", mat_in)
+    ex.input("in1", ncnn.Mat(np.full((300, 1), 0, dtype=np.int64)))
+
+    try:
+        ret, out0 = ex.extract("out0")
+        print("Output shape:", out0.shape)
+        return np.array(out0)
+    finally:
+        del ex, net
+
+
+def decode_tokens(logits, tokenizer_path="model"):
+    tkz = AutoTokenizer.from_pretrained(tokenizer_path)
+    token_ids = np.argmax(logits, axis=-1)
+    return tkz.decode(token_ids, skip_special_tokens=True)
+
+
+def post_process(text):
+    text = "".join(text.split()).replace("…", "...")
+    text = re.sub("[・.]{2,}", lambda x: (x.end() - x.start()) * ".", text)
+    return jaconv.h2z(text, ascii=True, digit=True)
+
+
+def main(image_path):
+    mat_in = preprocess_image(image_path)
+    logits = run_inference(mat_in)
+    raw_text = decode_tokens(logits)
+    return post_process(raw_text)
+
+
+if __name__ == "__main__":
+    print("PROCESSED OUT:", main("./manga_sample.jpg"))