File size: 3,432 Bytes
f24f2e7 c825110 f24f2e7 c825110 f24f2e7 c825110 f24f2e7 c825110 f24f2e7 c825110 f24f2e7 c825110 f24f2e7 c825110 f24f2e7 c825110 f24f2e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import sys
import torch
import onnx
import onnxruntime as rt
from torchvision import transforms as T
from tokenizer_base import Tokenizer
import pathlib
from PIL import Image
from huggingface_hub import Repository
class DocumentParserModel:
def __init__(
self,
repo_path,
model_subpath,
img_size,
charset,
repo_url="stevenchang/captcha",
token=None,
):
self.repo_path = pathlib.Path(repo_path).resolve()
self.model_path = self.repo_path / model_subpath
self.charset = charset
self.tokenizer_base = Tokenizer(self.charset)
self.initialize_repository(repo_url, token)
self.transform = self.create_transform_pipeline(img_size)
self.ort_session = self.initialize_onnx_model(str(self.model_path))
def initialize_repository(self, repo_url, token):
if not self.model_path.exists():
if not self.repo_path.exists():
print(
f"Repository does not exist. Cloning from {repo_url} into {self.repo_path}"
)
repo = Repository(
local_dir=str(self.repo_path),
clone_from=repo_url,
use_auth_token=token if token else True,
)
else:
print(
f"Model does not exist, but repository is already cloned. Pulling latest changes in {self.repo_path}"
)
repo = Repository(
local_dir=str(self.repo_path),
use_auth_token=token if token else True,
)
repo.git_pull()
else:
print(
f"Model {self.model_path} already exists, skipping repository update."
)
def create_transform_pipeline(self, img_size):
transforms = [
T.Resize(img_size, T.InterpolationMode.BICUBIC),
T.ToTensor(),
T.Normalize(0.5, 0.5),
]
return T.Compose(transforms)
def initialize_onnx_model(self, model_path):
onnx_model = onnx.load(model_path)
onnx.checker.check_model(onnx_model)
return rt.InferenceSession(model_path)
def predict_text(self, image_path):
try:
with Image.open(image_path) as img_org:
x = self.transform(img_org.convert("RGB")).unsqueeze(0)
ort_inputs = {self.ort_session.get_inputs()[0].name: x.cpu().numpy()}
logits = self.ort_session.run(None, ort_inputs)[0]
probs = torch.tensor(logits).softmax(-1)
preds, _ = self.tokenizer_base.decode(probs)
return preds[0]
except IOError:
print(f"Error: Cannot open image {image_path}")
return None
if __name__ == "__main__":
import sys
repo_path = "secret_models"
model_subpath = "captcha.onnx"
img_size = (32, 128)
charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
doc_parser = DocumentParserModel(
repo_path=repo_path,
model_subpath=model_subpath,
img_size=img_size,
charset=charset,
)
if len(sys.argv) > 1:
image_path = sys.argv[1]
result = doc_parser.predict_text(image_path)
print(result)
else:
print("Please provide an image path.")
|