Update inference_test.py

inference_test.py

@@ -1,34 +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
+import os
+from pathlib import Path
+from typing import Optional, Union
+
+import re
+import numpy as np
+
+def get_most_recent_file(
+    directory: Union[str, Path],
+    pattern: str = "*",
+    recursive: bool = False,
+    return_path: bool = True
+) -> Optional[Union[str, Path]]:
+   
+    dir_path = Path(directory)
+    
+    if not dir_path.exists():
+        raise FileNotFoundError(f"Directory not found: {directory}")
+    if not dir_path.is_dir():
+        raise ValueError(f"Path is not a directory: {directory}")
+
+    # Get all matching files
+    if recursive:
+        files = list(dir_path.rglob(pattern))
+    else:
+        files = list(dir_path.glob(pattern))
+    
+    # Filter out directories
+    files = [f for f in files if f.is_file()]
+    
+    if not files:
+        return None
+    
+    # Find most recent file by modification time
+    latest = max(files, key=lambda f: f.stat().st_mtime)
+    
+    return str(latest) if return_path else latest
+
+
+class TextProcessor:
+    def __init__(self, vocab_path: str):
+        self.vocab = self._load_vocab(vocab_path)
+        self.unk_token = "[UNK]"
+        self.special_token_pattern = re.compile(r'\[.*?\]')  # Pattern to match [ANYTHING]
+        
+    def _load_vocab(self, path: str) -> list:
+        """Load vocab with 1-based indexing"""
+        try:
+            with open(path, 'r', encoding='utf-8') as f:
+                return [line.strip() for line in f]
+        except FileNotFoundError:
+            raise RuntimeError(f"Vocab file not found: {path}")
+
+    def _remove_specials(self, text: str) -> str:
+        """Remove special tokens and consecutive duplicates"""
+        # Filter out special tokens
+        return self.special_token_pattern.sub('', text)
+        
+
+    def decode(self, indices) -> str:
+        """
+        Convert indices to cleaned text
+        """
+        if isinstance(indices, (int, np.integer)):
+            indices = [indices]
+            
+        raw_text = ''.join(
+            self.vocab[idx] if 1 <= idx <= len(self.vocab) else self.unk_token
+            for idx in indices
+        )
+        almost_done = self._remove_specials(raw_text)
+        final_text = post_process(almost_done)
+        return final_text
+
 
 
 def preprocess_image(image_path):
@@ -50,19 +99,21 @@ def run_inference(mat_in, model_path="model"):
 
     ex = net.create_extractor()
     ex.input("in0", mat_in)
-    ex.input("in1", ncnn.Mat(np.full((300, 1), 0, dtype=np.int64)))
+    ex.input("in1", ncnn.Mat(np.full((1, 300), 0, dtype=np.int64)))
 
     try:
         ret, out0 = ex.extract("out0")
         print("Output shape:", out0.shape)
-        return np.array(out0)
+        predictions = np.array(out0)
+        indexed = np.argmax(predictions, axis=-1)
+        print(indexed)
+        return indexed
     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)
 
 
@@ -75,9 +126,13 @@ def post_process(text):
 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)
-
+    processor = TextProcessor("model/vocab.txt")
+    result = processor.decode(logits)
+    print(f"Decoded text: {result}")
+#    raw_text = decode_tokens(logits)
+#    return post_process(raw_text)
+   
 
 if __name__ == "__main__":
-    print("PROCESSED OUT:", main("./manga_sample.jpg"))
+    recent_img = get_most_recent_file("/home/hashirama/screenshots/", pattern="*.png", recursive=True)
+    print("PROCESSED OUT:", main(recent_img))