decimer 2.7.1__tar.gz → 2.8.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (31) hide show
  1. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Efficient_Net_encoder.py +0 -2
  2. decimer-2.8.0/DECIMER/__init__.py +31 -0
  3. decimer-2.8.0/DECIMER/config.py +148 -0
  4. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/decimer.py +44 -4
  5. decimer-2.7.1/DECIMER/config.py → decimer-2.8.0/DECIMER/pre_process.py +8 -189
  6. decimer-2.8.0/DECIMER/utils.py +143 -0
  7. {decimer-2.7.1 → decimer-2.8.0}/PKG-INFO +85 -59
  8. {decimer-2.7.1 → decimer-2.8.0}/README.md +59 -50
  9. {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/PKG-INFO +85 -59
  10. {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/SOURCES.txt +4 -1
  11. {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/requires.txt +1 -1
  12. decimer-2.8.0/pyproject.toml +3 -0
  13. {decimer-2.7.1 → decimer-2.8.0}/setup.py +7 -10
  14. decimer-2.8.0/tests/test_functions.py +30 -0
  15. decimer-2.7.1/DECIMER/__init__.py +0 -29
  16. decimer-2.7.1/DECIMER/utils.py +0 -97
  17. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/DECIMER_EfficinetNetV2_Transfomer_Trainer.py +0 -0
  18. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Predictor_usingCheckpoints.py +0 -0
  19. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Repack_model.py +0 -0
  20. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Transformer_decoder.py +0 -0
  21. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Utils/Create_TFrecord_From_pickles.py +0 -0
  22. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Utils/Create_image_tokenizer.py +0 -0
  23. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/effnetv2_configs.py +0 -0
  24. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/effnetv2_model.py +0 -0
  25. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/hparams.py +0 -0
  26. {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/utils.py +0 -0
  27. {decimer-2.7.1 → decimer-2.8.0}/LICENSE +0 -0
  28. {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/dependency_links.txt +0 -0
  29. {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/entry_points.txt +0 -0
  30. {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/top_level.txt +0 -0
  31. {decimer-2.7.1 → decimer-2.8.0}/setup.cfg +0 -0
@@ -1,8 +1,6 @@
1
1
  # EfficientNet-V2 config
2
2
  import tensorflow as tf
3
3
 
4
- import DECIMER.efficientnetv2
5
- from DECIMER.efficientnetv2 import effnetv2_configs
6
4
  from DECIMER.efficientnetv2 import effnetv2_model
7
5
 
8
6
  BATCH_SIZE_DEBUG = 2
@@ -0,0 +1,31 @@
1
+ # -*- coding: utf-8 -*-
2
+ """
3
+ DECIMER V2.7.2 Python Package
4
+ =============================
5
+
6
+ Deep lEarning for Chemical ImagE Recognition (DECIMER) project
7
+ was launched to address the Optical Chemical Structure Recognition (OCSR)
8
+ problem using deep learning based methods,
9
+ providing an automated open-source software solution.
10
+
11
+ Typical usage example::
12
+
13
+ from decimer import predict_SMILES
14
+
15
+ # Chemical depiction to SMILES translation
16
+ image_path = "path/to/imagefile"
17
+ SMILES = predict_SMILES(image_path)
18
+ print(SMILES)
19
+
20
+ For comments, bug reports, or feature ideas,
21
+ please raise an issue on the GitHub repository.
22
+ """
23
+
24
+ __version__ = "2.7.2"
25
+
26
+ __all__ = [
27
+ "DECIMER",
28
+ ]
29
+
30
+
31
+ from .decimer import predict_SMILES
@@ -0,0 +1,148 @@
1
+ import tensorflow as tf
2
+
3
+ import DECIMER.Efficient_Net_encoder as Efficient_Net_encoder
4
+ import DECIMER.Transformer_decoder as Transformer_decoder
5
+
6
+ TARGET_DTYPE = tf.float32
7
+
8
+
9
+ class Config:
10
+ """Configuration class."""
11
+
12
+ def __init__(
13
+ self,
14
+ ):
15
+ self.encoder_config = {}
16
+ self.transformer_config = {}
17
+ self.lr_config = {}
18
+
19
+ def initialize_encoder_config(
20
+ self,
21
+ image_embedding_dim,
22
+ preprocessing_fn,
23
+ backbone_fn,
24
+ image_shape,
25
+ do_permute=False,
26
+ pretrained_weights=None,
27
+ ):
28
+ """This functions initializes the Efficient-Net V2 encoder with user
29
+ defined configurations.
30
+
31
+ Args:
32
+ image_embedding_dim (int): Embedding dimension of the input image
33
+ preprocessing_fn (method): Efficient Net preprocessing function for input image
34
+ backbone_fn (method): Calls Efficient-Net V2 as backbone for encoder
35
+ image_shape (int): Shape of the input image
36
+ do_permute (bool, optional): . Defaults to False.
37
+ pretrained_weights (keras weights, optional): Use pretrained efficient net weights or not. Defaults to None.
38
+ """
39
+ self.encoder_config = dict(
40
+ image_embedding_dim=image_embedding_dim,
41
+ preprocessing_fn=preprocessing_fn,
42
+ backbone_fn=backbone_fn,
43
+ image_shape=image_shape,
44
+ do_permute=do_permute,
45
+ pretrained_weights=pretrained_weights,
46
+ )
47
+
48
+ def initialize_transformer_config(
49
+ self,
50
+ vocab_len,
51
+ max_len,
52
+ n_transformer_layers,
53
+ transformer_d_dff,
54
+ transformer_n_heads,
55
+ image_embedding_dim,
56
+ rate=0.1,
57
+ ):
58
+ """This functions initializes the Transformer model as decoder with
59
+ user defined configurations.
60
+
61
+ Args:
62
+ vocab_len (int): Total number of words in the input vocabulary
63
+ max_len (int): Maximum length of the string found on the training dataset
64
+ n_transformer_layers (int): Number of layers present in the transformer model
65
+ transformer_d_dff (int): Transformer feed forward upwards projection size
66
+ transformer_n_heads (int): Number of heads present in the transformer model
67
+ image_embedding_dim (int): Total number of dimension the image gets embedded
68
+ dropout_rate (float, optional): Fraction of the input units to drop. Defaults to 0.1.
69
+ """
70
+ self.transformer_config = dict(
71
+ num_layers=n_transformer_layers,
72
+ d_model=image_embedding_dim,
73
+ num_heads=transformer_n_heads,
74
+ dff=transformer_d_dff,
75
+ target_vocab_size=vocab_len,
76
+ max_len=max_len,
77
+ rate=0.1,
78
+ )
79
+
80
+ def initialize_lr_config(self, warm_steps, n_epochs):
81
+ """This function sets the configuration to initialize learning rate.
82
+
83
+ Args:
84
+ warm_steps (int): Number of steps The learning rate is increased
85
+ n_epochs (int): Number of epochs
86
+ """
87
+ self.lr_config = dict(
88
+ warm_steps=warm_steps,
89
+ n_epochs=n_epochs,
90
+ )
91
+
92
+
93
+ class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
94
+ """Custom schedule for learning rate used during training.
95
+
96
+ Args:
97
+ tf (_type_): keras learning rate schedule
98
+ """
99
+
100
+ def __init__(self, d_model, warmup_steps=4000):
101
+ super(CustomSchedule, self).__init__()
102
+
103
+ self.d_model = d_model
104
+ self.d_model = tf.cast(self.d_model, tf.float32)
105
+
106
+ self.warmup_steps = warmup_steps
107
+
108
+ def __call__(self, step):
109
+ arg1 = tf.math.rsqrt(step)
110
+ arg2 = step * (self.warmup_steps**-1.5)
111
+
112
+ return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
113
+
114
+
115
+ def prepare_models(encoder_config, transformer_config, replica_batch_size, verbose=0):
116
+ """This function is used to initiate the Encoder and the Transformer with
117
+ appropriate configs set by the user. After initiating the models this
118
+ function returns the Encoder,Transformer and the optimizer.
119
+
120
+ Args:
121
+ encoder_config ([type]): Encoder configuration set by user in the config class.
122
+ transformer_config ([type]): Transformer configuration set by user in the config class.
123
+ replica_batch_size ([type]): Per replica batch size set by user(during distributed training).
124
+ verbose (int, optional): Defaults to 0.
125
+
126
+ Returns:
127
+ [type]: Optimizer, Encoder model and the Transformer
128
+ """
129
+
130
+ # Instantiate an optimizer
131
+ optimizer = tf.keras.optimizers.Adam(learning_rate=0.00051)
132
+
133
+ # Instantiate the encoder model
134
+
135
+ encoder = Efficient_Net_encoder.Encoder(**encoder_config)
136
+
137
+ # Instantiate the decoder model
138
+ transformer = Transformer_decoder.Decoder(**transformer_config)
139
+
140
+ # Show the model architectures and plot the learning rate
141
+ if verbose:
142
+ print("\nEncoder model\n")
143
+ print(encoder.summary())
144
+
145
+ print("\nTransformer model\n")
146
+ print(transformer.summary())
147
+
148
+ return optimizer, encoder, transformer
@@ -9,7 +9,7 @@ import numpy as np
9
9
  import pystow
10
10
  import tensorflow as tf
11
11
 
12
- import DECIMER.config as config
12
+ import DECIMER.pre_process as pre_process
13
13
  import DECIMER.utils as utils
14
14
 
15
15
  # Silence tensorflow model loading warnings.
@@ -18,6 +18,46 @@ logging.getLogger("absl").setLevel("ERROR")
18
18
  # Silence tensorflow errors - not recommended if your model is not working properly.
19
19
  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
20
20
 
21
+
22
+ # ============================================================================
23
+ # Keras 2 to Keras 3 Compatibility Layer
24
+ # ============================================================================
25
+ class KerasCompatibilityUnpickler(pickle.Unpickler):
26
+ """Custom unpickler to handle Keras 2.x tokenizers in Keras 3 environment.
27
+
28
+ This handles the module path changes between Keras 2 and Keras 3:
29
+ - keras.preprocessing.text -> tensorflow.keras.preprocessing.text
30
+ """
31
+
32
+ def find_class(self, module, name):
33
+ # Redirect old Keras 2.x imports to TensorFlow Keras
34
+ if module.startswith("keras."):
35
+ module = module.replace("keras.", "tensorflow.keras.")
36
+ return super().find_class(module, name)
37
+
38
+
39
+ def load_tokenizer(tokenizer_path: str) -> object:
40
+ """Load tokenizer with Keras 2/3 compatibility.
41
+
42
+ Args:
43
+ tokenizer_path (str): Path to the pickled tokenizer file
44
+
45
+ Returns:
46
+ tokenizer: Loaded tokenizer object
47
+ """
48
+ try:
49
+ # Try standard pickle load first (for Keras 3 tokenizers)
50
+ with open(tokenizer_path, "rb") as f:
51
+ return pickle.load(f)
52
+ except ModuleNotFoundError as e:
53
+ if "keras.preprocessing" in str(e):
54
+ # Use compatibility unpickler for Keras 2.x tokenizers
55
+ with open(tokenizer_path, "rb") as f:
56
+ return KerasCompatibilityUnpickler(f).load()
57
+ else:
58
+ raise
59
+
60
+
21
61
  # Set the absolute path
22
62
  HERE = os.path.dirname(os.path.abspath(__file__))
23
63
 
@@ -59,7 +99,7 @@ def get_models(model_urls: dict):
59
99
  tokenizer_path = os.path.join(
60
100
  model_paths["DECIMER"], "assets", "tokenizer_SMILES.pkl"
61
101
  )
62
- tokenizer = pickle.load(open(tokenizer_path, "rb"))
102
+ tokenizer = load_tokenizer(tokenizer_path)
63
103
 
64
104
  # Load DECIMER models
65
105
  DECIMER_V2 = tf.saved_model.load(model_paths["DECIMER"])
@@ -123,7 +163,7 @@ def detokenize_output_add_confidence(
123
163
 
124
164
 
125
165
  def predict_SMILES(
126
- image_input: [str, np.ndarray], confidence: bool = False, hand_drawn: bool = False
166
+ image_input: [str, np.ndarray], confidence: bool = False, hand_drawn: bool = False
127
167
  ) -> str:
128
168
  """Predicts SMILES representation of a molecule depicted in the given image.
129
169
 
@@ -135,7 +175,7 @@ def predict_SMILES(
135
175
  Returns:
136
176
  str: SMILES representation of the molecule in the input image, optionally with confidence values.
137
177
  """
138
- chemical_structure = config.decode_image(image_input)
178
+ chemical_structure = pre_process.decode_image(image_input)
139
179
 
140
180
  model = DECIMER_Hand_drawn if hand_drawn else DECIMER_V2
141
181
  predicted_tokens, confidence_values = model(tf.constant(chemical_structure))
@@ -1,24 +1,14 @@
1
- # Network configuration file
2
1
  import io
3
- import zipfile
4
- from pathlib import Path
5
-
6
2
  import cv2
7
3
  import efficientnet.tfkeras as efn
8
- import numpy as np
9
- import pystow
10
4
  import tensorflow as tf
5
+ import numpy as np
6
+
11
7
  from PIL import Image
12
8
  from PIL import ImageEnhance
13
9
  from pillow_heif import register_heif_opener
14
10
  from typing import Union
15
11
 
16
- import DECIMER.Efficient_Net_encoder as Efficient_Net_encoder
17
- import DECIMER.Transformer_decoder as Transformer_decoder
18
-
19
-
20
- TARGET_DTYPE = tf.float32
21
-
22
12
 
23
13
  def resize_byratio(image):
24
14
  """This function takes a Pillow Image object and will resize the image by
@@ -106,6 +96,7 @@ def remove_transparent(image: Union[str, np.ndarray]) -> Image.Image:
106
96
  Returns:
107
97
  PIL.Image.Image: The image with transparency removed.
108
98
  """
99
+
109
100
  def process_image(png: Image.Image) -> Image.Image:
110
101
  """
111
102
  Helper function to remove transparency from a single image.
@@ -183,14 +174,14 @@ def increase_contrast(image):
183
174
  Returns: PIL.Image
184
175
  """
185
176
 
186
- # Get brightness range
187
- min = np.min(image)
188
- max = np.max(image)
177
+ # Get brightness range - cast to int to avoid uint8 overflow
178
+ min_val = int(np.min(image))
179
+ max_val = int(np.max(image))
189
180
 
190
181
  # Use LUT to convert image values
191
182
  LUT = np.zeros(256, dtype=np.uint8)
192
- LUT[min : max + 1] = np.linspace(
193
- start=0, stop=255, num=(max - min) + 1, endpoint=True, dtype=np.uint8
183
+ LUT[min_val : max_val + 1] = np.linspace(
184
+ start=0, stop=255, num=(max_val - min_val) + 1, endpoint=True, dtype=np.uint8
194
185
  )
195
186
 
196
187
  # Apply LUT and return image
@@ -250,175 +241,3 @@ def decode_image(image_path: Union[str, np.ndarray]):
250
241
  img = tf.image.resize(img, (512, 512), method="gaussian", antialias=True)
251
242
  img = efn.preprocess_input(img)
252
243
  return img
253
-
254
-
255
- class Config:
256
- """Configuration class."""
257
-
258
- def __init__(
259
- self,
260
- ):
261
- self.encoder_config = {}
262
- self.transformer_config = {}
263
- self.lr_config = {}
264
-
265
- def initialize_encoder_config(
266
- self,
267
- image_embedding_dim,
268
- preprocessing_fn,
269
- backbone_fn,
270
- image_shape,
271
- do_permute=False,
272
- pretrained_weights=None,
273
- ):
274
- """This functions initializes the Efficient-Net V2 encoder with user
275
- defined configurations.
276
-
277
- Args:
278
- image_embedding_dim (int): Embedding dimension of the input image
279
- preprocessing_fn (method): Efficient Net preprocessing function for input image
280
- backbone_fn (method): Calls Efficient-Net V2 as backbone for encoder
281
- image_shape (int): Shape of the input image
282
- do_permute (bool, optional): . Defaults to False.
283
- pretrained_weights (keras weights, optional): Use pretrained efficient net weights or not. Defaults to None.
284
- """
285
- self.encoder_config = dict(
286
- image_embedding_dim=image_embedding_dim,
287
- preprocessing_fn=preprocessing_fn,
288
- backbone_fn=backbone_fn,
289
- image_shape=image_shape,
290
- do_permute=do_permute,
291
- pretrained_weights=pretrained_weights,
292
- )
293
-
294
- def initialize_transformer_config(
295
- self,
296
- vocab_len,
297
- max_len,
298
- n_transformer_layers,
299
- transformer_d_dff,
300
- transformer_n_heads,
301
- image_embedding_dim,
302
- rate=0.1,
303
- ):
304
- """This functions initializes the Transformer model as decoder with
305
- user defined configurations.
306
-
307
- Args:
308
- vocab_len (int): Total number of words in the input vocabulary
309
- max_len (int): Maximum length of the string found on the training dataset
310
- n_transformer_layers (int): Number of layers present in the transformer model
311
- transformer_d_dff (int): Transformer feed forward upwards projection size
312
- transformer_n_heads (int): Number of heads present in the transformer model
313
- image_embedding_dim (int): Total number of dimension the image gets embedded
314
- dropout_rate (float, optional): Fraction of the input units to drop. Defaults to 0.1.
315
- """
316
- self.transformer_config = dict(
317
- num_layers=n_transformer_layers,
318
- d_model=image_embedding_dim,
319
- num_heads=transformer_n_heads,
320
- dff=transformer_d_dff,
321
- target_vocab_size=vocab_len,
322
- max_len=max_len,
323
- rate=0.1,
324
- )
325
-
326
- def initialize_lr_config(self, warm_steps, n_epochs):
327
- """This function sets the configuration to initialize learning rate.
328
-
329
- Args:
330
- warm_steps (int): Number of steps The learning rate is increased
331
- n_epochs (int): Number of epochs
332
- """
333
- self.lr_config = dict(
334
- warm_steps=warm_steps,
335
- n_epochs=n_epochs,
336
- )
337
-
338
-
339
- class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
340
- """Custom schedule for learning rate used during training.
341
-
342
- Args:
343
- tf (_type_): keras learning rate schedule
344
- """
345
-
346
- def __init__(self, d_model, warmup_steps=4000):
347
- super(CustomSchedule, self).__init__()
348
-
349
- self.d_model = d_model
350
- self.d_model = tf.cast(self.d_model, tf.float32)
351
-
352
- self.warmup_steps = warmup_steps
353
-
354
- def __call__(self, step):
355
- arg1 = tf.math.rsqrt(step)
356
- arg2 = step * (self.warmup_steps**-1.5)
357
-
358
- return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
359
-
360
-
361
- def prepare_models(encoder_config, transformer_config, replica_batch_size, verbose=0):
362
- """This function is used to initiate the Encoder and the Transformer with
363
- appropriate configs set by the user. After initiating the models this
364
- function returns the Encoder,Transformer and the optimizer.
365
-
366
- Args:
367
- encoder_config ([type]): Encoder configuration set by user in the config class.
368
- transformer_config ([type]): Transformer configuration set by user in the config class.
369
- replica_batch_size ([type]): Per replica batch size set by user(during distributed training).
370
- verbose (int, optional): Defaults to 0.
371
-
372
- Returns:
373
- [type]: Optimizer, Encoder model and the Transformer
374
- """
375
-
376
- # Instantiate an optimizer
377
- optimizer = tf.keras.optimizers.Adam(learning_rate=0.00051)
378
-
379
- # Instantiate the encoder model
380
-
381
- encoder = Efficient_Net_encoder.Encoder(**encoder_config)
382
-
383
- # Instantiate the decoder model
384
- transformer = Transformer_decoder.Decoder(**transformer_config)
385
-
386
- # Show the model architectures and plot the learning rate
387
- if verbose:
388
- print("\nEncoder model\n")
389
- print(encoder.summary())
390
-
391
- print("\nTransformer model\n")
392
- print(transformer.summary())
393
-
394
- return optimizer, encoder, transformer
395
-
396
-
397
- # Downloads the model and unzips the file downloaded, if the model is not present on the working directory.
398
- def download_trained_weights(model_url: str, model_path: str, verbose=1):
399
- """This function downloads the trained models and tokenizers to a default
400
- location. After downloading the zipped file the function unzips the file
401
- automatically. If the model exists on the default location this function
402
- will not work.
403
-
404
- Args:
405
- model_url (str): trained model url for downloading.
406
- model_path (str): model default path to download.
407
-
408
- Returns:
409
- path (str): downloaded model.
410
- """
411
- # Download trained models
412
- if verbose > 0:
413
- print("Downloading trained model to " + str(model_path))
414
- model_path = pystow.ensure("DECIMER-V2", url=model_url)
415
- if verbose > 0:
416
- print(model_path)
417
- print("... done downloading trained model!")
418
-
419
- with zipfile.ZipFile(model_path.as_posix(), "r") as zip_ref:
420
- zip_ref.extractall(model_path.parent.as_posix())
421
-
422
- # Delete zipfile after downloading
423
- if Path(model_path).exists():
424
- Path(model_path).unlink()
@@ -0,0 +1,143 @@
1
+ import os
2
+ import re
3
+ import shutil
4
+
5
+ import pystow
6
+ import zipfile
7
+ from pathlib import Path
8
+
9
+ pattern = "R([0-9]*)|X([0-9]*)|Y([0-9]*)|Z([0-9]*)"
10
+ add_space_re = r"^(\W+)|(\W+)$"
11
+
12
+
13
+ def split_and_modify_atoms(SMILES):
14
+ splitted_SMILES = list(SMILES)
15
+ modified_SMILES = re.sub(r"\s+(?=[a-z])", "", " ".join(map(str, splitted_SMILES)))
16
+ return modified_SMILES
17
+
18
+
19
+ def replacer(match):
20
+ part = match.group(0)
21
+ text = (
22
+ part.replace("1", "!")
23
+ .replace("2", "$")
24
+ .replace("3", "^")
25
+ .replace("4", "<")
26
+ .replace("5", ">")
27
+ .replace("6", "?")
28
+ .replace("7", "£")
29
+ .replace("8", "¢")
30
+ .replace("9", "€")
31
+ .replace("0", "§")
32
+ )
33
+ return text
34
+
35
+
36
+ def add_space(match_):
37
+ if match_.group(1) is not None:
38
+ return "{} ".format(match_.group(1))
39
+ else:
40
+ return " {}".format(match_.group(2))
41
+
42
+
43
+ def encoder(SMILES):
44
+ replaced_SMILES = re.sub(pattern, replacer, SMILES)
45
+ splitted_SMILES = split_and_modify_atoms(replaced_SMILES)
46
+ modified_SMILES = " ".join(
47
+ [re.sub(add_space_re, add_space, word) for word in splitted_SMILES.split()]
48
+ )
49
+ return modified_SMILES
50
+
51
+
52
+ def decoder(predictions):
53
+ modified = (
54
+ predictions.replace("!", "1")
55
+ .replace("$", "2")
56
+ .replace("^", "3")
57
+ .replace("<", "4")
58
+ .replace(">", "5")
59
+ .replace("?", "6")
60
+ .replace("£", "7")
61
+ .replace("¢", "8")
62
+ .replace("€", "9")
63
+ .replace("§", "0")
64
+ )
65
+ return modified
66
+
67
+
68
+ # Downloads the model and unzips the file downloaded, if the model is not present on the working directory.
69
+ def download_trained_weights(model_url: str, model_path: str, verbose=1):
70
+ """This function downloads the trained models and tokenizers to a default
71
+ location. After downloading the zipped file the function unzips the file
72
+ automatically. If the model exists on the default location this function
73
+ will not work.
74
+
75
+ Args:
76
+ model_url (str): trained model url for downloading.
77
+ model_path (str): model default path to download.
78
+
79
+ Returns:
80
+ path (str): downloaded model.
81
+ """
82
+ # Download trained models
83
+ if verbose > 0:
84
+ print("Downloading trained model to " + str(model_path))
85
+ model_path = pystow.ensure("DECIMER-V2", url=model_url)
86
+ if verbose > 0:
87
+ print(model_path)
88
+ print("... done downloading trained model!")
89
+
90
+ with zipfile.ZipFile(model_path.as_posix(), "r") as zip_ref:
91
+ zip_ref.extractall(model_path.parent.as_posix())
92
+
93
+ # Delete zipfile after downloading
94
+ if Path(model_path).exists():
95
+ Path(model_path).unlink()
96
+
97
+
98
+ def ensure_models(default_path: str, model_urls: dict) -> dict:
99
+ """Function to ensure models are present locally.
100
+
101
+ Convenient function to ensure model downloads before usage.
102
+ Models are re-downloaded if the URL changes (e.g., new Zenodo record).
103
+
104
+ Args:
105
+ default_path (str): Default path for model data
106
+ model_urls (dict): Dictionary containing model names as keys and their corresponding URLs as values
107
+
108
+ Returns:
109
+ dict: A dictionary containing model names as keys and their local paths as values
110
+ """
111
+ model_paths = {}
112
+
113
+ for model_name, model_url in model_urls.items():
114
+ model_path = os.path.join(default_path, f"{model_name}_model")
115
+ saved_model_file = os.path.join(model_path, "saved_model.pb")
116
+ version_file = os.path.join(model_path, ".model_url")
117
+
118
+ # Check if model needs to be downloaded:
119
+ # 1. saved_model.pb doesn't exist, OR
120
+ # 2. The URL has changed (model was updated)
121
+ needs_download = not os.path.exists(saved_model_file)
122
+
123
+ if not needs_download and os.path.exists(version_file):
124
+ with open(version_file, "r") as f:
125
+ cached_url = f.read().strip()
126
+ if cached_url != model_url:
127
+ needs_download = True
128
+ print(f"Model {model_name} has been updated, re-downloading...")
129
+
130
+ if needs_download:
131
+ # Clean up incomplete/corrupted model directory if it exists
132
+ if os.path.exists(model_path):
133
+ shutil.rmtree(model_path)
134
+ download_trained_weights(model_url, default_path)
135
+
136
+ # Store the URL used for this download
137
+ os.makedirs(model_path, exist_ok=True)
138
+ with open(version_file, "w") as f:
139
+ f.write(model_url)
140
+
141
+ model_paths[model_name] = model_path
142
+
143
+ return model_paths