add squeeze 2/2
This commit is contained in:
parent
c23d9b04bc
commit
f77a83fca9
GPG key ID:
53E62470A86BC185
1 changed files with 162 additions and 168 deletions
|
|
@ -1,168 +1,162 @@
|
|||
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, List
|
||||
|
||||
|
||||
import re
|
||||
import numpy as np
|
||||
|
||||
|
||||
# we we're kinda forced to implement the argmax by hand because of two factors:
|
||||
# 1- the argmax from this model always takes the last axis, and this is not possible in numcpp , it only goes right or left.
|
||||
# 2- numpy's implementation of argmax was too obscure , so we couldn't replicate it in c++ .
|
||||
# well, now we can rewrite in c++, and also we dont need to give numpy any credits for this.
|
||||
def argmax_last_axis(mat: List[List[int]]) -> List[int]:
|
||||
"""Return first max index for each row in matrices of ANY size/shape."""
|
||||
return [
|
||||
max(
|
||||
((val, idx) for idx, val in enumerate(row)),
|
||||
default=(-1, -1)
|
||||
)[1]
|
||||
for row in mat
|
||||
]
|
||||
|
||||
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
|
||||
cleaned = self.special_token_pattern.sub('', text)
|
||||
|
||||
# Remove consecutive duplicates
|
||||
if not cleaned:
|
||||
return ''
|
||||
|
||||
result = [cleaned[0]]
|
||||
for char in cleaned[1:]:
|
||||
if char != result[-1]:
|
||||
result.append(char)
|
||||
return ''.join(result)
|
||||
|
||||
|
||||
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):
|
||||
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 = [0.5, 0.5, 0.5] * 2
|
||||
norm_vals = [1/255, 1/255, 1/255] * 2
|
||||
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 = 32
|
||||
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((1, 300), 0, dtype=np.int64)))
|
||||
|
||||
try:
|
||||
ret, out0 = ex.extract("out0")
|
||||
print("Output shape:", out0.shape)
|
||||
predictions = np.array(out0)
|
||||
# ids = np.argmax(predictions, axis=-1)
|
||||
ids = argmax_last_axis(predictions)
|
||||
print("argmax'ed: ", ids)
|
||||
return ids
|
||||
finally:
|
||||
del ex, net
|
||||
|
||||
|
||||
def decode_tokens(logits, tokenizer_path="model"):
|
||||
tkz = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
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)
|
||||
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__":
|
||||
recent_img = get_most_recent_file("/home/hashirama/screenshots/", pattern="*.png", recursive=True)
|
||||
print("PROCESSED OUT:", main(recent_img))
|
||||
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, List
|
||||
|
||||
|
||||
import re
|
||||
import numpy as np
|
||||
|
||||
|
||||
# we we're kinda forced to implement the argmax by hand because of two factors:
|
||||
# 1- the argmax from this model always takes the last axis, and this is not possible in numcpp , it only goes right or left.
|
||||
# 2- numpy's implementation of argmax was too obscure , so we couldn't replicate it in c++ .
|
||||
# well, now we can rewrite in c++, and also we dont need to give numpy any credits for this.
|
||||
def argmax_last_axis(mat: List[List[int]]) -> List[int]:
|
||||
"""Return first max index for each row in matrices of ANY size/shape."""
|
||||
return [
|
||||
max(
|
||||
((val, idx) for idx, val in enumerate(row)),
|
||||
default=(-1, -1)
|
||||
)[1]
|
||||
for row in mat
|
||||
]
|
||||
|
||||
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
|
||||
cleaned = self.special_token_pattern.sub('', text)
|
||||
|
||||
# Remove consecutive duplicates
|
||||
return cleaned
|
||||
|
||||
|
||||
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):
|
||||
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 = [0.5, 0.5, 0.5] * 3
|
||||
norm_vals = [1/255, 1/255, 1/255] * 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 = 32
|
||||
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((1, 300), 0, dtype=np.int64)))
|
||||
|
||||
try:
|
||||
ret, out0 = ex.extract("out0")
|
||||
print("Output shape:", out0.shape)
|
||||
predictions = np.array(out0)
|
||||
# ids = np.argmax(predictions, axis=-1)
|
||||
ids = argmax_last_axis(predictions)
|
||||
print("argmax'ed: ", ids)
|
||||
return np.squeeze(ids)
|
||||
finally:
|
||||
del ex, net
|
||||
|
||||
|
||||
def decode_tokens(logits, tokenizer_path="model"):
|
||||
tkz = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
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)
|
||||
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__":
|
||||
recent_img = get_most_recent_file("/home/hashirama/screenshots/", pattern="*.png", recursive=True)
|
||||
print("PROCESSED OUT:", main(recent_img))
|
||||
|
|
|
|||
Loading…
Reference in a new issue