First step refactoring

app.py

@@ -1,69 +1,101 @@
+import sys
 import torch
 import onnx
 import onnxruntime as rt
 from torchvision import transforms as T
 from tokenizer_base import Tokenizer
 import pathlib
-import os
-import sys
 from PIL import Image
-
 from huggingface_hub import Repository
 
-repo = Repository(
-    local_dir="secret_models",
-    repo_type="model",
-    clone_from="docparser/captcha",
-    token=True
-)
-repo.git_pull()
-
-cwd = pathlib.Path(__file__).parent.resolve()
-img_size = (32, 128)
-charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
-tokenizer_base = Tokenizer(charset)
-
 
-def get_transform(img_size):
-    transforms = []
-    transforms.extend([
-        T.Resize(img_size, T.InterpolationMode.BICUBIC),
-        T.ToTensor(),
-        T.Normalize(0.5, 0.5)
-    ])
-    return T.Compose(transforms)
+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 to_numpy(tensor):
-    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
+    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 initialize_model(model_file):
-    transform = get_transform(img_size)
-    onnx_model = onnx.load(model_file)
-    onnx.checker.check_model(onnx_model)
-    ort_session = rt.InferenceSession(model_file)
-    return transform, ort_session
+    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
 
 
-def get_text(image_path):
-    img_org = Image.open(image_path)
-    # Preprocess. Model expects a batch of images with shape: (B, C, H, W)
-    x = transform(img_org.convert('RGB')).unsqueeze(0)
-
-    # compute ONNX Runtime output prediction
-    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
-    logits = ort_session.run(None, ort_inputs)[0]
-    probs = torch.tensor(logits).softmax(-1)
-    preds, probs = tokenizer_base.decode(probs)
-    preds = preds[0]
-    return preds
-
+if __name__ == "__main__":
+    import sys
 
-model_file = os.path.join(cwd, "secret_models", "captcha.onnx")
-transform, ort_session = initialize_model(model_file=model_file)
+    repo_path = "secret_models"
+    model_subpath = "captcha.onnx"
+    img_size = (32, 128)
+    charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
 
-if __name__ == "__main__":
-    image_path = sys.argv[1]
-    res = get_text(image_path)
-    print(res)
+    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.")