yschneider commited on
Commit
6b1f545
1 Parent(s): 37527d6

Cache download model location, load images in L mode

Browse files
Files changed (1) hide show
  1. app.py +10 -2
app.py CHANGED
@@ -16,7 +16,7 @@ CONFIDENCE_PATTERN = r"(?P<confidence>[0-9.]+)" # For line
16
  TEXT_PATTERN = r"\s*(?P<text>.*)\s*"
17
  LINE_PREDICTION = re.compile(rf"{IMAGE_ID_PATTERN} {CONFIDENCE_PATTERN} {TEXT_PATTERN}")
18
  models_name = ["Teklia/pylaia-rimes"]
19
-
20
  DEFAULT_HEIGHT = 128
21
 
22
 
@@ -25,8 +25,14 @@ def get_width(image, height=DEFAULT_HEIGHT):
25
  return height * aspect_ratio
26
 
27
 
 
 
 
 
 
 
28
  def predict(model_name, input_img):
29
- model_dir = Path(snapshot_download(model_name))
30
 
31
  temperature = 2.0
32
  batch_size = 1
@@ -103,6 +109,7 @@ gradio_app = gr.Interface(
103
  type="pil",
104
  height=DEFAULT_HEIGHT,
105
  width=2000,
 
106
  ),
107
  ],
108
  outputs=[
@@ -114,6 +121,7 @@ gradio_app = gr.Interface(
114
  for filename in Path("examples").iterdir()
115
  ],
116
  title="Decode the transcription of an image using a PyLaia model",
 
117
  )
118
 
119
  if __name__ == "__main__":
 
16
  TEXT_PATTERN = r"\s*(?P<text>.*)\s*"
17
  LINE_PREDICTION = re.compile(rf"{IMAGE_ID_PATTERN} {CONFIDENCE_PATTERN} {TEXT_PATTERN}")
18
  models_name = ["Teklia/pylaia-rimes"]
19
+ MODELS = {}
20
  DEFAULT_HEIGHT = 128
21
 
22
 
 
25
  return height * aspect_ratio
26
 
27
 
28
+ def load_model(model_name):
29
+ if model_name not in MODELS:
30
+ MODELS[model_name] = Path(snapshot_download(model_name))
31
+ return MODELS[model_name]
32
+
33
+
34
  def predict(model_name, input_img):
35
+ model_dir = load_model(model_name)
36
 
37
  temperature = 2.0
38
  batch_size = 1
 
109
  type="pil",
110
  height=DEFAULT_HEIGHT,
111
  width=2000,
112
+ image_mode="L",
113
  ),
114
  ],
115
  outputs=[
 
121
  for filename in Path("examples").iterdir()
122
  ],
123
  title="Decode the transcription of an image using a PyLaia model",
124
+ cache_examples=True,
125
  )
126
 
127
  if __name__ == "__main__":