211 lines
6.9 KiB
Python
211 lines
6.9 KiB
Python
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 pyinotify
|
|
import fnmatch
|
|
import glob
|
|
import subprocess
|
|
from collections import Counter
|
|
import matplotlib.pyplot as plt
|
|
import re
|
|
import numpy as np
|
|
|
|
WATCH_DIR = "/tmp/mote-ocr-screenshots/"
|
|
|
|
|
|
|
|
def contrast_stretch(gray):
|
|
p2, p98 = np.percentile(gray, (2, 98))
|
|
return cv2.normalize(gray, None, p2, p98, cv2.NORM_MINMAX)
|
|
|
|
|
|
def copy_to_clipboard(text):
|
|
subprocess.run(["xsel", "--clipboard", "--input"], input=text.encode(), check=True)
|
|
|
|
|
|
# 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
|
|
]
|
|
|
|
|
|
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
|
|
)
|
|
text = self._remove_specials(raw_text)
|
|
|
|
return text
|
|
|
|
|
|
def bicubic(img):
|
|
h, w = img.shape[:2]
|
|
M = np.array([[1, 0, 0],
|
|
[0, 1, 0]], dtype=np.float32)
|
|
return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC)
|
|
|
|
|
|
def adjust_gamma(image, gamma=1.0, alpha=1.0, beta=0):
|
|
"""
|
|
Adjust the gamma and contrast of a grayscale image.
|
|
|
|
Parameters:
|
|
image (numpy.ndarray): Input grayscale image.
|
|
gamma (float): Gamma correction factor. Values < 1 darken the image, > 1 brighten it.
|
|
alpha (float): Contrast control. 1.0 means no change, < 1.0 reduces contrast, > 1.0 increases contrast.
|
|
beta (int): Brightness control. Positive values brighten the image, negative values darken it.
|
|
|
|
Returns:
|
|
numpy.ndarray: Adjusted image.
|
|
"""
|
|
# Apply gamma correction
|
|
invGamma = 1.0 / gamma
|
|
table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
|
|
gamma_corrected = cv2.LUT(image, table)
|
|
|
|
# Adjust contrast and brightness
|
|
adjusted = cv2.convertScaleAbs(gamma_corrected, alpha=alpha, beta=beta)
|
|
|
|
return adjusted
|
|
|
|
|
|
|
|
def preprocess_image(image_path):
|
|
# Load image in grayscale (1-channel)
|
|
gray = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
|
resampled = bicubic(gray)
|
|
# Convert grayscale to 3-channel RGB by duplicating the gray channel
|
|
img1 = contrast_stretch(adjust_gamma(cv2.cvtColor(resampled, cv2.COLOR_GRAY2RGB),gamma=0.3,alpha=1.6,beta=-5)) # perfect tunning
|
|
|
|
img = cv2.fastNlMeansDenoising(img1, None, 30, 10, 20)
|
|
|
|
plt.imshow(img)
|
|
plt.axis("off")
|
|
plt.savefig("/tmp/output.png", bbox_inches="tight", pad_inches=0)
|
|
|
|
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] * 3
|
|
norm_vals = [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.squeeze(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 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}")
|
|
copy_to_clipboard(result)
|
|
# raw_text = decode_tokens(logits)
|
|
# return post_process(raw_text)
|
|
|
|
|
|
def get_most_recent_file(directory, pattern="*.png", recursive=True):
|
|
search_path = os.path.join(directory, "**", pattern) if recursive else os.path.join(directory, pattern)
|
|
files = glob.glob(search_path, recursive=recursive)
|
|
return max(files, key=os.path.getctime) if files else None
|
|
|
|
|
|
class EventHandler(pyinotify.ProcessEvent):
|
|
def process_IN_CLOSE_WRITE(self, event):
|
|
if fnmatch.fnmatch(event.pathname, "*.png"):
|
|
recent_img = get_most_recent_file(WATCH_DIR, pattern="*.png", recursive=True)
|
|
print("PROCESSED OUT:", main(recent_img))
|
|
|
|
if __name__ == "__main__":
|
|
wm = pyinotify.WatchManager()
|
|
handler = EventHandler()
|
|
notifier = pyinotify.Notifier(wm, handler)
|
|
mask = pyinotify.IN_CLOSE_WRITE # triggers when a file is written and closed
|
|
|
|
wm.add_watch(WATCH_DIR, mask, rec=True, auto_add=True)
|
|
print(f"Watching {WATCH_DIR} for new PNG files...")
|
|
|
|
notifier.loop()
|