Update app.py

app.py

@@ -34,40 +34,30 @@ IMAGE_ID_PATTERN = r"(?P<image_id>[-a-z0-9]{36})"
 CONFIDENCE_PATTERN = r"(?P<confidence>[0-9.]+)"  # For line
 TEXT_PATTERN = r"\s*(?P<text>.*)\s*"
 LINE_PREDICTION = re.compile(rf"{IMAGE_ID_PATTERN} {CONFIDENCE_PATTERN} {TEXT_PATTERN}")
-models_name = ["johnlockejrr/pylaia-heb_sam_v1"]
-MODELS = {}
 
 def get_width(image, height=DEFAULT_HEIGHT):
     aspect_ratio = image.width / image.height
     return height * aspect_ratio
 
-def load_model(model_name):
-    if model_name not in MODELS:
-        MODELS[model_name] = Path(snapshot_download(model_name))
-    return MODELS[model_name]
-
 def predict(model_name, input_img):
-    model_dir = load_model(model_name)
-
+    model_dir = 'pylaia-mcdonald_v2'
     temperature = 2.0
     batch_size = 1
 
-    weights_path = model_dir / "weights.ckpt"
-    syms_path = model_dir / "syms.txt"
+    weights_path = f"{model_dir}/weights.ckpt"
+    syms_path = f"{model_dir}/syms.txt"
     language_model_params = {"language_model_weight": 1.0}
-    use_language_model = (model_dir / "tokens.txt").exists()
+    use_language_model = True
     if use_language_model:
-        language_model_params.update(
-            {
-	            "language_model_path": str(model_dir / "language_model.binary"),
-                "lexicon_path": str(model_dir / "lexicon.txt"),
-                "tokens_path": str(model_dir / "tokens.txt"),
-            }
-        )
+        language_model_params.update({
+            "language_model_path": f"{model_dir}/language_model.binary",
+            "lexicon_path": f"{model_dir}/lexicon.txt",
+            "tokens_path": f"{model_dir}/tokens.txt",
+        })
 
     common_args = CommonArgs(
-        checkpoint=str(weights_path.relative_to(model_dir)),
-        train_path=str(model_dir),
+        checkpoint="weights.ckpt",
+        train_path=f"{model_dir}",
         experiment_dirname="",
     )
 
@@ -128,7 +118,7 @@ def process_image(image):
         crop_pil = Image.fromarray(cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB))
 
         # Recognize text using PyLaia model
-        predicted = predict(models_name, crop_pil)
+        predicted = predict('johnlockejrr/pylaia-samaritan_v1', crop_pil)
         texts.append(predicted[1]["text"])
 
         bboxes.append((x1, y1, x2, y2))